| """D2-DATA d2.manifest.v1 reader and small synthetic internal block fixtures.""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import bisect |
| from pathlib import Path |
| from typing import Protocol, TypedDict |
|
|
| import torch |
| from torch import Tensor |
| from safetensors import safe_open |
|
|
| from .model import Config |
| from .synthetic import inputs |
|
|
| INPUT_NAMES = ("taps", "engine_mx", "has_taps", "positions", "mask", "anchor", "P") |
| TARGET_NAMES = ("previous", "top_ids", "top_logp", "tail_mass", "labels", "present", "position_mask") |
|
|
|
|
| class Block(TypedDict): |
| inputs: dict[str, Tensor] |
| targets: dict[str, Tensor] |
| request: str |
|
|
|
|
| class Blocks(Protocol): |
| fingerprint: str |
| def __len__(self) -> int: ... |
| def __getitem__(self, index: int) -> Block: ... |
|
|
|
|
| def collate(blocks: list[Block]) -> Block: |
| if not blocks: |
| raise ValueError("empty block batch") |
| width = max(b["inputs"]["positions"].shape[1] for b in blocks) |
| batches: dict[str, list[Tensor]] = {n: [] for n in INPUT_NAMES} |
| for block in blocks: |
| inp = block["inputs"] |
| pad = width-inp["positions"].shape[1] |
| for name in INPUT_NAMES: |
| t = inp[name] |
| if pad and name not in ("anchor", "P"): |
| if name == "positions": |
| before = torch.arange(int(t[0, 0])-pad, int(t[0, 0]), dtype=t.dtype)[None] |
| else: |
| before = torch.zeros(1, pad, *t.shape[2:], dtype=t.dtype) |
| t = torch.cat((before, t), 1) |
| batches[name].append(t) |
| return {"inputs": {n: torch.cat(v, 0) for n, v in batches.items()}, |
| "targets": {n: torch.cat([b["targets"][n] for b in blocks], 0) for n in TARGET_NAMES}, |
| "request": "batch"} |
|
|
|
|
| def sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with path.open("rb") as f: |
| for b in iter(lambda: f.read(4 << 20), b""): |
| h.update(b) |
| return h.hexdigest() |
|
|
|
|
| def source_families(split: str, units: dict[str, list[tuple[str, str]]]) -> dict[str, str]: |
| """Union requests sharing source documents (or a source-less prompt identity). |
| |
| Run after reading ALL shard sidecars of a split: a later request can bridge |
| two earlier components. Hash sorted component units so traversal order and |
| task labels cannot change the resampling key. Never expose source IDs. |
| """ |
| parent = {request: request for request in units} |
|
|
| def find(request: str) -> str: |
| while parent[request] != request: |
| parent[request] = parent[parent[request]] |
| request = parent[request] |
| return request |
|
|
| owners: dict[tuple[str, str], str] = {} |
| for request, sources in units.items(): |
| for unit in sources: |
| other = owners.setdefault(unit, request) |
| a, b = find(request), find(other) |
| parent[max(a, b)] = min(a, b) |
| components: dict[str, set[tuple[str, str]]] = {} |
| for request, sources in units.items(): |
| components.setdefault(find(request), set()).update(sources) |
| identities = {root: hashlib.sha256(json.dumps([split, sorted(sources)], separators=(",", ":")).encode()).hexdigest() |
| for root, sources in components.items()} |
| return {request: identities[find(request)] for request in units} |
|
|
|
|
| class SyntheticBlocks: |
| def __init__(self, c: Config, count: int = 8): |
| self.c, self.count = c, count |
| self.fingerprint = f"synthetic-v1-{c}-{count}" |
|
|
| def __len__(self) -> int: |
| return self.count |
|
|
| def __getitem__(self, index: int) -> Block: |
| inp = inputs(self.c, batch=1, P=12+index) |
| g = torch.Generator().manual_seed(111+index) |
| full = torch.randn(1, 5, self.c.vocab, generator=g).softmax(-1) |
| p, ids = full.topk(20, -1) |
| labels = torch.multinomial(full.flatten(0, 1), 1, generator=g).reshape(1, 5) |
| previous = torch.cat((inp["anchor"][:, None], labels[:, :-1]), -1) |
| target = {"previous": previous, "top_ids": ids, "top_logp": p.log(), "tail_mass": 1-p.sum(-1), |
| "labels": labels, "present": torch.ones(1, 5, dtype=torch.bool), |
| "position_mask": torch.ones(1, 5, dtype=torch.bool)} |
| return {"inputs": inp, "targets": target, "request": f"synthetic-{index}"} |
|
|
|
|
| class D2Blocks: |
| """Memory-map one bounded D2 shard at a time; hash all split metadata first. |
| |
| Blocks follow DATASET.md §4 exactly: query anchor P+1, targets P+2..P+6, |
| and teacher-forced Markov input P+1..P+5. Tail blocks are position masked. |
| """ |
| def __init__(self, manifest: Path, split: str): |
| self.root = manifest.parent.resolve() |
| self.manifest = json.loads(manifest.read_text()) |
| if self.manifest.get("schema") != "d2.manifest.v1": |
| raise ValueError("expected D2-DATA d2.manifest.v1") |
| checks = self.manifest.get("checks", {}) |
| if "plumbing_errors" in checks and checks["plumbing_errors"] != 0: |
| raise ValueError("dataset records plumbing errors") |
| self.sequences: list[tuple[Path, list[int], str]] = [] |
| self.request_families: dict[str, str] = {} |
| self.request_tasks: dict[str, str] = {} |
| self.sequence_b: list[int] = [] |
| self.ends: list[int] = [] |
| seen: dict[str, str] = {} |
| docs: dict[str, str] = {} |
| total = 0 |
| for split_name, spec in self.manifest["splits"].items(): |
| split_b = 0 |
| units: dict[str, list[tuple[str, str]]] = {} |
| for entry in spec["files"]: |
| path = self.verified(entry["file"], entry["sha256"]) |
| side = self.verified(entry["sidecar"], entry["sidecar_sha256"]) |
| meta = json.loads(side.read_text()) |
| constants = {"schema": "d2.shard.v1", "split": split_name, "W": 128, "K": 5, "topk": 20, |
| "tap_layers": [37, 38, 39], "hidden": 5120, "noise_token": 128799} |
| if any(meta.get(k) != v for k, v in constants.items()) or meta["sha256"] != entry["sha256"]: |
| raise ValueError(f"D2 shard sidecar mismatch {side}") |
| with safe_open(path, framework="pt", device="cpu") as f: |
| if f.metadata() != {"schema": "d2.shard.v1", "split": split_name}: |
| raise ValueError(f"D2 safetensors metadata mismatch {path}") |
| index = f.get_tensor("seq_index") |
| if index.dtype != torch.int64 or index.shape != (len(meta["sequences"]), 8): |
| raise ValueError(f"D2 sequence index mismatch {path}") |
| expected = (0, 0, 0) |
| for j, seq in enumerate(meta["sequences"]): |
| req = seq["seq_id"] |
| if req in seen: |
| raise ValueError("duplicate request across shards/splits") |
| seen[req] = split_name |
| if not isinstance(seq.get("family"), str) or not seq["family"]: |
| raise ValueError("D2 request missing task label") |
| self.request_tasks[req] = seq["family"] |
| source_docs = seq.get("source_docs") |
| if not isinstance(source_docs, list) or any(not isinstance(d, str) or not d for d in source_docs): |
| raise ValueError("D2 request invalid source documents") |
| if source_docs: |
| units[req] = [("document", doc) for doc in source_docs] |
| else: |
| prompt = seq.get("prompt_sha256") |
| if not isinstance(prompt, str) or not prompt: |
| raise ValueError("D2 source-less request missing prompt identity") |
| units[req] = [("prompt", prompt)] |
| for doc in source_docs: |
| if doc in docs and docs[doc] != split_name: |
| raise ValueError("source-document split leakage") |
| docs[doc] = split_name |
| if split_name in ("train", "dev") and seq["family"] in ("chronology", "motion_outline"): |
| raise ValueError("held-out family in train/dev") |
| row = index[j].tolist() |
| tok_off, L, gen, r0, row_off, nr, span_off, ns = row |
| if not 1 <= gen < L or (L, gen) != (seq["L"], seq["gen_start"]) or r0 != max(0, gen-128) or nr != L-r0 or ns != L-gen: |
| raise ValueError("D2 sequence offset/length mismatch") |
| if (tok_off, row_off, span_off) != expected: |
| raise ValueError("D2 sequence offsets are not cumulative") |
| tokens = f.get_slice("tokens")[tok_off:tok_off+L] |
| if tokens.dtype != torch.int32 or tokens.shape != (L,) or bool(((tokens < 0) | (tokens >= 129280)).any()): |
| raise ValueError("D2 token shape/dtype/vocabulary mismatch") |
| digest = hashlib.sha256(tokens.contiguous().numpy().astype("<i4", copy=False).tobytes()).hexdigest() |
| if digest != seq.get("tokens_sha256"): |
| raise ValueError("D2 per-sequence token hash differs") |
| rp = f.get_slice("row_pos")[row_off:row_off+nr] |
| if rp.dtype != torch.int32 or not torch.equal(rp, torch.arange(r0, L, dtype=torch.int32)): |
| raise ValueError("D2 row positions differ from sequence positions") |
| expected = tok_off+L, row_off+nr, span_off+ns |
| if split_name == split and L-gen > 1: |
| total += L-gen-1 |
| self.ends.append(total) |
| self.sequences.append((path, row, req)) |
| self.sequence_b.append(split_b) |
| split_b += 1 |
| shapes = [f.get_slice(n).get_shape()[0] for n in ("tokens", "row_pos", "top_lp")] |
| if tuple(shapes) != expected or f.get_slice("taps").get_shape() != [expected[1], 15360]: |
| raise ValueError("D2 aggregate tensor lengths differ from sequence offsets") |
| self.request_families.update(source_families(split_name, units)) |
| if not total: |
| raise ValueError(f"empty D2 block split {split}") |
| self.fingerprint = hashlib.sha256((sha256(manifest)+split).encode()).hexdigest() |
|
|
| def verified(self, relative: str, expected: str) -> Path: |
| path = (self.root/relative).resolve() |
| if not path.is_relative_to(self.root) or sha256(path) != expected: |
| raise ValueError(f"D2 file hash/path mismatch: {relative}") |
| return path |
|
|
| def __len__(self) -> int: |
| return self.ends[-1] |
|
|
| def supervised_tokens(self, index: int) -> int: |
| if index < 0 or index >= len(self): |
| raise IndexError(index) |
| seq_i = bisect.bisect_right(self.ends, index) |
| _, row, _ = self.sequences[seq_i] |
| _, length, gen = row[:3] |
| offset = index-(self.ends[seq_i-1] if seq_i else 0) |
| return min(5, length-(gen-1+offset)-2) |
|
|
| def __getitem__(self, index: int) -> Block: |
| if index < 0 or index >= len(self): |
| raise IndexError(index) |
| seq_i = bisect.bisect_right(self.ends, index) |
| path, row, req = self.sequences[seq_i] |
| tok_off, L, gen, r0, row_off, nr, span_off, ns = row |
| offset = index-(self.ends[seq_i-1] if seq_i else 0) |
| P = gen-1+offset |
| first = max(0, P-127) |
| positions = torch.arange(first, P+1) |
| with safe_open(path, framework="pt", device="cpu") as f: |
| tokens = f.get_slice("tokens")[tok_off:tok_off+L].long() |
| lo, hi = row_off+first-r0, row_off+P-r0+1 |
| taps = f.get_slice("taps")[lo:hi] |
| actual_pos = f.get_slice("row_pos")[lo:hi].long() |
| if taps.dtype != torch.bfloat16 or taps.shape != (positions.numel(), 15360) or not torch.equal(positions, actual_pos): |
| raise ValueError("tap shape/dtype/absolute positions differ") |
| target_pos = torch.arange(P+2, P+7) |
| mask = target_pos < L |
| n = int(mask.sum()) |
| start = span_off+P+2-gen |
| ids = torch.full((5, 20), -1, dtype=torch.long) |
| lp = torch.full((5, 20), -torch.inf) |
| resid = torch.ones(5) |
| ids[:n] = f.get_slice("top_ids")[start:start+n].long() |
| lp[:n] = f.get_slice("top_lp")[start:start+n] |
| resid[:n] = f.get_slice("resid")[start:start+n] |
| labels = torch.full((5,), -1, dtype=torch.long) |
| previous = torch.zeros(5, dtype=torch.long) |
| labels[:n] = tokens[P+2:P+2+n] |
| previous[:n] = tokens[P+1:P+1+n] |
| return {"inputs": {"taps": taps[None], "engine_mx": torch.zeros(1, positions.numel(), 5120), |
| "has_taps": torch.ones(1, positions.numel(), dtype=torch.bool), "positions": positions[None], |
| "mask": torch.ones(1, positions.numel(), dtype=torch.bool), "anchor": tokens[P+1:P+2], |
| "P": torch.tensor([P])}, |
| "targets": {"previous": previous[None], "top_ids": ids[None], "top_logp": lp[None], |
| "tail_mass": resid[None], "labels": labels[None], "present": mask[None], |
| "position_mask": mask[None]}, "request": req} |
|
|
