Add streaming, disable translation

* Also upgrade transformers, add sentencepiece

app.py

@@ -10,6 +10,7 @@ from transformers import (
     AutoModelForCausalLM,
     AutoModelForSeq2SeqLM,
     AutoTokenizer,
+    TextIteratorStreamer,
     pipeline,
     set_seed,
 )
@@ -41,6 +42,20 @@ def load_model(model_name, task):
     return tokenizer, model
 
 
+class StreamlitTextIteratorStreamer(TextIteratorStreamer):
+    def __init__(
+        self, output_placeholder, tokenizer, skip_prompt=False, **decode_kwargs
+    ):
+        super().__init__(tokenizer, skip_prompt, **decode_kwargs)
+        self.output_placeholder = output_placeholder
+        self.output_text = ""
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+        self.output_text += text
+        self.output_placeholder.markdown(self.output_text, unsafe_allow_html=True)
+        super().on_finalized_text(text, stream_end)
+
+
 class Generator:
     def __init__(self, model_name, task, desc):
         self.model_name = model_name
@@ -52,18 +67,38 @@ class Generator:
         self.load()
 
     def load(self):
-        if not self.pipeline:
+        if not self.model:
             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,
-            )
 
-    def get_text(self, text: str, **generate_kwargs) -> str:
-        return self.pipeline(text, **generate_kwargs)
+    def generate(self, text: str, streamer=None, **generate_kwargs) -> (str, dict):
+        batch_encoded = self.tokenizer(
+            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"],
+            streamer=streamer,
+            **generate_kwargs,
+        )
+        decoded_preds = self.tokenizer.batch_decode(
+            logits.cpu().numpy(), skip_special_tokens=False
+        )
+
+        def replace_tokens(pred):
+            pred = pred.replace("<pad> ", "").replace("<pad>", "").replace("</s>", "")
+            if hasattr(self.tokenizer, "newline_token"):
+                pred = pred.replace(self.tokenizer.newline_token, "\n")
+            return pred
+
+        decoded_preds = list(map(replace_tokens, decoded_preds))
+        return decoded_preds[0], generate_kwargs
 
 
 class GeneratorFactory:
@@ -82,11 +117,11 @@ class GeneratorFactory:
                 "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,
-            },
+            # {
+            #     "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']} ..."):
@@ -148,12 +183,13 @@ def main():
     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
-    )
+    num_return_sequences = 1
+    # 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")
+        print(f"Session state does not contain seed")
         st.session_state["seed"] = 4162549114
         print(f"Seed is set to: {st.session_state['seed']}")
 
@@ -218,69 +254,62 @@ and the [Huggingface text generation interface doc](https://huggingface.co/trans
     )
 
     if st.button("Run"):
-        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))
-        estimate = int(estimate)
-
-        with st.spinner(
-            text=f"Please wait ~ {estimate} second{'s' if estimate != 1 else ''} while getting results ..."
-        ):
-            memory = psutil.virtual_memory()
-            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*")
-                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"""
-            ---
-            *Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*        
-            *Text generated using seed {seed} in {time_diff:.5} seconds*
-            """
-            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()
+        st.subheader("Result")
+        container = st.container()
+        output_placeholder = container.empty()
+        streaming_enabled = True  # sampling_mode != "Beam Search" or num_beams == 1
+        generator = generators.get_generator(desc=model_desc)
+        streamer = (
+            StreamlitTextIteratorStreamer(output_placeholder, generator.tokenizer)
+            if streaming_enabled
+            else None
+        )
+        set_seed(seed)
+        time_start = time.time()
+        result = generator.generate(
+            text=st.session_state.text, streamer=streamer, **params
+        )
+        time_end = time.time()
+        time_diff = time_end - time_start
+
+        # for text in result:
+        # st.write(text.get("generated_text").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"""
+        ---
+        *Memory: {memory.total / 10**9:.2f}GB, used: {memory.percent}%, available: {memory.available / 10**9:.2f}GB*        
+        *Text generated using seed {seed} in {time_diff:.5} seconds*
+        """
+        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__":

requirements.txt

@@ -1,7 +1,10 @@
--f https://download.pytorch.org/whl/torch_stable.html
-streamlit==1.4.0
-torch==1.6.0+cpu
-torchvision==0.7.0+cpu
-transformers>=4.13.0
+#-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
+git+https://github.com/huggingface/transformers.git@1905384fd576acf4b645a8216907f980b4788d9b
 mtranslate
 psutil
+sentencepiece