Update app.py

app.py

@@ -7,9 +7,19 @@ import plotly.graph_objects as go
 model_id = "S-4-G-4-R/distilbert-base-uncased-finetuned-emotion"
 classifier = pipeline("text-classification", model=model_id)
 
-# Define emotion labels
+# Define emotion labels mapping (LABEL_0 to LABEL_5)
 EMOTION_LABELS = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
 
+# Label mapping from model output to emotion names
+LABEL_MAPPING = {
+    'LABEL_0': 'sadness',
+    'LABEL_1': 'joy',
+    'LABEL_2': 'love',
+    'LABEL_3': 'anger',
+    'LABEL_4': 'fear',
+    'LABEL_5': 'surprise'
+}
+
 # Emoji mapping for emotions
 EMOTION_EMOJIS = {
     'sadness': '😢',
@@ -30,12 +40,13 @@ def classify_emotion(text):
     # Get predictions
     preds = classifier(text, return_all_scores=True)[0]
     
-    # Create DataFrame
+    # Create DataFrame and map labels to emotion names
     df = pd.DataFrame(preds)
+    df['emotion'] = df['label'].map(LABEL_MAPPING)
     df['score'] = df['score'] * 100  # Convert to percentage
     
     # Add emojis to labels
-    df['display_label'] = df['label'].map(lambda x: f"{EMOTION_EMOJIS.get(x, '')} {x.capitalize()}")
+    df['display_label'] = df['emotion'].map(lambda x: f"{EMOTION_EMOJIS.get(x, '')} {x.capitalize()}")
     
     # Sort by score for better visualization
     df = df.sort_values('score', ascending=True)
@@ -79,14 +90,14 @@ def classify_emotion(text):
     sorted_df = df.sort_values('score', ascending=False)
     
     top_emotion = sorted_df.iloc[0]
-    results_text += f"**Top Emotion:** {EMOTION_EMOJIS.get(top_emotion['label'], '')} **{top_emotion['label'].capitalize()}** ({top_emotion['score']:.2f}%)\n\n"
+    results_text += f"**Top Emotion:** {EMOTION_EMOJIS.get(top_emotion['emotion'], '')} **{top_emotion['emotion'].capitalize()}** ({top_emotion['score']:.2f}%)\n\n"
     results_text += "---\n\n**All Emotions:**\n\n"
     
     for _, row in sorted_df.iterrows():
-        emoji = EMOTION_EMOJIS.get(row['label'], '')
+        emoji = EMOTION_EMOJIS.get(row['emotion'], '')
         bar_length = int(row['score'] / 5)
         bar = '█' * bar_length
-        results_text += f"{emoji} **{row['label'].capitalize()}**: {row['score']:.2f}% {bar}\n\n"
+        results_text += f"{emoji} **{row['emotion'].capitalize()}**: {row['score']:.2f}% {bar}\n\n"
     
     return fig, results_text