Update app.py

app.py

@@ -1,21 +1,36 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load classifier
+# Load the fine-tuned model from Hugging Face Hub
 classifier = pipeline("text-classification", model="Pisethan/khmer-classifier")
 
-# Map model label IDs to real names
+# Label mapping (match this to your training label order)
 label_map = {
-    "LABEL_0": "count_boys",
+    "LABEL_0": "most_students",
     "LABEL_1": "grade2_lesson",
-    "LABEL_2": "most_students"
+    "LABEL_2": "count_boys"
 }
 
+# Define prediction function
 def predict(text):
     output = classifier(text)[0]
     label_id = output["label"]
     label_name = label_map.get(label_id, label_id)
     return f"📚 Label: {label_name} (Score: {output['score']:.2f})"
 
-demo = gr.Interface(fn=predict, inputs="text", outputs="text", title="Khmer Prompt Classifier")
+# Build Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(label="Khmer Question"),
+    outputs=gr.Textbox(label="Predicted Label"),
+    title="Khmer Prompt Classifier",
+    description="🧠 Enter a Khmer question and get the predicted category.",
+    examples=[
+        ["សិស្សថ្នាក់ទី២ត្រូវរៀនអ្វី?"],
+        ["តើមានសិស្សប្រុសប៉ុន្មាននាក់?"],
+        ["សាលាណាមានសិស្សច្រើនជាងគេ?"]
+    ]
+)
+
+# Launch
 demo.launch()