feat: show batch generation progress

ml-explore · Aug 23, 2024 · 5105b31 · 5105b31
1 parent ef92993
commit 5105b31
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py
@@ -356,6 +356,8 @@ def generate(
             break
         if is_batch:
             output_toks.append(tokens)
+            if verbose:
+                print(".", end="", flush=True)
         else:
             token = tokens.item()
             logprobs = logprobs.squeeze(0)
@@ -385,6 +387,7 @@ def generate(
         if token_count <= 0:
             print("No tokens generated for this prompt")
         if is_batch:
+            print()
             for p, resp in zip(prompt, response):
                 print("=" * 10)
                 print("Prompt:", p)