code/bundle/run_jer_only.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "torch",
#   "torchvision",
#   "numpy",
#   "scipy",
#   "scikit-learn",
#   "pillow",
#   "timm",
#   "open-clip-torch",
#   "datasets",
#   "huggingface-hub",
# ]
# ///
"""Fill missing JER values for models that already have binding scores."""

from __future__ import annotations

import json
import os
import shutil
import sys
from pathlib import Path

ROOT = Path("/mnt/repro/code")
if not (ROOT / "metrics.py").exists():
    ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

from metrics import jacobian_effective_rank_noise  # noqa: E402
from models import load_encoder  # noqa: E402


def main() -> None:
    results_path = Path(os.environ.get("RESULTS", str(ROOT / "outputs" / "results_claim4.json")))
    n_jer = int(os.environ.get("N_JER", "20"))
    data = json.loads(results_path.read_text())
    for row in data["rows"]:
        if row.get("binding") is None:
            continue
        if row.get("JER") is not None:
            continue
        key = row["key"]
        print(f"[jer-only] {key}", flush=True)
        enc = load_encoder(key, device="cuda")
        try:
            jer = jacobian_effective_rank_noise(
                enc.encode_tensor, n_images=n_jer, k=32, device=enc.device, seed=42
            )
            row["JER"] = jer
            row["error"] = None
            print(f"[jer-only] {key} JER={jer:.4f}", flush=True)
        except Exception as e:
            row["error"] = f"JER failed: {type(e).__name__}: {e}"
            print(row["error"], flush=True)
        results_path.write_text(json.dumps(data, indent=2))
    out_dir = Path("/mnt/repro/outputs")
    if out_dir.parent.exists():
        out_dir.mkdir(parents=True, exist_ok=True)
        shutil.copy2(results_path, out_dir / results_path.name)
    print(json.dumps(data, indent=2), flush=True)


if __name__ == "__main__":
    main()