JJ Makely
sentiment analysis model's added
4abb8db
raw
history blame
1.41 kB
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"])