workspace/scripts/export_chart_observation_embeddings.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import hashlib
import io
import json
import subprocess
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))

import numpy as np  # noqa: E402

from cil.chart_features import OBSERVATION_EMBED_DIM  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Decode chart observation_ref JPEGs and write deployment-visible "
            "observation embeddings back into chart metadata."
        )
    )
    parser.add_argument(
        "--indexes",
        nargs="+",
        type=Path,
        default=[
            Path("data/cil_charts_rgb_refs/train/index.json"),
            Path("data/cil_charts_rgb_refs/val/index.json"),
            Path("data/cil_charts_rgb_refs/test/index.json"),
        ],
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("runs/chart_observation_embeddings_rgb_refs"),
    )
    parser.add_argument("--overwrite", action="store_true")
    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 = []
    for index_path in args.indexes:
        split_rows.append(_process_index(index_path, overwrite=args.overwrite))

    metrics = {
        "report_type": "chart_observation_embedding_export",
        "schema_version": 1,
        "embedding_dim": OBSERVATION_EMBED_DIM,
        "extractor": "rgb_jpeg_stats_v1",
        "indexes": [str(path) for path in args.indexes],
        "splits": split_rows,
        "data_hash": {
            row["split"]: row["content_hash_after"] for row in split_rows
        },
        "split_hash": {
            row["split"]: row["split_hash"] for row in split_rows
        },
        "leakage_contract": {
            "reads_outcomes": False,
            "reads_observation_ref": True,
            "writes_observation_embedding_path": True,
        },
    }
    (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; exported observation embeddings\n")
    (out_dir / "eval.log").write_text("decoded observation_ref JPEGs only; no outcomes read\n")
    print(json.dumps({"out_dir": str(out_dir), "splits": len(split_rows)}, indent=2))
    return 0


def _process_index(index_path: Path, *, overwrite: bool) -> dict[str, Any]:
    index = json.loads(index_path.read_text())
    split = str(index.get("split", index_path.parent.name))
    embed_path = index_path.parent / "observation_embeddings_rgb_stats.npz"
    if embed_path.exists() and not overwrite:
        raise FileExistsError(f"{embed_path} exists; pass --overwrite to replace it")

    ref_to_row: dict[tuple[str, str], int] = {}
    embeddings: list[np.ndarray] = []
    counters: Counter[str] = Counter()
    task_counts: Counter[str] = Counter()
    updated_shards: list[str] = []
    h5_cache: dict[Path, Any] = {}
    try:
        for shard in index.get("shards", []):
            shard_path = index_path.parent / str(shard["path"])
            with np.load(shard_path, allow_pickle=False) as data:
                arrays = {key: data[key] for key in data.files}
            metadata_values = arrays["metadata_json"]
            updated_metadata = []
            for raw in metadata_values:
                metadata = _json_loads(str(raw))
                task_counts[str(metadata.get("task_id", "unknown"))] += 1
                counters["rows"] += 1
                source_dataset = str(metadata.get("source_dataset", ""))
                observation_ref = str(metadata.get("observation_ref", ""))
                if not source_dataset or not observation_ref:
                    counters["missing_observation_ref"] += 1
                    updated_metadata.append(json.dumps(metadata, sort_keys=True))
                    continue
                key = (source_dataset, observation_ref)
                if key not in ref_to_row:
                    ref_to_row[key] = len(embeddings)
                    embeddings.append(_embedding_for_ref(key, h5_cache))
                metadata["observation_embedding_path"] = (
                    f"{embed_path.name}#embeddings/{ref_to_row[key]}"
                )
                metadata["observation_embedding_extractor"] = "rgb_jpeg_stats_v1"
                metadata["observation_embedding_dim"] = OBSERVATION_EMBED_DIM
                counters["rows_with_embedding"] += 1
                updated_metadata.append(json.dumps(metadata, sort_keys=True))
            arrays["metadata_json"] = np.asarray(updated_metadata)
            np.savez_compressed(shard_path, **arrays)
            updated_shards.append(str(shard_path))
    finally:
        for handle in h5_cache.values():
            handle.close()

    embedding_matrix = (
        np.stack(embeddings).astype(np.float32)
        if embeddings
        else np.zeros((0, OBSERVATION_EMBED_DIM), dtype=np.float32)
    )
    np.savez_compressed(
        embed_path,
        embeddings=embedding_matrix,
        extractor=np.asarray(["rgb_jpeg_stats_v1"]),
        observation_refs=np.asarray([ref for _, ref in ref_to_row]),
        source_datasets=np.asarray([source for source, _ in ref_to_row]),
    )

    index["observation_embedding_manifest"] = {
        "path": embed_path.name,
        "dataset": "embeddings",
        "dim": OBSERVATION_EMBED_DIM,
        "extractor": "rgb_jpeg_stats_v1",
        "num_embeddings": int(embedding_matrix.shape[0]),
        "reads_outcomes": False,
    }
    index["shard_content_hashes"] = {
        str(Path(path).name): _sha256(Path(path)) for path in updated_shards
    }
    index["embedding_content_hash"] = _sha256(embed_path)
    index["content_hash"] = _content_hash(index)
    index_path.write_text(json.dumps(index, indent=2, sort_keys=True) + "\n")

    return {
        "split": split,
        "index": str(index_path),
        "rows": int(counters["rows"]),
        "rows_with_embedding": int(counters["rows_with_embedding"]),
        "missing_observation_ref": int(counters["missing_observation_ref"]),
        "unique_observation_refs": int(embedding_matrix.shape[0]),
        "embedding_path": str(embed_path),
        "embedding_content_hash": index["embedding_content_hash"],
        "content_hash_after": index["content_hash"],
        "split_hash": index.get("split_hash"),
        "task_counts": dict(sorted(task_counts.items())),
    }


def _embedding_for_ref(key: tuple[str, str], h5_cache: dict[Path, Any]) -> np.ndarray:
    source_dataset, observation_ref = key
    archive_name, dataset_name, row_index = _parse_observation_ref(observation_ref)
    archive_path = Path(source_dataset) / archive_name
    if archive_path not in h5_cache:
        try:
            import h5py
        except ImportError as exc:  # pragma: no cover
            raise ImportError("export_chart_observation_embeddings.py requires h5py") from exc
        h5_cache[archive_path] = h5py.File(archive_path, "r")
    payload = np.asarray(h5_cache[archive_path][dataset_name][row_index], dtype=np.uint8)
    return _rgb_stats_embedding(payload.tobytes())


def _rgb_stats_embedding(jpeg_bytes: bytes) -> np.ndarray:
    try:
        from PIL import Image
    except ImportError as exc:  # pragma: no cover
        raise ImportError("export_chart_observation_embeddings.py requires Pillow") from exc
    image = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")
    arr = np.asarray(image, dtype=np.float32) / 255.0
    h, w, _ = arr.shape
    y0, y1 = h // 4, h - h // 4
    x0, x1 = w // 4, w - w // 4
    center = arr[y0:y1, x0:x1]
    grid_means = []
    for y_slice in (slice(0, h // 2), slice(h // 2, h)):
        for x_slice in (slice(0, w // 2), slice(w // 2, w)):
            grid_means.extend(arr[y_slice, x_slice].mean(axis=(0, 1)).tolist())
    luminance = arr.mean(axis=2)
    hist, _ = np.histogram(luminance, bins=8, range=(0.0, 1.0), density=False)
    hist = hist.astype(np.float32)
    hist = hist / max(float(hist.sum()), 1.0)
    feature = np.asarray(
        [
            *arr.mean(axis=(0, 1)).tolist(),
            *arr.std(axis=(0, 1)).tolist(),
            *center.mean(axis=(0, 1)).tolist(),
            *center.std(axis=(0, 1)).tolist(),
            *grid_means,
            *hist.tolist(),
        ],
        dtype=np.float32,
    )
    if feature.shape[0] != OBSERVATION_EMBED_DIM:
        raise AssertionError(f"expected {OBSERVATION_EMBED_DIM} dims, got {feature.shape[0]}")
    return feature


def _parse_observation_ref(value: str) -> tuple[str, str, int]:
    if "#" not in value:
        raise ValueError(f"invalid observation_ref: {value}")
    archive_name, ref = value.split("#", 1)
    parts = [part for part in ref.split("/") if part]
    if len(parts) != 2:
        raise ValueError(f"invalid observation_ref: {value}")
    return archive_name, parts[0], int(parts[1])


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 _metrics_by_task(rows: list[dict[str, Any]]) -> str:
    payload: dict[str, dict[str, int]] = {}
    for row in rows:
        for task, count in row["task_counts"].items():
            payload.setdefault(task, {})[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/export_chart_observation_embeddings.py",
        "\\begin{tabular}{lrrr}",
        "\\toprule",
        "Split & Rows & With embedding & Unique refs \\\\",
        "\\midrule",
    ]
    for row in rows:
        lines.append(
            f"{row['split']} & {row['rows']} & "
            f"{row['rows_with_embedding']} & {row['unique_observation_refs']} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _report(metrics: dict[str, Any]) -> str:
    lines = [
        "# Chart Observation Embedding Export",
        "",
        "Decoded deployment-visible RGB observation refs into 32D statistics embeddings. "
        "No outcome, label, or hidden-branch fields are read.",
        "",
        "| Split | Rows | With embedding | Missing refs | Unique refs |",
        "| --- | ---: | ---: | ---: | ---: |",
    ]
    for row in metrics["splits"]:
        lines.append(
            f"| {row['split']} | {row['rows']} | {row['rows_with_embedding']} | "
            f"{row['missing_observation_ref']} | {row['unique_observation_refs']} |"
        )
    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/export_chart_observation_embeddings.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:
        if index_path.exists():
            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 _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()


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())