← Walkthroughevals/neuroprobe/encode.py
Source snapshot · source SHA-256 6bdb0af6778b
1"""Encode the 12 Neuroprobe-Lite sessions into frozen per-region and per-contact features.23The three normalized 32 Hz magnitude bands receive Guard 3 at the model4input. Tap 0 uses the same clipped, decimated inputs as the encoder.5"""6from __future__ import annotations78import argparse9import hashlib10import os1112import numpy as np13import torch1415from mapa.models.stem import clip_band1617# The 15 Neuroprobe leaderboard tasks (upstream ``neuroprobe.config.NEUROPROBE_TASKS``).18# --tasks board re-labels the SAME task-agnostic features, so the only extra encode cost is19# materializing 15 label vectors instead of 4.20BOARD_TASKS: tuple[str, ...] = (21 "onset", "speech", "volume", "delta_volume", "pitch", "word_index",22 "word_gap", "gpt2_surprisal", "word_head_pos", "word_part_speech",23 "word_length", "global_flow", "local_flow", "frame_brightness", "face_num",24)25# The 12 Neuroprobe-Lite sessions (upstream ``NEUROPROBE_LITE_SUBJECT_TRIALS``): the board's26# CS anchor (2,4) + the 10 CS test cells + (2,0). --sessions board evaluates these.27BOARD_SESSIONS: tuple[tuple[int, int], ...] = (28 (1, 1), (1, 2), (2, 0), (2, 4), (3, 0), (3, 1),29 (4, 0), (4, 1), (7, 0), (7, 1), (10, 0), (10, 1),30)31FPS = 32.032CLIP_DUR_S = 1.033N_REGIONS = 7534GPU_TAPS: tuple[int, ...] = (3, 6, 9, 12) # raw block outputs read in one encoder forward353637def _load_ckpt(ckpt_path: str) -> tuple[dict, bool | None]:38 """The encoder weights and, when the file records it, whether relative position is on.3940 Reads both layouts through the same helper the public loader uses, so there is one41 implementation of what counts as a checkpoint.42 """43 from mapa.encoder import _tower_state4445 return _tower_state(torch.load(ckpt_path, map_location="cpu", weights_only=True))464748def _freeze(m, *, device):49 m.eval().to(device)50 for p in m.parameters():51 p.requires_grad_(False)52 return m535455def _load_encoder(tsd: dict, *, device: torch.device, space_rope: bool):56 """Build the encoder shell and load the weights into it.5758 Everything the shell needs is read off the weights rather than passed in, so a wrong guess59 fails the strict check below instead of being absorbed. Deep supervision shows up as60 ``norms_block.*`` with no ``norm_out``, and a single tap as the reverse. Width is read off a61 LayerNorm, and a wrong width cannot load at all, whatever ``strict`` says.6263 Relative position is the one exception, handled by the caller: see ``--no-relative-position``.64 """65 from mapa.models.tower import MaeEncoder6667 peek = [v.shape[0] for kk, v in tsd.items() if kk.endswith("region_embed.embed.weight")]68 if peek and int(peek[0]) != N_REGIONS:69 raise ValueError(f"ckpt region table {peek[0]} != expected {N_REGIONS}")70 deep_sup = any(kk.startswith("encoder.norms_block.") for kk in tsd)71 # region_embed is READ OFF the ckpt for the same reason deep_sup is: the --no-region-embed72 # arm ships a tower with no region_embed submodule at all, so a shell built with one puts73 # region_embed.embed.weight in `missing` and the check below raises. Inferring it keeps that74 # check as the verifier — a wrong inference still fails loud, it is never silently absorbed.75 region_embed = bool(peek)76 dkey = "encoder.blocks.0.norm1.weight"77 if dkey not in tsd:78 raise RuntimeError(f"ckpt has no '{dkey}'; cannot infer encoder width")79 d_model = int(tsd[dkey].shape[0])80 print(f"[encode] encoder deep_sup={deep_sup} region_embed={region_embed} d_model={d_model} "81 f"(read off the weights) relative_position={space_rope}")82 tower = MaeEncoder(n_regions=N_REGIONS, deep_sup=deep_sup, region_embed=region_embed,83 space_rope=space_rope, d_model=d_model)84 missing, unexpected = tower.load_state_dict(tsd, strict=False)85 bad = [m for m in missing if "num_batches_tracked" not in m]86 if bad or unexpected:87 raise RuntimeError(f"encoder state_dict mismatch: missing={bad[:8]} unexpected={unexpected[:8]}")88 return _freeze(tower, device=device)899091def _load_targets(session, bt_root, tasks=BOARD_TASKS):92 from mapa.probe.labels import build_label_events as _label_events93 from mapa.probe.labels import build_session_targets9495 subject_id, trial_id = session96 events = _label_events(subject_id, trial_id, f"btbank{subject_id}_trial{trial_id}",97 tasks, bt_root, lite_cap=True)98 # build_session_targets derives its task list from events["task"].unique(), so passing99 # `tasks` to _label_events is what widens 4 -> 15.100 return build_session_targets(events, subject_id=subject_id, trial_id=trial_id)101102103def _lite_keep_labels_fn(bt_root):104 """``keep_labels_fn`` restricting a session to its Neuroprobe-Lite montage.105106 Injected into ``load_v3_sessions`` (NOT applied to the built spec) so keep_idx / region_id107 / sidecar / geom / band_stats are all constructed on the Lite axis by the normal code path108 — ``geom`` cannot be masked post-hoc, its gather_idx stores indices into the survivor axis.109 The realized montage is a SET intersection (``voltage_order ∩ lite_labels``), matching110 upstream ``datasets.py`` ``[full.index(e) for e in lite if e in full]``; we keep OUR voltage111 order, which is free for the encoder (the per-region pool is permutation-invariant within a112 region) and is what index-RoPE expects."""113 from mapa.data.anatomy import lite_electrode_set114115 def fn(subject_id, trial_id, labels):116 return set(lite_electrode_set(subject_id))117118 return fn119120121def _window_bands(spec, starts, clip_frames):122 """Slice normalized evaluation windows from three continuous 32 Hz band caches.123124 Load selected contacts once per band to avoid repeated scattered disk reads."""125 keep = spec.keep_idx.numpy()126 starts = np.asarray(starts, dtype=float)127 n_frames_native = spec.n_frames128 t0 = np.rint(starts * FPS).astype(np.int64)129 end = t0 + clip_frames130 oob = np.where((t0 < 0) | (end > n_frames_native))[0]131 if len(oob):132 raise RuntimeError(133 f"{spec.session_key}: {len(oob)} union windows out of cache bounds "134 f"(n_frames_native={n_frames_native}, first bad start={float(starts[oob[0]]):.4f}s)"135 )136 bands = []137 for path, norm in zip(spec.band_paths, spec.band_norms):138 mm = np.load(path, mmap_mode="r")139 full = np.asarray(mm[keep], dtype=np.float32) # (N, F, T_total) bulk → RAM140 del mm141 clips = np.stack([full[:, :, a:b] for a, b in zip(t0.tolist(), end.tolist())], axis=0)142 t = torch.from_numpy(clips)143 bands.append(norm.transform(t)) # (n, N, F_b, T32) vectorized144 return bands # 3 × (n, N, F_b, T32)145146147def _canon_regions(grid, region_id):148 """Canonical (grid-order) contact indices + their region ids + present region atlas ids.149150 build_r4_grid lays tokens contact-major (k_full block per contact); the first token of151 each block carries that contact's index (``grid.contact``), so reshaping to (n, k_full)152 and taking column 0 recovers the n canonical contacts and their regions."""153 k = grid.k_full154 canon = grid.contact.reshape(-1, k)[:, 0].cpu().numpy() # (n,) contact index into N155 region_canon = region_id.cpu().numpy()[canon] # (n,) DKT tag per canon contact156 present = np.unique(region_canon) # sorted present atlas ids157 return canon, region_canon, present158159160def _pool_regions(x, region_canon, present):161 """Pool electrodes→regions: x (B, n, *feat) → (B, |P|, prod(feat)) flattened last dim.162163 Per present region, MEAN over its electrodes. Returns (B, |P|, F) fp16, region order == present."""164 B = x.shape[0]165 blocks = []166 for p in present:167 cols = np.where(region_canon == p)[0]168 sub = x[:, cols] # (B, |cols|, *feat)169 blocks.append(sub.mean(1).reshape(B, -1)) # (B, prod(feat))170 return torch.stack(blocks, dim=1).to(torch.float16) # (B, |P|, F)171172173def _enc0_band_lengths(bands, strides):174 """Frame counts after the model’s per-band decimation."""175 st = _enc0_strides(bands, strides)176 return tuple(int(-(-b.shape[-1] // s)) for b, s in zip(bands, st)) # len of x[..., ::s]177178179def _enc0_strides(bands, strides):180 """Resolve/validate the per-band decimation strides shared by enc0's pooled + elec paths."""181 from mapa.models.pack_r4 import BAND_STRIDES182183 if strides is None:184 strides = BAND_STRIDES185 if len(strides) != len(bands):186 raise ValueError(f"enc0 got {len(bands)} bands but {len(strides)} strides")187 return strides188189190def _enc0_bands_canon(bands, canon, strides):191 """Per band: decimate ``x[..., ::stride]`` (the model's own input frames), reorder to192 canonical contacts, time-major. Yields (n_win, n_canon, T_b, F_b) tensors, one per band —193 the shared prefix of enc0's pooled and unpooled (elec) paths."""194 for b, (x, st) in enumerate(zip(bands, strides)): # x (n, N, F_b, T_clock)195 xd = clip_band(x[..., ::st], b) # (n, N, F_b, T_b)196 xd = xd.transpose(-1, -2).contiguous() # (n, N, T_b, F_b) time-major197 yield xd[:, canon] # (n, n_canon, T_b, F_b)198199200def _enc0_pooled(bands, canon, region_canon, present, strides=None):201 """Model-clipped input baseline, decimated and pooled by atlas region."""202 strides = _enc0_strides(bands, strides)203 per_band = [_pool_regions(xd, region_canon, present)204 for xd in _enc0_bands_canon(bands, canon, strides)]205 return torch.cat(per_band, dim=-1) # (n, |P|, F0)206207208def _enc0_elec(bands, canon, strides=None):209 """Unpooled enc0: same canonical-contact axis and band-concat order as ``_enc0_pooled``'s210 pre-pool input — the depth-0 sibling of ``_encode_taps``' ``enc{t}_elec`` (GPU taps211 3/6/12), which stores the equivalent unpooled tensor for those taps."""212 strides = _enc0_strides(bands, strides)213 per_band = []214 for xd in _enc0_bands_canon(bands, canon, strides):215 B = xd.shape[0]216 per_band.append(xd.reshape(B, xd.shape[1], -1).to(torch.float16))217 return torch.cat(per_band, dim=-1) # (n, n_canon, F0)218219220@torch.no_grad()221def _encode_taps(encoder, bands, grid, region_packed, region_canon, present,222 *, device, batch_size, elec_taps=()):223 """One forward of the encoder over all windows → per-tap region-pooled keep-time features.224225 Cache stores the raw region-mean feature (n,|P|,k_full·d) — the most flexible storage: a226 readout can standardize columns on train stats (the FM linear-probe convention) or feed it227 raw (r1/M9-comparable), but neither is recoverable from a baked per-token LN. So we keep raw228 only. Returns {tap: {'raw': ...}}."""229 n = bands[0].shape[0]230 k = grid.k_full231 # Per-electrode keep-time: WS keeps ALL electrodes; the region-mean is the232 # comparison. Stored UNPOOLED on the canonical-contact axis (same order as region_canon), so233 # it is the pooled tap's exact pre-mean input — the diff is the pooling and nothing else.234 #235 # Each tap is preallocated at full length and written slice-by-slice. Accumulating a list of236 # per-batch tensors and torch.cat-ing at the end held the list AND its concatenation alive237 # simultaneously — a 2x peak on the ~40 GB enc12_elec tap, which is what forced --mem=300G238 # (measured 230-245 GiB on the 12-session board encode). Values, dtypes and row order are239 # unchanged; only the allocation pattern differs.240 acc: dict = {key: None for key in list(GPU_TAPS) + [f"elec{t}" for t in elec_taps]}241242 def _write(key, lo, hi, x):243 if acc[key] is None:244 acc[key] = torch.empty((n, *x.shape[1:]), dtype=x.dtype)245 acc[key][lo:hi] = x246247 for s in range(0, n, batch_size):248 e = min(s + batch_size, n)249 bb = [b[s:e].to(device) for b in bands]250 Bb = e - s251 with torch.autocast(device_type=device.type, dtype=torch.bfloat16,252 enabled=(device.type == "cuda")):253 _z, taps = encoder.forward(bb, grid, region_packed, tap_blocks=GPU_TAPS)254 for t in GPU_TAPS:255 enc = taps[t].float().reshape(Bb, -1, k, taps[t].shape[-1]).cpu() # (Bb, n, k, d)256 _write(t, s, e, _pool_regions(enc, region_canon, present))257 if t in elec_taps:258 # (Bb, n_contacts, k·d) fp16 — the SAME tensor, just unpooled.259 _write(f"elec{t}", s, e, enc.reshape(Bb, enc.shape[1], -1).to(torch.float16))260 return {t: {"raw": v} for t, v in acc.items()}261262263def main() -> None:264 p = argparse.ArgumentParser()265 p.add_argument("--ckpt", default=None,266 help="required unless --enc0-only (enc0 never reads weights)")267 p.add_argument("--tag", required=True)268 p.add_argument("--out-dir", required=True)269 p.add_argument("--band-cache-dir", dest="band_cache_dirs", action="append", required=True,270 help="3× in v3 concat order: slow, mid, hga")271 p.add_argument("--span-dir", required=True)272 p.add_argument("--bt-root", required=True)273 p.add_argument("--batch-size", type=int, default=64)274 p.add_argument("--sessions", choices=("board",), default="board",275 help="the 12 Neuroprobe-Lite sessions")276 p.add_argument("--session-index", type=int, default=None,277 help="encode ONLY cohort[i], one Slurm array task per session. "278 "Default = all.")279 p.add_argument("--tasks", choices=("board15",), default="board15",280 help="board15 = the 15 leaderboard tasks (re-labels the same features)")281 p.add_argument("--electrode-set", choices=("lite",), default="lite",282 help="lite = the Neuroprobe-Lite montage (leaderboard parity)")283 p.add_argument("--no-relative-position", "--no-space-rope", dest="no_space_rope",284 action="store_true",285 help="the checkpoint was trained without the relative positional encoding. "286 "A released checkpoint records this itself and needs no flag. A raw "287 "training checkpoint records nothing, and the setting leaves no trace in "288 "the weights, so it has to be passed by hand: get it wrong and the encode "289 "rotates by a contact index the trained model never saw, with no error.")290 p.add_argument("--enc0-only", action="store_true",291 help="compute ONLY the enc0 input floor (raw decimated band bins pooled to "292 "regions). No ckpt, no model, no GPU — enc0 never touched weights. Lets "293 "the published frontend baseline run without loading a checkpoint.")294 p.add_argument("--elec-taps", default="0,12",295 help="comma-separated taps to ALSO write per-electrode (unpooled), e.g. "296 "'0,12' -> feats['enc0_elec'], feats['enc12_elec']. 0 = the |STFT| "297 "frontend (enc0); 3/6/9/12 route through the encoder forward (GPU_TAPS). "298 "WS keeps all electrodes by default; each costs "299 "~N/|P| (~5x) the pooled tap on disk.")300 args = p.parse_args()301302 from mapa.data.region_fn import make_bt_region_fn303 from mapa.data.session_loader import load_v3_sessions304 from mapa.models.pack_r4 import build_r4_grid305306 n_cache = 3307 if len(args.band_cache_dirs) != n_cache:308 raise SystemExit(309 f"need {n_cache} --band-cache-dir (slow, mid, hga), got {len(args.band_cache_dirs)}")310 os.makedirs(args.out_dir, exist_ok=True)311 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")312 clip_frames = round(CLIP_DUR_S * FPS)313 region_fn = make_bt_region_fn(args.bt_root)314 cohort = BOARD_SESSIONS315 if args.session_index is not None:316 if not 0 <= args.session_index < len(cohort):317 raise SystemExit(f"--session-index {args.session_index} out of range [0,{len(cohort)})")318 cohort = (cohort[args.session_index],)319 for s, t in cohort:320 path = os.path.join(args.out_dir, f"enc_s{s}_t{t}_{args.tag}.pt")321 if os.path.exists(path):322 raise SystemExit(323 f"output already exists: {path}. Use a new tag/directory, or --session-index "324 "to encode only a missing session; existing files are not verified or resumed."325 )326 tasks = BOARD_TASKS327 keep_labels_fn = _lite_keep_labels_fn(args.bt_root)328 elec_taps = tuple(int(t) for t in args.elec_taps.split(",") if t.strip())329 bad = [t for t in elec_taps if t not in (0,) + GPU_TAPS]330 if bad:331 raise SystemExit(f"--elec-taps {bad} not in {(0,) + GPU_TAPS} (0 = enc0_elec, off the "332 f"|STFT| frontend directly; the rest route through GPU_TAPS {GPU_TAPS})")333 # tap 0 never goes through the encoder forward (_encode_taps/GPU_TAPS) — it's the frontend334 # input, handled by _enc0_elec below instead of _encode_taps' elec_taps loop.335 gpu_elec_taps = tuple(t for t in elec_taps if t != 0)336 want_enc0_elec = 0 in elec_taps337 checkpoint = None338 if args.enc0_only:339 # enc0 is a pure function of the CACHED BANDS (no weights ever touched), so there is340 # nothing to load and nothing to run on a GPU.341 encoder = None342 tower_note = "ENC0-ONLY (no ckpt, no model, CPU)"343 else:344 if not args.ckpt:345 raise SystemExit("--ckpt is required unless --enc0-only")346 sd, recorded = _load_ckpt(args.ckpt)347 if args.no_space_rope:348 space_rope = False349 elif recorded is not None:350 space_rope = recorded351 else:352 raise SystemExit(353 f"{args.ckpt} does not record whether the relative positional encoding is "354 "active, and it cannot be read off the weights. Pass --no-relative-position "355 "for an ablation checkpoint, or re-export the checkpoint with its metadata.")356 encoder = _load_encoder(sd, device=device, space_rope=space_rope)357 digest = hashlib.sha256()358 with open(args.ckpt, "rb") as f:359 for block in iter(lambda: f.read(1024 * 1024), b""):360 digest.update(block)361 checkpoint = {362 "sha256": digest.hexdigest(), "space_rope": space_rope,363 "region_embed": any(k.endswith("region_embed.embed.weight") for k in sd),364 "d_model": encoder.d_model,365 }366 del sd367 tower_note = "encoder"368 print(f"[encode-r4] tag={args.tag} device={device} gpu_taps={GPU_TAPS} + enc0 "369 f"{tower_note}, sessions={args.sessions}({len(cohort)}) "370 f"tasks={args.tasks}({len(tasks)}) electrodes={args.electrode_set}", flush=True)371372 for session in cohort:373 subject_id, trial_id = session374 path = os.path.join(args.out_dir, f"enc_s{subject_id}_t{trial_id}_{args.tag}.pt")375 spec = load_v3_sessions(376 sessions=[session], band_cache_dirs=args.band_cache_dirs, span_dir=args.span_dir,377 region_fn=region_fn, lof_report_path=None,378 keep_labels_fn=keep_labels_fn,379 )[0]380 if keep_labels_fn is not None:381 # The montage is the WHOLE parity claim — print what was realized, per session.382 from mapa.data.anatomy import lite_electrode_set383 lite = lite_electrode_set(subject_id)384 kept = spec.setup.sidecar.labels385 ok = set(kept) <= set(lite)386 print(f"[check] lite montage s{subject_id}_t{trial_id}: kept={len(kept)} "387 f"of lite-list {len(lite)} subset-of-lite={ok} -> {'OK' if ok else 'VIOLATED'}",388 flush=True)389 if not ok:390 raise RuntimeError("lite montage kept a non-Lite electrode — refusing to write")391 targets = _load_targets(session, args.bt_root, tasks)392 bands = _window_bands(spec, targets.clip_starts, clip_frames)393394 geom = spec.setup.geom.to(device)395 region_id = spec.setup.region_id.to(device)396 grid = build_r4_grid(geom, n_time=clip_frames)397 region_true = region_fn(subject_id, trial_id, list(spec.setup.sidecar.labels)).long()398 if not torch.equal(region_true, spec.setup.region_id.cpu()):399 raise RuntimeError(400 f"region tag parity FAILED for s{subject_id}_t{trial_id}: the recomputed atlas "401 f"tag disagrees with the session spec on "402 f"{int((region_true != spec.setup.region_id.cpu()).sum())} of "403 f"{region_true.numel()} contacts. Refusing to encode."404 )405 region_packed = region_id[grid.contact] # MODEL-side tag406 canon, region_canon, present = _canon_regions(grid, region_true) # TRUE atlas ids407408 feats = {"enc0": {"raw": _enc0_pooled(bands, canon, region_canon, present)}}409 if want_enc0_elec:410 feats["enc0_elec"] = {"raw": _enc0_elec(bands, canon)}411 enc0_lengths = _enc0_band_lengths(bands, None)412 if enc0_lengths != tuple(int(x) for x in grid.band_lengths):413 raise RuntimeError(f"enc0 band lengths {enc0_lengths} != model grid {grid.band_lengths}")414 if args.enc0_only:415 tap_pooled = {}416 else:417 tap_pooled = _encode_taps(encoder, bands, grid, region_packed, region_canon,418 present, device=device, batch_size=args.batch_size,419 elec_taps=gpu_elec_taps)420 for t in GPU_TAPS:421 feats[f"enc{t}"] = tap_pooled[t]422 for t in gpu_elec_taps:423 feats[f"enc{t}_elec"] = tap_pooled[f"elec{t}"]424425 payload = {426 "subject_id": subject_id, "trial_id": trial_id, "ckpt_tag": args.tag,427 "checkpoint": checkpoint,428 "present_parcels": np.asarray(present, dtype=np.int64), # (|P|,) atlas ids, feature order429 # (n_contacts,) atlas id per CANONICAL contact — the same axis the ``enc*_elec`` taps430 # are stored on, so a readout can re-pool electrodes→regions itself instead of being431 # stuck with the mean baked in here. present_regions alone is not enough: it gives the432 # region ORDER but not the membership.433 "parcel_canon": np.asarray(region_canon, dtype=np.int64),434 "elec_labels": np.asarray(spec.setup.sidecar.labels)[canon],435 "band_lengths": enc0_lengths,436 # Frequency bins per band. band_lengths alone does NOT let a consumer slice enc0 by437 # band: the encoder taps are (k_full, d) so band_lengths is enough there, but enc0 is438 # the raw spectrogram at Σ_b F_b·T_b, and F_b is not recoverable from that total.439 "band_fdims": tuple(int(b.shape[-2]) for b in bands),440 "feats": {k: {v: t for v, t in d.items()} for k, d in feats.items()},441 "clip_starts": np.asarray(targets.clip_starts),442 "labels": {lt: np.asarray(v) for lt, v in targets.labels.items()},443 "ws_split": targets.ws_split,444 "cs_split": targets.cs_split,445 "n_windows": int(bands[0].shape[0]),446 }447 with open(path, "xb") as f:448 torch.save(payload, f)449 shp = {k: tuple(next(iter(d.values())).shape) for k, d in feats.items()}450 print(f"[encode-r4] {session}: |P|={len(present)} n={payload['n_windows']} "451 f"shapes={shp} -> {path}", flush=True)452 del bands, feats, tap_pooled, payload453454455if __name__ == "__main__":456 main()