Update app.py

app.py

@@ -13,6 +13,7 @@ st.markdown("<img width=400px src='https://img.freepik.com/free-photo/young-pret
 
 title = st.text_area("Enter the title of the article", height=20)
 abstract = st.text_area("Enter the abstract of the article", height=350)
+top_k = st.slider('How many topics from top to show?', 1, 10, 3)
 text = title + " " + abstract
 
 tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-cased")
@@ -22,5 +23,5 @@ n_classes = 40
 my_linear = nn.Linear(in_features=768, out_features=n_classes, bias=True)
 my_linear.load_state_dict(torch.load(MY_LINEAR_NAME, map_location=torch.device('cpu')))
 
-for ms in get_answer_with_desc(text, model, tokenizer, my_linear):
-    st.markdown("### " + ms)
+for ms in get_answer_with_desc(text, model, tokenizer, my_linear, top_k=top_k):
+    st.markdown("#### " + ms)