| """Strict prepared readers and explicit without-replacement question orders.""" |
|
|
| import collections |
| import json |
| import random |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| from mdprep.common import sha256_file |
| from mdprep.native import digest_json, validate_contract |
| from mdprep.records import TrainingExample, parse_record |
| from mdprep.verify import validate_verification_contract |
|
|
| EVALUATION_SPLITS = ("dev", "calibration", "policy", "test_id", "test_ood") |
|
|
| DEFAULT_SCHEDULE = "paired" |
|
|
|
|
| @dataclass(frozen=True) |
| class TrainingSchedule: |
| name: str |
| phase_components: tuple[tuple[str, tuple[str, ...]], ...] |
|
|
| @property |
| def phases(self) -> tuple[str, ...]: |
| return tuple(phase for phase, _ in self.phase_components) |
|
|
| def compose(self, components: dict[str, list[str]]) -> dict[str, list[str]]: |
| orders = {} |
| for phase, names in self.phase_components: |
| order = [] |
| for name in names: |
| order.extend(components[name]) |
| orders[phase] = order |
| return orders |
|
|
|
|
| SCHEDULES = { |
| "paired": TrainingSchedule("paired", (("head", ("warmup",)), ("joint", ("training",)))), |
| "pointer-only": TrainingSchedule("pointer-only", (("head", ("warmup", "training")),)), |
| } |
| OPTIMIZER_PHASES = tuple(dict.fromkeys( |
| phase for schedule in SCHEDULES.values() for phase in schedule.phases)) |
|
|
|
|
| def training_schedule(name: str) -> TrainingSchedule: |
| try: |
| return SCHEDULES[name] |
| except KeyError as error: |
| raise ValueError(f"Unknown training schedule: {name}") from error |
|
|
|
|
| def bound_schedule(payload: dict) -> str: |
| """Resolve historical paired identities while requiring new schedules to name themselves.""" |
| name = payload.get("schedule", DEFAULT_SCHEDULE) |
| training_schedule(name) |
| return name |
|
|
|
|
| def qualified_inputs(root: Path, *, processor: dict | None = None, |
| verification: dict | None = None) -> dict: |
| """Commit successful qualification, processor implementation, rows, and media.""" |
| processor = processor if processor is not None else validate_contract(root) |
| verification = verification if verification is not None else validate_verification_contract(root) |
| for split in set(processor["prepared_hashes"]) & set(verification["prepared_hashes"]): |
| if processor["prepared_hashes"][split] != verification["prepared_hashes"][split]: |
| raise ValueError("Processor and data qualification disagree on prepared rows") |
| for relative in set(processor["media_hashes"]) & set(verification["media_hashes"]): |
| if processor["media_hashes"][relative] != verification["media_hashes"][relative]: |
| raise ValueError("Processor and data qualification disagree on media") |
| return {"processor_id": processor["processor_id"], |
| "processor_contract_id": digest_json(processor), |
| "qualification_id": verification["qualification_sha256"], |
| "prepared_hashes": processor["prepared_hashes"], |
| "media_hashes": processor["media_hashes"]} |
|
|
|
|
| @dataclass(frozen=True) |
| class DatasetView: |
| examples: tuple[TrainingExample, ...] |
| rows: tuple[dict, ...] |
| split: str |
| recipe: str |
| identity: str |
|
|
| @classmethod |
| def open(cls, root: Path, *, split: str, recipe: str = "mixed"): |
| if split not in ("train", *EVALUATION_SPLITS) or recipe not in ("real", "mixed"): |
| raise ValueError("Unknown dataset split or recipe") |
| processor = validate_contract(root) |
| verification = validate_verification_contract(root) |
| if split not in verification["scope"]: |
| raise ValueError(f"Verified scope does not include {split}") |
| if split not in processor["prepared_hashes"]: |
| raise ValueError(f"Processor scope does not include {split}") |
| path = root / "data/prepared" / f"{split}.jsonl" |
| rows, examples = [], [] |
| with path.open() as stream: |
| for line in stream: |
| row = json.loads(line) |
| if row["split"] != split: |
| raise ValueError("Prepared row split mismatch") |
| example = parse_record(row) |
| if recipe == "real" and example.origin != "real_human_label": |
| continue |
| rows.append(row) |
| examples.append(example) |
| ids = [example.input.id for example in examples] |
| if not ids or len(ids) != len(set(ids)): |
| raise ValueError("Dataset view is empty or contains duplicate IDs") |
| prepared_hash = sha256_file(path) |
| if prepared_hash != processor["prepared_hashes"][split]: |
| raise ValueError("Prepared rows changed while constructing the qualified view") |
| media = {relative for example in examples |
| for relative in (*example.input.images, *(video.path for video in example.input.videos))} |
| commitment = qualified_inputs(root, processor=processor, verification=verification) |
| identity = digest_json({"prepared_sha256": prepared_hash, "recipe": recipe, "selected_ids": ids, |
| "media_hashes": {name: processor["media_hashes"][name] for name in sorted(media)}, |
| "qualified_inputs": commitment}) |
| return cls(tuple(examples), tuple(rows), split, recipe, identity) |
|
|
| def by_id(self) -> dict[str, TrainingExample]: |
| return {example.input.id: example for example in self.examples} |
|
|
|
|
| def ordered_ids(view: DatasetView, *, seed: int, lengths: dict[str, int]) -> list[str]: |
| rng = random.Random(seed) |
| strata = collections.defaultdict(lambda: collections.defaultdict(list)) |
| for example, row in zip(view.examples, view.rows): |
| strata[(example.origin, row["modality"], example.input.kind)][lengths[example.input.id] // 256].append(example.input.id) |
| queues = [] |
| for stratum in sorted(strata): |
| buckets = strata[stratum] |
| keys = sorted(buckets) |
| rng.shuffle(keys) |
| queue = [] |
| for key in keys: |
| rng.shuffle(buckets[key]) |
| queue.extend(buckets[key]) |
| queues.append(collections.deque(queue)) |
| rng.shuffle(queues) |
| result = [] |
| while any(queues): |
| for queue in queues: |
| if queue: |
| result.append(queue.popleft()) |
| if len(result) != len(view.examples) or len(set(result)) != len(result): |
| raise ValueError("Without-replacement order is not a permutation") |
| return result |
|
|
|
|
| def plan_orders(view: DatasetView, stage: str, lengths: dict[str, int], |
| schedule: str = DEFAULT_SCHEDULE) -> dict: |
| limits = {"diagnostic": (64, 8), "pilot": (1000, 16), "main": (len(view.examples), 100)} |
| if stage not in limits: |
| raise ValueError("Unknown training stage") |
| selected_schedule = training_schedule(schedule) |
| size, head_updates = limits[stage] |
| if len(view.examples) < size: |
| raise ValueError(f"{stage} requires {size} rows without recipe fallback") |
| order = ordered_ids(view, seed=1729, lengths=lengths)[:size] |
| warmup = [order[index % size] for index in range(head_updates * 16)] |
| training = ([order[index % size] for index in range(64 * 16)] |
| if stage == "diagnostic" else order) |
| phase_orders = selected_schedule.compose({"warmup": warmup, "training": training}) |
| orders = {phase: phase_orders.get(phase, []) for phase in OPTIMIZER_PHASES} |
| updates = {phase: (len(orders[phase]) + 15) // 16 for phase in OPTIMIZER_PHASES} |
| return {**orders, "selection": order, |
| "head_updates": updates["head"], "joint_updates": updates["joint"], |
| "without_replacement_joint": stage != "diagnostic" and "joint" in selected_schedule.phases} |
|
|
