refactored inference.py

inference.py

@@ -1,8 +1,35 @@
+# 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  # You may need to rename or include this class here
+from fin_tinybert_pytorch import TinyFinBERTRegressor
 
-# 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))
@@ -11,19 +38,24 @@ 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
+
+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 = [
@@ -32,8 +64,5 @@ def predict(texts):
 #         "There was no noticeable change in performance."
 #     ]
 #
-#     predictions = predict(texts)
-#     for pred in predictions:
-#         print(f"Text: {pred['text']}")
-#         print(f"Score: {pred['score']:.3f}")
-#         print(f"Sentiment: {pred['sentiment']}\n")
+#     predictions = pipeline("The stock price soared after the earnings report.")[0]
+#     print(f"sentiment: {predictions['label']}, score: {predictions['score']}")