Update app.py

app.py

@@ -25,7 +25,7 @@ def decode(tokenizer, pred_idx, top_clean):
   tokens = []
   for w in pred_idx:
     token = ''.join(tokenizer.decode(w).split())
-    if token not in ignore_tokens:
+    if token not in ignore_tokens and len(token) > 1 and not token.startswith('.') and not token.startswith('['):
       #tokens.append(token.replace('##', ''))
       tokens.append(token)
   return '\n'.join(tokens[:top_clean])
@@ -48,8 +48,8 @@ def get_all_predictions(text_sentence, top_clean=5):
   input_ids, mask_idx = encode(bert_tokenizer, text_sentence)
   with torch.no_grad():
     predict = bert_model(input_ids)[0]
-  bert = decode(bert_tokenizer, predict[0, mask_idx, :].topk(top_k*2).indices.tolist(), top_clean)
-  cls = decode(bert_tokenizer, predict[0, 0, :].topk(top_k*2).indices.tolist(), top_clean)
+  bert = decode(bert_tokenizer, predict[0, mask_idx, :].topk(top_k*5).indices.tolist(), top_clean)
+  cls = decode(bert_tokenizer, predict[0, 0, :].topk(top_k*5).indices.tolist(), top_clean)
   return {'bert': bert,'[CLS]':cls}
 
 def get_bert_prediction(input_text,top_k):
@@ -67,7 +67,7 @@ st.markdown("""
 st.write("Incomplete. Work in progress...")
   #st.write("https://ajitrajasekharan.github.io/2021/01/02/my-first-post.html")
 st.write("CLS vectors as well as the model prediction for a blank position are examined")
-top_k = 10
+top_k = 20
 print(top_k)