update for API

app.py

@@ -2,32 +2,27 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
+# Load model
 model_id = "Rerandaka/Cild_safety_bigbird"
-
 tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
 model = AutoModelForSequenceClassification.from_pretrained(model_id)
 
-label_map = {
-    0: "Safe / Normal",
-    1: "Inappropriate / Unsafe"
-}
-
-def classify_text(text: str):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256)
+# Inference function
+def classify(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
     with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-        predicted = torch.argmax(probs, dim=1).item()
-        confidence = probs[0][predicted].item()
-    return f"{label_map.get(predicted, predicted)} (Confidence: {confidence:.2f})"
+        logits = model(**inputs).logits
+        predicted_class = torch.argmax(logits, dim=1).item()
+    return str(predicted_class)
 
+# Gradio interface with a named API
 demo = gr.Interface(
-    fn=classify_text,
-    inputs=gr.Textbox(label="Enter text to classify"),
+    fn=classify,
+    inputs=gr.Textbox(label="Text Input"),
     outputs=gr.Textbox(label="Prediction"),
     title="Child-Safety Text Classifier",
-    description="This model detects unsafe or inappropriate text for children.",
-    flagging_mode="never"
+    description="This model detects unsafe or inappropriate text for children."
 )
 
-demo.launch()  # 🚫 DO NOT include api_name
+# ✅ Launch with API endpoint
+demo.launch(api_name="predict")