implement stream generating

src/cli_demo.py

@@ -3,12 +3,13 @@
 # Usage: python cli_demo.py --checkpoint_dir path_to_checkpoint
 
 
-import torch
 from utils import (
     load_pretrained,
     prepare_infer_args,
     get_logits_processor
 )
+from threading import Thread
+from transformers import TextIteratorStreamer
 
 
 def main():
@@ -34,25 +35,32 @@ def main():
         return prompt
 
     format_example = format_example_alpaca if data_args.prompt_template == "alpaca" else format_example_ziya
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
-    def predict(query, history: list):
+    def predict_and_print(query, history: list):
         input_ids = tokenizer([format_example(query, history)], return_tensors="pt")["input_ids"]
         input_ids = input_ids.to(model.device)
         gen_kwargs = {
+            "input_ids": input_ids,
             "do_sample": True,
             "top_p": 0.7,
             "temperature": 0.95,
             "num_beams": 1,
             "max_new_tokens": 256,
             "repetition_penalty": 1.0,
-            "logits_processor": get_logits_processor()
+            "logits_processor": get_logits_processor(),
+            "streamer": streamer
         }
-        with torch.no_grad():
-            generation_output = model.generate(input_ids=input_ids, **gen_kwargs)
-        outputs = generation_output.tolist()[0][len(input_ids[0]):]
-        response = tokenizer.decode(outputs, skip_special_tokens=True)
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+        response = ""
+        print("{}: ".format(model_name), end="")
+        for new_text in streamer:
+            print(new_text, end="", flush=True)
+            response += new_text
+        print()
         history = history + [(query, response)]
-        return response, history
+        return history
 
     history = []
     print("欢迎使用 {} 模型，输入内容即可对话，clear清空对话历史，stop终止程序".format(model_name))
@@ -73,8 +81,7 @@ def main():
             print("History has been removed.")
             continue
 
-        response, history = predict(query, history)
-        print("{}:".format(model_name), response)
+        history = predict_and_print(query, history)
 
 
 if __name__ == "__main__":