Add translation. Keep GeneratorFactory in the session cache.

app.py

@@ -6,11 +6,18 @@ from random import randint
 import psutil
 import streamlit as st
 import torch
-from transformers import (AutoModelForCausalLM, AutoModelForSeq2SeqLM,
-                          AutoTokenizer, pipeline, set_seed)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    pipeline,
+    set_seed,
+)
 
 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):
@@ -63,43 +70,45 @@ class GeneratorFactory:
     def __init__(self):
         self.generators = []
 
+    def instantiate_generators(self):
+        GENERATOR_LIST = [
+            {
+                "model_name": "yhavinga/gpt-neo-125M-dutch-nedd",
+                "desc": "GPT-Neo Small Dutch(book finetune)",
+                "task": "text-generation",
+            },
+            {
+                "model_name": "yhavinga/gpt2-medium-dutch-nedd",
+                "desc": "GPT2 Medium Dutch (book finetune)",
+                "task": "text-generation",
+            },
+            {
+                "model_name": "yhavinga/t5-small-24L-ccmatrix-multi",
+                "desc": "Dutch<->English T5 small 24 layers",
+                "task": TRANSLATION_NL_TO_EN,
+            },
+        ]
+        for g in GENERATOR_LIST:
+            with st.spinner(text=f"Loading the model {g['desc']} ..."):
+                self.add_generator(**g)
+
+        return self
+
     def add_generator(self, model_name, task, desc):
-        g = Generator(model_name, task, desc)
-        g.load()
-        self.generators.append(g)
+        # 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, model_desc):
+    def get_generator(self, **kwargs):
         for g in self.generators:
-            if g.desc == model_desc:
+            if all([g.__dict__.get(k) == v for k, v in kwargs.items()]):
                 return g
         return None
 
-
-GENERATORS = [
-    {
-        "model_name": "yhavinga/gpt-neo-125M-dutch-nedd",
-        "desc": "GPT-Neo Small Dutch(book finetune)",
-        "task": "text-generation",
-    },
-    {
-        "model_name": "yhavinga/gpt2-medium-dutch-nedd",
-        "desc": "GPT2 Medium Dutch (book finetune)",
-        "task": "text-generation",
-    },
-    {
-        "model_name": "yhavinga/t5-small-24L-ccmatrix-multi",
-        "desc": "Dutch<->English T5 small 24 layers",
-        "task": "translation_nl_to_en",
-    },
-]
-
-generators = GeneratorFactory()
-
-
-def instantiate_generators():
-    for g in GENERATORS:
-        with st.spinner(text=f"Loading the model {g['desc']} ..."):
-            generators.add_generator(**g)
+    def gpt_descs(self):
+        return [g.desc for g in self.generators if g.task == "text-generation"]
 
 
 def main():
@@ -109,7 +118,11 @@ def main():
         initial_sidebar_state="expanded",  # Can be "auto", "expanded", "collapsed"
         page_icon="📚",  # String, anything supported by st.image, or None.
     )
-    instantiate_generators()
+
+    if "generators" not in st.session_state:
+        st.session_state["generators"] = GeneratorFactory().instantiate_generators()
+
+    generators = st.session_state["generators"]
 
     with open("style.css") as f:
         st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
@@ -119,18 +132,11 @@ def main():
         """# Netherator
     Nederlandse verhalenverteller"""
     )
-
-    model_desc = st.sidebar.selectbox(
-        "Model", [p["desc"] for p in GENERATORS if "generation" in p["task"]], index=1
-    )
-
+    model_desc = st.sidebar.selectbox("Model", generators.gpt_descs(), index=1)
     st.sidebar.title("Parameters:")
-
     if "prompt_box" not in st.session_state:
         st.session_state["prompt_box"] = "Het was een koude winterdag"
-
     st.session_state["text"] = st.text_area("Enter text", st.session_state.prompt_box)
-
     max_length = st.sidebar.number_input(
         "Lengte van de tekst",
         value=200,
@@ -145,7 +151,6 @@ def main():
     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 {st.session_state} does not contain seed")
@@ -231,20 +236,37 @@ and the [Huggingface text generation interface doc](https://huggingface.co/trans
             text=f"Please wait ~ {estimate} second{'s' if estimate != 1 else ''} while getting results ..."
         ):
             memory = psutil.virtual_memory()
-            generator = generators.get_generator(model_desc)
+            generator = generators.get_generator(desc=model_desc)
             set_seed(seed)
             time_start = time.time()
             result = generator.get_text(text=st.session_state.text, **params)
             time_end = time.time()
             time_diff = time_end - time_start
-
             st.subheader("Result")
+
             for text in result:
                 st.write(text.get("generated_text").replace("\n", "  \n"))
-
-            # st.text("*Translation*")
-            # translation = translate(result, "en", "nl")
-            # st.write(translation.replace("\n", "  \n"))
+                st.text("*Translation*")
+                translate_params = {
+                    "num_return_sequences": 1,
+                    "num_beams": 4,
+                    "early_stopping": True,
+                    "length_penalty": 1.1,
+                    "max_length": 200,
+                }
+                text_lines = [
+                    "translate Dutch to English: " + t
+                    for t in text.get("generated_text").splitlines()
+                ]
+                translated_lines = [
+                    t["translation_text"]
+                    for t in generators.get_generator(
+                        task=TRANSLATION_NL_TO_EN
+                    ).get_text(text_lines, **translate_params)
+                ]
+                translation = "  \n".join(translated_lines)
+                st.write(translation)
+                st.write("---")
             #
             info = f"""
             ---