Update emotion_detection.py

emotion_detection.py

@@ -1,15 +1,16 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from transformers_interpret import SequenceClassificationExplainer
 import torch
+import os
 
 
 class EmotionDetection:
     """
     Emotion Detection on text data.
     Attributes:
-        tokenizer: An instance of Hugging Face Tokenizer
-        model: An instance of Hugging Face Model
-        explainer: An instance of SequenceClassificationExplainer from Transformers interpret
+        tokenizer: Hugging Face Tokenizer instance
+        model: Hugging Face Sequence Classification model
+        explainer: SequenceClassificationExplainer instance for model interpretability
     """
 
     def __init__(self):
@@ -18,7 +19,7 @@ class EmotionDetection:
         self.model = AutoModelForSequenceClassification.from_pretrained(hub_location)
         self.explainer = SequenceClassificationExplainer(self.model, self.tokenizer)
 
-        # Friendly emoji mapping
+        # Emoji map for friendly display
         self.emoji_map = {
             "joy": "😊",
             "anger": "😠",
@@ -26,7 +27,7 @@ class EmotionDetection:
             "sadness": "😢"
         }
 
-        # Friendly explanation mapping
+        # Simple explanation map
         self.explanation_map = {
             "joy": "The person is happy or excited.",
             "anger": "The person is upset or angry.",
@@ -36,23 +37,31 @@ class EmotionDetection:
 
     def justify(self, text):
         """
-        Get html annotation for displaying emotion justification over text.
+        Generate HTML visualization of word attributions for emotion.
         Parameters:
-            text (str): The user input string for emotion justification
+            text (str): Input text
         Returns:
-            html (str): html string for plotting emotion prediction justification
+            html (str): HTML string with justification visualization
         """
         word_attributions = self.explainer(text)
-        html = self.explainer.visualize(return_html=True)  # Changed to return HTML string
+        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()
+
+        # Clean up file
+        os.remove(html_path)
         return html
 
     def classify(self, text):
         """
-        Recognize Emotion in text (simplified).
+        Classify the main emotion in the input text.
         Parameters:
-            text (str): The user input string to perform emotion classification on
+            text (str): Input text
         Returns:
-            result (str): User-friendly emotion label with emoji and explanation
+            result (str): Friendly output with emoji and short explanation
         """
         tokens = self.tokenizer.encode_plus(text, return_tensors='pt')
         outputs = self.model(**tokens)
@@ -72,12 +81,12 @@ class EmotionDetection:
 
     def run(self, text):
         """
-        Classify and Justify Emotion in text.
+        Perform classification and justification.
         Parameters:
-            text (str): The user input string to perform emotion classification on
+            text (str): Input text
         Returns:
-            result (str): Friendly emotion classification output
-            html (str): HTML visualization string for model justification
+            result (str): Emotion classification result
+            html (str): Justification HTML
         """
         result = self.classify(text)
         html = self.justify(text)