import streamlit as st
from transformers import (AutoTokenizer, TFAutoModelForSequenceClassification,
                          pipeline)

st.title("CS-GY-6613 Project Milestone 2")
model_choices = (
    "distilbert-base-uncased-finetuned-sst-2-english",
    "j-hartmann/emotion-english-distilroberta-base",
    "joeddav/distilbert-base-uncased-go-emotions-student",
)

with st.form("Input Form"):
    text = st.text_area("Write your text here:", "CS-GY-6613 is a great course!")
    model_name = st.selectbox("Select a model:", model_choices)
    submitted = st.form_submit_button("Submit")

if submitted:
    model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    classifier = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
    res = classifier(text)
    label = res[0]["label"].upper()
    score = res[0]["score"]
    st.markdown(
        f"This text was classified as **{label}** with a confidence score of **{score}**."
    )