FS-sentiment / app.py
Anthos23's picture
Create app.py
9d41f87
raw
history blame
1.13 kB
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TextClassificationPipeline
import operator
import matplotlib.pyplot as plt
import pandas as pd
def get_sentiment(out):
d = dict()
for k in out:
print(k)
label = k['label']
score = k['score']
d[label] = score
winning_lab = max(d.items(), key=operator.itemgetter(1))[0]
winning_score = d[winning_lab]
df = pd.DataFrame.from_dict(d, orient = 'index')
return df #winning_lab, winning_score
model_name = "Anthos23/FS-distilroberta-fine-tuned"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer, return_all_scores=True)
text = st.text_area(f'Ciao! This app uses {model_name}.\nEnter your text to test it ❤️')
if text:
out = pipe(text)
df = get_sentiment(out[0])
fig, ax = plt.subplots()
c = ['#C34A36', '#FFC75F', '#008F7A']
ax.bar(df.index, df[0], color=c, width=0.4)
st.pyplot(fig)
#st.json(get_sentiment(out[0][0]))