Upload handler.py

handler.py

@@ -1,37 +1,33 @@
 import torch
 from transformers import BertTokenizer, BertForSequenceClassification
+from typing import List
 
 class SentimentModel:
     def __init__(self, model_path):
         self.tokenizer = BertTokenizer.from_pretrained(model_path)
         self.model = BertForSequenceClassification.from_pretrained(model_path)
         self.model.eval()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
 
-    def predict(self, texts):
-        max_length = 250
-        batch_size = 64
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(device)
-
-        probabilities = []
+    def predict(self, texts: List[str]) -> List[float]:
+        max_length = 128
         encodings = self.tokenizer(texts, truncation=True, padding=True, max_length=max_length, return_tensors="pt")
-        input_ids = encodings['input_ids']
-        attention_masks = encodings['attention_mask']
+        input_ids = encodings['input_ids'].to(self.device)
+        attention_masks = encodings['attention_mask'].to(self.device)
         
-        dataset = torch.utils.data.TensorDataset(input_ids, attention_masks)
-        dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
-
         with torch.no_grad():
-            for batch in dataloader:
-                input_ids_batch, attention_masks_batch = batch
-                input_ids_batch = input_ids_batch.to(device)
-                attention_masks_batch = attention_masks_batch.to(device)
-                
-                outputs = self.model(input_ids_batch, attention_mask=attention_masks_batch)
-                logits = outputs.logits
-                
-                probs = torch.nn.functional.softmax(logits, dim=-1)
-                strong_probs = probs[:, 1].tolist()
-                probabilities.extend(strong_probs)
+            outputs = self.model(input_ids, attention_mask=attention_masks)
+            logits = outputs.logits
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            strong_probs = probs[:, 1].tolist()  # Probability of the "strong" class (class 1)
+
+        return strong_probs
+
+# Initialize the model
+model = SentimentModel(".")
 
-        return probabilities
+def infer(inputs):
+    texts = [inputs[i]['data'] for i in range(len(inputs))]
+    probabilities = model.predict(texts)
+    return [{"sentiment intensity": prob} for prob in probabilities]