Add app

.gitignore

@@ -0,0 +1,4 @@
+venv
+.idea
+__pycache__
+*~

README.md

@@ -1,13 +1,39 @@
 ---
-title: Babel
-emoji: 📊
-colorFrom: indigo
-colorTo: gray
+title: Babel - translate between Dutch and English
+emoji: 🧙
+colorFrom: gray
+colorTo: indigo
 sdk: streamlit
-sdk_version: 1.10.0
 app_file: app.py
 pinned: false
+sdk_version: 1.0.0
 license: postgresql
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Configuration
+
+`title`: _string_
+Display title for the Space
+
+`emoji`: _string_
+Space emoji (emoji-only character allowed)
+
+`colorFrom`: _string_
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`colorTo`: _string_
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`sdk`: _string_
+Can be either `gradio`, `streamlit`, or `static`
+
+`sdk_version` : _string_
+Only applicable for `streamlit` SDK.  
+See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
+
+`app_file`: _string_
+Path to your main application file (which contains either `gradio` or `streamlit` Python code, or `static` html code).
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,211 @@
+import json
+import os
+import time
+from random import randint
+
+import psutil
+import streamlit as st
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    pipeline,
+    set_seed,
+)
+
+from generator import GeneratorFactory
+
+device = torch.cuda.device_count() - 1
+
+TRANSLATION_NL_TO_EN = "translation_en_to_nl"
+
+GENERATOR_LIST = [
+    {
+        "model_name": "yhavinga/longt5-local-eff-large-nl8-voc8k-ddwn-512beta-512l-nedd-256ccmatrix-en-nl",
+        "desc": "longT5 large nl8 256cc/512beta/512l en->nl",
+        "task": TRANSLATION_NL_TO_EN,
+    },
+    {
+        "model_name": "yhavinga/longt5-local-eff-large-nl8-voc8k-ddwn-512beta-512-nedd-en-nl",
+        "desc": "longT5 large nl8 512beta/512l en->nl",
+        "task": TRANSLATION_NL_TO_EN,
+    },
+    {
+        "model_name": "yhavinga/t5-small-24L-ccmatrix-multi",
+        "desc": "T5 small nl24 ccmatrix en->nl",
+        "task": TRANSLATION_NL_TO_EN,
+    },
+]
+
+
+def main():
+    st.set_page_config(  # Alternate names: setup_page, page, layout
+        page_title="Babel",  # String or None. Strings get appended with "• Streamlit".
+        layout="wide",  # Can be "centered" or "wide". In the future also "dashboard", etc.
+        initial_sidebar_state="expanded",  # Can be "auto", "expanded", "collapsed"
+        page_icon="📚",  # String, anything supported by st.image, or None.
+    )
+
+    if "generators" not in st.session_state:
+        st.session_state["generators"] = GeneratorFactory(GENERATOR_LIST)
+
+    generators = st.session_state["generators"]
+
+    with open("style.css") as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+    st.sidebar.image("babel.png", width=200)
+    st.sidebar.markdown(
+        """# Babel
+    Vertaal van en naar Engels"""
+    )
+    model_desc = st.sidebar.selectbox("Model", generators.gpt_descs(), index=1)
+    st.sidebar.title("Parameters:")
+    if "prompt_box" not in st.session_state:
+        # Text is from https://www.gutenberg.org/files/35091/35091-h/35091-h.html
+        st.session_state[
+            "prompt_box"
+        ] = """It was a wet, gusty night and I had a lonely walk home. By taking the river road, though I hated it, I saved two miles, so I sloshed ahead trying not to think at all. Through the barbed wire fence I could see the racing river. Its black swollen body writhed along with extraordinary swiftness, breathlessly silent, only occasionally making a swishing ripple. I did not enjoy looking at it. I was somehow afraid.
+
+And there, at the end of the river road where I swerved off, a figure stood waiting for me, motionless and enigmatic. I had to meet it or turn back.
+
+It was a quite young girl, unknown to me, with a hood over her head, and with large unhappy eyes.
+
+“My father is very ill,” she said without a word of introduction. “The nurse is frightened. Could you come in and help?”"""
+    st.session_state["text"] = st.text_area(
+        "Enter text", st.session_state.prompt_box, height=300
+    )
+    max_length = st.sidebar.number_input(
+        "Lengte van de tekst",
+        value=200,
+        max_value=4096,
+    )
+    no_repeat_ngram_size = st.sidebar.number_input(
+        "No-repeat NGram size", min_value=1, max_value=5, value=3
+    )
+    repetition_penalty = st.sidebar.number_input(
+        "Repetition penalty", min_value=0.0, max_value=5.0, value=1.2, step=0.1
+    )
+    num_return_sequences = st.sidebar.number_input(
+        "Num return sequences", min_value=1, max_value=5, value=1
+    )
+    seed_placeholder = st.sidebar.empty()
+    if "seed" not in st.session_state:
+        print(f"Session state does not contain seed")
+        st.session_state["seed"] = 4162549114
+        print(f"Seed is set to: {st.session_state['seed']}")
+
+    seed = seed_placeholder.number_input(
+        "Seed", min_value=0, max_value=2**32 - 1, value=st.session_state["seed"]
+    )
+
+    def set_random_seed():
+        st.session_state["seed"] = randint(0, 2**32 - 1)
+        seed = seed_placeholder.number_input(
+            "Seed", min_value=0, max_value=2**32 - 1, value=st.session_state["seed"]
+        )
+        print(f"New random seed set to: {seed}")
+
+    if st.button("Set new random seed"):
+        set_random_seed()
+
+    if sampling_mode := st.sidebar.selectbox(
+        "select a Mode", index=0, options=["Top-k Sampling", "Beam Search"]
+    ):
+        if sampling_mode == "Beam Search":
+            num_beams = st.sidebar.number_input(
+                "Num beams", min_value=1, max_value=10, value=4
+            )
+            length_penalty = st.sidebar.number_input(
+                "Length penalty", min_value=0.0, max_value=2.0, value=1.0, step=0.1
+            )
+            params = {
+                "max_length": max_length,
+                "no_repeat_ngram_size": no_repeat_ngram_size,
+                "repetition_penalty": repetition_penalty,
+                "num_return_sequences": num_return_sequences,
+                "num_beams": num_beams,
+                "early_stopping": True,
+                "length_penalty": length_penalty,
+            }
+        else:
+            top_k = st.sidebar.number_input(
+                "Top K", min_value=0, max_value=100, value=50
+            )
+            top_p = st.sidebar.number_input(
+                "Top P", min_value=0.0, max_value=1.0, value=0.95, step=0.05
+            )
+            temperature = st.sidebar.number_input(
+                "Temperature", min_value=0.05, max_value=1.0, value=1.0, step=0.05
+            )
+            params = {
+                "max_length": max_length,
+                "no_repeat_ngram_size": no_repeat_ngram_size,
+                "repetition_penalty": repetition_penalty,
+                "num_return_sequences": num_return_sequences,
+                "do_sample": True,
+                "top_k": top_k,
+                "top_p": top_p,
+                "temperature": temperature,
+            }
+
+    st.sidebar.markdown(
+        """For an explanation of the parameters, head over to the [Huggingface blog post about text generation](https://huggingface.co/blog/how-to-generate)
+and the [Huggingface text generation interface doc](https://huggingface.co/transformers/main_classes/model.html?highlight=generate#transformers.generation_utils.GenerationMixin.generate).
+"""
+    )
+
+    def estimate_time():
+        """Estimate the time it takes to generate the text."""
+        estimate = max_length / 18
+        if device == -1:
+            ## cpu
+            estimate = estimate * (1 + 0.7 * (num_return_sequences - 1))
+            if sampling_mode == "Beam Search":
+                estimate = estimate * (1.1 + 0.3 * (num_beams - 1))
+        else:
+            ## gpu
+            estimate = estimate * (1 + 0.1 * (num_return_sequences - 1))
+            estimate = 0.5 + estimate / 5
+            if sampling_mode == "Beam Search":
+                estimate = estimate * (1.0 + 0.1 * (num_beams - 1))
+        return int(estimate)
+
+    if st.button("Run"):
+        estimate = estimate_time()
+
+        with st.spinner(
+            text=f"Please wait ~ {estimate} second{'s' if estimate != 1 else ''} while getting results ..."
+        ):
+            memory = psutil.virtual_memory()
+
+            for generator in generators:
+                st.subheader(f"Result from {generator}")
+                set_seed(seed)
+                time_start = time.time()
+                result = generator.generate(text=st.session_state.text, **params)
+                time_end = time.time()
+                time_diff = time_end - time_start
+
+                for text in result:
+                    st.write(text.replace("\n", "  \n"))
+                    st.write(f"--- generated in {time_diff:.2f} seconds ---")
+
+            info = f"""
+            ---
+            *Memory: {memory.total / 10**9:.2f}GB, used: {memory.percent}%, available: {memory.available / 10**9:.2f}GB*
+            *Text generated using seed {seed}*
+            """
+            st.write(info)
+
+            params["seed"] = seed
+            params["prompt"] = st.session_state.text
+            params["model"] = generator.model_name
+            params_text = json.dumps(params)
+            print(params_text)
+            st.json(params_text)
+
+
+if __name__ == "__main__":
+    main()

