#!/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")) parser.add_argument( "--no-markdown-report", action="store_true", help="Do not write report.md; persistent prose is consolidated in README.md.", ) 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", "object_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), "object_embedding_path_present": int(counts["object_embedding_path_present"]), "object_embedding_path_rate": _rate( counts["object_embedding_path_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, "data_hash": {row["split"]: _index_hash(row, "content_hash") for row in split_rows}, "split_hash": {row["split"]: _index_hash(row, "split_hash") for row in 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") _write_markdown_report(out_dir, metrics, no_markdown_report=args.no_markdown_report) (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 _index_hash(row: dict[str, Any], key: str) -> Any: path = Path(str(row["index"])) if not path.exists(): return None return json.loads(path.read_text()).get(key) def _conclusion(rows: list[dict[str, Any]]) -> str: if all(float(row["observation_embedding_path_rate"]) > 0.0 for row in rows): if all(float(row["object_embedding_path_rate"]) > 0.0 for row in rows): return ( "Current chart exports expose observation and object-layout embeddings " "in every split; visual/object-layout chart features can be evaluated " "with leakage-audited indexes." ) return ( "Current chart exports expose observation embeddings in every split; " "visual chart features can be evaluated with leakage-audited indexes." ) if all(float(row["observation_ref_rate"]) > 0.0 for row in rows): return ( "Current chart exports expose raw observation refs in every split but " "not observation embeddings; run scripts/export_chart_observation_embeddings.py " "before claiming embedded visual chart tokens." ) if all(float(row["observation_embedding_path_rate"]) == 0.0 for row in rows) and all( float(row["observation_ref_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 ( "Observation feature availability is partial across splits; inspect per-split " "embedding/ref rates before running visual chart-token experiments." ) 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}{lrrrrr}", "\\toprule", "Split & Rows & ObsEmbed & ObjEmbed & ObsRef & Instruction \\\\", "\\midrule", ] for row in rows: lines.append( f"{row['split']} & {row['rows']} & " f"{row['observation_embedding_path_present']} & " f"{row['object_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 | Object 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['object_embedding_path_present']} ({row['object_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_markdown_report( out_dir: Path, metrics: dict[str, Any], *, no_markdown_report: bool, ) -> None: report_path = out_dir / "report.md" if no_markdown_report: report_path.unlink(missing_ok=True) return report_path.write_text(_report(metrics) + "\n") 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())