Spaces:
Runtime error
Runtime error
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"]) | |