Use session state instead of global vars, and add progress spinners for loading models

app.py

@@ -17,10 +17,6 @@ EDIT_ALGS = [
     "LU: Lookup Cache",
 ]
 
-tokenizer = None
-model = None
-editable_model = None
-
 def reset():
     st.session_state.edits.drop(st.session_state.edits.index, inplace=True)
     st.session_state.model_outputs.drop(st.session_state.edits.index, inplace=True)
@@ -28,28 +24,16 @@ def reset():
     selected_alg = st.session_state.alg_selector
     selected_alg_idx = EDIT_ALGS.index(selected_alg)
 
-    ############# TODO: show progress spinner
-    global tokenizer
-    global model
-    global editable_model
-
-    if tokenizer is None:
-        tokenizer = AutoTokenizer.from_pretrained("google/t5-large-ssm-nq")
-    if model is None:
-        model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-large-ssm-nq").eval()
-    del editable_model
-
-    alg_name = st.session_state.alg_selector
-    alg_abbrv = alg_name[:alg_name.index(":")]
-    alg_module = importlib.import_module(f"algs.{alg_abbrv.lower()}")
-    alg_class = getattr(alg_module, alg_abbrv.upper())
-    cfg = getattr(config, f"{alg_abbrv.lower()}_config")
-    editable_model = alg_class(
-        model,
-        cfg,
-        lambda: copy.deepcopy(model),
-    ).eval()
-
+    with st.spinner('Loading model...'):
+        alg_abbrv = selected_alg[:selected_alg.index(":")]
+        alg_module = importlib.import_module(f"algs.{alg_abbrv.lower()}")
+        alg_class = getattr(alg_module, alg_abbrv.upper())
+        cfg = getattr(config, f"{alg_abbrv.lower()}_config")
+        st.session_state.editable_model = alg_class(
+            st.session_state.model,
+            cfg,
+            lambda: copy.deepcopy(st.session_state.model),
+        ).eval()
 
 def apply_edit():
     st.session_state.edits.loc[len(st.session_state.edits)] = [str(edit_input), str(edit_label)]
@@ -71,8 +55,10 @@ if "init" not in st.session_state:
     st.session_state.edits = pd.DataFrame([], columns=["Edit input", "Edit label"])
     st.session_state.model_outputs = pd.DataFrame([], columns=["Input", "Output", "N edits", "Alg"])
     st.session_state.init = True
-    st.session_state.model = None  ############## 
-
+    with st.spinner('Loading model...'):
+        st.session_state.tokenizer = AutoTokenizer.from_pretrained("google/t5-large-ssm-nq")
+        st.session_state.model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-large-ssm-nq").eval()
+        st.session_state.editable_model = None
 
 ########################
 #### Interface code ####