Refactor model+task code using a factory. Run black

app.py

@@ -6,13 +6,8 @@ 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
 
@@ -39,11 +34,11 @@ def load_model(model_name, task):
     return tokenizer, model
 
 
-class ModelTask:
-    def __init__(self, p):
-        self.model_name = p["model_name"]
-        self.task = p["task"]
-        self.desc = p["desc"]
+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.pipeline = None
@@ -64,27 +59,47 @@ class ModelTask:
         return self.pipeline(text, **generate_kwargs)
 
 
-PIPELINES = [
+class GeneratorFactory:
+    def __init__(self):
+        self.generators = []
+
+    def add_generator(self, model_name, task, desc):
+        g = Generator(model_name, task, desc)
+        g.load()
+        self.generators.append(g)
+
+    def get_generator(self, model_desc):
+        for g in self.generators:
+            if g.desc == model_desc:
+                return g
+        return None
+
+
+GENERATORS = [
     {
         "model_name": "yhavinga/gpt-neo-125M-dutch-nedd",
-        "desc": "Dutch GPTNeo Small",
+        "desc": "GPT-Neo Small Dutch(book finetune)",
         "task": "text-generation",
-        "pipeline": None,
     },
     {
         "model_name": "yhavinga/gpt2-medium-dutch-nedd",
-        "desc": "Dutch GPT2 Medium",
+        "desc": "GPT2 Medium Dutch (book finetune)",
         "task": "text-generation",
-        "pipeline": None,
+    },
+    {
+        "model_name": "yhavinga/t5-small-24L-ccmatrix-multi",
+        "desc": "Dutch<->English T5 small 24 layers",
+        "task": "translation_nl_to_en",
     },
 ]
 
+generators = GeneratorFactory()
+
 
-def instantiate_models():
-    for p in PIPELINES:
-        p["pipeline"] = ModelTask(p)
-        with st.spinner(text=f"Loading the model {p['desc']} ..."):
-            p["pipeline"].load()
+def instantiate_generators():
+    for g in GENERATORS:
+        with st.spinner(text=f"Loading the model {g['desc']} ..."):
+            generators.add_generator(**g)
 
 
 def main():
@@ -94,7 +109,7 @@ def main():
         initial_sidebar_state="expanded",  # Can be "auto", "expanded", "collapsed"
         page_icon="📚",  # String, anything supported by st.image, or None.
     )
-    instantiate_models()
+    instantiate_generators()
 
     with open("style.css") as f:
         st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
@@ -106,7 +121,7 @@ def main():
     )
 
     model_desc = st.sidebar.selectbox(
-        "Model", [p["desc"] for p in PIPELINES], index=1
+        "Model", [p["desc"] for p in GENERATORS if "generation" in p["task"]], index=1
     )
 
     st.sidebar.title("Parameters:")
@@ -138,13 +153,13 @@ def main():
         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"]
+        "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)
+        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"]
+            "Seed", min_value=0, max_value=2**32 - 1, value=st.session_state["seed"]
         )
         print(f"New random seed set to: {seed}")
 
@@ -152,7 +167,7 @@ def main():
         set_random_seed()
 
     if sampling_mode := st.sidebar.selectbox(
-            "select a Mode", index=0, options=["Top-k Sampling", "Beam Search"]
+        "select a Mode", index=0, options=["Top-k Sampling", "Beam Search"]
     ):
         if sampling_mode == "Beam Search":
             num_beams = st.sidebar.number_input(
@@ -171,7 +186,9 @@ def main():
                 "length_penalty": length_penalty,
             }
         else:
-            top_k = st.sidebar.number_input("Top K", min_value=0, max_value=100, value=50)
+            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
             )
@@ -211,17 +228,10 @@ and the [Huggingface text generation interface doc](https://huggingface.co/trans
         estimate = int(estimate)
 
         with st.spinner(
-                text=f"Please wait ~ {estimate} second{'s' if estimate != 1 else ''} while getting results ..."
+            text=f"Please wait ~ {estimate} second{'s' if estimate != 1 else ''} while getting results ..."
         ):
             memory = psutil.virtual_memory()
-            generator = next(
-                (
-                    x["pipeline"]
-                    for x in PIPELINES
-                    if x["desc"] == model_desc
-                ),
-                None,
-            )
+            generator = generators.get_generator(model_desc)
             set_seed(seed)
             time_start = time.time()
             result = generator.get_text(text=st.session_state.text, **params)