import gradio as gr
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import torch

description = "Sentiment Analysis :) && :("
title = "SentBERT"
examples = [["That ice cream was really bad"], ["Great to meet you!"], ["Hey, there's a snake there"]]

class2interpret = {
    0: 'Positive/Neutral',
    1: 'Negative'
}

def classify(example):
  tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
  model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")  
  inputs = tokenizer(example, return_tensors="pt")
  with torch.no_grad():
      logits = model(**inputs).logits
  
  probs = torch.nn.Softmax(dim=1)(logits).tolist()[0]

  return {class2interpret[0]: probs[0], class2interpret[1]: probs[1]}, {class2interpret[0]: probs[0], class2interpret[1]: probs[1]}

interface = gr.Interface(fn=classify, inputs='text', outputs=['label', 'json'], examples=examples, description=description, title=title)
interface.launch()