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)