hummus / app.py
Anthos23's picture
Update app.py
3d7ef34
raw
history blame
1.14 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 = ['red', 'yellow', 'green']
ax.bar(df.index, df[0], color=c)
st.pyplot(fig)
#st.json(get_sentiment(out[0][0]))