File size: 2,434 Bytes
92b3bd3 feb2463 92b3bd3 feb2463 92b3bd3 feb2463 92b3bd3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
# 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']}") |