Update app.py

app.py

@@ -2,17 +2,25 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import gradio as gr
 
+# Load tokenizer and model from your saved folder
 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 classification 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 = ["toxic", "non-toxic"]  # <-- corrected label order
+    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()
 
-