generator.py

@@ -0,0 +1,124 @@
+import os
+import streamlit as st
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+)
+
+device = torch.cuda.device_count() - 1
+
+TRANSLATION_NL_TO_EN = "translation_en_to_nl"
+
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_model(model_name, task):
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    try:
+        if not os.path.exists(".streamlit/secrets.toml"):
+            raise FileNotFoundError
+        access_token = st.secrets.get("netherator")
+    except FileNotFoundError:
+        access_token = os.environ.get("HF_ACCESS_TOKEN", None)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, from_flax=True, use_auth_token=access_token
+    )
+    if tokenizer.pad_token is None:
+        print("Adding pad_token to the tokenizer")
+        tokenizer.pad_token = tokenizer.eos_token
+    auto_model_class = (
+        AutoModelForSeq2SeqLM if "translation" in task else AutoModelForCausalLM
+    )
+    model = auto_model_class.from_pretrained(
+        model_name, from_flax=True, use_auth_token=access_token
+    )
+    if device != -1:
+        model.to(f"cuda:{device}")
+    return tokenizer, model
+
+
+class Generator:
+    def __init__(self, model_name, task, desc):
+        self.model_name = model_name
+        self.task = task
+        self.desc = desc
+        self.tokenizer = None
+        self.model = None
+        self.prefix = ""
+        self.load()
+
+    def load(self):
+        if not self.model:
+            print(f"Loading model {self.model_name}")
+            self.tokenizer, self.model = load_model(self.model_name, self.task)
+
+            try:
+                if self.task in self.model.config.task_specific_params:
+                    task_specific_params = self.model.config.task_specific_params[
+                        self.task
+                    ]
+                    if "prefix" in task_specific_params:
+                        self.prefix = task_specific_params["prefix"]
+            except TypeError:
+                pass
+
+    def generate(self, text: str, **generate_kwargs) -> str:
+        #
+        # import pydevd_pycharm
+        # pydevd_pycharm.settrace('10.1.0.144', port=12345, stdoutToServer=True, stderrToServer=True)
+        #
+        batch_encoded = self.tokenizer(
+            self.prefix + text,
+            max_length=generate_kwargs["max_length"],
+            padding=False,
+            truncation=False,
+            return_tensors="pt",
+        )
+        if device != -1:
+            batch_encoded.to(f"cuda:{device}")
+        logits = self.model.generate(
+            batch_encoded["input_ids"],
+            attention_mask=batch_encoded["attention_mask"],
+            **generate_kwargs,
+        )
+        decoded_preds = self.tokenizer.batch_decode(
+            logits.cpu().numpy(), skip_special_tokens=False
+        )
+        decoded_preds = [
+            pred.replace("<pad> ", "").replace("<pad>", "").replace("</s>", "")
+            for pred in decoded_preds
+        ]
+        return decoded_preds
+
+        # return self.pipeline(text, **generate_kwargs)
+
+    def __str__(self):
+        return self.desc
+
+
+class GeneratorFactory:
+    def __init__(self, generator_list):
+        self.generators = []
+        for g in generator_list:
+            with st.spinner(text=f"Loading the model {g['desc']} ..."):
+                self.add_generator(**g)
+
+    def add_generator(self, model_name, task, desc):
+        # If the generator is not yet present, add it
+        if not self.get_generator(model_name=model_name, task=task, desc=desc):
+            g = Generator(model_name, task, desc)
+            g.load()
+            self.generators.append(g)
+
+    def get_generator(self, **kwargs):
+        for g in self.generators:
+            if all([g.__dict__.get(k) == v for k, v in kwargs.items()]):
+                return g
+        return None
+
+    def __iter__(self):
+        return iter(self.generators)
+
+    def gpt_descs(self):
+        return [g.desc for g in self.generators if g.task == TRANSLATION_NL_TO_EN]

