zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-30 10:56:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: show batch generation progress

Browse Source
This commit is contained in:
L Lllvvuu 2024-08-23 16:27:50 +09:00
parent 280b3784d4
commit 2caa8329c0
No known key found for this signature in database
GPG Key ID: CFAD5A25056DDD0F
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
llms/mlx_lm/utils.py
View File

@ -375,6 +375,8 @@ def generate(
break break
if is_batch: if is_batch:
output_toks.append(tokens) output_toks.append(tokens)
if verbose:
print(".", end="", flush=True)
else: else:
token = tokens.item() token = tokens.item()
logprobs = logprobs.squeeze(0) logprobs = logprobs.squeeze(0)
@ -404,6 +406,7 @@ def generate(
if token_count <= 0: if token_count <= 0:
print("No tokens generated for this prompt") print("No tokens generated for this prompt")
if is_batch: if is_batch:
print()
for p, resp in zip(prompt, response): for p, resp in zip(prompt, response):
print("=" * 10) print("=" * 10)
print("Prompt:", p) print("Prompt:", p)