Syntactic changes

app.py

@@ -1,20 +1,24 @@
 import json
 import os
-import pprint
 import time
 from random import randint
 
 import psutil
 import streamlit as st
 import torch
-from transformers import (AutoModelForCausalLM, AutoTokenizer, pipeline,
-                          set_seed)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    pipeline,
+    set_seed,
+)
 
 device = torch.cuda.device_count() - 1
 
 
 @st.cache(suppress_st_warning=True, allow_output_mutation=True)
-def load_model(model_name):
+def load_model(model_name, task):
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
     try:
         if not os.path.exists(".streamlit/secrets.toml"):
@@ -23,70 +27,68 @@ def load_model(model_name):
     except FileNotFoundError:
         access_token = os.environ.get("HF_ACCESS_TOKEN", None)
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 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, use_auth_token=access_token)
     if device != -1:
         model.to(f"cuda:{device}")
     return tokenizer, model
 
 
-class StoryGenerator:
-    def __init__(self, model_name):
-        self.model_name = model_name
+class ModelTask:
+    def __init__(self, p):
+        self.model_name = p["model_name"]
+        self.task = p["task"]
+        self.desc = p["desc"]
         self.tokenizer = None
         self.model = None
-        self.generator = None
-        self.model_loaded = False
+        self.pipeline = None
+        self.load()
 
     def load(self):
-        if not self.model_loaded:
-            self.tokenizer, self.model = load_model(self.model_name)
-            self.generator = pipeline(
-                "text-generation",
+        if not self.pipeline:
+            print(f"Loading model {self.model_name}")
+            self.tokenizer, self.model = load_model(self.model_name, self.task)
+            self.pipeline = pipeline(
+                task=self.task,
                 model=self.model,
                 tokenizer=self.tokenizer,
                 device=device,
             )
-            self.model_loaded = True
 
     def get_text(self, text: str, **generate_kwargs) -> str:
-        return self.generator(text, **generate_kwargs)
+        return self.pipeline(text, **generate_kwargs)
 
 
-STORY_GENERATORS = [
+PIPELINES = [
     {
         "model_name": "yhavinga/gpt-neo-125M-dutch-nedd",
         "desc": "Dutch GPTNeo Small",
-        "story_generator": None,
+        "task": "text-generation",
+        "pipeline": None,
     },
     {
         "model_name": "yhavinga/gpt2-medium-dutch-nedd",
         "desc": "Dutch GPT2 Medium",
-        "story_generator": None,
+        "task": "text-generation",
+        "pipeline": None,
     },
-    # {
-    #     "model_name": "yhavinga/gpt-neo-125M-dutch",
-    #     "desc": "Dutch GPTNeo Small",
-    #     "story_generator": None,
-    # },
-    # {
-    #     "model_name": "yhavinga/gpt2-medium-dutch",
-    #     "desc": "Dutch GPT2 Medium",
-    #     "story_generator": None,
-    # },
 ]
 
 
 def instantiate_models():
-    for sg in STORY_GENERATORS:
-        sg["story_generator"] = StoryGenerator(sg["model_name"])
-        with st.spinner(text=f"Loading the model {sg['desc']} ..."):
-            sg["story_generator"].load()
+    for p in PIPELINES:
+        p["pipeline"] = ModelTask(p)
+        with st.spinner(text=f"Loading the model {p['desc']} ..."):
+            p["pipeline"].load()
 
 
 def set_new_seed():
-    seed = randint(0, 2 ** 32 - 1)
+    seed = randint(0, 2**32 - 1)
     set_seed(seed)
     return seed
 
@@ -104,14 +106,13 @@ def main():
         st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
 
     st.sidebar.image("demon-reading-Stewart-Orr.png", width=200)
-
     st.sidebar.markdown(
         """# Netherator
-Teller of tales from the Netherlands"""
+    Nederlandse verhalenverteller"""
     )
 
     model_desc = st.sidebar.selectbox(
-        "Model", [sg["desc"] for sg in STORY_GENERATORS], index=1
+        "Model", [p["desc"] for p in PIPELINES], index=1
     )
 
     st.sidebar.title("Parameters:")
@@ -126,7 +127,7 @@ Teller of tales from the Netherlands"""
     # )
     max_length = st.sidebar.number_input(
         "Lengte van de tekst",
-        value=300,
+        value=200,
         max_value=512,
     )
     no_repeat_ngram_size = st.sidebar.number_input(
@@ -147,7 +148,7 @@ Teller of tales from the Netherlands"""
                 "Num beams", min_value=1, max_value=10, value=4
             )
             length_penalty = st.sidebar.number_input(
-                "Length penalty", min_value=0.0, max_value=5.0, value=1.5, step=0.1
+                "Length penalty", min_value=0.0, max_value=2.0, value=1.0, step=0.1
             )
             params = {
                 "max_length": max_length,
@@ -159,14 +160,12 @@ Teller of tales from the Netherlands"""
                 "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
             )
             temperature = st.sidebar.number_input(
-                "Temperature", min_value=0.05, max_value=1.0, value=0.8, step=0.05
+                "Temperature", min_value=0.05, max_value=1.0, value=1.0, step=0.05
             )
             params = {
                 "max_length": max_length,
@@ -204,17 +203,17 @@ 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()
-            story_generator = next(
+            generator = next(
                 (
-                    x["story_generator"]
-                    for x in STORY_GENERATORS
+                    x["pipeline"]
+                    for x in PIPELINES
                     if x["desc"] == model_desc
                 ),
                 None,
             )
             seed = set_new_seed()
             time_start = time.time()
-            result = story_generator.get_text(text=st.session_state.text, **params)
+            result = generator.get_text(text=st.session_state.text, **params)
             time_end = time.time()
             time_diff = time_end - time_start
 
@@ -235,7 +234,7 @@ and the [Huggingface text generation interface doc](https://huggingface.co/trans
 
             params["seed"] = seed
             params["prompt"] = st.session_state.text
-            params["model"] = story_generator.model_name
+            params["model"] = generator.model_name
             params_text = json.dumps(params)
             print(params_text)
             st.json(params_text)