import gradio as gr
from transformers import pipeline

# 1) Initialize zero-shot classifier
classifier = pipeline(
    task="zero-shot-classification",
    model="facebook/bart-large-mnli"
)

# 2) Define candidate labels
LABELS = ["linear algebra", "calculus", "probability", "geometry"]

# 3) Gradio interface function
def tag_question(question):
    result = classifier(question, candidate_labels=LABELS)
    return {lbl: round(score, 3) for lbl, score in zip(result["labels"], result["scores"])}

# 4) Build UI
iface = gr.Interface(
    fn=tag_question,
    inputs=gr.Textbox(lines=3, placeholder="Enter your MCQ here..."),
    outputs=gr.Label(num_top_classes=3),
    title="Zero-Shot Question Tagger",
    description="Classify questions into math topics without any training data."
)

if __name__ == "__main__":
    iface.launch()