GitHub

Original file line numberDiff line numberDiff line change

@@ -365,10 +365,9 @@ def prefill(model, toks, start_pos=0):

365365

last_seen_toks = toks

366366

toks = toks[i:]

367367
368-

# prefill the model

369-

for tok in tqdm(toks):

370-

model(Tensor([[tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).realize()

371-

start_pos += 1

368+

# batch prefill the model

369+

model(Tensor([toks], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).realize()

370+

start_pos += len(toks)

372371

return start_pos

373372
374373

if __name__ == "__main__":

Read the original on github.com ↗