fix streamer problems (#3601)

* fix streamer problems

* format correction

---------

Co-authored-by: Chunyuan WU <[email protected]>

Author: ZailiWang

examples/cpu/llm/inference/distributed/run_generation_tp.py

@@ -248,11 +248,16 @@
 model = model.to(memory_format=torch.channels_last)
 
 num_beams = 1 if args.greedy else 4
-# generate args
+streamer = None
 if args.streaming:
-    streamer = TextStreamer(tokenizer)
-else:
-    streamer = None
+    if num_beams != 1 or args.batch_size != 1:
+        logger.warning(
+            "--streaming only supported in greedy search mode (--greedy) with --batch-size 1. Disabling streaming output."
+        )
+    else:
+        streamer = TextStreamer(tokenizer)
+
+# generate args
 generate_kwargs = dict(
     do_sample=False,
     temperature=0.9,

examples/cpu/llm/inference/distributed/run_generation_with_deepspeed.py

@@ -713,10 +713,15 @@ def write_checkpoints_json():
 # Generate
 print_rank0(f"*** Starting to generate {num_tokens} tokens with bs={args.batch_size}")
 
+streamer = None
 if args.streaming:
-    streamer = TextStreamer(tokenizer)
-else:
-    streamer = None
+    if num_beams != 1 or args.batch_size != 1:
+        logger.warning(
+            "--streaming only supported in greedy search mode (--greedy) with --batch-size 1. Disabling streaming output."
+        )
+    elif local_rank == 0:
+        streamer = TextStreamer(tokenizer)
+
 generate_kwargs = dict(
     do_sample=False,
     num_beams=num_beams,

examples/cpu/llm/inference/single_instance/run_generation.py

@@ -251,11 +251,16 @@
 model = model.eval()
 model = model.to(memory_format=torch.channels_last)
 num_beams = 1 if args.greedy else 4
-# generate args
+streamer = None
 if args.streaming:
-    streamer = TextStreamer(tokenizer)
-else:
-    streamer = None
+    if num_beams != 1 or args.batch_size != 1:
+        logger.warning(
+            "--streaming only supported in greedy search mode (--greedy) with --batch-size 1. Disabling streaming output."
+        )
+    else:
+        streamer = TextStreamer(tokenizer)
+
+# generate args
 generate_kwargs = dict(
     do_sample=False,
     temperature=0.9,

examples/cpu/llm/inference/single_instance/run_quantization.py

@@ -579,10 +579,15 @@ def download_and_open(url: str) -> Image.Image:
 
 tokenizer = model.get_tokenizer()
 print("Data type of the model:", user_model.dtype)
+streamer = None
 if args.streaming:
-    streamer = TextStreamer(tokenizer)
-else:
-    streamer = None
+    if num_beams != 1 or args.batch_size != 1:
+        print(
+            "--streaming only supported in greedy search mode (--greedy) with --batch-size 1. Disabling streaming output."
+        )
+    else:
+        streamer = TextStreamer(tokenizer)
+
 generate_kwargs = dict(
     do_sample=False,
     temperature=0.9,