Create codebook_contributions.py

codebook_contributions.py

@@ -0,0 +1,215 @@
+"""
+battery_ablation.py — test contribution signals across batteries.
+
+For each battery: load it frozen, extract its projective codebook, compute the
+contribution signals (codebook_contributions), and pull its recon MSE as the
+target. Then rank every signal by:
+  * std across batteries  — does it vary at all, or is it a dead signal?
+  * |corr| with recon MSE — does it track downstream quality?
+
+This is the "run N trains, test each contribution as a whole" pass: each
+battery is one data point; the ablation table says which contributions earn a
+slot in the omega-phase classifier before we hardwire any of them.
+
+Cell workflow: paste codebook_contributions cell first, then this. Edit
+BATTERIES to your set (≥3 needed for correlation). `pip install ripser` for the
+H1/H2 void signals; without it they self-exclude as NaN.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+# cell-tolerant: from the codebook_contributions cell (or installed)
+try:
+    from codebook_contributions import (
+        collect_signatures, ablation_table, SIGNAL_SPECS, HAVE_RIPSER,
+    )
+except ModuleNotFoundError:
+    pass
+
+
+# ── edit this to your battery set ───────────────────────────────────
+BATTERIES: List[str] = [
+    "h2_linear_tiny_imagenet_64",
+    # add your other battery folder names here, e.g.:
+    # "h2_linear_imagenet_128",
+    # "byte_trigram_proto_64_patch_2_v1",
+    # "v40_freckles_noise", "v50_fresnel_64", ...
+]
+REPO_ID = "AbstractPhil/geolip-SVAE"
+
+
+def discover_batteries(repo_id: str = REPO_ID) -> List[str]:
+    """List every battery folder in the repo that has a checkpoints/best.pt.
+    Saves you maintaining BATTERIES by hand — `run_ablation(discover_batteries())`
+    ablates over the whole zoo (mixed classes/D are fine; signals are D-normalized)."""
+    from huggingface_hub import HfApi
+    files = HfApi().list_repo_files(repo_id)
+    vers = sorted({f.split("/")[0] for f in files if f.endswith("/checkpoints/best.pt")})
+    print(f"  discovered {len(vers)} batteries in {repo_id}")
+    return vers
+
+
+def _load_model_safe(ver: str, device: str, repo_id: str):
+    """load_model, with a fallback for torch.compile checkpoints whose state-dict
+    keys carry an '_orig_mod.' prefix. On that specific failure: re-download, strip
+    the prefix (and backfill config from final_report.json the way load_model would,
+    since checkpoint_path loads skip hf_version backfill), re-save, re-enter via
+    checkpoint_path so all of load_model's construction logic is reused."""
+    from geolip_svae.inference.loading import load_model
+    try:
+        return load_model(hf_version=ver, device=device, repo_id=repo_id)
+    except RuntimeError as e:
+        if "_orig_mod." not in str(e):
+            raise
+        import torch, os, tempfile, json
+        from huggingface_hub import hf_hub_download
+        path = hf_hub_download(repo_id=repo_id, filename=f"{ver}/checkpoints/best.pt",
+                               repo_type="model")
+        ckpt = torch.load(path, map_location="cpu", weights_only=False)
+        pref = "_orig_mod."
+        ckpt["model_state_dict"] = {
+            (k[len(pref):] if k.startswith(pref) else k): v
+            for k, v in ckpt["model_state_dict"].items()
+        }
+        # mirror load_model's final_report backfill into the temp config
+        cfg0 = dict(ckpt.get("config", {}))
+        backfillable = ("n_heads", "smooth_mid", "linear_readout",
+                        "svd_mode", "match_params", "channels")
+        if any(k not in cfg0 for k in backfillable):
+            try:
+                rp = hf_hub_download(repo_id=repo_id, filename=f"{ver}/final_report.json",
+                                     repo_type="model")
+                rc = json.load(open(rp)).get("config", {})
+                for k in backfillable:
+                    if k not in cfg0 and rc.get(k) is not None:
+                        cfg0[k] = rc[k]
+                ckpt["config"] = cfg0
+            except Exception:
+                pass
+        tmp = os.path.join(tempfile.gettempdir(), f"{ver.replace('/', '_')}_stripped.pt")
+        torch.save(ckpt, tmp)
+        model, cfg = load_model(checkpoint_path=tmp, device=device, repo_id=repo_id)
+        print(f"  (recovered {ver}: stripped _orig_mod. torch.compile prefix)")
+        return model, cfg
+
+
+def extract_row(ver: str, device: str) -> Dict[str, Any]:
+    """Load a frozen battery, extract its codebook, return an ablation row
+    {id, axes, D, n_pairs, n_unpaired, target=recon_mse, class}."""
+    from geolip_svae.inference.calibration import make_calibration
+    from geolip_svae.inference.codebook import extract_codebook
+    from geolip_svae.inference.train_codebook import (
+        infer_class_from_cfg, DEFAULT_CALIBRATIONS,
+    )
+    import torch
+
+    model, cfg = _load_model_safe(ver, device, REPO_ID)
+    cls = infer_class_from_cfg(cfg)
+    cal = DEFAULT_CALIBRATIONS.get(cls, DEFAULT_CALIBRATIONS["unknown"])
+    size = cfg.get("img_size") or cal["size"]
+
+    calib = make_calibration(cal["name"], n=cal["n"], size=size)
+    if not isinstance(calib, torch.Tensor):
+        calib = torch.as_tensor(calib)
+    ch = int(cfg.get("channels", 3))                   # match model input channels
+    if calib.shape[1] != ch:
+        if ch < calib.shape[1]:
+            calib = calib[:, :ch]
+        else:
+            reps = (ch + calib.shape[1] - 1) // calib.shape[1]
+            calib = calib.repeat(1, reps, 1, 1)[:, :ch]
+
+    cb = extract_codebook(model, calib.to(device), model_id=ver,
+                          model_class=cls, calibration_name=cal["name"])
+    axes = cb.axes.detach().cpu().numpy()
+    n_pairs = getattr(cb.metadata, "n_pairs", None)
+    n_unpaired = getattr(cb.metadata, "n_unpaired", None)
+    if n_pairs is None:
+        n_pairs, n_unpaired = len(cb.pairs), len(cb.unpaired)
+
+    return {
+        "id": ver,
+        "class": cls,
+        "axes": axes,
+        "D": int(cfg.get("D") or axes.shape[1]),
+        "n_pairs": int(n_pairs),
+        "n_unpaired": int(n_unpaired),
+        "target": cfg.get("_test_mse"),               # recon MSE (None if absent)
+        "n_axes": int(axes.shape[0]),
+    }
+
+
+def run_ablation(batteries: Optional[List[str]] = None, device: Optional[str] = None,
+                 enabled=None) -> Dict[str, Any]:
+    """Extract every battery's codebook, compute signatures, rank contributions."""
+    import torch
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    batteries = batteries or BATTERIES
+    print(f"[battery_ablation] {len(batteries)} batteries on {device} | ripser={HAVE_RIPSER}")
+
+    cb_rows: List[Dict[str, Any]] = []
+    for ver in batteries:
+        try:
+            row = extract_row(ver, device)
+            cb_rows.append(row)
+            print(f"  ok   {ver:42s} class={row['class']:12s} "
+                  f"n_axes={row['n_axes']:3d} target_mse={row['target']}")
+        except Exception as e:
+            print(f"  SKIP {ver:42s} {type(e).__name__}: {e}")
+
+    if not cb_rows:
+        print("  no batteries loaded — check BATTERIES / network")
+        return {}
+
+    rows = collect_signatures(cb_rows, enabled=enabled)
+
+    # per-battery signature table
+    names = [s[0] for s in SIGNAL_SPECS if (enabled is None or s[0] in enabled)]
+    print("\n── per-battery contribution values ──")
+    header = "battery".ljust(42) + "".join(f"{n[:11]:>13s}" for n in names)
+    print(header)
+    for r in rows:
+        line = r["id"][:40].ljust(42)
+        for n in names:
+            v = r["values"].get(n, float("nan"))
+            line += f"{v:>13.4f}"
+        print(line)
+
+    # ablation ranking
+    table = ablation_table(rows)
+    n_target = max((s["n_target"] for s in table.values()), default=0)
+    classes_present = sorted({r.get("class") for r in rows if r.get("class") is not None})
+    print(f"\n── contribution informativeness ──")
+    print(f"   cv = scale-free spread | |rho| = |Spearman| w/ recon MSE (n={n_target}, detects BROKEN)")
+    print(f"   eta2 = variance explained by class (detects CLASS SEPARATION) | classes: {classes_present}")
+    def _key(it):
+        e = it[1]["eta2_by_class"]
+        rho = it[1]["abs_spearman_with_target"]
+        return (-(e if e == e else -1), -(rho if rho == rho else -1))
+    for name, stats in sorted(table.items(), key=_key):
+        rho = stats["abs_spearman_with_target"]; rho_s = f"{rho:.3f}" if rho == rho else "  -- "
+        eta = stats["eta2_by_class"];            eta_s = f"{eta:.3f}" if eta == eta else "  -- "
+        cv = stats["cv"];                        cv_s = f"{cv:6.2f}" if cv == cv else "   -- "
+        print(f"  {name:26s} eta2={eta_s}  |rho|={rho_s}  cv={cv_s}  n={stats['n_valid']}")
+
+    # per-class means for the strongest class separators
+    top = sorted(table.items(), key=_key)[:4]
+    print(f"\n── per-class means (top {len(top)} class-separating signals) ──")
+    hdr = "class".ljust(16) + "".join(f"{n[:11]:>13s}" for n, _ in top)
+    print(hdr)
+    for c in classes_present:
+        line = str(c).ljust(16)
+        for _, stats in top:
+            mv = stats["class_means"].get(str(c))
+            line += (f"{mv:>13.3f}" if mv is not None else f"{'--':>13s}")
+        print(line)
+    return {"rows": rows, "table": table}
+
+
+if __name__ == "__main__":
+    # If BATTERIES is left at the lone default, ablate the whole discovered zoo.
+    bats = BATTERIES if len(BATTERIES) > 1 else discover_batteries()
+    run_ablation(bats)