AI-project / app.py
corykhal's picture
Update app.py
a057f42
raw
history blame
2.79 kB
import streamlit as st
from transformers import pipeline
import pandas as pd
import numpy as np
from transformers import BertTokenizerFast
st.set_page_config(page_title="AI Project", page_icon=":tada:", layout="wide")
pre_trained_models = ["finiteautomata/bertweet-base-sentiment-analysis",
'bhadresh-savani/distilbert-base-uncased-emotion',
"nlptown/bert-base-multilingual-uncased-sentiment"]
fine_tuned_model = "corykhal/twitter-finetuned"
fine_tuned_labels = {"LABEL_0": "Toxic",
"LABEL_1": "Severe Toxic",
"LABEL_2": "Obscene",
"LABEL_3": "Threat",
"LABEL_4": "Insult",
"LABEL_5": "Identity Hate"}
models = pre_trained_models + [fine_tuned_model]
with st.container():
st.title("Hello! Welcome to the Sentiment Analysis App :wave:")
st.header("By: Cory Khalilollahi")
with st.container():
st.write("---")
text = st.text_input("Please enter the text of a tweet to use for the sentiment analysis:",
value="Hello! It is a pleasure to meet you!")
st.write("---")
st.write("The first 3 models in the select box are for any sentiment analysis.")
st.write("The last model (which is finetuned) in the select box is specifically for toxicity analysis.")
model = st.selectbox("Please select one of the following pre-trained models:", models)
with st.container():
st.write("---")
if st.button("Analyze!"):
# If the user selects a pre-trained model
if model in pre_trained_models:
analysis = pipeline("sentiment-analysis", model=model)
result = analysis(text)
sentiment = result[0]["label"]
score = result[0]["score"]
# Create a table for the label and score
data = pd.DataFrame({"Tweet": [text], "Sentiment": [sentiment], "Confidence Score": [score]})
st.table(data)
# If the user selects the finetuned model
else:
analysis = pipeline("sentiment-analysis", model=model, top_k=2)
result = analysis(text)
# Get the top two labels and scores
sentiment1 = fine_tuned_labels[result[0][0]["label"]]
score1 = result[0][0]["score"]
sentiment2 = fine_tuned_labels[result[0][1]["label"]]
score2 = result[0][1]["score"]
# Create a table for the labels and scores
data = pd.DataFrame({"Tweet": [text], "Toxic": [sentiment1], "Toxic Confidence Score": [score1],
"Type of Toxicity": [sentiment2], "Toxicity Type Confidence Score": [score2]})
st.table(data)