import gradio as gr
import torch.nn.functional as F
import torch
from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast

def translate(text):
  model_name = 'sbenel/emotion-distilbert'
  tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
  model= DistilBertForSequenceClassification.from_pretrained(model_name)

  input = tokenizer(text, return_tensors="pt")
  labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
  output = model(**input, labels=labels)
  logits = output.logits 
  prediction = F.softmax(logits, dim=1)
  y_pred = torch.argmax(prediction).numpy()
  class_names = ['sad','joy','love','anger','fear','surprise']
  return class_names[y_pred]
#  output = model.generate(input["input_ids"], max_length=40, num_beams=4, early_stopping=True)
  
title = "Text Emotion Classification"
inputs = gr.inputs.Textbox(lines=1, label="Text")
outputs = [gr.outputs.Textbox(label="Emotions")]

description = "Here use the [emotion-distilbert](https://huggingface.co/sbenel/emotion-distilbert) that was trained with [emotion dataset](https://huggingface.co/datasets/emotion)."

iface = gr.Interface(fn=translate, inputs=inputs, outputs=outputs, theme="grass", title=title, description=description)
iface.launch(enable_queue=True)