Update app.py

app.py

@@ -36,7 +36,10 @@ def encode(tokenizer, text_sentence, add_special_tokens=True):
     text_sentence += ' .'
 
   input_ids = torch.tensor([tokenizer.encode(text_sentence, add_special_tokens=add_special_tokens)])
-  mask_idx = torch.where(input_ids == tokenizer.mask_token_id)[1].tolist()[0]
+  if (tokenizer.mask_token in text_sentence.split()):
+    mask_idx = torch.where(input_ids == tokenizer.mask_token_id)[1].tolist()[0]
+  else:
+    mask_idx = 0
   return input_ids, mask_idx
 
 def get_all_predictions(text_sentence, top_clean=5):