Update app.py

app.py

@@ -1,18 +1,21 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TextClassificationPipeline
 import operator
+import plotly.express as px
 
 def get_sentiment(out):
-  print(out)
   d = dict()
-  for k in out.keys():
-    label = out[k]['label']
-    score = out[k]['score']
-    d[label] = score
-    
-  winning_lab = max(d.iteritems(), key=operator.itemgetter(1))[0]
+  for k in out:
+    print(k)
+    label = k['label']
+    score = k['score']
+    d[label] = score.reset_index().rename(columns={'index': 'sentiment', 0:'score'})
+
+  winning_lab = max(d.items(), key=operator.itemgetter(1))[0]
   winning_score = d[winning_lab]
-  return winning_lab, winning_score
+
+  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)
@@ -24,5 +27,9 @@ text = st.text_area(f'Ciao! This app uses {model_name}.\nEnter your text  to tes
 
 if text:
   out = pipe(text)
+  df = get_sentiment(out[0])
+  fig = px.pie(df, values='score', names='sentiment')
+  
+  st.pyplot(fig)
   
-  st.json(get_sentiment(out[0][0]))
+  #st.json(get_sentiment(out[0][0]))