Update app.py

app.py

@@ -3,10 +3,6 @@ import streamlit as st
 import torch
 import string
 
-bert_tokenizer = None
-bert_model = None
-global top_k
-model_name = "ajitrajaskharan/biomedical"
 
 
 from transformers import BertTokenizer, BertForMaskedLM
@@ -75,14 +71,12 @@ def get_bert_prediction(input_text,top_k,model_name):
  
 def run_test(sent,top_k,model_name):
   start = None
-  global bert_tokenizer
-  global bert_model
-  if (bert_tokenizer is None):
-        bert_tokenizer, bert_model  = load_bert_model(model_name)
+  if (st.session_state['bert_tokenizer'] is None):
+        st.session_state['bert_tokenizer'], st.session_state['bert_model']  = load_bert_model(st.session_state['model_name'])
   with st.spinner("Computing"):
           start = time.time()
           try:
-            res = get_bert_prediction(sent,top_k,model_name)
+            res = get_bert_prediction(sent,st.session_state['top_k'],st.session_state['model_name'])
             st.caption("Results in JSON")
             st.json(res)
             
@@ -93,47 +87,48 @@ def run_test(sent,top_k,model_name):
     st.text(f"prediction took {time.time() - start:.2f}s")
     
 def on_text_change():
-  global top_k,model_name
+
   text = st.session_state.my_text
-  run_test(text,top_k,model_name)
+  run_test(text,st.session_state['top_k']),st.session_state['model_name'])
 
 def on_option_change():
-  global top_k,model_name
+
   text = st.session_state.my_choice
-  run_test(text,top_k,model_name)
+  run_test(text,st.session_state['top_k']),st.session_state['model_name'])
   
 def on_results_count_change():
-   global top_k
-   top_k = int(st.session_state.my_slider)
-   st.info("Results count changed " + str(top_k))
+
+   st.session_state['top_k'] = int(st.session_state.my_slider)
+   st.info("Results count changed " + str(st.session_state['top_k']))
 
 def on_model_change1():
-  global model_name
-  global bert_tokenizer
-  global bert_model
-  model_name = st.session_state.my_model1
-  st.info("Pre-selected model chosen: " + model_name)
-  bert_tokenizer, bert_model  = load_bert_model(model_name)
+  st.session_state['model_name'] = st.session_state.my_model1
+  st.info("Pre-selected model chosen: " + st.session_state['model_name'])
+  st.session_state['bert_tokenizer'], st.session_state['bert_model']  = load_bert_model(st.session_state['model_name'])
 
 def on_model_change2(): 
-  global model_name
-  global bert_tokenizer
-  global bert_model
-  model_name = st.session_state.my_model2
-  st.info("Custom model chosen: " + model_name)
-  bert_tokenizer, bert_model  = load_bert_model(model_name)
+  st.session_state['model_name'] = st.session_state.my_model2
+  st.info("Custom model chosen: " + st.session_state['model_name'])
+  st.session_state['bert_tokenizer'], st.session_state['bert_model']  = load_bert_model(st.session_state['model_name'])
   
 def init_selectbox():
   st.selectbox(
      'Choose any of these sentences or type any text below',
      ('', "[MASK] who lives in New York and works for XCorp suffers from Parkinson's", "Lou Gehrig who lives in [MASK] and works for XCorp suffers from Parkinson's","Lou Gehrig who lives in New York and works       for [MASK] suffers from Parkinson's","Lou Gehrig who lives in New York and works for XCorp suffers from [MASK]","[MASK] who lives in New York and works for XCorp suffers from Lou Gehrig's", "Parkinson who      lives in [MASK] and works for XCorp suffers from Lou Gehrig's","Parkinson who lives in New York and works for [MASK] suffers from Lou Gehrig's","Parkinson who lives in New York and works for XCorp suffers      from [MASK]","Lou Gehrig","Parkinson","Lou Gehrigh's is a [MASK]","Parkinson is a [MASK]","New York is a [MASK]","New York","XCorp","XCorp is a [MASK]","acute lymphoblastic leukemia","acute lymphoblastic       leukemia is a [MASK]"),on_change=on_option_change,key='my_choice') 
   
-
+def init_session_states():
+  if 'top_k' not in st.session_state:
+    st.session_state['top_k'] = 20
+  if 'bert_tokenizer' not in st.session_state:
+    st.session_state['bert_tokenizer'] = None
+  if 'bert_model' not in st.session_state:
+    st.session_state['bert_model'] = None
+  if 'model_name' not in st.session_state:
+    st.session_state['model_name'] = "ajitrajasekharan/biomedical"
 
 def main():
-  global top_k 
-  global bert_tokenizer
-  global bert_model
+  init_session_states()
+  
 
   
   st.markdown("<h3 style='text-align: center;'>Qualitative evaluation of  any pretrained BERT model</h3>", unsafe_allow_html=True)
@@ -144,8 +139,8 @@ def main():
   st.write("This app can be used to examine both model prediction for a masked  position as well as the neighborhood of CLS vector")
   st.write("   - To examine model prediction for a position, enter the token [MASK] or <mask>")
   st.write("   - To examine just the [CLS] vector, enter a word/phrase or sentence. Example: eGFR or EGFR or non small cell lung cancer")
-  top_k = st.sidebar.slider("Select how many predictions do you need", 1 , 50, 20,key='my_slider',on_change=on_results_count_change) #some times it is possible to have less words
-  print(top_k)
+  st.sidebar.slider("Select how many predictions do you need", 1 , 50, 20,key='my_slider',on_change=on_results_count_change) #some times it is possible to have less words
+
 
 
 
@@ -153,13 +148,13 @@ def main():
     
   #  with st.spinner("Computing"):
   try:
-      model_name = st.sidebar.selectbox(label='Select Model to Apply',  options=['ajitrajasekharan/biomedical', 'bert-base-cased','bert-large-cased','microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext','allenai/scibert_scivocab_cased','dmis-lab/biobert-v1.1'], index=0,  key = "my_model1",on_change=on_model_change1)
+      st.sidebar.selectbox(label='Select Model to Apply',  options=['ajitrajasekharan/biomedical', 'bert-base-cased','bert-large-cased','microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext','allenai/scibert_scivocab_cased','dmis-lab/biobert-v1.1'], index=0,  key = "my_model1",on_change=on_model_change1)
       init_selectbox()
       st.text_input("Enter text below", "",on_change=on_text_change,key='my_text')
-      custom_model_name = st.text_input("Model not listed on left? Type the model name (fill-mask BERT models only)", "",key="my_model2",on_change=on_model_change2)
-      if (len(custom_model_name) > 0):
-        model_name = custom_model_name
-        st.info("Custom model selected: " + model_name)
+      st.text_input("Model not listed on left? Type the model name (fill-mask BERT models only)", "",key="my_model2",on_change=on_model_change2)
+      #if (len(custom_model_name) > 0):
+      #  model_name = custom_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,model_name)
@@ -167,10 +162,10 @@ def main():
        # if len(option) > 0:
        #   run_test(option,top_k,model_name)
 
-      st.info("Top k = " + str(top_k))
-      st.info("Model name = " + model_name)
-      if (bert_tokenizer is None):
-        bert_tokenizer, bert_model  = load_bert_model(model_name)
+      st.info("Top k = " + str(st.session_state['top_k']))
+      st.info("Model name = " + st.session_state['model_name'])
+      if (st.session_state['bert_tokenizer'] is None):
+        st.session_state['bert_tokenizer'], st.session_state['bert_model']  = load_bert_model(st.session_state['model_name'])