import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from torch.nn.functional import softmax


class TextEmotionDetector:
    def __init__(self, model_path):
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.label_key = {0: 'surprise', 1: 'anger', 2: 'love', 3: 'joy', 4: 'fear', 5: 'sadness'}
    
    def inference(self, text):
        embedding = self.tokenizer(text, return_tensors="pt")
        logits = self.model(input_ids=embedding['input_ids'], attention_mask=embedding['attention_mask'], 
                      token_type_ids=embedding['token_type_ids'])['logits']

        probabilities = softmax(logits)
        predication = { self.label_key[i]: float(prob) for i, prob in enumerate(probabilities[0]) }

        return predication
        
detector = TextEmotionDetector("emotion-nlp-model")

text = gr.inputs.Textbox(lines=1, placeholder="Please enter text")
label = gr.outputs.Label()
interface = gr.Interface(fn=detector.inference, inputs=text, outputs=label, title="Text Emotion Classifier")

interface.launch(inline=False)