File size: 1,337 Bytes
33310e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
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) |