Update app.py

app.py

@@ -1,5 +1,3 @@
-import time
-
 import streamlit as st
 
 from torch.nn import Softmax
@@ -7,34 +5,30 @@ from torch.nn import Softmax
 from model import ArxivModel, load_model
 from tokenizer import get_tokenizer
 
-from lables import num_to_classes
+from lables import num_to_classes, taxonomy
 
 
-start_time = time.time()
 model = load_model()
-end_time = time.time()
-
-print("Model:", (end_time - start_time))
-
-start_time = time.time()
 tokenizer = get_tokenizer()
-end_time = time.time()
-
-print("Tokenizer:", (end_time - start_time))
 
 arxiv_model = ArxivModel(model, tokenizer)
 softmax = Softmax(dim=1)
 
 st.markdown("### Classification of article topics")
-# st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
 
-text = st.text_area("Write title and (optional) summary of article")
+title_text = st.text_area("Write title of article")
+summary_text = st.text_area("Write summary of article (optional)")
+text = title_text.strip() + " " + summary_text.strip()
 text = text.strip()
 
 if text != "":
     idxs = arxiv_model.get_idx_class(text, thr=0.95)
+    idxs = idxs[:10]
 
-    for idx in idxs:
-        st.markdown(num_to_classes[idx])
+    for idx, prob in idxs:
+        for tax in taxonomy:
+            if num_to_classes[idx] in tax[0]:
+                st.markdown("{} \t {}%".format(tax[1], round(prob*100, 1)))
+                break
 else:
     st.markdown("")