hummus / app.py
Anthos23's picture
Update app.py
28f45b5
raw history blame
No virus
1.16 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 = "mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis"
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]))