fix: bug fix

improve the UX.
- disable all components while a process is in progress, such as loading a model or generating output.

app.py

@@ -100,12 +100,15 @@ def init_session_state():
             'current_config': None
         }
 
-    if 'first_run' not in st.session_state:
-        st.session_state.first_run = True
-
     if 'translate_in_progress' not in st.session_state:
         st.session_state.translate_in_progress = False
 
+    if "load_model_in_progress" not in st.session_state:
+        st.session_state.load_model_in_progress = False
+
+    if "select_model_button" in st.session_state and st.session_state.select_model_button == True:
+        st.session_state.load_model_in_progress = True
+
     if 'translate_button' in st.session_state and st.session_state.translate_button == True:
         st.session_state.translate_in_progress = True
 
@@ -115,6 +118,20 @@ def init_session_state():
             'output': None
         }
 
+
+def get_translation_result():
+    if "translation_result" in st.session_state and st.session_state.translation_result['input'] is not None:
+        input_text_content = st.session_state.translation_result['input']
+    else:
+        input_text_content = "Enter Text Here"
+
+    if "translation_result" in st.session_state and st.session_state.translation_result['output'] is not None:
+        output_text_content = st.session_state.translation_result['output']
+    else:
+        output_text_content = None
+    return input_text_content, output_text_content
+
+
 def main():
     hold_deterministic(SEED)
     config = load_json(DATASETS_MODEL_INFO_PATH)
@@ -123,19 +140,51 @@ def main():
 
     init_session_state()
 
+    st.write(st.session_state)
+
+    with right:
+        right_placeholder = st.empty()
+
+    if st.session_state.load_model_in_progress:
+
+        # Placeholder for right column, to display the input text area and translation result. If do not overwrite the
+        # right column from previous run, the translate button and input text area will be available for user to interace
+        # during the loading of model.
+        disable = True
+        with right_placeholder.container():
+            input_text_content, output_text_content = get_translation_result()
+            input_text = st.text_area("Enter Text", input_text_content, max_chars=MAX_INPUT_LEN, disabled=disable)
+
+            msg_model = "Please confirm model selection via the \'Select Model\' Button first!" \
+                if st.session_state.model_state['current_config'] is None \
+                else f"Current Model: {st.session_state.model_state['current_config']}"
+
+            st.write(msg_model)
+
+            btn_translate = st.button("Translate",
+                                      disabled=disable,
+                                      use_container_width=True,
+                                      key="translate_button")
+
+
     with left:
-        dataset = st.selectbox("Choose a dataset used for fine-tuning", list(DATASETS_MODEL_INFO.keys()))
+        disable = st.session_state.translate_in_progress or st.session_state.load_model_in_progress
+        dataset = st.selectbox("Choose a dataset used for fine-tuning", list(DATASETS_MODEL_INFO.keys()), disabled=disable)
         language_pairs_list = list(DATASETS_MODEL_INFO[dataset]["languages pairs"].keys())
-        language_pair = st.selectbox("Language pair for translation", language_pairs_list)
+        language_pair = st.selectbox("Language pair for translation", language_pairs_list, disabled=disable)
         src_lang, tgt_lang = language_pair.split("-")
         epsilon_options = list(DATASETS_MODEL_INFO[dataset]['languages pairs'][language_pair]['epsilon'].keys())
-        epsilon = st.radio("Select a privacy budget epsilon", epsilon_options, horizontal=True)
-        btn_select_model = st.button("Select Model", disabled=st.session_state.translate_in_progress,
-                                     use_container_width=True, key="select_model_button")
+        epsilon = st.radio("Select a privacy budget epsilon", epsilon_options, horizontal=True, disabled=disable)
+        btn_select_model = st.button(
+            "Select Model",
+            disabled=disable,
+            use_container_width=True,
+            key="select_model_button")
         model_status_box = st.empty()
 
     # Load model to cache, if the user has selected a model for the first time
     if btn_select_model:
+        st.session_state.load_model_in_progress = True
         current_config = f"{dataset}_{language_pair}_{epsilon}"
 
         st.session_state.model_state['loaded'] = False
@@ -144,49 +193,51 @@ def main():
             model_path = load_selected_model(config, dataset, language_pair, epsilon)
             model, tokenizer = load_model(model_path, tokenizer_name=DEFAULT_MODEL)
         model_status_box.success('Model loaded!')
-        st.session_state.model_state['loaded'] = True
-        st.session_state.model_state['current_config'] = current_config
-        st.session_state.first_run = False
 
-    with right:
-        if "translation_result" in st.session_state and st.session_state.translation_result['input'] is not None:
-            input_text_content = st.session_state.translation_result['input']
-        else:
-            input_text_content = "Enter Text Here"
+        st.session_state.model_state['current_config'] = current_config
+        st.session_state.load_model_in_progress = False
+        st.rerun()
 
-        if "translation_result" in st.session_state and st.session_state.translation_result['output'] is not None:
-            output_text_content = st.session_state.translation_result['output']
-        else:
-            output_text_content = None
+    with right_placeholder.container():
+        disable = st.session_state.load_model_in_progress or st.session_state.translate_in_progress
+        input_text_content, output_text_content = get_translation_result()
 
-        input_text = st.text_area("Enter Text", input_text_content, max_chars=MAX_INPUT_LEN)
+        input_text = st.text_area("Enter Text", input_text_content, max_chars=MAX_INPUT_LEN, disabled=disable, key="input_text")
 
         msg_model = "Please confirm model selection via the \'Select Model\' Button first!" \
-            if st.session_state.model_state['current_config'] is None \
-            else f"Current Model: {st.session_state.model_state['current_config']}"
+                    if st.session_state.model_state['current_config'] is None \
+                    else f"Current Model: {st.session_state.model_state['current_config']}"
 
         st.write(msg_model)
 
         btn_translate = st.button("Translate",
-                                  disabled=not st.session_state.model_state['loaded'],
+                                  disabled=(disable or st.session_state.translate_in_progress),
                                   use_container_width=True,
                                   key="translate_button")
         result_container = st.empty()
 
         if output_text_content is not None and not st.session_state.translate_in_progress:
-            result_container.write("**Translation:**")
-            output_container = result_container.container(border=True)
-            output_container.write("".join([postprocess(output_text_content)]))
+            with result_container.container():
+                st.write("**Translation:**")
+                output_container = result_container.container(border=True)
+                output_container.write("".join([postprocess(output_text_content)]))
 
     # Load model from cache when click translate button, if the user has selected a model previously
     if not st.session_state.select_model_button and st.session_state.translate_button:
         model_config = st.session_state.model_state['current_config']
+        if model_config is None:
+
+            # If the user click translate button without selecting a model, set st.session_state.translate_in_progress to False,
+            # to avoid death of program and then refresh the page
+            st.session_state.translate_in_progress = False
+            st.rerun()
+
         dataset, language_pair, epsilon = model_config.split("_")
         model_path = load_selected_model(config, dataset, language_pair, epsilon)
         model, tokenizer = load_model(model_path, tokenizer_name=DEFAULT_MODEL)
         st.session_state.model_state['loaded'] = True
 
-    if btn_translate and st.session_state.model_state['loaded']:
+    if btn_translate:
         st.session_state.translate_in_progress = True
         with right:
             with st.spinner("Translating..."):