Update app.py

app.py

@@ -1,4 +1,4 @@
-
+!pip install transformers torch
 
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -17,15 +17,19 @@ def preprocess(text):
     inputs = tokenizer(text, padding=True, truncation=True, max_length=128, return_tensors='pt')
     return inputs['input_ids'], inputs['attention_mask']
 
-# Define a function to classify a text input
+# Define a function to classify a text input and return the predicted categories with probabilities
 def classify(text):
     input_ids, attention_mask = preprocess(text)
     with torch.no_grad():
         logits = model(input_ids, attention_mask=attention_mask).logits
-    preds = torch.sigmoid(logits) > 0.5
-    return [labels[i] for i, pred in enumerate(preds.squeeze().tolist()) if pred]
+    preds = torch.sigmoid(logits).squeeze().tolist()
+    return {labels[i]: preds[i] for i in range(len(labels))}
 
-# Example usage
-text = "You are a stupid idiot"
+# Prompt the user to input text and classify it
+text = input("Enter text to check toxicity: ")
 preds = classify(text)
-print(preds)  # Output: ['toxic', 'insult']
+
+# Print the predicted categories with probabilities
+print("Predicted toxicity categories and probabilities:")
+for label, prob in preds.items():
+    print(f"{label}: {prob:.2f}")