import gradio as gr
from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="bhadresh-savani/distilbert-base-uncased-emotion",
    return_all_scores=True,
)

EMOTIONS = ["sadness", "joy", "love", "anger", "fear", "surprise"]

def predict_emotion(text):
    results = classifier(text)[0]
    return {result["label"]: result["score"] for result in results if result["label"] in EMOTIONS}


iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(lines=3, placeholder="Enter text here..."),
    outputs=gr.Label(num_top_classes=6),
    title="Creative Machines: Sentiment Analysis",
    description="Enter some text and see the predicted emotions.",
)

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