ctt artifacts 2026-07-02 workspace/scripts/build_data_accounting.py

workspace/scripts/build_data_accounting.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import sys
+from collections import Counter, defaultdict
+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.generation.tangent_targets import (  # noqa: E402
+    DEFAULT_BASE_CANDIDATE_TYPES,
+    action_matrix,
+    action_shape,
+    choose_base_record,
+    label_from_delta_utility,
+    record_utility,
+)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Build CIL chart data accounting tables.")
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/data_accounting"))
+    parser.add_argument("--epsilon", type=float, default=0.05)
+    parser.add_argument("--split-fractions", default="0.70,0.15,0.15")
+    parser.add_argument("--split-seed", type=int, default=0)
+    parser.add_argument(
+        "--base-candidate-types",
+        default=",".join(DEFAULT_BASE_CANDIDATE_TYPES),
+    )
+    args = parser.parse_args(argv)
+
+    fractions = _parse_split_fractions(args.split_fractions)
+    base_candidate_types = _parse_csv(args.base_candidate_types)
+    dataset = CILDataset(args.dataset)
+    rows: dict[str, dict[str, Any]] = {
+        split: _empty_row(split) for split in ("train", "val", "test")
+    }
+    skip_counts: Counter[str] = Counter()
+
+    for group in dataset.iter_groups():
+        if not group:
+            skip_counts["empty_group"] += 1
+            continue
+        split = _assign_split(group[0].group_id, fractions=fractions, seed=args.split_seed)
+        row = rows[split]
+        base = choose_base_record(group, base_candidate_types=base_candidate_types)
+        if base is None:
+            skip_counts["missing_base"] += 1
+            continue
+        base_action = action_matrix(base)
+        if not base_action:
+            skip_counts["empty_base_action"] += 1
+            continue
+        base_shape = action_shape(base_action)
+        base_utility = record_utility(base)
+        row["chart_ids"].add(group[0].group_id)
+        row["state_hashes"].add(group[0].state_hash)
+        row["task_ids"].add(group[0].task_id)
+        seed = group[0].metadata.get("episode_seed", group[0].metadata.get("random_seed"))
+        if seed is not None:
+            row["seeds"].add(str(seed))
+        for record in group:
+            action = action_matrix(record)
+            if not action or action_shape(action) != base_shape:
+                skip_counts["unusable_branch"] += 1
+                continue
+            row["num_branches"] += 1
+            candidate_type = str(record.candidate_type)
+            row["candidate_type_counts"][candidate_type] += 1
+            if record.record_id == base.record_id:
+                row["num_base_branches"] += 1
+            if candidate_type == "expert" or candidate_type.endswith("_expert"):
+                row["num_expert"] += 1
+            if record.record_id == base.record_id:
+                label = "neutral"
+            else:
+                label = label_from_delta_utility(
+                    record_utility(record) - base_utility,
+                    epsilon=args.epsilon,
+                )
+            row[f"num_{label}"] += 1
+
+    materialized = [_finalize_row(row) for row in rows.values()]
+    payload = {
+        "schema_version": 1,
+        "dataset": str(args.dataset),
+        "epsilon": args.epsilon,
+        "split_fractions": fractions,
+        "split_seed": args.split_seed,
+        "split_hash": _split_hash(materialized),
+        "skip_counts": dict(sorted(skip_counts.items())),
+        "rows": materialized,
+    }
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    (out_dir / "table.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    (out_dir / "table.tex").write_text(_latex_table(materialized) + "\n")
+    (out_dir / "report.md").write_text(_markdown_report(payload) + "\n")
+    print(json.dumps({"out_dir": str(out_dir), "split_hash": payload["split_hash"]}, indent=2))
+    return 0
+
+
+def _empty_row(split: str) -> dict[str, Any]:
+    return {
+        "split": split,
+        "task_ids": set(),
+        "seeds": set(),
+        "chart_ids": set(),
+        "state_hashes": set(),
+        "num_branches": 0,
+        "num_base_branches": 0,
+        "num_positive": 0,
+        "num_neutral": 0,
+        "num_negative": 0,
+        "num_expert": 0,
+        "candidate_type_counts": Counter(),
+    }
+
+
+def _finalize_row(row: dict[str, Any]) -> dict[str, Any]:
+    split = row["split"]
+    num_branches = int(row["num_branches"])
+    return {
+        "split": split,
+        "num_tasks": len(row["task_ids"]),
+        "num_seeds": len(row["seeds"]),
+        "num_charts": len(row["chart_ids"]),
+        "num_branches": num_branches,
+        "num_base_branches": int(row["num_base_branches"]),
+        "num_positive": int(row["num_positive"]),
+        "num_neutral": int(row["num_neutral"]),
+        "num_negative": int(row["num_negative"]),
+        "num_expert": int(row["num_expert"]),
+        "num_train_only_retrieval_rows": num_branches if split == "train" else 0,
+        "num_eval_only_rows": num_branches if split != "train" else 0,
+        "candidate_type_counts": dict(sorted(row["candidate_type_counts"].items())),
+        "chart_hashes": sorted(_hash_id(value) for value in row["chart_ids"]),
+        "state_hashes": sorted(_hash_id(value) for value in row["state_hashes"]),
+    }
+
+
+def _latex_table(rows: list[dict[str, Any]]) -> str:
+    lines = [
+        "% Auto-generated by scripts/build_data_accounting.py",
+        "\\begin{tabular}{lrrrrrrrrrrr}",
+        "\\toprule",
+        "Split & Tasks & Seeds & Charts & Branches & Base & Positive & Neutral & Negative & Expert & TrainRows & EvalRows \\\\",
+        "\\midrule",
+    ]
+    for row in rows:
+        lines.append(
+            f"{row['split']} & {row['num_tasks']} & {row['num_seeds']} & "
+            f"{row['num_charts']} & {row['num_branches']} & {row['num_base_branches']} & "
+            f"{row['num_positive']} & {row['num_neutral']} & {row['num_negative']} & "
+            f"{row['num_expert']} & {row['num_train_only_retrieval_rows']} & "
+            f"{row['num_eval_only_rows']} \\\\"
+        )
+    lines.extend(["\\bottomrule", "\\end{tabular}"])
+    return "\n".join(lines)
+
+
+def _markdown_report(payload: dict[str, Any]) -> str:
+    lines = [
+        "# CIL Chart Data Accounting",
+        "",
+        f"Dataset: `{payload['dataset']}`",
+        f"Split seed: `{payload['split_seed']}`",
+        f"Split hash: `{payload['split_hash']}`",
+        "",
+        "| Split | Tasks | Seeds | Charts | Branches | Base | Positive | Neutral | Negative | Expert | Train rows | Eval rows |",
+        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
+    ]
+    for row in payload["rows"]:
+        lines.append(
+            f"| {row['split']} | {row['num_tasks']} | {row['num_seeds']} | "
+            f"{row['num_charts']} | {row['num_branches']} | {row['num_base_branches']} | "
+            f"{row['num_positive']} | {row['num_neutral']} | {row['num_negative']} | "
+            f"{row['num_expert']} | {row['num_train_only_retrieval_rows']} | "
+            f"{row['num_eval_only_rows']} |"
+        )
+    if payload["skip_counts"]:
+        lines.extend(["", "Skip counts:", ""])
+        lines.extend(f"- `{key}`: {value}" for key, value in payload["skip_counts"].items())
+    return "\n".join(lines)
+
+
+def _parse_csv(value: str) -> tuple[str, ...]:
+    return tuple(item.strip() for item in value.split(",") if item.strip())
+
+
+def _parse_split_fractions(value: str) -> dict[str, float]:
+    parts = [float(item.strip()) for item in value.split(",") if item.strip()]
+    if len(parts) != 3:
+        raise ValueError("--split-fractions must contain train,val,test values")
+    total = sum(parts)
+    if total <= 0.0 or any(part < 0.0 for part in parts):
+        raise ValueError("split fractions must be non-negative with positive sum")
+    train, val, test = [part / total for part in parts]
+    return {"train": train, "val": val, "test": test}
+
+
+def _assign_split(group_id: str, *, fractions: dict[str, float], seed: int) -> str:
+    value = _stable_uniform(group_id, seed=seed)
+    if value < fractions["train"]:
+        return "train"
+    if value < fractions["train"] + fractions["val"]:
+        return "val"
+    return "test"
+
+
+def _stable_uniform(value: str, *, seed: int) -> float:
+    digest = hashlib.sha256(f"{seed}:{value}".encode("utf-8")).digest()
+    return int.from_bytes(digest[:8], "big") / float(2**64)
+
+
+def _hash_id(value: str) -> str:
+    return hashlib.sha256(str(value).encode()).hexdigest()
+
+
+def _split_hash(rows: list[dict[str, Any]]) -> str:
+    digest_rows = [
+        {
+            "split": row["split"],
+            "chart_hashes": row["chart_hashes"],
+            "state_hashes": row["state_hashes"],
+        }
+        for row in rows
+    ]
+    return hashlib.sha256(json.dumps(digest_rows, sort_keys=True).encode()).hexdigest()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())