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

model_name = "ciro-c/emobot"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

mapping = {
    3:'raiva',
    4:'medo',
    1:'alegria',
    2:'amor',
    0:'tristeza',
    5:'surpresa'
}

def predict(text):
    inputs = tokenizer(text, return_tensors="pt")

    outputs = model(**inputs)
    predictions = outputs.logits
    # predictionItem = predictions.item()
    result = torch.sum(predictions)
    result = max(result, 0)
    result = min(result, 5)
    return mapping[round(result)]

iface = gr.Interface(
    title="Identificador de emoções em uma frase",
    description="Esse modelo foi treinado com o objetivo de identificar a emoção de uma frase escrita em inglês.  Ele identifica 6 emoções raiva, medo, alegria , amor, tristeza e surpresa ",
    fn=predict,
    inputs="text",
    layout="vertical",
    outputs="text"
)

iface.launch(share=True)