import gradio as gr
import numpy as np
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch


target_list = ["Playful", "Infuriating", "Sentimental", "Cynical", "Depressing", "Awe-inspiring", "Patriotic", "Begrudging", "Educational", "Hopeful",
"Sarcastic", "Disrespectful", "Disparaging"]
device = torch.device('cuda' if torch.cuda.is_available() else 'CPU')

model_name = "valurank/finetuned-distilbert-multi-label-emotion"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def get_probs(logits, threshold=0.5):
    sigm = 1 / (1 + np.exp(-logits))
    return sigm 

def multi_label_emotions(text):
  inputs = tokenizer(text, return_tensors="pt", truncation=True)
  model.to(device)
    
  with torch.no_grad():
      logits = model(**inputs).logits
      #probs = logits.int().numpy()[0]
      log_probs = get_probs(logits)

      final_log_probs = []
      for log in log_probs:
        final_log_probs.append(log.numpy())

      final_output = []
      for i in zip(final_log_probs[0], target_list):
        final_output.append(i)

      final_output.sort(reverse=True)
      
      final_dict = {}
      for k,v in final_output:
        final_dict[v] = float(k)

      return final_dict


demo = gr.Interface(multi_label_emotions, inputs=gr.Textbox(),
                    outputs = gr.Label(num_top_classes=16),
                    title="Multi-label-emotion-classification")

if __name__ == "__main__":
  demo.launch(debug=True)