Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 
-from torch import nn
+from torch.nn import Softmax
 
 from model import ArxivModel, load_model
 from tokenizer import get_tokenizer
@@ -13,7 +13,7 @@ model = load_model()
 tokenizer = get_tokenizer()
 
 arxiv_model = ArxivModel(model, tokenizer)
-softmax = nn.Softmax(dim=1)
+softmax = Softmax(dim=1)
 
 st.markdown("### Classification of article topics")
 
@@ -27,9 +27,13 @@ with col2:
     summary_text = st.text_area("Write summary of article (optional)", key='arxiv_sum_input')
 click_button_text = st.button('Submit title and summary', key=1)
 
-if click_button_text and summary_text.strip() != "":
+if click_button_text and title_text.strip() == "":
+    text = ""
+    if summary_text.strip() != "":
+        st.markdown(f'<p style="color:#FF2D00;font-size:18px">Please, input title</p>', unsafe_allow_html=True)
+elif click_button_text and title_text.strip() != "" and summary_text.strip() != "":
     text = title_text.strip() + '\t' + summary_text.strip()
-else:
+elif click_button_text and title_text.strip() != "":
     text = title_text.strip()
 text = text.strip()
 
@@ -48,10 +52,15 @@ if click_button_url and id_url != "":
 print(text)
 
 if text != "":
-    idxs = arxiv_model.get_idx_class(text, thr=0.95)[:10]
-
-    for idx, prob in idxs:
-        if taxonomy.get(num_to_classes[idx], -1) != -1:
-            st.markdown("{} \t {}%".format(taxonomy.get(num_to_classes[idx], -1), round(prob * 100, 1)))
+    idxs = arxiv_model.get_idx_class(text, thr=0.95)
+    print(len(idxs))
+    if len(idxs) > 85:
+        st.markdown("#### Sorry, model can't classify the article with high confidence")
+    else:
+        idxs = idxs[:10]
+        st.markdown("#### The model have defined:")
+        for idx, prob in idxs:
+            if taxonomy.get(num_to_classes[idx], -1) != -1:
+                st.markdown("{} \t {}%".format(taxonomy.get(num_to_classes[idx], -1), round(prob * 100, 1)))
 else:
     st.markdown("")