import gradio as gr
from huggingface_hub import from_pretrained_fastai

learn = from_pretrained_fastai("sadie27/E3-classifier")
LABELS = learn.dls.vocab[1]

def classify_prompt(text):
    if not text.strip():
        return "Escribe un prompt para clasificarlo.", ""
    pred, _, probs = learn.predict(text)
    top = sorted(zip(LABELS, probs), key=lambda x: x[1], reverse=True)[:3]
    details = "\n".join([f"{l}: {float(p)*100:.1f}%" for l, p in top])
    return str(pred), details

with gr.Blocks(theme=gr.themes.Soft(), title="Clasificador de prompts") as demo:
    gr.Markdown("# Clasificador de prompts por categoria")
    gr.Markdown("Introduce un prompt y el modelo predecirá a qué categoría pertenece.")

    with gr.Row():
        with gr.Column(scale=2):
            txt = gr.Textbox(lines=5, placeholder="Escribe tu prompt aquí...", label="Prompt")
            btn = gr.Button("Clasificar", variant="primary")
        with gr.Column(scale=1):
            out_pred  = gr.Textbox(label="Categoría predicha", interactive=False)
            out_probs = gr.Textbox(label="Top 3 categorías", interactive=False, lines=4)

    btn.click(fn=classify_prompt, inputs=txt, outputs=[out_pred, out_probs])

    gr.Examples(
        examples=[
            ["Write a Python function that sorts a list of integers."],
            ["What are the main causes of climate change?"],
            ["Explain the symptoms of type 2 diabetes."],
            ["What is the proof of the Pythagorean theorem?"]
        ],
        inputs=txt
    )

demo.launch()