import gradio as gr
from transformers import pipeline

MODEL_ID = "dk409/emotion-roberta"
classifier = pipeline("text-classification", model=MODEL_ID, top_k=None)

# Prediction function
def classify_emotion(text):
    """
    Run the classifier and return a dict of {label: score} for Gradio's Label component.
    The Label component automatically sorts and displays as a bar chart.
    """
    if not text or not text.strip():
        return {}

    results = classifier(text)[0]  # list of {"label": ..., "score": ...}
    return {r["label"]: r["score"] for r in results}


# Gradio interface
examples = [
    "I'm so happy to see you after all these years!",
    "This is absolutely terrifying, I can't watch.",
    "I can't believe they cancelled the show. So angry right now.",
    "She looked at him with so much love in her eyes.",
    "I feel so alone and empty inside.",
    "Wait, you got promoted? I had no idea! That's amazing!",
]

demo = gr.Interface(
    fn=classify_emotion,
    inputs=gr.Textbox(
        label="Enter text",
        placeholder="Type a sentence and I'll detect the emotion...",
        lines=3,
    ),
    outputs=gr.Label(label="Emotion Probabilities", num_top_classes=6),
    title="Emotion Text Classifier",
    description=(
        "Detects 6 emotions in text: **sadness, joy, love, anger, fear, surprise**. "
        "Fine-tuned RoBERTa model trained on the dair-ai/emotion dataset."
    ),
    examples=examples,
)

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