update app layout

app.py

@@ -1,6 +1,7 @@
 import json
 import os
 import random
+import re
 
 import numpy as np
 import streamlit as st
@@ -51,38 +52,65 @@ def hold_deterministic(seed):
     random.seed(seed)
 
 
-def main():
-    hold_deterministic(SEED)
-
-    st.title("Neural Machine Translation with DP-SGD")
-
-    st.write("This is a demo for private neural machine translation with DP-SGD. More detail can be found in the [repository](https://github.com/trusthlt/dp-nmt)")
-
-    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())
+def postprocess(output_text):
+    output = re.sub(r"<extra_id[^>]*>", "", output_text)
+    return output
 
-    language_pair = st.selectbox("Language pair for translation", language_pairs_list)
 
-    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)
+def main():
+    hold_deterministic(SEED)
 
-    st_model_load = st.text(f'Loading model trained on {dataset} with epsilon {epsilon}...')
+    st.set_page_config(page_title="DP-NMT DEMO", layout="wide")
 
-    model_path = DATASETS_MODEL_INFO[dataset]['languages pairs'][language_pair]['epsilon'][str(epsilon)]
+    st.title("Neural Machine Translation with DP-SGD")
 
-    model, tokenizer = load_model(model_path, tokenizer_name=DEFAULT_MODEL)
-    st.success('Model loaded!')
-    st_model_load.text("")
+    st.write(
+        "[![Star](https://img.shields.io/github/stars/trusthlt/dp-nmt.svg?logo=github&style=social)](https://github.com/trusthlt/dp-nmt)"
+        "&nbsp;&nbsp;&nbsp;"
+        "[![ACL](https://img.shields.io/badge/ACL-Link-red.svg?logo=data:image/svg%2bxml;base64,PHN2ZyBoZWlnaHQ9IjI2MC4wOTA0ODIiIHZpZXdCb3g9IjAgMCA2OCA0NiIgd2lkdGg9IjM4NC40ODE1ODIiIHhtbG5zPSJodHRwOi8vd3d3LnczLm9yZy8yMDAwL3N2ZyI+PHBhdGggZD0ibTQxLjk3NzU1MyAwdjMuMDE1OGgtMzQuNDkwNjQ3Ni03LjQ4NjkwNTR2Ny40ODQ5OSAyNy45Nzc4OCA3LjUyMTMzaDcuNDg2OTA1NCA0Mi4wMTM4OTY2IDcuNDg2OTA2IDExLjAxMjI5MnYtMTUuMDA2MzJoLTExLjAxMjI5MnYtMjAuNDkyODktNy40ODQ5OWMwLTEuNTczNjkgMC0xLjI1NDAyIDAtMy4wMTU4em0tMjYuOTY3Mzk4IDE3Ljk4NTc4aDI2Ljk2NzM5OHYxMy4wMDc5aC0yNi45NjczOTh6IiBmaWxsPSIjZWQxYzI0IiBmaWxsLXJ1bGU9ImV2ZW5vZGQiLz48L3N2Zz4=&link=https%3A%2F%2Faclanthology.org%2F2024.eacl-demo.11%2F)](https://aclanthology.org/2024.eacl-demo.11/)"
+    )
 
-    input_text = st.text_area("Enter Text", "Enter Text Here", max_chars=200)
+    st.write("This is a demo for private neural machine translation with DP-SGD.")
+
+    left, right = st.columns(2)
+    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())
+        language_pair = st.selectbox("Language pair for translation", language_pairs_list)
+        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)
+
+    with right:
+        input_text = st.text_area("Enter Text", "Enter Text Here", max_chars=MAX_INPUT_LEN)
+        btn_translate = st.button("Translate")
+
+    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
+
+    with left:
+        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!')
+
+    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)]))
 
-    if st.button("Translate"):
-        st.write("Translation")
-        prediction = translate(input_text, model, tokenizer, src_lang, tgt_lang)
-        st.success("".join([prediction]))
 
 if __name__ == '__main__':
     DATASETS_MODEL_INFO_PATH = os.path.join(os.getcwd(), "dataset_and_model_info.json")
@@ -93,5 +121,5 @@ if __name__ == '__main__':
     MAX_SEQ_LEN = 512
     NUM_BEAMS = 3
     SEED = 2023
+    MAX_INPUT_LEN = 500
     main()
-

dataset_and_model_info.json

@@ -5,7 +5,7 @@
                 "epsilon": {
                     "1": "MarcoYuTono/privalingo-WMT-de-en-eps1",
                     "5": "MarcoYuTono/privalingo-WMT-de-en-eps5",
-                    "non": "MarcoYuTono/privalingo-WMT-de-en-infinite"
+                    "infinite": "MarcoYuTono/privalingo-WMT-de-en-infinite"
                 }
             }
         }
@@ -18,7 +18,7 @@
                     "2": "2023_09_04-16_23_41",
                     "5": "2023_09_04-16_51_06",
                     "10": "2023_09_04-17_17_44",
-                    "non": "2023_10_22-19_08_23"
+                    "infinite": "2023_10_22-19_08_23"
                 }
             }
         }
@@ -29,7 +29,7 @@
                 "epsilon": {
                     "1": "2023_10_01-00_38_27",
                     "5": "2023_10_01-01_11_49",
-                    "non": "2023_10_23-15_46_22"
+                    "infinite": "2023_10_23-15_46_22"
                 }
             }
         }