| """Canonical fixed-prefix conditional generation matrix construction.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| from dataclasses import dataclass |
| from dataclasses import asdict |
| import os |
| from pathlib import Path |
| from typing import Literal |
|
|
| from dlb.conditional_prompts import ConditionalProtocol, load_protocol, verify_prompts |
| from dlb.io import sha256_file |
| from dlb.matrix import build_matrix, unsupported_inventory, write_unsupported_inventory |
| from dlb.registry import ExperimentRegistry, load_registry |
|
|
|
|
| CONDITIONAL_MATRIX_SCHEMA = "dlb-conditional-generation-matrix-v1" |
| CONDITIONAL_MATRIX_COLUMNS = ( |
| "task_id", |
| "category", |
| "model", |
| "dataset", |
| "steps", |
| "sample_count", |
| "seed", |
| "environment", |
| "adapter", |
| "source", |
| "provenance", |
| "protocol", |
| "conditioning_manifest", |
| "conditioning_manifest_sha256", |
| "sample_dir", |
| "metrics_path", |
| "timing_path", |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class ConditionalMatrixTask: |
| """One registry-backed conditional request using the C64 protocol.""" |
|
|
| task_id: str |
| category: str |
| model: str |
| dataset: str |
| steps: int |
| sample_count: Literal[2048] |
| seed: int |
| environment: str |
| adapter: str |
| source: str |
| provenance: str |
| protocol: Literal["c64_zs_v1"] |
| conditioning_manifest: str |
| conditioning_manifest_sha256: str |
| sample_dir: str |
| metrics_path: str |
| timing_path: str |
|
|
| @property |
| def model_id(self) -> str: |
| return self.model |
|
|
| @property |
| def dataset_id(self) -> str: |
| return self.dataset |
|
|
| @property |
| def step_count(self) -> int: |
| return self.steps |
|
|
| def row(self) -> dict[str, str]: |
| value = asdict(self) |
| return {key: str(value[key]) for key in CONDITIONAL_MATRIX_COLUMNS} |
|
|
|
|
| def _root_path(root: Path | None) -> Path: |
| return (root or Path(".")).resolve() |
|
|
|
|
| def _default_protocol(root: Path) -> ConditionalProtocol: |
| """Prefer a supplied project config while keeping isolated matrix tests usable.""" |
|
|
| config = root / "configs" / "conditional.yaml" |
| if not config.is_file(): |
| config = Path(__file__).resolve().parents[2] / "configs" / "conditional.yaml" |
| return load_protocol(config) |
|
|
|
|
| def _sidecar_path(root: Path, dataset: str) -> Path: |
| return root / "data" / "manifests" / f"conditional-{dataset}-c64.json" |
|
|
|
|
| def build_conditional_matrix( |
| registry: ExperimentRegistry, |
| *, |
| root: Path | None = None, |
| protocol: ConditionalProtocol | None = None, |
| ) -> list[ConditionalMatrixTask]: |
| """Expand the supported registry coverage into the isolated C64 result tree. |
| |
| Verify both prompt artifacts before expanding tasks so each matrix row is |
| already bound to its canonical sidecar manifest. |
| """ |
|
|
| root_path = _root_path(root) |
| protocol = protocol or _default_protocol(root_path) |
| verified_sidecars: dict[str, tuple[Path, str]] = {} |
| for dataset in sorted(protocol.datasets): |
| manifest = verify_prompts(root_path, dataset, protocol) |
| if manifest.dataset != dataset: |
| raise ValueError("verified conditional prompt manifest has the wrong dataset") |
| sidecar = _sidecar_path(root_path, dataset) |
| if sidecar.is_symlink() or not sidecar.is_file(): |
| raise ValueError(f"conditional prompt manifest is missing or unsafe: {sidecar}") |
| verified_sidecars[dataset] = (sidecar, sha256_file(sidecar)) |
| ordinary = build_matrix(registry, root=root_path, sample_count=2048, seed=protocol.sampling_seed) |
| tasks: list[ConditionalMatrixTask] = [] |
| for task in ordinary: |
| sidecar, manifest_sha256 = verified_sidecars[task.dataset] |
| sample_dir = root_path / "results" / "conditional" / "samples" / task.dataset / task.model / f"steps_{task.steps}" |
| metrics_path = root_path / "results" / "conditional" / "metrics" / task.dataset / task.model / f"steps_{task.steps}" / "metrics.json" |
| timing_path = root_path / "results" / "conditional" / "timing" / task.dataset / task.model / f"steps_{task.steps}" / "timing.json" |
| tasks.append( |
| ConditionalMatrixTask( |
| task_id=task.task_id, |
| category=task.category, |
| model=task.model, |
| dataset=task.dataset, |
| steps=task.steps, |
| sample_count=2048, |
| seed=protocol.sampling_seed, |
| environment=task.environment, |
| adapter=task.adapter, |
| source=task.source, |
| provenance=task.provenance, |
| protocol=protocol.protocol, |
| conditioning_manifest=sidecar.as_posix(), |
| conditioning_manifest_sha256=manifest_sha256, |
| sample_dir=sample_dir.as_posix(), |
| metrics_path=metrics_path.as_posix(), |
| timing_path=timing_path.as_posix(), |
| ) |
| ) |
| return tasks |
|
|
|
|
| def conditional_unsupported_inventory(registry: ExperimentRegistry) -> list[dict[str, str]]: |
| """Return the same explicit unsupported inventory as the ordinary matrix.""" |
|
|
| records = unsupported_inventory(registry) |
| |
| |
| for record in records: |
| if record["model"] == "rdlm" and record["dataset"] == "owt": |
| record["category"] = "few" |
| return records |
|
|
|
|
| def write_conditional_unsupported_inventory( |
| path: Path, records: list[dict[str, str]] |
| ) -> Path: |
| """Write unsupported cells with the conditional matrix schema header.""" |
|
|
| return write_unsupported_inventory(path, records, schema=CONDITIONAL_MATRIX_SCHEMA) |
|
|
|
|
| def _validate_tasks(tasks: list[ConditionalMatrixTask]) -> None: |
| seen: set[str] = set() |
| previous: tuple[str, str, int] | None = None |
| for task in tasks: |
| if task.task_id in seen: |
| raise ValueError(f"duplicate conditional matrix task ID: {task.task_id}") |
| seen.add(task.task_id) |
| key = (task.model, task.dataset, task.steps) |
| if previous is not None and key < previous: |
| raise ValueError("conditional matrix tasks are not in stable order") |
| previous = key |
| if task.protocol != "c64_zs_v1" or task.sample_count != 2048: |
| raise ValueError("conditional matrix task has the wrong C64 protocol") |
| if task.steps <= 0: |
| raise ValueError("conditional matrix task steps must be positive") |
| if any("\t" in value or "\n" in value or "\r" in value for value in task.row().values()): |
| raise ValueError(f"conditional matrix task contains a control character: {task.task_id}") |
|
|
|
|
| def write_conditional_matrix(path: Path, tasks: list[ConditionalMatrixTask]) -> Path: |
| """Write a deterministic, versioned TSV matrix for C64 conditional runs.""" |
|
|
| path = Path(path).absolute() |
| _validate_tasks(tasks) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.partial") |
| try: |
| with temporary.open("w", encoding="utf-8", newline="") as handle: |
| handle.write(f"# schema={CONDITIONAL_MATRIX_SCHEMA}\n") |
| writer = csv.DictWriter( |
| handle, |
| fieldnames=CONDITIONAL_MATRIX_COLUMNS, |
| delimiter="\t", |
| lineterminator="\n", |
| ) |
| writer.writeheader() |
| for task in tasks: |
| writer.writerow(task.row()) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temporary, path) |
| except BaseException: |
| temporary.unlink(missing_ok=True) |
| raise |
| return path |
|
|
|
|
| def _main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--root", type=Path, default=Path.cwd()) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--unsupported-output", type=Path) |
| arguments = parser.parse_args(argv) |
| root = arguments.root.resolve() |
| registry = load_registry(root / "configs" / "experiments.yaml") |
| tasks = build_conditional_matrix(registry, root=root) |
| write_conditional_matrix(arguments.output, tasks) |
| write_conditional_unsupported_inventory( |
| arguments.unsupported_output or arguments.output.with_name("unsupported.tsv"), |
| conditional_unsupported_inventory(registry), |
| ) |
| print(f"conditional_matrix={arguments.output} tasks={len(tasks)}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(_main()) |
|
|
