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

model_name = "rafaelcarvalhoj/emotion-classifier"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

mapping = {
    0: 'anger',
    1: 'joy',
    2: 'fear'
}

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

    outputs = model(**inputs)
    predictions = outputs.logits

    return mapping[round(predictions.item())]

iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    layout="vertical",
    title="Classificador de emoções em uma frase",
    description="Este modelo analisa uma frase em inglês e diz qual sentimento mais se aproxima da frase apresentada. A frase pode ser classificada em Joy, Anger e Fear"
)

iface.launch(share=True)