zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-31 03:19:23 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

nit

Browse Source
This commit is contained in:
Awni Hannun 2024-12-02 12:19:53 -08:00
parent 1036798dac
commit 21a05d14d2
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
llms/mlx_lm/utils.py
View File

@ -345,9 +345,7 @@ def stream_generate(
with wired_limit(model, [generation_stream]): with wired_limit(model, [generation_stream]):
detokenizer.reset() detokenizer.reset()
tic = time.perf_counter() tic = time.perf_counter()
for n, (token, logprobs) in enumerate( for n, (token, logprobs) in enumerate(generate_step(prompt, model, **kwargs)):
generate_step(prompt, model, **kwargs),
):
if n == 0: if n == 0:
prompt_time = time.perf_counter() - tic prompt_time = time.perf_counter() - tic
prompt_tps = prompt.size / prompt_time prompt_tps = prompt.size / prompt_time