requirements.txt

@@ -0,0 +1,13 @@
+#-f https://download.pytorch.org/whl/torch_stable.html
+-f https://download.pytorch.org/whl/cu116
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+protobuf<3.20
+streamlit>=1.4.0,<=1.10.0
+torch
+transformers>=4.13.0
+mtranslate
+psutil
+jax[cuda]==0.3.16
+chex>=0.1.4
+##jaxlib==0.1.67
+flax>=0.5.3

style.css

@@ -0,0 +1,42 @@
+body {
+    background-color: #eee;
+}
+/*.fullScreenFrame > div {*/
+/*    display: flex;*/
+/*    justify-content: center;*/
+/*}*/
+/*.stButton>button {*/
+/*    color: #4F8BF9;*/
+/*    border-radius: 50%;*/
+/*    height: 3em;*/
+/*    width: 3em;*/
+/*}*/
+
+.stTextInput>div>div>input {
+    color: #4F8BF9;
+}
+.stTextArea>div>div>input {
+    color: #4F8BF9;
+    min-height: 300px;
+}
+
+
+/*.st-cj {*/
+/*    min-height: 500px;*/
+/*    spellcheck="false";*/
+/*    color: #4F8BF9;*/
+/*}*/
+/*.st-ch {*/
+/*    min-height: 500px;*/
+/*    spellcheck="false";*/
+/*    color: #4F8BF9;*/
+/*}*/
+/*.st-bb {*/
+/*    min-height: 500px;*/
+/*    spellcheck="false";*/
+/*    color: #4F8BF9;*/
+/*}*/
+
+/*body {*/
+/*    background-color: #f1fbff*/
+/*}*/