import streamlit as st
from transformers import pipeline, AutoTokenizer

# Define a list of pretrained models
sentimentAnalysisModels = {
    "Roberta": "deepset/roberta-base-squad2",
    "RESPIN": "RESPIN/Telugu_LanguageModels_for_ASR",
    "GPT2-n": "DunnBC22/gpt2-Causal_Language_Model-AG_News",
    "Roberta-Base (ADDITIONAL MODEL!)": "achimoraites/TextClassification-roberta-base_ag_news",
}

# Display a selection box for the user to choose a model
sentimentAnalysisLanguageSelections = st.selectbox("Please select one of these finetuned models from the dropdown", list(sentimentAnalysisModels.keys()))

# roBERTa specific label map
sentimentAnalysisLanguageMap = {"LABEL_0": "NEGATIVE",
                     "LABEL_1": "POSITIVE", "LABEL_2": "NEUTRAL"}

# Load the selected model and tokenizer
sentimentAnalysisModelName = sentimentAnalysisModels[sentimentAnalysisLanguageSelections]
sentimentAnalysisTokenizer = AutoTokenizer.from_pretrained(sentimentAnalysisModelName)
sentiment_pipeline = pipeline(
    "sentiment-analysis", model=sentimentAnalysisModelName, tokenizer=sentimentAnalysisTokenizer)

# Get user input and perform sentiment analysis
sentimentAnalysisTextInput = st.text_input("Please enter sample text for finetuned language model below:",
                           "I am grateful for Data Science Program at NJIT and to have amazing faculty to learn under!")
sentimentAnalysisSubmitButton = st.button("Press the submit button for final grading")

if sentimentAnalysisSubmitButton and sentimentAnalysisTextInput:
    sentimentAnalysisFinalOutput = sentiment_pipeline(sentimentAnalysisTextInput)
    if sentimentAnalysisLanguageSelections == "roBERTa":
        st.write("Roberta Sentiment Analysis Resultant Value:", sentimentAnalysisLanguageMap[sentimentAnalysisFinalOutput[0]["Roberta Label"]])
        st.write("Roberta Sentiment Analysis Resultant Score:", sentimentAnalysisFinalOutput[0]["Probability Assigned with Sentiment Analysis Label"])
    else:
        st.write("Sentiment Analysis Resultant Value:", sentimentAnalysisFinalOutput[0]["Sentiment Analysis Label"])
        st.write("Sentiment Analysis Resultant Score:", sentimentAnalysisFinalOutput[0]["Probability Assigned with Sentiment Analysis Label"])