import gradio as gr
from transformers import pipeline

# Load trained model
classifier = pipeline(
    "text-classification",
    model="King-8/request-classifier",
    tokenizer="King-8/request-classifier",
    return_all_scores=True
)

# Label mapping (MUST match training)
label_map = {
    "LABEL_0": "administrative_action",
    "LABEL_1": "attendance",
    "LABEL_2": "check_in",
    "LABEL_3": "clarification",
    "LABEL_4": "general_chat",
    "LABEL_5": "technical_help"
}

def classify_request(role, context, request):
    text = f"Role: {role} | Context: {context} | Request: {request}"
    
    outputs = classifier(text)[0]
    
    # Get highest scoring label
    best = max(outputs, key=lambda x: x["score"])

    readable_label = label_map.get(best["label"], best["label"])
    
    return {
        "Predicted intent": readable_label,
        "Confidence": round(best["score"], 3)
    }

with gr.Blocks() as demo:
    gr.Markdown("## Internship Request Classifier")
    gr.Markdown(
        "This demo uses a fine-tuned transformer model to classify internship-related requests "
        "into routing categories such as attendance, check-ins, technical help, and more."
    )

    role = gr.Dropdown(
        ["student", "parent", "supervisor", "admin"],
        label="User Role"
    )

    context = gr.Textbox(
        label="Conversation Context",
        placeholder="e.g., No prior context or Earlier discussion about check-ins"
    )

    request = gr.Textbox(
        label="User Request",
        placeholder="e.g., I missed the Zoom meeting this morning"
    )

    output = gr.JSON(label="Model Output")

    submit = gr.Button("Classify Request")

    submit.click(
        classify_request,
        inputs=[role, context, request],
        outputs=output
    )

demo.launch()