zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 18:26:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

optional kwarg pop

Browse Source
This commit is contained in:
Awni Hannun 2024-12-13 21:03:30 -08:00
parent 8c0b4ee7f3
commit a8da484c8b
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
llms/mlx_lm/utils.py
View File

@ -503,10 +503,10 @@ def stream_generate(
detokenizer = tokenizer.detokenizer detokenizer = tokenizer.detokenizer
if draft_model is None: if draft_model is None:
kwargs.pop("num_draft_tokens") kwargs.pop("num_draft_tokens", None)
token_generator = generate_step(prompt, model, **kwargs) token_generator = generate_step(prompt, model, **kwargs)
else: else:
kwargs.pop("max_kv_size") kwargs.pop("max_kv_size", None)
token_generator = speculative_generate_step( token_generator = speculative_generate_step(
prompt, model, draft_model, **kwargs prompt, model, draft_model, **kwargs
) )