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

labels = labels = ['Comment (Expert / Leadership)', 'Personal News','Event Participation', 'Obituary', 'Award / Recognition', 'Company achievement',
          'Financial Insight of stockholding', 'Job Updates', 'Philanthropy', 'Negative News', 'Achievement / Highlighting']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


model_name = 'almalabs/finetuned-distilbert-article-emotions-categorization'
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def get_category(text):
  #text = read_in_text(file.name)

  input_tensor = tokenizer.encode(text, return_tensors='pt', truncation=True)
  logits = model(input_tensor).logits

  softmax = torch.nn.Softmax(dim=1)
  probs = softmax(logits)[0]
  probs = probs.cpu().detach().numpy()
  max_index = np.argmax(probs)
  sentiment = labels[max_index]
  return sentiment

demo = gr.Interface(get_category, inputs=gr.inputs.Textbox(),
                    outputs = 'text',
                    title='Articles emotion Categorization')

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