# Import statements
import streamlit as st
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Get Input
models = {"Distilbert": "distilbert-base-uncased-finetuned-sst-2-english", "roBERTa" : "cardiffnlp/twitter-roberta-base-sentiment", "Bert" : "finiteautomata/bertweet-base-sentiment-analysis"}
text = st.text_area('Text to analyze', "On the review aggregation website Rotten Tomatoes, The Super Mario Bros Movie has a critics score of just 54 per cent positive, marking it as Rotten.")
model_name = st.selectbox("What model do you wanna use?",  ("Distilbert", "roBERTa", "Bert"))

# Default Model - DistilBert      
if text and st.button('Submit!'):
  model = AutoModelForSequenceClassification.from_pretrained(models[model_name])
  tokenizer = AutoTokenizer.from_pretrained(models[model_name])
  classifier = pipeline('sentiment-analysis', model = model, tokenizer = tokenizer)
  out = classifier(text)
  st.json(out)