Update app.py

app.py

@@ -83,18 +83,27 @@ def run_sentiment_with_selected_model(text, model_name):
         outputs = model(**inputs)
 
     logits = outputs.logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)
+    probs = F.softmax(logits, dim=-1)
     pred = torch.argmax(probs, dim=-1).item()
 
-    # Get label from model config if available
-    if model.config.id2label:
-        label = model.config.id2label[pred]
+    # Custom label mapping
+    label_map = {
+        "assemblyai/bert-large-uncased-sst2": {0: "Negative", 1: "Positive"},
+        "sohan-ai/sentiment-analysis-model-amazon-reviews": {0: "Negative", 1: "Positive"},
+    }
+
+    if model_name in label_map:
+        label = label_map[model_name][pred]
+    elif model.config.id2label:
+        label = model.config.id2label.get(pred, f"LABEL_{pred}")
     else:
-        label = "Positive" if pred == 1 else "Negative"
+        label = f"LABEL_{pred}"
 
     emoji = "✅" if "positive" in label.lower() else "❌" if "negative" in label.lower() else "⚠️"
-    return f"{emoji} '{text}' -> {label}"
-
+    
+    # Add confidence score
+    confidence = probs[0][pred].item() * 100
+    return f"{emoji} '{text}' -> {label} ({confidence:.1f}%)"
 
 # ---------------- Gradio UI ----------------
 background_css = """