AR / ar /eval /batch_gen.py

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 18 days ago

1.99 kB

"""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

	"""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