ctt chart export script

workspace/scripts/export_cil_charts.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import sys
+from collections import Counter
+from pathlib import Path
+from typing import Any
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.data.datasets import CILDataset  # noqa: E402
+from dovla_cil.data.schema import CILRecord, OutcomeVector  # noqa: E402
+from dovla_cil.generation.tangent_targets import (  # noqa: E402
+    DEFAULT_BASE_CANDIDATE_TYPES,
+    action_matrix,
+    action_shape,
+    choose_base_record,
+    label_from_delta_utility,
+    record_utility,
+    subtract_actions,
+)
+
+
+LABEL_TO_ID = {"negative": -1, "neutral": 0, "positive": 1, "hidden": 9}
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Export same-state CIL charts to NPZ shards for train-only CIL-Atlas "
+            "retrieval/generator training, with an index that can be leakage-audited."
+        )
+    )
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument(
+        "--out-dir",
+        type=Path,
+        default=Path("data/cil_charts/train"),
+        help="Output split directory, e.g. data/cil_charts/train.",
+    )
+    parser.add_argument("--split", default="train")
+    parser.add_argument("--epsilon", type=float, default=0.0)
+    parser.add_argument("--max-groups", type=int, default=None)
+    parser.add_argument("--shard-size", type=int, default=50000)
+    parser.add_argument(
+        "--base-candidate-types",
+        default=",".join(DEFAULT_BASE_CANDIDATE_TYPES),
+    )
+    parser.add_argument(
+        "--include-outcomes",
+        action="store_true",
+        help=(
+            "Include measured utilities/outcome vectors. Defaults to true only for "
+            "the train split; non-train outcome exports are evaluator-only."
+        ),
+    )
+    parser.add_argument(
+        "--no-outcomes",
+        action="store_true",
+        help="Force-hide utilities/outcomes even for train export.",
+    )
+    args = parser.parse_args(argv)
+
+    if args.epsilon < 0.0:
+        parser.error("--epsilon must be non-negative")
+    if args.max_groups is not None and args.max_groups <= 0:
+        parser.error("--max-groups must be positive when provided")
+    if args.shard_size <= 0:
+        parser.error("--shard-size must be positive")
+
+    try:
+        import numpy as np
+    except ImportError as exc:  # pragma: no cover - environment guard
+        raise SystemExit("export_cil_charts.py requires numpy to write .npz shards") from exc
+
+    include_outcomes = (args.split == "train" or args.include_outcomes) and not args.no_outcomes
+    dataset = CILDataset(args.dataset)
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    rows: list[dict[str, Any]] = []
+    shard_paths: list[dict[str, Any]] = []
+    counters: Counter[str] = Counter()
+    task_counts: Counter[str] = Counter()
+    label_counts: Counter[str] = Counter()
+    group_hashes: set[str] = set()
+    state_hashes: set[str] = set()
+    max_flat_dim = 0
+    group_count = 0
+    base_candidate_types = _parse_csv(args.base_candidate_types)
+
+    for group in dataset.iter_groups():
+        if args.max_groups is not None and group_count >= args.max_groups:
+            break
+        group_count += 1
+        if not group:
+            counters["empty_group"] += 1
+            continue
+        base = choose_base_record(group, base_candidate_types=base_candidate_types)
+        if base is None:
+            counters["missing_base"] += 1
+            continue
+        base_action = action_matrix(base)
+        if not base_action:
+            counters["empty_base_action"] += 1
+            continue
+        base_shape = action_shape(base_action)
+        base_utility = record_utility(base) if include_outcomes else math.nan
+        emitted_for_group = 0
+        for record in group:
+            action = action_matrix(record)
+            if not action:
+                counters["empty_action"] += 1
+                continue
+            if action_shape(action) != base_shape:
+                counters["action_shape_mismatch"] += 1
+                continue
+            utility = record_utility(record) if include_outcomes else math.nan
+            delta_utility = utility - base_utility if include_outcomes else math.nan
+            label = (
+                label_from_delta_utility(delta_utility, epsilon=args.epsilon)
+                if include_outcomes and record.record_id != base.record_id
+                else ("neutral" if include_outcomes else "hidden")
+            )
+            delta_action = subtract_actions(action, base_action)
+            row = _row_from_record(
+                record,
+                base=base,
+                split=args.split,
+                action=action,
+                base_action=base_action,
+                delta_action=delta_action,
+                utility=utility,
+                base_utility=base_utility,
+                delta_utility=delta_utility,
+                label=label,
+                include_outcomes=include_outcomes,
+            )
+            max_flat_dim = max(max_flat_dim, len(row["action_flat"]))
+            rows.append(row)
+            emitted_for_group += 1
+            task_counts[str(record.task_id)] += 1
+            label_counts[label] += 1
+        if emitted_for_group:
+            group_hashes.add(_hash_id(group[0].group_id))
+            state_hashes.add(_hash_id(group[0].state_hash))
+        if len(rows) >= args.shard_size:
+            shard_paths.append(_write_shard(np, out_dir, rows, len(shard_paths)))
+            rows = []
+    if rows:
+        shard_paths.append(_write_shard(np, out_dir, rows, len(shard_paths)))
+
+    index = {
+        "schema_version": 1,
+        "format": "cil_charts_npz",
+        "dataset": str(args.dataset),
+        "split": args.split,
+        "epsilon": args.epsilon,
+        "include_outcomes": include_outcomes,
+        "audience": "train_retrieval" if args.split == "train" else "evaluator_only",
+        "retrieval_index_allowed": args.split == "train",
+        "deployment_clean": args.split == "train",
+        "base_candidate_types": list(base_candidate_types),
+        "num_groups_scanned": group_count,
+        "num_groups_exported": len(group_hashes),
+        "num_rows": sum(int(item["num_rows"]) for item in shard_paths),
+        "max_flat_action_dim": max_flat_dim,
+        "label_counts": dict(sorted(label_counts.items())),
+        "task_counts": dict(sorted(task_counts.items())),
+        "skip_counts": dict(sorted(counters.items())),
+        "group_hashes": sorted(group_hashes),
+        "state_hashes": sorted(state_hashes),
+        "shards": shard_paths,
+        "leakage_contract": {
+            "train_split_only_for_retrieval": True,
+            "nontrain_outcomes_are_evaluator_only": True,
+            "deployment_must_not_read_outcomes": args.split != "train",
+        },
+    }
+    index["content_hash"] = _content_hash(index)
+    (out_dir / "index.json").write_text(json.dumps(index, indent=2, sort_keys=True) + "\n")
+    print(json.dumps({k: index[k] for k in ("split", "num_groups_exported", "num_rows", "content_hash")}, indent=2))
+    return 0
+
+
+def _row_from_record(
+    record: CILRecord,
+    *,
+    base: CILRecord,
+    split: str,
+    action: list[list[float]],
+    base_action: list[list[float]],
+    delta_action: list[list[float]],
+    utility: float,
+    base_utility: float,
+    delta_utility: float,
+    label: str,
+    include_outcomes: bool,
+) -> dict[str, Any]:
+    outcome = OutcomeVector.from_reward(record.reward, failure=record.failure)
+    metadata = {
+        "group_id": record.group_id,
+        "state_hash": record.state_hash,
+        "task_id": record.task_id,
+        "split_id": split,
+        "instruction": record.instruction,
+        "observation_ref": record.observation_ref,
+        "record_id": record.record_id,
+        "candidate_type": record.candidate_type,
+        "base_record_id": base.record_id,
+        "base_candidate_type": base.candidate_type,
+        "action_id": record.action_chunk.action_id,
+        "base_action_id": base.action_chunk.action_id,
+        "scene_id": record.scene_id,
+        "rank_within_group": record.rank_within_group,
+        "failure_type": record.failure.type if record.failure else None,
+        "source_dataset": record.metadata.get("source_dataset"),
+    }
+    return {
+        "group_id": record.group_id,
+        "state_hash": record.state_hash,
+        "task_id": record.task_id,
+        "record_id": record.record_id,
+        "candidate_type": record.candidate_type,
+        "label": label,
+        "label_id": LABEL_TO_ID[label],
+        "shape": list(action_shape(action)),
+        "action_flat": _flatten(action),
+        "base_action_flat": _flatten(base_action),
+        "delta_action_flat": _flatten(delta_action),
+        "utility": utility,
+        "base_utility": base_utility,
+        "delta_utility": delta_utility,
+        "outcome_vector": [
+            outcome.success,
+            outcome.progress,
+            outcome.contact_quality,
+            outcome.safety_violation,
+            outcome.task_stage_quality,
+            outcome.smoothness,
+            outcome.recovery,
+        ]
+        if include_outcomes
+        else [math.nan] * 7,
+        "metadata_json": json.dumps(metadata, sort_keys=True),
+    }
+
+
+def _write_shard(np: Any, out_dir: Path, rows: list[dict[str, Any]], shard_index: int) -> dict[str, Any]:
+    max_dim = max((len(row["action_flat"]) for row in rows), default=0)
+    path = out_dir / f"charts_{shard_index:05d}.npz"
+    np.savez_compressed(
+        path,
+        action=_pad(np, [row["action_flat"] for row in rows], max_dim),
+        base_action=_pad(np, [row["base_action_flat"] for row in rows], max_dim),
+        delta_action=_pad(np, [row["delta_action_flat"] for row in rows], max_dim),
+        action_shape=np.asarray([row["shape"] for row in rows], dtype=np.int32),
+        utility=np.asarray([row["utility"] for row in rows], dtype=np.float32),
+        base_utility=np.asarray([row["base_utility"] for row in rows], dtype=np.float32),
+        delta_utility=np.asarray([row["delta_utility"] for row in rows], dtype=np.float32),
+        label_id=np.asarray([row["label_id"] for row in rows], dtype=np.int8),
+        outcome_vector=np.asarray([row["outcome_vector"] for row in rows], dtype=np.float32),
+        group_id=np.asarray([row["group_id"] for row in rows]),
+        state_hash=np.asarray([row["state_hash"] for row in rows]),
+        task_id=np.asarray([row["task_id"] for row in rows]),
+        record_id=np.asarray([row["record_id"] for row in rows]),
+        candidate_type=np.asarray([row["candidate_type"] for row in rows]),
+        label=np.asarray([row["label"] for row in rows]),
+        metadata_json=np.asarray([row["metadata_json"] for row in rows]),
+    )
+    return {
+        "path": path.name,
+        "num_rows": len(rows),
+        "sha256": _sha256(path),
+    }
+
+
+def _pad(np: Any, vectors: list[list[float]], width: int) -> Any:
+    array = np.full((len(vectors), width), np.nan, dtype=np.float32)
+    for index, vector in enumerate(vectors):
+        array[index, : len(vector)] = np.asarray(vector, dtype=np.float32)
+    return array
+
+
+def _flatten(values: list[list[float]]) -> list[float]:
+    return [float(value) for row in values for value in row]
+
+
+def _parse_csv(value: str) -> tuple[str, ...]:
+    return tuple(item.strip() for item in value.split(",") if item.strip())
+
+
+def _hash_id(value: str) -> str:
+    return hashlib.sha256(str(value).encode()).hexdigest()
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _content_hash(index: dict[str, Any]) -> str:
+    payload = dict(index)
+    payload.pop("content_hash", None)
+    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())