sync chart feature audit script 2026-07-03

workspace/scripts/audit_chart_feature_sources.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import subprocess
+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))
+
+import numpy as np  # noqa: E402
+
+from cil.chart_features import CHART_FEATURE_MODES, chart_feature_dim  # noqa: E402
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description="Audit deployment-visible chart feature sources in CIL chart shards."
+    )
+    parser.add_argument(
+        "--indexes",
+        nargs="+",
+        type=Path,
+        default=[
+            Path("data/cil_charts/train/index.json"),
+            Path("data/cil_charts/val/index.json"),
+            Path("data/cil_charts/test/index.json"),
+        ],
+    )
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/chart_feature_audit"))
+    args = parser.parse_args(argv)
+
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    _write_provenance(out_dir, args)
+
+    split_rows = []
+    row_details = []
+    for index_path in args.indexes:
+        index = json.loads(index_path.read_text())
+        split = str(index.get("split", index_path.parent.name))
+        counts: Counter[str] = Counter()
+        task_counts: Counter[str] = Counter()
+        feature_dims: dict[str, int] = {}
+        for shard in index.get("shards", []):
+            shard_path = index_path.parent / shard["path"]
+            with np.load(shard_path, allow_pickle=False) as data:
+                metadata_values = data["metadata_json"]
+                base_actions = data["base_action"]
+                action_shapes = data["action_shape"]
+                task_ids = data["task_id"]
+                chart_ids = data["chart_id"]
+                for row in range(metadata_values.shape[0]):
+                    metadata = _json_loads(str(metadata_values[row]))
+                    task_id = str(task_ids[row])
+                    task_counts[task_id] += 1
+                    counts["rows"] += 1
+                    for key in (
+                        "observation_embedding_path",
+                        "observation_ref",
+                        "scene_id",
+                        "instruction",
+                        "source_dataset",
+                    ):
+                        if metadata.get(key):
+                            counts[f"{key}_present"] += 1
+                    if counts["sample_details"] < 5:
+                        shape = tuple(int(value) for value in action_shapes[row])
+                        flat_count = int(math.prod(shape))
+                        base = np.asarray(base_actions[row][:flat_count], dtype=np.float32).reshape(shape)
+                        for mode in CHART_FEATURE_MODES:
+                            feature_dims[mode] = chart_feature_dim(base, mode=mode)
+                        row_details.append(
+                            {
+                                "split": split,
+                                "chart_id": str(chart_ids[row]),
+                                "task_id": task_id,
+                                "has_observation_embedding_path": bool(
+                                    metadata.get("observation_embedding_path")
+                                ),
+                                "has_observation_ref": bool(metadata.get("observation_ref")),
+                                "has_scene_id": bool(metadata.get("scene_id")),
+                                "has_instruction": bool(metadata.get("instruction")),
+                            }
+                        )
+                        counts["sample_details"] += 1
+        rows = int(counts["rows"])
+        split_rows.append(
+            {
+                "split": split,
+                "index": str(index_path),
+                "rows": rows,
+                "charts": int(index.get("num_groups_exported", 0)),
+                "retrieval_index_allowed": bool(index.get("retrieval_index_allowed")),
+                "include_outcomes": bool(index.get("include_outcomes")),
+                "observation_embedding_path_present": int(
+                    counts["observation_embedding_path_present"]
+                ),
+                "observation_embedding_path_rate": _rate(
+                    counts["observation_embedding_path_present"], rows
+                ),
+                "observation_ref_present": int(counts["observation_ref_present"]),
+                "observation_ref_rate": _rate(counts["observation_ref_present"], rows),
+                "scene_id_present": int(counts["scene_id_present"]),
+                "scene_id_rate": _rate(counts["scene_id_present"], rows),
+                "instruction_present": int(counts["instruction_present"]),
+                "instruction_rate": _rate(counts["instruction_present"], rows),
+                "source_dataset_present": int(counts["source_dataset_present"]),
+                "source_dataset_rate": _rate(counts["source_dataset_present"], rows),
+                "feature_dims": feature_dims,
+                "task_counts": dict(sorted(task_counts.items())),
+            }
+        )
+
+    metrics = {
+        "report_type": "chart_feature_source_audit",
+        "schema_version": 1,
+        "indexes": [str(path) for path in args.indexes],
+        "splits": split_rows,
+        "sample_details": row_details,
+        "conclusion": _conclusion(split_rows),
+    }
+    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")
+    (out_dir / "metrics_by_task.json").write_text(_metrics_by_task(split_rows) + "\n")
+    (out_dir / "metrics_by_seed.json").write_text("{}\n")
+    (out_dir / "table.tex").write_text(_table(split_rows) + "\n")
+    (out_dir / "report.md").write_text(_report(metrics) + "\n")
+    (out_dir / "train.log").write_text("not a training run; audited chart feature sources\n")
+    (out_dir / "eval.log").write_text("audited chart feature sources in exported chart indexes\n")
+    print(json.dumps({"out_dir": str(out_dir), "splits": len(split_rows)}, indent=2))
+    return 0
+
+
+def _json_loads(value: str) -> dict[str, Any]:
+    try:
+        payload = json.loads(value)
+    except json.JSONDecodeError:
+        return {}
+    return payload if isinstance(payload, dict) else {}
+
+
+def _rate(count: int, total: int) -> float:
+    return float(count) / float(total) if total else 0.0
+
+
+def _conclusion(rows: list[dict[str, Any]]) -> str:
+    if all(float(row["observation_embedding_path_rate"]) == 0.0 for row in rows):
+        return (
+            "Current chart exports do not contain observation embeddings or raw observation refs; "
+            "visual/object-centric chart tokens require a new export or embedding pass."
+        )
+    return "At least one split exposes observation references; visual chart features can be wired next."
+
+
+def _metrics_by_task(rows: list[dict[str, Any]]) -> str:
+    payload: dict[str, dict[str, int]] = defaultdict(dict)
+    for row in rows:
+        for task, count in row["task_counts"].items():
+            payload[task][str(row["split"])] = int(count)
+    return json.dumps(payload, indent=2, sort_keys=True)
+
+
+def _table(rows: list[dict[str, Any]]) -> str:
+    lines = [
+        "% Auto-generated by scripts/audit_chart_feature_sources.py",
+        "\\begin{tabular}{lrrrr}",
+        "\\toprule",
+        "Split & Rows & ObsEmbed & ObsRef & Instruction \\\\",
+        "\\midrule",
+    ]
+    for row in rows:
+        lines.append(
+            f"{row['split']} & {row['rows']} & "
+            f"{row['observation_embedding_path_present']} & "
+            f"{row['observation_ref_present']} & "
+            f"{row['instruction_present']} \\\\"
+        )
+    lines.extend(["\\bottomrule", "\\end{tabular}"])
+    return "\n".join(lines)
+
+
+def _report(metrics: dict[str, Any]) -> str:
+    lines = [
+        "# Chart Feature Source Audit",
+        "",
+        metrics["conclusion"],
+        "",
+        "| Split | Rows | Obs embedding path | Obs ref | Scene id | Instruction | Feature dims |",
+        "| --- | ---: | ---: | ---: | ---: | ---: | --- |",
+    ]
+    for row in metrics["splits"]:
+        lines.append(
+            f"| {row['split']} | {row['rows']} | "
+            f"{row['observation_embedding_path_present']} ({row['observation_embedding_path_rate']:.2%}) | "
+            f"{row['observation_ref_present']} ({row['observation_ref_rate']:.2%}) | "
+            f"{row['scene_id_present']} ({row['scene_id_rate']:.2%}) | "
+            f"{row['instruction_present']} ({row['instruction_rate']:.2%}) | "
+            f"{json.dumps(row['feature_dims'], sort_keys=True)} |"
+        )
+    return "\n".join(lines)
+
+
+def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None:
+    (out_dir / "config.yaml").write_text(
+        "\n".join(f"{key}: {value}" for key, value in sorted(vars(args).items())) + "\n"
+    )
+    (out_dir / "command.txt").write_text(
+        "python scripts/audit_chart_feature_sources.py " + " ".join(sys.argv[1:]) + "\n"
+    )
+    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
+    hashes = {}
+    for index_path in args.indexes:
+        index = json.loads(index_path.read_text())
+        hashes[str(index_path)] = {
+            "content_hash": index.get("content_hash"),
+            "split_hash": index.get("split_hash"),
+        }
+    (out_dir / "data_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n")
+    (out_dir / "split_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n")
+
+
+def _run(command: list[str]) -> str:
+    try:
+        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return ""
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())