Update app.py

app.py

@@ -22,7 +22,7 @@ def load_bert_model(model_name):
 
   
 def decode(tokenizer, pred_idx, top_clean):
-  ignore_tokens = string.punctuation + '[PAD]'
+  ignore_tokens = string.punctuation
   tokens = []
   for w in pred_idx:
     token = ''.join(tokenizer.decode(w).split())
@@ -51,7 +51,10 @@ def get_all_predictions(text_sentence, top_clean=5):
     predict = bert_model(input_ids)[0]
   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 {'Input sentence':text_sentence,'Masked position': bert,'[CLS]':cls}
+  if ("[MASK]" in text_sentence or "<mask>" in text_sentecence):
+    return {'Input sentence':text_sentence,'Masked position': bert,'[CLS]':cls}
+  else:
+    return {'Input sentence':text_sentence,'[CLS]':cls}
 
 def get_bert_prediction(input_text,top_k):
   try: