ciro-c's picture
Fixing result error
31e70fa
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)