import gradio as gr
from transformers import pipeline

# Load the fine-tuned model from Hugging Face Hub
classifier = pipeline("text-classification", model="Pisethan/khmer-classifier")

# Label mapping (match this to your training label order)
label_map = {
    "LABEL_0": "most_students",
    "LABEL_1": "grade2_lesson",
    "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})"

# 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()