Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
import torch | |
model_id = "Rerandaka/child-safety-01" | |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False) | |
model = AutoModelForSequenceClassification.from_pretrained(model_id) | |
label_map = { | |
0: "Safe / Normal", | |
1: "Inappropriate / Unsafe" | |
} | |
def classify_text(text: str): | |
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256) | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
probs = torch.nn.functional.softmax(outputs.logits, dim=1) | |
predicted = torch.argmax(probs, dim=1).item() | |
confidence = probs[0][predicted].item() | |
return f"{label_map.get(predicted, predicted)} (Confidence: {confidence:.2f})" | |
demo = gr.Interface( | |
fn=classify_text, | |
inputs=gr.Textbox(label="Enter text to classify"), | |
outputs=gr.Textbox(label="Prediction"), | |
title="Child-Safety Text Classifier", | |
description="This model detects unsafe or inappropriate text for children.", | |
flagging_mode="never" | |
) | |
demo.launch() # π« DO NOT include api_name | |