Update app.py

app.py

@@ -47,21 +47,24 @@ def analyse_comment(comment):
     sentiment_logits = outputs["sentiment_logits"]
     toxicity_logits = outputs["toxicity_logits"]
     
-    # Process sentiment (single label classification)
-    sentiment_probs = F.softmax(sentiment_logits, dim=1)
-    sentiment_idx = torch.argmax(sentiment_probs, dim=1).item()
-    sentiment_prediction = sentiment_labels[sentiment_idx]
-    
+    # Process sentiment (multi-class classification)
+    sentiment_probs = F.softmax(sentiment_logits, dim=1).squeeze(0)  # shape: (3,)
+    sentiment_predictions = {}
+
+    for idx, label in enumerate(sentiment_labels):
+        prob = sentiment_probs[idx].item()
+        sentiment_predictions[label] = round(prob, 4)
+
     # Process toxicity (multi-label classification)
     toxicity_probs = torch.sigmoid(toxicity_logits).squeeze(0)  # shape: (6,)
     toxicity_predictions = {}
 
     for idx, label in enumerate(toxicity_labels):
         prob = toxicity_probs[idx].item()
-        toxicity_predictions[label] = round(prob, 2) 
+        toxicity_predictions[label] = round(prob, 4)
 
     return {
-        "Sentiment": sentiment_prediction,
+        "Sentiment Probabilities": sentiment_predictions,
         "Toxicity Probabilities": toxicity_predictions
     }