Repository Public Research preview
1"""Shared dataset primitives: examples, vocabularies, split containers and hashing.23Every example is a token sequence ``<bos> <prompt tokens> = <a_1 ... a_k>`` where the ``k`` answer4tokens form a *scratchpad* (one intermediate result per execution step; k = 1 for single-answer5tasks). The primary metric is exact match of the final answer token under greedy decoding; the6trace metric is exact match of all ``k`` answer tokens. ``cluster`` identifies the7generating template (e.g. the operator program) and is the unit of resampling in the cluster8bootstrap: items sharing a cluster are not treated as independent.9"""10from __future__ import annotations1112import hashlib13import json14from dataclasses import asdict, dataclass, field15from pathlib import Path16from typing import Iterable1718import numpy as np1920PAD, BOS, EQ = "<pad>", "<bos>", "="212223@dataclass(frozen=True)24class Example:25 tokens: tuple[str, ...] # full sequence including the answer tokens26 n_answer: int # number of answer (scratchpad) tokens at the end of ``tokens``27 cluster: str # template / program identifier (bootstrap unit)28 meta: dict = field(default_factory=dict)2930 @property31 def prompt(self) -> tuple[str, ...]:32 return self.tokens[: -self.n_answer]3334 @property35 def answers(self) -> tuple[str, ...]:36 return self.tokens[-self.n_answer :]3738 @property39 def answer(self) -> str:40 return self.tokens[-1]4142 @property43 def item_hash(self) -> str:44 return hashlib.sha256(" ".join(self.tokens).encode()).hexdigest()[:16]454647class Vocab:48 def __init__(self, tokens: Iterable[str]) -> None:49 base = [PAD, BOS, EQ]50 extra = [t for t in dict.fromkeys(tokens) if t not in base]51 self.itos: list[str] = base + extra52 self.stoi: dict[str, int] = {t: i for i, t in enumerate(self.itos)}5354 def __len__(self) -> int:55 return len(self.itos)5657 @property58 def pad_id(self) -> int:59 return self.stoi[PAD]6061 def encode(self, tokens: Iterable[str]) -> list[int]:62 return [self.stoi[t] for t in tokens]6364 def decode(self, ids: Iterable[int]) -> list[str]:65 return [self.itos[i] for i in ids]6667 def to_json(self) -> str:68 return json.dumps(self.itos)6970 @classmethod71 def from_json(cls, text: str) -> "Vocab":72 return cls(json.loads(text))737475@dataclass76class Split:77 name: str78 examples: list[Example]79 generator: str80 generator_version: str81 seed: int82 config: dict8384 def __len__(self) -> int:85 return len(self.examples)8687 @property88 def hash(self) -> str:89 return split_hash(self.examples)9091 def item_hashes(self) -> set[str]:92 return {e.item_hash for e in self.examples}9394 def answer_distribution(self) -> dict[str, float]:95 vals, counts = np.unique([e.answer for e in self.examples], return_counts=True)96 return {str(v): float(c) / len(self.examples) for v, c in zip(vals, counts)}9798 def save(self, path: Path) -> None:99 path.parent.mkdir(parents=True, exist_ok=True)100 payload = {101 "name": self.name,102 "generator": self.generator,103 "generator_version": self.generator_version,104 "seed": self.seed,105 "config": self.config,106 "hash": self.hash,107 "n": len(self.examples),108 "examples": [asdict(e) for e in self.examples],109 }110 path.write_text(json.dumps(payload, sort_keys=True))111112 @classmethod113 def load(cls, path: Path) -> "Split":114 d = json.loads(path.read_text())115 ex = [Example(tuple(e["tokens"]), int(e["n_answer"]), e["cluster"], e.get("meta", {})) for e in d["examples"]]116 s = cls(d["name"], ex, d["generator"], d["generator_version"], d["seed"], d["config"])117 if s.hash != d["hash"]:118 raise ValueError(f"split hash mismatch for {path}: stored {d['hash']} computed {s.hash}")119 return s120121122def split_hash(examples: Iterable[Example]) -> str:123 h = hashlib.sha256()124 for e in examples:125 h.update((" ".join(e.tokens) + "\n").encode())126 return h.hexdigest()[:16]127128129def assert_disjoint(*splits: Split) -> None:130 seen: dict[str, str] = {}131 for s in splits:132 for e in s.examples:133 if e.item_hash in seen:134 raise AssertionError(f"item leakage: {e.tokens} present in {seen[e.item_hash]} and {s.name}")135 seen[e.item_hash] = s.name136137138def encode_split(split: Split, vocab: Vocab, max_len: int) -> dict[str, np.ndarray]:139 """Encode a split for teacher-forced training and greedy evaluation.140141 Returns a dict of arrays:142 ``ids`` [N,T] full sequences (prompt + answers), right-padded with ``<pad>``;143 ``loss_mask`` [N,T] 1 at positions whose *next* token is an answer token;144 ``prompt_len`` [N] length of the prompt (index of the first answer token);145 ``n_answer`` [N] number of answer tokens;146 ``answers`` [N,K] answer token ids right-padded with ``<pad>`` (K = max n_answer).147 """148 n = len(split.examples)149 K = max(e.n_answer for e in split.examples)150 ids = np.full((n, max_len), vocab.pad_id, dtype=np.int64)151 loss_mask = np.zeros((n, max_len), dtype=np.float32)152 prompt_len = np.zeros(n, dtype=np.int64)153 n_answer = np.zeros(n, dtype=np.int64)154 answers = np.full((n, K), vocab.pad_id, dtype=np.int64)155 for i, e in enumerate(split.examples):156 seq = vocab.encode(e.tokens)157 if len(seq) > max_len:158 raise ValueError(f"sequence longer than max_len={max_len}: {e.tokens}")159 ids[i, : len(seq)] = seq160 pl = len(e.prompt)161 assert e.prompt[-1] == EQ162 prompt_len[i] = pl163 n_answer[i] = e.n_answer164 loss_mask[i, pl - 1 : pl - 1 + e.n_answer] = 1.0165 answers[i, : e.n_answer] = vocab.encode(e.answers)166 return {"ids": ids, "loss_mask": loss_mask, "prompt_len": prompt_len, "n_answer": n_answer, "answers": answers}167