Update emotion_detection.py

emotion_detection.py

@@ -1,13 +1,11 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from transformers_interpret import SequenceClassificationExplainer
 import torch
-import pandas as pd
 
 
 class EmotionDetection:
     """
     Emotion Detection on text data.
-
     Attributes:
         tokenizer: An instance of Hugging Face Tokenizer
         model: An instance of Hugging Face Model
@@ -20,55 +18,67 @@ class EmotionDetection:
         self.model = AutoModelForSequenceClassification.from_pretrained(hub_location)
         self.explainer = SequenceClassificationExplainer(self.model, self.tokenizer)
 
+        # Friendly emoji mapping
+        self.emoji_map = {
+            "joy": "😊",
+            "anger": "😠",
+            "optimism": "😎",
+            "sadness": "😢"
+        }
+
+        # Friendly explanation mapping
+        self.explanation_map = {
+            "joy": "The person is happy or excited.",
+            "anger": "The person is upset or angry.",
+            "optimism": "The person is feeling hopeful or positive.",
+            "sadness": "The person is feeling low or unhappy."
+        }
+
     def justify(self, text):
         """
         Get html annotation for displaying emotion justification over text.
-
         Parameters:
-            text (str): The user input string to emotion justification
-
+            text (str): The user input string for emotion justification
         Returns:
-            html (hmtl): html object for plotting emotion prediction justification
+            html (str): html string for plotting emotion prediction justification
         """
-
         word_attributions = self.explainer(text)
-        html = self.explainer.visualize("example.html")
-
+        html = self.explainer.visualize(return_html=True)  # Changed to return HTML string
         return html
 
     def classify(self, text):
         """
-        Recognize Emotion in text.
-
+        Recognize Emotion in text (simplified).
         Parameters:
             text (str): The user input string to perform emotion classification on
-
         Returns:
-            predictions (str): The predicted probabilities for emotion classes
+            result (str): User-friendly emotion label with emoji and explanation
         """
-
-        tokens = self.tokenizer.encode_plus(text, add_special_tokens=False, return_tensors='pt')
+        tokens = self.tokenizer.encode_plus(text, 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')
+        max_index = probs.argmax()
+        emotion = labels[max_index]
+        confidence = probs[max_index]
 
-        return preds
+        emoji = self.emoji_map.get(emotion, "")
+        explanation = self.explanation_map.get(emotion, "")
+
+        result = f"{emoji} **{emotion.capitalize()}** ({confidence:.1%})\n{explanation}"
+        return result
 
     def run(self, text):
         """
         Classify and Justify Emotion in text.
-
         Parameters:
             text (str): The user input string to perform emotion classification on
-
         Returns:
-            predictions (str): The predicted probabilities for emotion classes
-            html (hmtl): html object for plotting emotion prediction justification
+            result (str): Friendly emotion classification output
+            html (str): HTML visualization string for model justification
         """
-
-        preds = self.classify(text)
+        result = self.classify(text)
         html = self.justify(text)
-
-        return preds, html
+        return result, html