optional kwarg pop

This commit is contained in:
Awni Hannun 2024-12-13 21:03:30 -08:00
parent 8c0b4ee7f3
commit a8da484c8b

View File

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