Update emotion_detection.py

emotion_detection.py

@@ -3,6 +3,7 @@ from transformers_interpret import SequenceClassificationExplainer
 import torch
 import pandas as pd
 
+
 class EmotionDetection:
     """
     Emotion Detection on text data.
@@ -24,10 +25,12 @@ class EmotionDetection:
         Parameters:
             text (str): The user input string to emotion justification
         Returns:
-            html (html): html object for plotting emotion prediction justification
+            html (hmtl): html object for plotting emotion prediction justification
         """
+
         word_attributions = self.explainer(text)
         html = self.explainer.visualize("example.html")
+
         return html
 
     def classify(self, text):
@@ -43,28 +46,9 @@ class EmotionDetection:
         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')
 
-        # 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])
-
-        preds = pd.Series(filtered_probs, index=filtered_labels, name='Predicted Probability')
         return preds
 
     def run(self, text):
@@ -74,8 +58,10 @@ class EmotionDetection:
             text (str): The user input string to perform emotion classification on
         Returns:
             predictions (str): The predicted probabilities for emotion classes
-            html (html): html object for plotting emotion prediction justification
+            html (hmtl): html object for plotting emotion prediction justification
         """
+
         preds = self.classify(text)
         html = self.justify(text)
-        return preds, html
+
+        return preds, html