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

fix stream generate

Browse Source
This commit is contained in:
Awni Hannun 2024-11-05 17:18:00 -08:00
parent f5cd03c64d
commit c9994f80e6
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
llms/mlx_lm/utils.py
View File

@ -300,10 +300,9 @@ def stream_generate(
range(max_tokens), range(max_tokens),
generate_step(prompt_tokens, model, **kwargs), generate_step(prompt_tokens, model, **kwargs),
): ):
if token == tokenizer.eos_token_id:
break
detokenizer.add_token(token) detokenizer.add_token(token)
if n == (max_tokens - 1) or token == tokenizer.eos_token_id:
break
# Yield the last segment if streaming # Yield the last segment if streaming
yield detokenizer.last_segment, token, logits yield detokenizer.last_segment, token, logits