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()