Update app.py

app.py

@@ -46,7 +46,7 @@ def encode(tokenizer, text_sentence, add_special_tokens=True):
     mask_idx = 0
   return input_ids, mask_idx
 
-def get_all_predictions(text_sentence, top_clean=5):
+def get_all_predictions(text_sentence, model_name,top_clean=5):
     # ========================= BERT =================================
   input_ids, mask_idx = encode(bert_tokenizer, text_sentence)
   with torch.no_grad():
@@ -55,20 +55,20 @@ def get_all_predictions(text_sentence, top_clean=5):
   cls = decode(bert_tokenizer, predict[0, 0, :].topk(top_k*5).indices.tolist(), top_clean)
   
   if ("[MASK]" in text_sentence or "<mask>" in text_sentence):
-    return {'Input sentence':text_sentence,'Masked position': bert,'[CLS]':cls}
+    return {'Input sentence':text_sentence,'Model':model_name,'Masked position': bert,'[CLS]':cls}
   else:
-    return {'Input sentence':text_sentence,'[CLS]':cls}
+    return {'Input sentence':text_sentence,'Model':model_name,'[CLS]':cls}
 
-def get_bert_prediction(input_text,top_k):
+def get_bert_prediction(input_text,top_k,model_name):
   try:
     #input_text += ' <mask>'
-    res = get_all_predictions(input_text, top_clean=int(top_k))
+    res = get_all_predictions(input_text,model_name, top_clean=int(top_k))
     return res
   except Exception as error:
     pass
     
  
-def run_test(sent,top_k):
+def run_test(sent,top_k,model_name):
   start = None
   global bert_tokenizer
   global bert_model
@@ -77,7 +77,7 @@ def run_test(sent,top_k):
   with st.spinner("Computing"):
           start = time.time()
           try:
-            res = get_bert_prediction(sent,top_k)
+            res = get_bert_prediction(sent,top_k,model_name)
             st.caption("Results in JSON")
             st.json(res)
             
@@ -115,13 +115,13 @@ try:
       custom_model_name = st.text_input("Model not listed on left? Type the model name (fill-mask models only)", "")
       if (len(custom_model_name) > 0):
         model_name = custom_model_name
-        st.info("Custom model selected:" + model_name)
+        st.info("Custom model selected: " + model_name)
         bert_tokenizer, bert_model  = load_bert_model(model_name)
       if len(input_text) > 0:
-        run_test(input_text,top_k)
+        run_test(input_text,top_k,model_name)
       else:
         if len(option) > 0:
-          run_test(option,top_k)
+          run_test(option,top_k,model_name)
       if (bert_tokenizer is None):
         bert_tokenizer, bert_model  = load_bert_model(model_name)