import streamlit as st
from transformers import pipeline

st.title('Sentiment Analyser')
st.image("https://www.pngall.com/wp-content/uploads/5/Emotion-Transparent.png")

st.write("---")
st.header('Step 1 - Select a language model')
col1, col2 = st.columns(2)

selection = "N/A"
with col1:
    selection = st.radio("Pick one of the four pre-trained models below:",
        key = "modelChoice",
        options = ["DistilBERT", "FinBERT", "SiEBERT", "Twitter-roBERTa"],
    )

pipe = pipeline('sentiment-analysis')
if selection == "DistilBERT":
    pipe = pipeline(model = "distilbert-base-uncased-finetuned-sst-2-english")
if selection == "FinBERT":
    pipe = pipeline(model = "yiyanghkust/finbert-tone")
if selection == "Twitter-roBERTa":
    pipe = pipeline(model = "cardiffnlp/twitter-roberta-base-sentiment-latest")
if selection == "SiEBERT":
    pipe = pipeline(model = "siebert/sentiment-roberta-large-english")

with col2:
    st.caption('DistilBERT - One of the most popular and widely-used language models. Labels text as POSITIVE or NEGATIVE. Developed by Hugging Face.')
    st.caption('FinBERT - A finance-oriented model trained with analysis reports. Labels text as POSITIVE, NEGATIVE or NEUTRAL. Developed by yiyanghkust.')
    st.caption('SiEBERT - A model trained on diverse text sources to improve generalization. Labels text as POSITIVE or NEGATIVE. Developed by siebert.')
    st.caption('Twitter-roBERTa - A model trained on over 124M tweets. Labels text as POSITIVE, NEGATIVE or NEUTRAL. Developed by cardiffnlp.')

st.write("---")
st.header('Step 2 - Enter some text')
text = st.text_area('The sentiment of the text entered here will be determined based on the model you chose above.',
                   value = "It was the best of times, it was the worst of times.")

st.write("---")
st.header('Step 3 - View your results')

if text:
    st.write('Model used: ', selection)
    out = pipe(text)
    st.json(out)