File size: 1,989 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """Batched greedy generation for benchmarks (no KV cache: simple, not fast)."""
import torch
from ar.model.mask import build_doc_block_mask
@torch.no_grad()
def generate_batch(model, tok, prompts, seq_len, max_new, stop, device):
"""Greedy-decode a batch of prompts; each output is cut at its first stop string."""
B = len(prompts)
# fixed-size buffer: each row has its own prompt length and write position;
# the causal mask hides the placeholder slots to the right of each position
buf = torch.full((B, seq_len), tok.eot_id, dtype=torch.long, device=device)
pos = []
for b, p in enumerate(prompts):
x = tok.encode(p)[-(seq_len - max_new):] # keep the end if a prompt is too long
buf[b, :len(x)] = torch.tensor(x, device=device)
pos.append(len(x))
start = list(pos)
seg = torch.zeros(B, seq_len, dtype=torch.int32, device=device) # one doc per row
mask = build_doc_block_mask(seg, seq_len, device=device)
done, outs = [False] * B, [""] * B
rows = torch.arange(B, device=device)
for _ in range(max_new):
logits = model(buf, mask)
# each row's next token = most likely token at its own current position
idx = torch.tensor(pos, device=device) - 1
nxt = logits[rows, idx].argmax(-1).tolist()
for b in range(B):
if done[b]:
continue
if nxt[b] == tok.eot_id: # model ended the text
done[b] = True
continue
buf[b, pos[b]] = nxt[b]
pos[b] += 1
outs[b] = tok.decode(buf[b, start[b]:pos[b]].tolist())
if pos[b] >= seq_len or any(s in outs[b] for s in stop):
done[b] = True
if all(done):
break
# cut each output at its first stop string
for b in range(B):
hits = [outs[b].find(s) for s in stop if s in outs[b]]
outs[b] = outs[b][: min(hits)] if hits else outs[b]
return outs |