Update emotion_detection.py

emotion_detection.py

@@ -34,22 +34,42 @@ class EmotionDetection:
         return html
 
     def classify(self, text):
-        """
-        Recognize Emotion in text.
-        Parameters:
-            text (str): The user input string to perform emotion classification on
-        Returns:
-            predictions (str): The predicted probabilities for emotion classes
-        """
+    """
+    Recognize Emotion in text.
+    Parameters:
+        text (str): The user input string to perform emotion classification on
+    Returns:
+        predictions (str): The predicted probabilities for emotion classes
+    """
+
+    tokens = self.tokenizer.encode_plus(text, add_special_tokens=False, return_tensors='pt')
+    outputs = self.model(**tokens)
+    probs = torch.nn.functional.softmax(outputs[0], dim=-1)
+    probs = probs.mean(dim=0).detach().numpy()
+
+    # Original labels from model
+    original_labels = list(self.model.config.id2label.values())
+
+    # Only keep the 4 specific emotions and map custom names
+    desired_labels = ['joy', 'anger', 'sadness', 'optimism']
+    custom_labels = {
+        'joy': 'Happiness 😊',
+        'anger': 'Anger 😡',
+        'sadness': 'Sadness 😢',
+        'optimism': 'Hopeful ✨'
+    }
+
+    filtered_probs = []
+    filtered_labels = []
+
+    for label, prob in zip(original_labels, probs):
+        if label in desired_labels:
+            filtered_probs.append(prob)
+            filtered_labels.append(custom_labels[label])
 
-        tokens = self.tokenizer.encode_plus(text, add_special_tokens=False, return_tensors='pt')
-        outputs = self.model(**tokens)
-        probs = torch.nn.functional.softmax(outputs[0], dim=-1)
-        probs = probs.mean(dim=0).detach().numpy()
-        labels = list(self.model.config.id2label.values())
-        preds = pd.Series(probs, index=labels, name='Predicted Probability')
+    preds = pd.Series(filtered_probs, index=filtered_labels, name='Predicted Probability')
 
-        return preds
+    return preds
 
     def run(self, text):
         """