checking for toxicchat-t5-large

app.py

@@ -102,7 +102,12 @@ def predict_toxicity(text, model, tokenizer, device, model_name):
             outputs = model.generate(inputs, max_new_tokens=5)
         
         prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).strip().lower()
-        prediction = "Toxic" if prediction == "positive" else "Not Toxic"
+        
+        # Print raw output for debugging
+        print(f"Raw model output: {prediction}")
+        
+        # Adjust the condition to check for both "positive" and potential variations
+        prediction = "Toxic" if prediction in ["positive", "pos", "toxic", "yes"] else "Not Toxic"
     else:
         inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding="max_length").to(device)