From 5105b31cf75c81369ba21ed2692e547317bb5bdf Mon Sep 17 00:00:00 2001
From: L Lllvvuu <git@llllvvuu.dev>
Date: Fri, 23 Aug 2024 16:27:50 +0900
Subject: [PATCH] feat: show batch generation progress

---
 llms/mlx_lm/utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/llms/mlx_lm/utils.py b/llms/mlx_lm/utils.py
index 10609ecc3..58ab084fb 100644
--- 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)