implement streaming

app.py

@@ -1,15 +1,16 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from transformers.cache_utils import DynamicCache
 import torch
 import itertools
+from threading import Thread
 
 DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 MODEL_IDS = [
     "70M",
     "160M",
     "410M",
-    "610M",
+    # "610M",
 ]
 MODEL_MAPPING = {
     model_id: f"LinguaCustodia/multilingual-multidomain-fin-mt-{model_id}"
@@ -211,31 +212,28 @@ def translate_with_model(model_name, text, tgt_lang, src_lang, domain):
     src_lang_token_pos = domain_token_pos - 1
     _tgt_lang_token_pos = src_lang_token_pos - 1
 
-    outputs = model.generate(
+    streamer = TextIteratorStreamer(TOKENIZER, skip_prompt=True)
+    generation_kwargs = dict(
         input_ids=inputs["input_ids"],
         attention_mask=inputs["attention_mask"],
         num_beams=1,
         max_new_tokens=500,
-        pad_token_id=TOKENIZER.pad_token_id,
-        eos_token_id=TOKENIZER.eos_token_id,
         past_key_values=past_key_values,
+        streamer=streamer,
     )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
 
-    generated_translation = TOKENIZER.decode(
-        outputs[0, domain_token_pos + 1 :], skip_special_tokens=True
-    )
-
-    source_language_token = TOKENIZER.convert_ids_to_tokens(
-        outputs[0, src_lang_token_pos].item()
-    )
-    dom_token = TOKENIZER.convert_ids_to_tokens(outputs[0, domain_token_pos].item())
+    generated_translation = ""
+    for new_text in streamer:
+        generated_translation += new_text.replace(TOKENIZER.eos_token, "")
 
-    return {
-        "model": model_name,
-        "source_lang": CODE2LANG[language_token_to_str(source_language_token)],
-        "domain": DOMAIN_MAPPING_REVERSED[domain_token_to_str(dom_token)],
-        "translation": generated_translation,
-    }
+        yield {
+            "model": model_name,
+            "source_lang": CODE2LANG[src_lang],
+            "domain": DOMAIN_MAPPING_REVERSED[domain],
+            "translation": generated_translation,
+        }
 
 
 def translate_with_all_models(selected_models, text, tgt_lang, src_lang, domain):
@@ -257,11 +255,13 @@ def translate_with_all_models(selected_models, text, tgt_lang, src_lang, domain)
 
     for model_id in selected_models:
         i = MODEL_INDEX[model_id]
-        model_output = translate_with_model(model_id, text, tgt_lang, src_lang, domain)
-        outputs[i * 3] = model_output["translation"]
-        outputs[i * 3 + 1] = model_output["source_lang"]
-        outputs[i * 3 + 2] = model_output["domain"]
-        yield outputs
+        for model_output in translate_with_model(
+            model_id, text, tgt_lang, src_lang, domain
+        ):
+            outputs[i * 3] = model_output["translation"]
+            outputs[i * 3 + 1] = model_output["source_lang"]
+            outputs[i * 3 + 2] = model_output["domain"]
+            yield outputs
 
 
 with gr.Blocks() as demo: