Update app.py

app.py

@@ -1,14 +1,14 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from torch.nn.functional import softmax
+from torch.nn.functional import sigmoid
 import torch
 
-# Load model and tokenizer
-model_name = "bhadresh-savani/distilbert-base-uncased-emotion"
+# Load multi-label GoEmotions model
+model_name = "joeddav/distilbert-base-uncased-go-emotions-group"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Emotion label to icon mapping
+# Emotion label to icon mapping (subset)
 emotion_icons = {
     "admiration": "😍",
     "amusement": "😅",
@@ -40,15 +40,18 @@ emotion_icons = {
     "neutral": "😐"
 }
 
-# Prediction function
-def get_emotion(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    probs = softmax(outputs.logits, dim=1)
-    predicted_class = torch.argmax(probs).item()
-    label = model.config.id2label[predicted_class]
-    icon = emotion_icons.get(label, "")
-    return f"{icon} {label.capitalize()}"
+# Prediction function for multi-label output
+def get_emotions(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    probs = sigmoid(logits)[0]
+
+    threshold = 0.3  # You can adjust this threshold
+    labels = [model.config.id2label[i] for i, p in enumerate(probs) if p > threshold]
+    icons = [emotion_icons.get(label, '') + ' ' + label.capitalize() for label in labels]
+
+    return ", ".join(icons) if icons else "No strong emotion detected."
 
 # Gradio UI
 custom_css = """
@@ -82,11 +85,11 @@ body {
 """
 
 demo = gr.Interface(
-    fn=get_emotion,
-    inputs=gr.Textbox(lines=3, placeholder="What's on your mind today?", label="Your Text"),
-    outputs=gr.Textbox(label="Detected Emotion", elem_classes=["output-textbox"]),
-    title="🥰 Emotion Detector",
-    description="Type a sentence below and hit Submit to reveal the emotion behind your words.",
+    fn=get_emotions,
+    inputs=gr.Textbox(lines=5, placeholder="Write a sentence or a full paragraph...", label="Your Text"),
+    outputs=gr.Textbox(label="Detected Emotions", elem_classes=["output-textbox"]),
+    title="🥰 Multi-Label Emotion Detector",
+    description="Enter a sentence or paragraph to detect multiple emotions present in the text.",
     theme="default",
     css=custom_css
 )