Update model.py

model.py

@@ -1,57 +1,20 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-import numpy as np
-
-# Load model and tokenizer
-model_name = "SamLowe/roberta-base-go_emotions"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-
-#Define distress weights for emotions 
-DISTRESS_WEIGHTS = {
-    "grief": 1.0,
-    "sadness": 0.9,
-    "fear": 0.9,
-    "disgust": 0.9,
-    "anger": 0.9,
-    "remorse": 0.8,
-    "disappointment": 0.8,
-    "nervousness": 0.8,
-    "disapproval": 0.7,
-    "embarrassment": 0.7,
-    "annoyance": 0.6,
-    "confusion": 0.6,
-    "surprise": 0.4,
-    "desire": 0.4,
-    "love": 0.3,
-    "excitement": 0.3,
-    "pride": 0.3,
-    "optimism": 0.3,
-    "admiration": 0.2,
-    "gratitude": 0.2,
-    "relief": 0.2,
-    "joy": 0.2,
-    "amusement": 0.2,
-    "neutral": 0.1,
-}
-# Labels from GoEmotions
-EMOTION_LABELS = tokenizer.convert_ids_to_tokens(list(range(model.config.num_labels)))
-
-# Ensure the model is in evaluation mode
-model.eval()    
-
-def predict_emotions(text: str):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True)
-    outputs = model(**inputs)
-    probs = torch.sigmoid(outputs.logits)[0].detach().numpy()
-
-    threshold = 0.3  # You can tune this
-    predicted = {label: float(prob) for label, prob in zip(model.config.id2label.values(), probs) if prob > threshold}
-    return predicted
-
-def calculate_distress(emotions: dict):
-    distress_score = sum(
-        emotions.get(emotion, 0) * DISTRESS_WEIGHTS.get(emotion, 0)
-        for emotion in emotions
-    )
-    return round(distress_score, 3)
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
+class EmotionModel:
+    def __init__(self):
+        self.model_name = "SamLowe/roberta-base-go_emotions"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+        self.labels = self.model.config.id2label
+
+    def predict(self, text):
+        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            logits = self.model(**inputs).logits
+        probs = torch.sigmoid(logits)[0]
+        return {
+            self.labels[i]: float(probs[i])
+            for i in range(len(probs)) if probs[i] > 0.3
+        }
+