Update app.py

app.py

@@ -24,9 +24,9 @@ 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:
-    #  tokens.append(token.replace('##', ''))
-    tokens.append(token)
+    if token not in ignore_tokens:
+      #tokens.append(token.replace('##', ''))
+      tokens.append(token)
   return '\n'.join(tokens[:top_clean])
 
 def encode(tokenizer, text_sentence, add_special_tokens=True):
@@ -44,8 +44,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).indices.tolist(), top_clean)
-  cls = decode(bert_tokenizer, predict[0, 0, :].topk(top_k).indices.tolist(), top_clean)
+  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)
   return {'bert': bert,'[CLS]':cls}
 
 def get_bert_prediction(input_text,top_k):