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

model = AutoModelForSequenceClassification.from_pretrained("vangmayy/emotion")
tokenizer = AutoTokenizer.from_pretrained("vangmayy/emotion")
labels = ["sadness", "Joy", "love", "anger", "fear", "surprise"]
label_map={
    'LABEL_0':'Sadness',
    'LABEL_1':'Joy',
    'LABEL_2':'love',
    'LABEL_3':'anger',
    'LABEL_4':'fear',
    'LABEL_5':'surprise'
}
pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer)

def run_inference(text):
    emotion = label_map[pipe(text)[0]['label']]
    return "Emotion detected: " + emotion

intf = gr.Interface(fn = run_inference, inputs =["text"], outputs = ["text"], examples = ["Woah this is so cool", "I feel sad about what happened", "That room is so dark! I am not going inside"])
intf.launch()