Update emotion_detection.py

emotion_detection.py

@@ -1,16 +1,16 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from transformers_interpret import SequenceClassificationExplainer
 import torch
-import os
+import pandas as pd
 
 
 class EmotionDetection:
     """
     Emotion Detection on text data.
     Attributes:
-        tokenizer: Hugging Face Tokenizer instance
-        model: Hugging Face Sequence Classification model
-        explainer: SequenceClassificationExplainer instance for model interpretability
+        tokenizer: An instance of Hugging Face Tokenizer
+        model: An instance of Hugging Face Model
+        explainer: An instance of SequenceClassificationExplainer from Transformers interpret
     """
 
     def __init__(self):
@@ -19,75 +19,49 @@ class EmotionDetection:
         self.model = AutoModelForSequenceClassification.from_pretrained(hub_location)
         self.explainer = SequenceClassificationExplainer(self.model, self.tokenizer)
 
-        # Emoji map for friendly display
-        self.emoji_map = {
-            "joy": "😊",
-            "anger": "😠",
-            "optimism": "😎",
-            "sadness": "😢"
-        }
-
-        # Simple explanation map
-        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):
         """
-        Generate HTML visualization of word attributions for emotion.
+        Get html annotation for displaying emotion justification over text.
         Parameters:
-            text (str): Input text
+            text (str): The user input string to emotion justification
         Returns:
-            html (str): HTML string with justification visualization
+            html (hmtl): html object for plotting emotion prediction justification
         """
-        word_attributions = self.explainer(text)
-        html_path = "justification_output.html"
-        self.explainer.visualize(html_path)
 
-        # Read from file
-        with open(html_path, "r", encoding="utf-8") as f:
-            html = f.read()
+        word_attributions = self.explainer(text)
+        html = self.explainer.visualize("example.html")
 
-        # Clean up file
-        os.remove(html_path)
         return html
 
     def classify(self, text):
         """
-        Classify the main emotion in the input text.
+        Recognize Emotion in text.
         Parameters:
-            text (str): Input text
+            text (str): The user input string to perform emotion classification on
         Returns:
-            result (str): Friendly output with emoji and short explanation
+            predictions (str): The predicted probabilities for emotion classes
         """
-        tokens = self.tokenizer.encode_plus(text, return_tensors='pt')
+
+        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())
-        max_index = probs.argmax()
-        emotion = labels[max_index]
-        confidence = probs[max_index]
-
-        emoji = self.emoji_map.get(emotion, "")
-        explanation = self.explanation_map.get(emotion, "")
+        preds = pd.Series(probs, index=labels, name='Predicted Probability')
 
-        result = f"{emoji} **{emotion.capitalize()}** ({confidence:.1%})\n{explanation}"
-        return result
+        return preds
 
     def run(self, text):
         """
-        Perform classification and justification.
+        Classify and Justify Emotion in text.
         Parameters:
-            text (str): Input text
+            text (str): The user input string to perform emotion classification on
         Returns:
-            result (str): Emotion classification result
-            html (str): Justification HTML
+            predictions (str): The predicted probabilities for emotion classes
+            html (hmtl): html object for plotting emotion prediction justification
         """
-        result = self.classify(text)
+
+        preds = self.classify(text)
         html = self.justify(text)
-        return result, html
+
+        return preds, html