Update app.py

app.py

@@ -1,60 +1,62 @@
 import torch
-import torch.nn as nn
-from torch.nn.functional import softmax
 import gradio as gr
+import torch.nn.functional as F
 
-# ---------- Load Your Model ----------
-class SentimentModel(nn.Module):
+# --------------------
+# Load Model
+# --------------------
+class YourModelClass(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.model = nn.Linear(768, 2)  # example
-
+        self.linear = torch.nn.Linear(768, 5)  # example
     def forward(self, x):
-        return self.model(x)
+        return self.linear(x)
 
-model = SentimentModel()
-model.load_state_dict(torch.load("model.pth", map_location="cpu"))
+model = YourModelClass()
+model.load_state_dict(torch.load("model.pt", map_location="cpu"))
 model.eval()
 
-# ---------- Prediction Function ----------
-def predict_sentiment(text):
-    # Convert text → embedding (example)
-    # replace with your own preprocessing
+labels = ["anger", "fear", "joy", "sadness", "surprise"]
+
+# --------------------
+# Prediction Function
+# --------------------
+def predict(text):
+    # TODO: replace with your real preprocessing:
+    x = torch.rand(1, 768)  # dummy input
     with torch.no_grad():
-        x = torch.rand(1, 768)  # dummy embedding, replace with real one
-        output = model(x)
-        probs = softmax(output, dim=1)
-
-    labels = ["Negative", "Positive"]
-    result = {
-        labels[i]: float(probs[0][i])
-        for i in range(2)
-    }
+        logits = model(x)
+        probs = F.softmax(logits, dim=1)[0]
+
+    result = {labels[i]: float(probs[i]) for i in range(len(labels))}
     return result
 
-# ---------- Gradio UI ----------
-with gr.Blocks() as demo:
-    gr.Markdown("""
-    ## 🔍 Sentiment Analysis  
-    Paste text on the right → Get prediction on the left.
-    """)
+# --------------------
+# Gradio UI
+# --------------------
+with gr.Blocks() as app:
+    gr.Markdown("## 🌈 Emotion Classifier — Multi-label Text Analysis")
 
     with gr.Row():
-        info_box = gr.Markdown("""
-        ### About This Model  
-        - Custom PyTorch sentiment classifier  
-        - Trained on your dataset  
-        - Upload text to analyze
+        left = gr.Markdown("""
+        ### 📄 About  
+        Upload or type text → get emotional predictions:
+        - Anger  
+        - Fear  
+        - Joy  
+        - Sadness  
+        - Surprise  
         """)
 
-        input_text = gr.Textbox(
+        text_input = gr.Textbox(
             label="Enter Text",
-            placeholder="Type your sentence here..."
+            placeholder="Type here...",
+            lines=6
         )
 
-    output = gr.Label(label="Prediction")
+    output = gr.Label(label="Emotion Probabilities")
 
-    submit_btn = gr.Button("Analyze")
-    submit_btn.click(predict_sentiment, inputs=input_text, outputs=output)
+    btn = gr.Button("Analyze")
+    btn.click(predict, inputs=text_input, outputs=output)
 
-demo.launch()
+app.launch()