import gradio as gr
from transformers import pipeline

bert_clf = pipeline("text-classification",
                    model="magomerob/clasificador-emotion",
                    top_k=None)

label_map = {
    "LABEL_0": "sadness",
    "LABEL_1": "joy",
    "LABEL_2": "love",
    "LABEL_3": "anger",
    "LABEL_4": "fear",
    "LABEL_5": "surprise"
}

def classify(text, model_choice):
    clf = bert_clf
    results = clf(text)[0]
    return {label_map[r["label"]]: round(r["score"], 4) for r in results}

demo = gr.Interface(
    fn=classify,
    inputs=[
        gr.Textbox(label="Introduce un texto en inglés", lines=3,
                   placeholder="I feel so happy today!"),
        gr.Radio(choices=["BERT"], value="BERT", label="Modelo")
    ],
    outputs=gr.Label(num_top_classes=6, label="Emoción detectada"),
    title="Emotion Classifier",
    description="Clasifica el sentimiento de un texto en inglés en 6 emociones: "
                "sadness, joy, love, anger, fear, surprise.",
    examples=[
        ["I am so happy and excited about this!", "BERT"],
        ["I feel devastated and heartbroken.", "BERT"],
        ["This made me so angry I could scream.", "BERT"],
        ["I am terrified of what might happen.", "BERT"],
    ]
)

demo.launch()