Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 import tokenizers
 
+import pandas as pd
 import streamlit as st
 import torch.nn as nn
 from transformers import RobertaTokenizer, RobertaModel
@@ -29,19 +30,26 @@ cats = ['Computer Science', 'Economics', 'Electrical Engineering',
 def predict(outputs):
     top = 0
     probs = nn.functional.softmax(outputs, dim=1).tolist()[0]
+    
+    top_cats = []
+    top_probs = []
 
     for prob, cat in sorted(zip(probs, cats), reverse=True):
         if top < 95:
             percent = prob * 100
             top += percent
-            st.write(f'{cat}: {round(percent, 1)}')
+            top_cats.append(cat)
+            top_probs.append(prob)
+            
+    chart_data = pd.DataFrame(top_probs, columns=top_cats)
+    st.bar_chart(chart_data)
 
 tokenizer = RobertaTokenizer.from_pretrained("roberta-large-mnli")
 model = init_model()
         
 st.markdown("### Title")
 
-title = st.text_area("Enter title", height=20)
+title = st.text_area("* Enter title (required)", height=20)
 
 st.markdown("### Abstract")
 
@@ -50,6 +58,7 @@ abstract = st.text_area("Enter abstract", height=200)
 if not title:
     st.warning("Please fill out so required fields")
 else:    
+    st.markdown("### Result")
     encoded_input = tokenizer(title + '. ' + abstract, return_tensors='pt', padding=True, 
                           max_length = 512, truncation=True)
     with torch.no_grad():