import gradio as gr
from transformers import pipeline

model_id = "GV05/distilbert-base-uncased-finetuned-emotion"
classifier = pipeline("text-classification", model=model_id)

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

def classify_emotion(text):
    preds = classifier(text, return_all_scores=True)
    res = {}
    for x in preds[0]:
        res[label_to_emotion[x['label']]] = x['score']
    return res

image = gr.Textbox()
label = gr.Label()
examples = ["you are not too sensitive. you are not overreacting",
            "Thinking of you keeps me awake. Dreaming of you keeps me asleep. Being with you keeps me alive."]

title = "Emotion Detector"
description = "This model is a fine-tuned version of distilbert-base-uncased on the emotion dataset"

intf = gr.Interface(fn=classify_emotion, inputs=image, outputs=label, examples=examples, title=title,
                    description=description)

intf.launch(inline=False)