| """Batched sampling for many prompts — used by DPO (building preference pairs) and evaluation.""" |
| from __future__ import annotations |
|
|
| from collections import defaultdict |
|
|
| import torch |
|
|
| from common import autocast_ctx, build_prompt |
| from tokenizer import BPETokenizer |
|
|
|
|
| @torch.no_grad() |
| def sample_responses(model, tok: BPETokenizer, word_lists: list[list[str]], n_samples: int, max_tokens: int, |
| temperature: float, top_k: int, device: str, batch_rows: int = 64): |
| """ |
| Returns: list (per prompt) of list (per sample) of dict {ids, text, finished}. |
| Prompts are grouped by token length so they can be batched without padding. |
| """ |
| model.eval() |
| ctx, _ = autocast_ctx(device) |
| prompts = [tok.encode(build_prompt(w)) for w in word_lists] |
| groups = defaultdict(list) |
| for i, p in enumerate(prompts): |
| groups[len(p)].append(i) |
|
|
| results: list[list[dict]] = [[] for _ in prompts] |
| for plen, idxs in groups.items(): |
| rows = [i for i in idxs for _ in range(n_samples)] |
| budget = min(max_tokens, model.cfg.context_len - plen) |
| for s in range(0, len(rows), batch_rows): |
| chunk = rows[s:s + batch_rows] |
| x = torch.tensor([prompts[i] for i in chunk], device=device) |
| with ctx: |
| out, finished = model.generate(x, budget, temperature=temperature, top_k=top_k, |
| eos_id=tok.eos_id, valid_vocab=tok.vocab_size) |
| for r, i in enumerate(chunk): |
| ids = out[r].tolist() |
| if tok.eos_id in ids: |
| ids = ids[: ids.index(tok.eos_id) + 1] |
| text = tok.decode([t for t in ids if t != tok.eos_id]) |
| results[i].append({"ids": ids, "text": text, "finished": bool(finished[r])}) |
| return prompts, results |
|
|
