Update app.py

app.py

@@ -2,25 +2,34 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import gradio as gr
 
-# Load tokenizer and model from your saved folder
+# Load your saved model and tokenizer
 model_dir = "saved_model"
 tokenizer = AutoTokenizer.from_pretrained(model_dir)
+model = AutoModelForSequenceClassification.from_pretrained(model_dir)
 
-# Load the model with explicit label mappings
-model = AutoModelForSequenceClassification.from_pretrained(
-    model_dir,
-    id2label={0: "non-toxic", 1: "toxic"},
-    label2id={"non-toxic": 0, "toxic": 1}
-)
+# Define all 6 labels (Jigsaw-style multi-label toxic comment classification)
+labels = [
+    "toxic",
+    "severe_toxic",
+    "obscene",
+    "threat",
+    "insult",
+    "identity_hate"
+]
 
-# Define classification function
+# Inference function
 def classify(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    outputs = model(**inputs)
-    probs = torch.softmax(outputs.logits, dim=1)
-    labels = ["non-toxic", "toxic"]  # must match id2label order
-    return {labels[i]: float(probs[0][i]) for i in range(len(labels))}
-
-# Launch Gradio app
-gr.Interface(fn=classify, inputs="text", outputs="label").launch()
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.sigmoid(outputs.logits)[0]  # Sigmoid for multi-label
+        result = {label: float(probs[i]) for i, label in enumerate(labels)}
+    return result
 
+# Gradio interface
+gr.Interface(
+    fn=classify,
+    inputs=gr.Textbox(placeholder="Enter your comment..."),
+    outputs=gr.Label(num_top_classes=6),
+    title="Toxic Comment Classifier"
+).launch()