Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
tokenizer = AutoTokenizer.from_pretrained("Sk1306/student_chat_toxicity_classifier_model") | |
model = AutoModelForSequenceClassification.from_pretrained("Sk1306/student_chat_toxicity_classifier_model") | |
def predict_toxicity(text): | |
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128) | |
outputs = model(**inputs) | |
logits = outputs.logits | |
# Apply softmax to get probabilities | |
probabilities = torch.nn.functional.softmax(logits, dim=-1) | |
# Get the predicted class (index 0 for non-toxic, index 1 for toxic) | |
predicted_class = torch.argmax(probabilities, dim=-1).item() | |
# Map the prediction to the label (0 = Non-toxic, 1 = Toxic) | |
if predicted_class == 0: | |
return "Non-toxic" | |
else: | |
return "Toxic" | |
interface = gr.Interface( | |
fn=predict_toxicity, | |
inputs="text", # Text input from the user | |
outputs="text", # Text output for the prediction | |
title="Student Chat Toxicity Classifier", | |
description="Enter a message", | |
theme="dark", | |
examples=[ | |
"You can copy in exam to pass!", | |
"Study well.Hardwork pays off!", | |
"Take these drugs.It will boost your memory", | |
], | |
) | |
interface.launch(inline=False) |