Spaces:
Running
Running
import gradio as gr | |
import torch | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
# Load model | |
model_id = "Rerandaka/Cild_safety_bigbird" | |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False) | |
model = AutoModelForSequenceClassification.from_pretrained(model_id) | |
# Classification 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 | |
prediction = torch.argmax(logits, dim=1).item() | |
return prediction # 0 = safe, 1 = unsafe | |
# β API-compatible Interface with explicit name | |
demo = gr.Interface( | |
fn=classify, | |
inputs=gr.Textbox(label="Enter paragraph..."), | |
outputs=gr.Number(label="Prediction (0=safe, 1=unsafe)"), | |
api_name="/classify" # π₯ This only works in gr.Interface (not Blocks) | |
) | |
demo.queue() | |
demo.launch(show_api=True) | |