from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import streamlit as st

@st.cache
def sentiment_analysis(inp):

    model_name = 'distilbert-base-uncased-finetuned-sst-2-english'
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    classifier = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
    res = classifier(inp)
    return res

def main():
    st.header("Check the sentiment of your text")
    my_string = 'Hello, it was fun making this streamlit application.'
    user_input = st.text_input("Enter text here:", value=my_string)
    st.write("You entered:", user_input)
    res = sentiment_analysis(user_input)
    print(res)
    with st.form("my_form"):
        submit_button = st.form_submit_button(label='Submit')
    sentiment = res[0]['label']
    conf = res[0]['score']
    if submit_button:
        st.write("Sentiment of the input: ", sentiment)
        st.write("Confidence of the predicted sentiment: ", conf)

if __name__ == '__main__':
    main()