Remove unneeded generator settings. Get model generator settings if exist

app.py

@@ -1,18 +1,8 @@
-import json
-import os
 import time
-from random import randint
+import torch
 
 import psutil
 import streamlit as st
-import torch
-from transformers import (
-    AutoModelForCausalLM,
-    AutoModelForSeq2SeqLM,
-    AutoTokenizer,
-    pipeline,
-    set_seed,
-)
 
 from generator import GeneratorFactory
 
@@ -25,16 +15,19 @@ 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,
+        "split_sentences": False,
     },
     {
         "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,
+        "split_sentences": False,
     },
     {
         "model_name": "yhavinga/t5-small-24L-ccmatrix-multi",
         "desc": "T5 small nl24 ccmatrix en->nl",
         "task": TRANSLATION_NL_TO_EN,
+        "split_sentences": True,
     },
 ]
 
@@ -49,18 +42,15 @@ def main():
 
     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
@@ -76,135 +66,47 @@ It was a quite young girl, unknown to me, with a hood over her head, and with la
     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
+    num_beams = st.sidebar.number_input("Num beams", min_value=1, max_value=10, value=6)
+    num_beam_groups = st.sidebar.number_input(
+        "Num beam groups", min_value=1, max_value=10, value=3
     )
-    repetition_penalty = st.sidebar.number_input(
-        "Repetition penalty", min_value=0.0, max_value=5.0, value=1.2, step=0.1
+    length_penalty = st.sidebar.number_input(
+        "Length penalty", min_value=0.0, max_value=2.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)
+    params = {
+        "num_beams": num_beams,
+        "num_beam_groups": num_beam_groups,
+        "length_penalty": length_penalty,
+    }
 
     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)
+        memory = psutil.virtual_memory()
+
+        for generator in generators:
+            st.markdown(f"🧮 **Model `{generator}`**")
+            time_start = time.time()
+            result, params_used = generator.generate(
+                text=st.session_state.text, **params
+            )
+            time_end = time.time()
+            time_diff = time_end - time_start
+
+            st.write(result[0].replace("\n", "  \n"))
+            text_line = ", ".join([f"{k}={v}" for k, v in params_used.items()])
+            st.markdown(f"    🕙 *generated in {time_diff:.2f}s, `{text_line}`*")
+
+        st.write(
+            f"""
+        ---
+        *Memory: {memory.total / 10**9:.2f}GB, used: {memory.percent}%, available: {memory.available / 10**9:.2f}GB*
+        """
+        )
 
 
 if __name__ == "__main__":

generator.py

@@ -1,37 +1,37 @@
 import os
+import re
+
 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"
+def get_access_token():
     try:
         if not os.path.exists(".streamlit/secrets.toml"):
             raise FileNotFoundError
         access_token = st.secrets.get("babel")
     except FileNotFoundError:
         access_token = os.environ.get("HF_ACCESS_TOKEN", None)
+    return access_token
+
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_model(model_name):
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
     tokenizer = AutoTokenizer.from_pretrained(
-        model_name, from_flax=True, use_auth_token=access_token
+        model_name, from_flax=True, use_auth_token=get_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
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_name, from_flax=True, use_auth_token=get_access_token()
     )
     if device != -1:
         model.to(f"cuda:{device}")
@@ -39,35 +39,61 @@ def load_model(model_name, task):
 
 
 class Generator:
-    def __init__(self, model_name, task, desc):
+    def __init__(self, model_name, task, desc, split_sentences):
         self.model_name = model_name
         self.task = task
         self.desc = desc
         self.tokenizer = None
         self.model = None
         self.prefix = ""
+        self.gen_kwargs = {
+            "max_length": 128,
+            "num_beams": 6,
+            "num_beam_groups": 3,
+            "no_repeat_ngram_size": 0,
+            "early_stopping": True,
+            "num_return_sequences": 1,
+            "length_penalty": 1.0,
+        }
         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)
-
+            self.tokenizer, self.model = load_model(self.model_name)
+
+            for key in self.gen_kwargs:
+                if key in self.model.config.__dict__:
+                    self.gen_kwargs[key] = self.model.config.__dict__[key]
+                    print(
+                        "Setting",
+                        key,
+                        "to",
+                        self.gen_kwargs[key],
+                        "for model",
+                        self.model_name,
+                    )
             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"]
+                    self.prefix = (
+                        task_specific_params["prefix"]
+                        if "prefix" in task_specific_params
+                        else ""
+                    )
+                    for key in self.gen_kwargs:
+                        if key in task_specific_params:
+                            self.gen_kwargs[key] = task_specific_params[key]
             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)
-        #
+    def generate(self, text: str, **generate_kwargs) -> (str, dict):
+        # Replace two or more newlines with a single newline in text
+        text = re.sub(r"\n{2,}", "\n", text)
+
+        generate_kwargs = {**self.gen_kwargs, **generate_kwargs}
         batch_encoded = self.tokenizer(
             self.prefix + text,
             max_length=generate_kwargs["max_length"],
@@ -89,9 +115,7 @@ class Generator:
             pred.replace("<pad> ", "").replace("<pad>", "").replace("</s>", "")
             for pred in decoded_preds
         ]
-        return decoded_preds
-
-        # return self.pipeline(text, **generate_kwargs)
+        return decoded_preds, generate_kwargs
 
     def __str__(self):
         return self.desc
@@ -104,10 +128,10 @@ class GeneratorFactory:
             with st.spinner(text=f"Loading the model {g['desc']} ..."):
                 self.add_generator(**g)
 
-    def add_generator(self, model_name, task, desc):
+    def add_generator(self, model_name, task, desc, split_sentences):
         # 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 = Generator(model_name, task, desc, split_sentences)
             g.load()
             self.generators.append(g)
 
@@ -120,5 +144,5 @@ class GeneratorFactory:
     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]
+    def get_descriptions(self, task=None):
+        return [g.desc for g in self.generators if task is None or task == g.task]