from openmind import AutoTokenizer, AutoModelForSequenceClassification
import torch
from openmind_hub import snapshot_download
import torch_npu

device = torch.device('npu')
import argparse

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="model_path",
        default='Jinan_AICC/FinTwitBERT',
    )
    args = parser.parse_args()
    return args

args = parse_args()
model_name = args.model_name_or_path

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.to(device)


text = "Nice 9% pre market move for $para, pump my calls Uncle Buffett"


inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
inputs = {key: value.to(device) for key, value in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)

logits = outputs.logits
predicted_class_idx = torch.argmax(logits, dim=-1).item()


labels = ["negative", "neutral", "positive"]
predicted_label = labels[predicted_class_idx]

print(f"Predicted sentiment: {predicted_label}")