Spaces:
Running
Running
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
import torch | |
# Load model and tokenizer | |
model_id = "Rerandaka/Child_safty_bigbird_1" | |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False) | |
model = AutoModelForSequenceClassification.from_pretrained(model_id) | |
# Inference function | |
def classify(text): | |
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512) | |
with torch.no_grad(): | |
logits = model(**inputs).logits | |
predicted_class = torch.argmax(logits, dim=1).item() | |
return str(predicted_class) | |
# Use Blocks to define interface and API | |
with gr.Blocks() as demo: | |
gr.Markdown("## Child-Safety Text Classifier\nThis model detects unsafe or inappropriate text for children.") | |
with gr.Row(): | |
input_box = gr.Textbox(label="Enter text to classify") | |
output_box = gr.Textbox(label="Prediction") | |
btn = gr.Button("Submit") | |
btn.click(fn=classify, inputs=input_box, outputs=output_box) | |
# β Named API endpoint | |
demo.load(fn=classify, inputs=input_box, outputs=output_box, api_name="predict") | |
# Launch space | |
demo.launch() | |