Update emotion_detection.py

emotion_detection.py

@@ -3,7 +3,6 @@ from transformers_interpret import SequenceClassificationExplainer
 import torch
 import pandas as pd
 
-
 class EmotionDetection:
     """
     Emotion Detection on text data.
@@ -25,51 +24,48 @@ class EmotionDetection:
         Parameters:
             text (str): The user input string to emotion justification
         Returns:
-            html (hmtl): html object for plotting emotion prediction justification
+            html (html): html object for plotting emotion prediction justification
         """
-
         word_attributions = self.explainer(text)
         html = self.explainer.visualize("example.html")
-
         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
-    """
-
-    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()
+        """
+        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
+        """
 
-    # Original labels from model
-    original_labels = list(self.model.config.id2label.values())
+        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()
 
-    # 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 ✨'
-    }
+        # Original labels from model
+        original_labels = list(self.model.config.id2label.values())
 
-    filtered_probs = []
-    filtered_labels = []
+        # 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 ✨'
+        }
 
-    for label, prob in zip(original_labels, probs):
-        if label in desired_labels:
-            filtered_probs.append(prob)
-            filtered_labels.append(custom_labels[label])
+        filtered_probs = []
+        filtered_labels = []
 
-    preds = pd.Series(filtered_probs, index=filtered_labels, name='Predicted Probability')
+        for label, prob in zip(original_labels, probs):
+            if label in desired_labels:
+                filtered_probs.append(prob)
+                filtered_labels.append(custom_labels[label])
 
-    return preds
+        preds = pd.Series(filtered_probs, index=filtered_labels, name='Predicted Probability')
+        return preds
 
     def run(self, text):
         """
@@ -78,10 +74,8 @@ class EmotionDetection:
             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
+            html (html): html object for plotting emotion prediction justification
         """
-
         preds = self.classify(text)
         html = self.justify(text)
-
-        return preds, html
+        return preds, html