import streamlit as st
from transformers import pipeline, AutoTokenizer

# Define a list of pretrained models
models = {
    "BERTweet": "finiteautomata/bertweet-base-sentiment-analysis",
    "roBERTa": "cardiffnlp/twitter-roberta-base-sentiment",
    "Distilbert": "bhadresh-savani/distilbert-base-uncased-emotion",
    "BERT (BONUS!: Item Review)": "nlptown/bert-base-multilingual-uncased-sentiment",
}

# Display a selection box for the user to choose a model
selected_model = st.selectbox("Select a pretrained model", list(models.keys()))

# roBERTa specific label map
roberta_label_map = {"LABEL_0": "negative", "LABEL_1": "neutral", "LABEL_2": "positive"}

# Load the selected model and tokenizer
model_name = models[selected_model]
tokenizer = AutoTokenizer.from_pretrained(model_name)
sentiment_pipeline = pipeline("sentiment-analysis", model=model_name, tokenizer=tokenizer)

# Get user input and perform sentiment analysis
text_input = st.text_input("Enter text for sentiment analysis", "Anish is a very awesome and dedicated TA! 🤗")
submit_btn = st.button("Submit")

if submit_btn and text_input:
    result = sentiment_pipeline(text_input)
    if selected_model == "roBERTa":
      st.write("Sentiment:", roberta_label_map[result[0]["label"]])
      st.write("Score:", result[0]["score"])
    else:
      st.write("Sentiment:", result[0]["label"])
      st.write("Score:", result[0]["score"])