import streamlit as st
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
from transformers import pipeline

#st.title("Enter Phrase: ")
uInput = st.text_input("Enter Phrase: ")
data = [uInput]

pipe = pipeline(model="distilbert-base-uncased-finetuned-sst-2-english")

if st.button("Bertweet"):
    pipe = pipeline(model="finiteautomata/bertweet-base-sentiment-analysis")
if st.button("Distilbert"):
    pipe = pipeline(model="distilbert-base-uncased-finetuned-sst-2-english")

col1, col2 = st.columns(2)
predictions = pipe(data)

col2.header("Probabilities")
for p in predictions:
    col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")