Update app.py

app.py

@@ -33,15 +33,13 @@ def decode(tokenizer, pred_idx, top_clean):
   return '\n'.join(tokens[:top_clean])
 
 def encode(tokenizer, text_sentence, add_special_tokens=True):
-  bert_tokenizer = st.session_state['bert_tokenizer']
-  bert_model = st.session_state['bert_model']
   
   text_sentence = text_sentence.replace('<mask>', tokenizer.mask_token)
     # if <mask> is the last token, append a "." so that models dont predict punctuation.
   #if tokenizer.mask_token == text_sentence.split()[-1]:
   #  text_sentence += ' .'
 
-  tokenized_text = bert_tokenizer.tokenize(text_sentence) 
+  tokenized_text = tokenizer.tokenize(text_sentence) 
   input_ids = torch.tensor([tokenizer.encode(text_sentence, add_special_tokens=add_special_tokens)])
   if (tokenizer.mask_token in text_sentence.split()):
     mask_idx = torch.where(input_ids == tokenizer.mask_token_id)[1].tolist()[0]
@@ -52,6 +50,7 @@ def encode(tokenizer, text_sentence, add_special_tokens=True):
 def get_all_predictions(text_sentence, model_name,top_clean=5):
   bert_tokenizer = st.session_state['bert_tokenizer']
   bert_model = st.session_state['bert_model']
+  top_k = st.session_state['top_k']
   
     # ========================= BERT =================================
   input_ids, mask_idx,tokenized_text = encode(bert_tokenizer, text_sentence)