Flush the output after each print in streaming mode (#213)

IBM · Dec 12, 2024 · ad048b8 · ad048b8
1 parent 4a4759e
commit ad048b8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/pdl/pdl_scheduler.py b/src/pdl/pdl_scheduler.py
@@ -115,7 +115,7 @@ def schedule(
                             text = stringify(result)
                         else:
                             text = colored(stringify(result), msg.color)
-                        print(text, end="")
+                        print(text, end="", flush=True)
                         todo_next.append((i, gen, None))
                     case YieldBackgroundMessage(background=background):
                         if len(background) > 0 and background[0]["role"] == _LAST_ROLE:
@@ -127,7 +127,7 @@ def schedule(
                         s += "\n".join(
                             [f"{msg['role']}: {msg['content']}" for msg in background]
                         )
-                        print(s, end="")
+                        print(s, end="", flush=True)
                         todo_next.append((i, gen, None))
                     case ModelCallMessage():
                         text_msg = BamModel.generate_text(