import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import numpy as np


model_name = "Ngit/MiniLMv2-L6-H384-goemotions-v2"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def evaluate(text):
    text = text.strip()
    proba = [0]*28
    if text:
        input_ids = tokenizer(text, return_tensors="pt").input_ids
        output = model(input_ids)
        proba = 1 / (1 + np.exp(-output.logits.detach().numpy()[0]))
        proba = [int(v*1000)/10 for v in proba]
    return proba


with gr.Blocks() as demo:
    text = gr.Textbox(label="Text to evaluate", lines=6)
    with gr.Row():
        with gr.Group():
            t_adm = gr.Slider(label="admiration", value=0, maximum=100)
            t_amu = gr.Slider(label="amusement", value=0, maximum=100)
            t_ang = gr.Slider(label="anger", value=0, maximum=100)
            t_ann = gr.Slider(label="annoyance", value=0, maximum=100)
            t_app = gr.Slider(label="approval", value=0, maximum=100)
            t_car = gr.Slider(label="caring", value=0, maximum=100)
            t_con = gr.Slider(label="confusion", value=0, maximum=100)
        with gr.Group():
            t_cur = gr.Slider(label="curiosity", value=0, maximum=100)
            t_des = gr.Slider(label="desire", value=0, maximum=100)
            t_dis = gr.Slider(label="disappointment", value=0, maximum=100)
            t_dip = gr.Slider(label="disapproval", value=0, maximum=100)
            t_dit = gr.Slider(label="disgust", value=0, maximum=100)
            t_emb = gr.Slider(label="embarrassment", value=0, maximum=100)
            t_exc = gr.Slider(label="excitement", value=0, maximum=100)
        with gr.Group():
            t_fea = gr.Slider(label="fear", value=0, maximum=100)
            t_gra = gr.Slider(label="gratitude", value=0, maximum=100)
            t_gri = gr.Slider(label="grief", value=0, maximum=100)
            t_joy = gr.Slider(label="joy", value=0, maximum=100)
            t_lov = gr.Slider(label="love", value=0, maximum=100)
            t_ner = gr.Slider(label="nervousness", value=0, maximum=100)
            t_opt = gr.Slider(label="optimism", value=0, maximum=100)
        with gr.Group():
            t_pri = gr.Slider(label="pride", value=0, maximum=100)
            t_rea = gr.Slider(label="realization", value=0, maximum=100)
            t_rel = gr.Slider(label="relief", value=0, maximum=100)
            t_rem = gr.Slider(label="remorse", value=0, maximum=100)
            t_sad = gr.Slider(label="sadness", value=0, maximum=100)
            t_sur = gr.Slider(label="surprise", value=0, maximum=100)
            t_neu = gr.Slider(label="neutral", value=0, maximum=100)

    btn = gr.Button(value="Evaluate Emotion")
    btn.click(
            evaluate,
            inputs=[text],
            outputs=[
                t_adm,
                t_amu,
                t_ang,
                t_ann,
                t_app,
                t_car,
                t_con,
                t_cur,
                t_des,
                t_dis,
                t_dip,
                t_dit,
                t_emb,
                t_exc,
                t_fea,
                t_gra,
                t_gri,
                t_joy,
                t_lov,
                t_ner,
                t_opt,
                t_pri,
                t_rea,
                t_rel,
                t_rem,
                t_sad,
                t_sur,
                t_neu,
            ],
        )

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