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

tokenizer = AutoTokenizer.from_pretrained("daspartho/text-emotion")
model = AutoModelForSequenceClassification.from_pretrained("daspartho/text-emotion") # i've uploaded the model on HuggingFace :)

pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer, top_k=6)

label_map={
    'LABEL_0':'🙁',
    'LABEL_1':'😃',
    'LABEL_2':'🥰',
    'LABEL_3':'😠',
    'LABEL_4':'😬',
    'LABEL_5':'😳'
  }

def classify_text(text):
    predictions = pipe(text)[0]
    return {label_map[pred['label']]: float(pred['score']) for pred in predictions}

iface = gr.Interface(
    title='Text Emotion',
    description = "enter a text and the model will attempt to predict the emotion.",
    article = "<p style='text-align: center'><a href='https://github.com/daspartho/text-emotion' target='_blank'>Github</a></p>",
    fn=classify_text, 
    inputs=gr.inputs.Textbox(label="type the text here"), 
    outputs=gr.outputs.Label(label='what the model thinks'),
    )
iface.launch()