| import torch |
| from transformers import AutoTokenizer |
| from fin_tinybert_pytorch import TinyFinBERTRegressor |
|
|
|
|
| class InferenceAPI: |
| def __init__(self): |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.model = TinyFinBERTRegressor() |
| self.model.load_state_dict(torch.load("./saved_model/pytorch_model.bin", map_location=self.device)) |
| self.model.to(self.device) |
| self.model.eval() |
| self.tokenizer = AutoTokenizer.from_pretrained("./saved_model") |
|
|
| def __call__(self, inputs): |
| if not isinstance(inputs, list): |
| inputs = [inputs] |
|
|
| results = [] |
| for text in inputs: |
| encoded = self.tokenizer(text, return_tensors="pt", truncation=True, |
| padding='max_length', max_length=128) |
| encoded = {k: v.to(self.device) for k, v in encoded.items() if k != "token_type_ids"} |
|
|
| with torch.no_grad(): |
| score = self.model(**encoded)["score"].item() |
|
|
| sentiment = "positive" if score > 0.3 else "negative" if score < -0.3 else "neutral" |
|
|
| results.append({ |
| "label": sentiment, |
| "score": round(score, 4) |
| }) |
|
|
| if len(results) == 1: |
| return results[0] |
| return results |
