File size: 785 Bytes
76eebae
5423fc0
 
 
ff14db7
 
 
5423fc0
ff14db7
 
 
5423fc0
76eebae
 
 
 
 
 
5423fc0
76eebae
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
d = "Rerandaka/Cild_safety_bigbird"
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)

# βœ… Create API-ready interface
demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="Enter text"),
    outputs=gr.Textbox(label="Prediction")
)

# βœ… Enable queue and API support
demo.queue()              # <-- Enables /queue/join
demo.launch(show_api=True)  # <-- Enables gradio_client + /predict