Rebase and fix streaming

neuralmagic · Oct 11, 2023 · 8c2e74f · 8c2e74f
1 parent 9886574
commit 8c2e74f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/deepsparse/transformers/inference/infer.py b/src/deepsparse/transformers/inference/infer.py
@@ -212,10 +212,10 @@ def _display_generation_speed(prompt_sequence_length, pipeline):
 def _display_bot_response(stream: bool, response):
     # print response from pipeline, streaming or not
 
-    print("Bot:", end="")
+    print("Bot:", end="", flush=True)
     if stream:
         for generation in response:
-            print(generation.generations[0].text, end="")
+            print(generation.generations[0].text, end="", flush=True)
         print()
     else:
         print(response.generations[0].text)