haitaoh's picture
Update app.py
0746424
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
import tensorflow as tf
import gradio as gr
# Load the model
tokenizer = AutoTokenizer.from_pretrained("haitaoh/emotion_classification_toy")
model = TFAutoModelForSequenceClassification.from_pretrained("haitaoh/emotion_classification_toy")
# Emotion dictionary
EMOTION_DICT = {0: ("anger", '😠'),
1: ("disapproval", '❌'),
2: ("disgust", '🤢'),
3: ("fear", '😨'),
4: ("happy", '😄'),
5: ("sadness", '😫'),
6: ("surprise", '🤯'),
7: ("neutral", '😑')}
def get_emotion(userInput):
tokenized_input = tokenizer(userInput, padding=True, truncation=True, return_tensors="tf")
output = model(tokenized_input).logits[0]
index = tf.keras.backend.get_value(tf.math.argmax(output))
score = tf.keras.backend.get_value(tf.nn.softmax(output)[index])
emotion, emoji = EMOTION_DICT[index]
result = f"With probability {score}, the emotion is {emotion}, {emoji}"
return result
textbox = gr.Textbox(label="Enter your text here", placeholder="Hello, how are you doing?")
gr.Interface(fn=get_emotion, inputs=textbox, outputs="text").launch()