Update app.py

app.py

@@ -1,15 +1,17 @@
-import mlflow.pyfunc
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 import gradio as gr
 
-# Load the model from MLflow registry
-model = mlflow.pyfunc.load_model("models:/bert-toxic-classifier/1")
+model_dir = "saved_model"
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+model = AutoModelForSequenceClassification.from_pretrained(model_dir)
 
 def classify(text):
-    return model.predict([text])[0]
-
-gr.Interface(
-    fn=classify,
-    inputs=gr.Textbox(label="Enter your comment"),
-    outputs=gr.Textbox(label="Prediction"),
-    title="Toxic Comment Classifier"
-).launch()
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    outputs = model(**inputs)
+    probs = torch.softmax(outputs.logits, dim=1)
+    labels = ["non-toxic", "toxic"]  # Adjust if needed
+    return {labels[i]: float(probs[0][i]) for i in range(len(labels))}
+
+gr.Interface(fn=classify, inputs="text", outputs="label").launch()
+