danwath weda krpn

app.py

@@ -1,28 +1,27 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification  # ✅ required
+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)
 
-# Inference function
+# 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
-        predicted_class = torch.argmax(logits, dim=1).item()
-    return str(predicted_class)
+        prediction = torch.argmax(logits, dim=1).item()
+    return prediction  # 0 = safe, 1 = unsafe
 
-# API-ready Gradio Interface
+# ✅ API-compatible Interface with explicit name
 demo = gr.Interface(
     fn=classify,
-    inputs=gr.Textbox(label="Enter text"),
-    outputs=gr.Textbox(label="Prediction"),
-    api_name="/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)
 )
 
-# ✅ Enable API and queue
 demo.queue()
 demo.launch(show_api=True)