Update app.py

app.py

@@ -39,7 +39,7 @@ toxicity_labels = ["Toxic", "Severe Toxic", "Obscene", "Threat", "Insult", "Iden
 # Define the prediction function
 def analyse_comment(comment):
     inputs = tokenizer(comment, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+    inputs = {k: v.to(device) for k, v in inputs.items() if k in ['input_ids', 'attention_mask', 'token_type_ids']}
     
     with torch.no_grad():
         outputs = model(**inputs)
@@ -47,19 +47,22 @@ def analyse_comment(comment):
     sentiment_logits = outputs["sentiment_logits"]
     toxicity_logits = outputs["toxicity_logits"]
     
-    # Process sentiment
+    # 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 toxicity
-    toxicity_probs = F.softmax(toxicity_logits, dim=1)
-    toxicity_idx = torch.argmax(toxicity_probs, dim=1).item()
-    toxicity_prediction = toxicity_labels[toxicity_idx]
-    
+    # 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) 
+
     return {
         "Sentiment": sentiment_prediction,
-        "Toxicity": toxicity_prediction
+        "Toxicity Probabilities": toxicity_predictions
     }
 
 # Create Gradio interface