bugs fixed:

- added more instructions and slightly modified the layout
- disable translate while loading model and vice versa.
- input text would be disappeared while clicking either translate button or model select button.

app.py

@@ -7,6 +7,9 @@ import numpy as np
 import streamlit as st
 import torch
 from transformers import FlaxAutoModelForSeq2SeqLM, AutoTokenizer
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 @st.cache_resource(show_spinner=False)
@@ -14,12 +17,15 @@ def load_model(model_name, tokenizer_name):
     try:
         model = FlaxAutoModelForSeq2SeqLM.from_pretrained(model_name)
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-    except Exception as e:
+    except OSError as e:
         st.error(f"Error loading model: {e}")
         st.error(f"Model not found. Use {DEFAULT_MODEL} instead")
         model_path = DEFAULT_MODEL
         model = FlaxAutoModelForSeq2SeqLM.from_pretrained(model_path)
         tokenizer = AutoTokenizer.from_pretrained(DEFAULT_MODEL)
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        raise RuntimeError("Error loading model")
     return model, tokenizer
 
 
@@ -75,6 +81,7 @@ def display_ui():
 
 def load_selected_model(config, dataset, language_pair, epsilon):
     ckpt = config[dataset]['languages pairs'][language_pair]['epsilon'][str(epsilon)]
+    logger.info(f"Loading model from {ckpt}")
     if "privalingo" in ckpt:
         model_path = ckpt  # load model from huggingface hub
     else:
@@ -86,12 +93,36 @@ def load_selected_model(config, dataset, language_pair, epsilon):
     return model_path
 
 
+def init_session_state():
+    if 'model_state' not in st.session_state:
+        st.session_state.model_state = {
+            'loaded': False,
+            '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 'translate_button' in st.session_state and st.session_state.translate_button == True:
+        st.session_state.translate_in_progress = True
+
+    if 'translation_result' not in st.session_state:
+        st.session_state.translation_result = {
+            'input': None,
+            'output': None
+        }
+
 def main():
     hold_deterministic(SEED)
     config = load_json(DATASETS_MODEL_INFO_PATH)
 
     left, right = display_ui()
 
+    init_session_state()
+
     with left:
         dataset = st.selectbox("Choose a dataset used for fine-tuning", list(DATASETS_MODEL_INFO.keys()))
         language_pairs_list = list(DATASETS_MODEL_INFO[dataset]["languages pairs"].keys())
@@ -99,28 +130,74 @@ def main():
         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")
         model_status_box = st.empty()
 
-    with right:
-        input_text = st.text_area("Enter Text", "Enter Text Here", max_chars=MAX_INPUT_LEN)
-        btn_translate = st.button("Translate")
-        result_container = st.empty()
-
-    model_path = load_selected_model(config, dataset, language_pair, epsilon)
+    # Load model to cache, if the user has selected a model for the first time
+    if btn_select_model:
+        current_config = f"{dataset}_{language_pair}_{epsilon}"
 
-    with left:
+        st.session_state.model_state['loaded'] = False
         model_status_box.write("")
         with st.spinner(f'Loading model trained on {dataset} with epsilon {epsilon}...'):
+            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
 
-    if btn_translate:
+    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"
+
+        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
+
+        input_text = st.text_area("Enter Text", input_text_content, max_chars=MAX_INPUT_LEN)
+
+        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=not st.session_state.model_state['loaded'],
+                                  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)]))
+
+    # 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']
+        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']:
+        st.session_state.translate_in_progress = True
         with right:
             with st.spinner("Translating..."):
+                input_text = st.session_state.model_state['current_config'] + input_text
                 prediction = translate(input_text, model, tokenizer, src_lang, tgt_lang)
-            result_container.write("**Translation:**")
-            output_container = result_container.container(border=True)
-            output_container.write("".join([postprocess(prediction)]))
+
+            st.session_state.translation_result['input'] = input_text
+            st.session_state.translation_result['output'] = prediction
+
+        st.session_state.translate_in_progress = False
+        st.rerun()
 
 
 if __name__ == '__main__':