mirror/exo
1
0
Fork 0
mirror of https://github.com/exo-explore/exo.git synced 2025-12-23 22:27:50 -05:00
Code Issues Packages Projects Releases Wiki Activity

Revert dumb merge mistake

Browse Source
This commit is contained in:
rltakashige
2025-11-06 18:39:08 -08:00
committed by GitHub
parent 6bcac37d98
commit 612f58c78d
1 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/exo/worker/runner/generate.py
View File

@@ -215,11 +215,9 @@ def generate_step(
y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
mx.async_eval(y, logprobs)
next_y, next_logprobs = _step(y)
mx.async_eval(next_y, next_logprobs)
next_y: array | None = None
next_logprobs: array | None = None
n = 0
while True:
@@ -236,11 +234,8 @@ def generate_step(
yield int(y.item()), logprobs
if n % 256 == 0:
mx.clear_cache()
if n == max_tokens:
break
y, logprobs = next_y, logprobs
next_y, next_logprobs = _step(y)
mx.async_eval(next_y, next_logprobs)
y, logprobs = next_y, next_logprobs
n += 1
def stream_generate(