update model paths. improve readability of code.

app.py

@@ -57,11 +57,8 @@ def postprocess(output_text):
     return output
 
 
-def main():
-    hold_deterministic(SEED)
-
+def display_ui():
     st.set_page_config(page_title="DP-NMT DEMO", layout="wide")
-
     st.title("Neural Machine Translation with DP-SGD")
 
     st.write(
@@ -73,6 +70,27 @@ def main():
     st.write("This is a demo for private neural machine translation with DP-SGD.")
 
     left, right = st.columns(2)
+    return left, right
+
+
+def load_selected_model(config, dataset, language_pair, epsilon):
+    ckpt = config[dataset]['languages pairs'][language_pair]['epsilon'][str(epsilon)]
+    if "privalingo" in ckpt:
+        model_path = ckpt  # load model from huggingface hub
+    else:
+        model_name = DEFAULT_MODEL.split('/')[-1]
+        model_path = os.path.join(CHECKPOINTS_DIR, ckpt, model_name)
+        if not os.path.exists(model_path):
+            st.error(f"Model not found. Using default model: {DEFAULT_MODEL}")
+            model_path = DEFAULT_MODEL
+    return model_path
+
+
+def main():
+    hold_deterministic(SEED)
+
+    left, right = display_ui()
+
     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())
@@ -80,36 +98,28 @@ 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)
+        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()
 
-    ckpt = DATASETS_MODEL_INFO[dataset]['languages pairs'][language_pair]['epsilon'][str(epsilon)]
-
-    if "privalingo" in ckpt:
-        # that means the model is loaded from huggingface hub rather checkpoints locally
-        model_path = ckpt
-    else:
-        model_name = DEFAULT_MODEL.split('/')[-1]
-        model_path = os.path.join(CHECKPOINTS_DIR, ckpt, model_name)
-        if not os.path.exists(model_path):
-            with left:
-                st.error(f"Model not found. Use {DEFAULT_MODEL} instead")
-            model_path = DEFAULT_MODEL
+    model_path = load_selected_model(config, dataset, language_pair, epsilon)
 
     with left:
+        model_status_box.write("")
         with st.spinner(f'Loading model trained on {dataset} with epsilon {epsilon}...'):
             model, tokenizer = load_model(model_path, tokenizer_name=DEFAULT_MODEL)
-        sc_box_model_loaded = st.success('Model loaded!')
+        model_status_box.success('Model loaded!')
 
     if btn_translate:
         with right:
             with st.spinner("Translating..."):
                 prediction = translate(input_text, model, tokenizer, src_lang, tgt_lang)
-            st.write("**Translation:**")
-            result_container = st.container(border=True)
-            result_container.write("".join([postprocess(prediction)]))
+            result_container.write("**Translation:**")
+            output_container = result_container.container(border=True)
+            output_container.write("".join([postprocess(prediction)]))
 
 
 if __name__ == '__main__':

dataset_and_model_info.json

@@ -1,11 +1,11 @@
 {
-    "WMT": {
+    "WMT-16": {
         "languages pairs": {
             "German-English": {
                 "epsilon": {
                     "1": "TrustHLT/privalingo-WMT-de-en-eps1",
                     "5": "TrustHLT/privalingo-WMT-de-en-eps5",
-                    "infinite": "TrustHLT/privalingo-WMT-de-en-Infinity"
+                    "Infinity": "TrustHLT/privalingo-WMT-de-en-Infinity"
                 }
             }
         }
@@ -18,7 +18,7 @@
                     "2": "TrustHLT/privalingo-BSD-jp-en-eps2",
                     "5": "TrustHLT/privalingo-BSD-jp-en-eps5",
                     "10": "TrustHLT/privalingo-BSD-jp-en-eps10",
-                    "infinite": "TrustHLT/privalingo-BSD-jp-en-Infinity"
+                    "Infinity": "TrustHLT/privalingo-BSD-jp-en-Infinity"
                 }
             }
         }
@@ -29,7 +29,7 @@
                 "epsilon": {
                     "1": "TrustHLT/privalingo-ClinSpEn-es-en-eps1",
                     "5": "TrustHLT/privalingo-ClinSpEn-es-en-eps5",
-                    "infinite": "TrustHLT/privalingo-ClinSpEn-es-en-Infinity"
+                    "Infinity": "TrustHLT/privalingo-ClinSpEn-es-en-Infinity"
                 }
             }
         }