# import torch
# from transformers import AutoTokenizer
# from fin_tinybert_pytorch import TinyFinBERTRegressor  # You may need to rename or include this class here
#
# # Load model
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model = TinyFinBERTRegressor()
# model.load_state_dict(torch.load("./saved_model/pytorch_model.bin", map_location=device))
# model.to(device)
# model.eval()
#
# tokenizer = AutoTokenizer.from_pretrained("./saved_model")
#
# def predict(texts):
#     if isinstance(texts, str):
#         texts = [texts]
#
#     results = []
#     for text in texts:
#         inputs = tokenizer(text, return_tensors="pt", truncation=True, padding='max_length', max_length=128)
#         inputs = {k: v.to(device) for k, v in inputs.items() if k != "token_type_ids"}
#         with torch.no_grad():
#             score = model(**inputs)["score"].item()
#         sentiment = "positive" if score > 0.3 else "negative" if score < -0.3 else "neutral"
#         results.append({"text": text, "score": score, "sentiment": sentiment})
#     return results


import torch
from transformers import AutoTokenizer
from fin_tinybert_pytorch import TinyFinBERTRegressor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TinyFinBERTRegressor()
model.load_state_dict(torch.load("./saved_model/pytorch_model.bin", map_location=device))
model.to(device)
model.eval()

tokenizer = AutoTokenizer.from_pretrained("./saved_model")


def pipeline(text):
    if not isinstance(text, str):
        raise ValueError("Input must be a string")

    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding='max_length', max_length=128)
    inputs = {k: v.to(device) for k, v in inputs.items() if k != "token_type_ids"}

    with torch.no_grad():
        score = model(**inputs)["score"].item()

    sentiment = "positive" if score > 0.3 else "negative" if score < -0.3 else "neutral"

    return [{
        "label": sentiment,
        "score": round(score, 4)
    }]

#
# if __name__ == "__main__":
#     texts = [
#         "The stock price soared after the earnings report.",
#         "The company reported significant losses this quarter.",
#         "There was no noticeable change in performance."
#     ]
#
#     predictions = pipeline("The stock price soared after the earnings report.")[0]
#     print(f"sentiment: {predictions['label']}, score: {predictions['score']}")