cs-gy-6613-project-final / milestone-2.py
dbleek
removed unnecessary files
71ee167
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}**."
)