Adapt to streaming interface (only when num_beams is equal to 1)

app.py

@@ -4,6 +4,7 @@ import psutil
 import streamlit as st
 import torch
 from langdetect import detect
+from transformers import TextIteratorStreamer
 
 from default_texts import default_texts
 from generator import GeneratorFactory
@@ -60,6 +61,20 @@ GENERATOR_LIST = [
 ]
 
 
+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)
+
+
 def main():
     st.set_page_config(  # Alternate names: setup_page, page, layout
         page_title="Rosetta en/nl",  # String or None. Strings get appended with "• Streamlit".
@@ -132,16 +147,28 @@ and the [Huggingface text generation interface doc](https://huggingface.co/trans
             left.error("Num beams should be a multiple of num beam groups")
             return
 
+        streaming_enabled = num_beams == 1
+        if not streaming_enabled:
+            left.markdown("*`num_beams > 1` so streaming is disabled*")
+
         for generator in generators.filter(task=task):
-            right.markdown(f"🧮 **Model `{generator}`**")
+            model_container = right.container()
+            model_container.markdown(f"🧮 **Model `{generator}`**")
+            output_placeholder = model_container.empty()
+            streamer = (
+                StreamlitTextIteratorStreamer(output_placeholder, generator.tokenizer)
+                if streaming_enabled
+                else None
+            )
             time_start = time.time()
             result, params_used = generator.generate(
-                text=st.session_state.text, **params
+                text=st.session_state.text, streamer=streamer, **params
             )
             time_end = time.time()
             time_diff = time_end - time_start
 
-            right.write(result.replace("\n", "  \n"))
+            if not streaming_enabled:
+                right.write(result.replace("\n", "  \n"))
             text_line = ", ".join([f"{k}={v}" for k, v in params_used.items()])
             right.markdown(f"    🕙 *generated in {time_diff:.2f}s, `{text_line}`*")
 

generator.py

@@ -20,7 +20,7 @@ def get_access_token():
 
 @st.cache(suppress_st_warning=True, allow_output_mutation=True)
 def load_model(model_name):
-    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
     tokenizer = AutoTokenizer.from_pretrained(
         model_name,
         from_flax=True,
@@ -30,19 +30,18 @@ def load_model(model_name):
     if tokenizer.pad_token is None:
         print("Adding pad_token to the tokenizer")
         tokenizer.pad_token = tokenizer.eos_token
-    try:
-        model = AutoModelForSeq2SeqLM.from_pretrained(
-            model_name, use_auth_token=get_access_token()
-        )
-    except EnvironmentError:
+    for framework in [None, "flax", "tf"]:
         try:
             model = AutoModelForSeq2SeqLM.from_pretrained(
-                model_name, from_flax=True, use_auth_token=get_access_token()
+                model_name,
+                from_flax=(framework == "flax"),
+                from_tf=(framework == "tf"),
+                use_auth_token=get_access_token(),
             )
+            break
         except EnvironmentError:
-            model = AutoModelForSeq2SeqLM.from_pretrained(
-                model_name, from_tf=True, use_auth_token=get_access_token()
-            )
+            if framework == "tf":
+                raise
     if device != -1:
         model.to(f"cuda:{device}")
     return tokenizer, model
@@ -89,7 +88,7 @@ class Generator:
             except TypeError:
                 pass
 
-    def generate(self, text: str, **generate_kwargs) -> (str, dict):
+    def generate(self, text: str, streamer=None, **generate_kwargs) -> (str, dict):
         # Replace two or more newlines with a single newline in text
         text = re.sub(r"\n{2,}", "\n", text)
 
@@ -98,7 +97,9 @@ class Generator:
         # if there are newlines in the text, and the model needs line-splitting, split the text and recurse
         if re.search(r"\n", text) and self.split_sentences:
             lines = text.splitlines()
-            translated = [self.generate(line, **generate_kwargs)[0] for line in lines]
+            translated = [
+                self.generate(line, streamer, **generate_kwargs)[0] for line in lines
+            ]
             return "\n".join(translated), generate_kwargs
 
         # if self.tokenizer has a newline_token attribute, replace \n with it
@@ -117,6 +118,7 @@ class Generator:
         logits = self.model.generate(
             batch_encoded["input_ids"],
             attention_mask=batch_encoded["attention_mask"],
+            streamer=streamer,
             **generate_kwargs,
         )
         decoded_preds = self.tokenizer.batch_decode(

requirements.txt

@@ -4,7 +4,7 @@
 protobuf<3.20
 streamlit>=1.4.0,<=1.10.0
 torch
-transformers>=4.13.0
+git+https://github.com/huggingface/transformers.git@1905384fd576acf4b645a8216907f980b4788d9b
 langdetect
 psutil
 jax[cuda]==0.3.16