token_probe.py

"""
Token probe — the GO / NO-GO experiment before touching the LLM
==============================================================
Question: do the 3x3x3 = 27 spatial positions sitting before the encoder's
AdaptiveAvgPool3d(1) carry information worth handing to an LLM as "visual
tokens"?

Why this comes first: if the 27 tokens say nothing beyond the pooled 512-d
vector, wiring the encoder into the LLM as a token sequence is pointless — you
would just be moving the pipeline inside a transformer. Learning that in an hour
beats learning it after a week of projector training.

Method:
  1. Load the checkpoint and FREEZE the entire encoder.
  2. For every scan, compute the ASPP output (B, 512, 3, 3, 3) once and cache it
     as a (B, 27, 512) token sequence. The encoder is frozen, so it cannot change
     between epochs and the 3D CNN never has to run twice.
  3. Train two small probes on those tokens:
       - mean-pool probe : discards spatial information (what the model does now)
       - attention probe : weighted pooling with a learned query over 27 tokens
  4. Compare on the test split, and measure whether the attention is degenerate
     (uniform vs selective) through its entropy.

Reading the result:
  attention probe clearly ahead   → the token grid is meaningful, wire it in.
  probes equal and entropy ~max   → no spatial information; fix the encoder's
                                     strides / ASPP dilations first.

NOTE: the normalizer is loaded from the checkpoint and never re-fitted on the
evaluation data — re-fitting would leak.

This script uses the model modules (config, model, dataset); the path bootstrap
is below. YOU MUST SET YOUR OWN PATH via VBAI_MODEL_DIR if they are elsewhere.

Run:
    python token_probe.py --ckpt Vbai-2.6AD.pt
    python token_probe.py --ckpt ... --epochs 40 --batch-size 8
"""
from __future__ import annotations
import argparse
import os
import sys

# ----------------------------------------------------------------------
# The modality must be chosen BEFORE config is imported: config decides at
# import time which visit manifest to read (USE_TBM).
#
# CRITICAL: the checkpoint and the input modality must match. If they do not,
# the encoder drops to chance and the probe result is meaningless.
# ----------------------------------------------------------------------
_ap = argparse.ArgumentParser(add_help=False)
_ap.add_argument("--tbm", action="store_true")
_ap.add_argument("--t1", action="store_true")
_known, _ = _ap.parse_known_args()
if _known.tbm == _known.t1:
    sys.exit("ERROR: pass exactly one of --tbm / --t1 "
             "(match whichever modality the checkpoint was trained on).")
os.environ["VBAI_USE_TBM"] = "1" if _known.tbm else "0"
MODALITY = "TBM" if _known.tbm else "raw T1"


# ----------------------------------------------------------------------
# Make the model modules importable.
# ----------------------------------------------------------------------
def _bootstrap_model_path() -> str:
    env = os.environ.get("VBAI_MODEL_DIR")
    candidates = []
    if env:
        candidates.append(env)
    here = os.path.dirname(os.path.abspath(__file__))
    candidates.append(here)                                          # next to this file
    candidates.append(os.path.join(here, "Vbai-2.6AD"))              # ./Vbai-2.6AD
    for c in candidates:
        if os.path.isfile(os.path.join(c, "config.py")):
            if c not in sys.path:
                sys.path.insert(0, c)
            return c
    raise ImportError(
        "Could not locate the model modules (config.py).\n"
        f"Tried: {candidates}\n"
        "YOU MUST SET YOUR OWN PATH: point VBAI_MODEL_DIR at the directory "
        "holding config.py / model.py / dataset.py."
    )


MODEL_DIR = _bootstrap_model_path()

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, f1_score

import config as C
from model import Vbai26ADModel
from dataset import (PairedVisitDataset, TabularNormalizer, collate_pad,
                     subject_split, load_paired)


# ----------------------------------------------------------------------
# Token extraction — the encoder's forward pass without the pooling step
# ----------------------------------------------------------------------
@torch.no_grad()
def encode_tokens(encoder, mri: torch.Tensor):
    """(B, 1, 96, 96, 96) → (B, 27, 512) sequence of spatial tokens."""
    x = encoder.stem(mri)
    x = encoder.stage1(x)
    x = encoder.stage2(x)
    x = encoder.stage3(x)
    x = encoder.stage4(x)
    x = encoder.aspp(x)                      # (B, 512, 3, 3, 3)
    ch = x.shape[1]
    return x.flatten(2).transpose(1, 2).contiguous(), (ch, tuple(x.shape[2:]))


@torch.no_grad()
def build_token_cache(model, loader, device):
    """The encoder is frozen, so compute the tokens once and keep them in RAM."""
    toks, labels, shape_info = [], [], None
    for batch in loader:
        if "mri" not in batch:
            continue
        mri = batch["mri"].to(device, non_blocking=True)
        t, shape_info = encode_tokens(model.mri_encoder, mri)
        toks.append(t.float().cpu())
        labels.append(batch["label"].clone())
    if not toks:
        raise RuntimeError("No MRI batch could be loaded — check your data paths.")
    return torch.cat(toks), torch.cat(labels), shape_info


# ----------------------------------------------------------------------
# Probes
# ----------------------------------------------------------------------
class MeanPoolProbe(nn.Module):
    """Discards spatial information — equivalent to the current AdaptiveAvgPool3d(1)."""
    def __init__(self, dim, n_cls=3):
        super().__init__()
        self.head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, 128),
                                  nn.GELU(), nn.Dropout(0.2), nn.Linear(128, n_cls))

    def forward(self, tok):                       # (B, N, D)
        return self.head(tok.mean(dim=1)), None


class AttnPoolProbe(nn.Module):
    """Weighted pooling over the 27 tokens with a single learned query."""
    def __init__(self, dim, n_cls=3):
        super().__init__()
        self.q = nn.Parameter(torch.randn(1, 1, dim) * 0.02)
        self.norm = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads=8, batch_first=True)
        self.head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, 128),
                                  nn.GELU(), nn.Dropout(0.2), nn.Linear(128, n_cls))

    def forward(self, tok):                       # (B, N, D)
        x = self.norm(tok)
        q = self.q.expand(x.size(0), -1, -1)
        pooled, w = self.attn(q, x, x, need_weights=True, average_attn_weights=True)
        return self.head(pooled.squeeze(1)), w.squeeze(1)   # w: (B, N)


def train_probe(probe, tr_tok, tr_y, te_tok, te_y, device, epochs, bs, lr=3e-4):
    probe = probe.to(device)
    opt = torch.optim.AdamW(probe.parameters(), lr=lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

    # Compensate for class imbalance — MCI is already the weakest class
    counts = torch.bincount(tr_y, minlength=3).float().clamp(min=1)
    w = (counts.sum() / (3 * counts)).to(device)

    n = tr_tok.size(0)
    for _ in range(epochs):
        probe.train()
        perm = torch.randperm(n)
        for i in range(0, n, bs):
            idx = perm[i:i + bs]
            xb = tr_tok[idx].to(device)
            yb = tr_y[idx].to(device)
            logits, _ = probe(xb)
            loss = F.cross_entropy(logits, yb, weight=w)
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
        sched.step()

    probe.eval()
    preds, attns = [], []
    with torch.no_grad():
        for i in range(0, te_tok.size(0), bs):
            logits, a = probe(te_tok[i:i + bs].to(device))
            preds.append(logits.argmax(-1).cpu())
            if a is not None:
                attns.append(a.cpu())
    preds = torch.cat(preds).numpy()
    y = te_y.numpy()
    attn = torch.cat(attns) if attns else None
    return {
        "acc": accuracy_score(y, preds),
        "f1_macro": f1_score(y, preds, average="macro"),
        "f1_per": f1_score(y, preds, average=None, labels=[0, 1, 2]),
        "attn": attn,
    }


def remap_nifti_paths(df):
    """
    The visit manifest stores absolute volume paths from the machine that built
    it, which will not resolve anywhere else. The tail of each path (after the
    dataset root) is re-attached to the roots configured here. Harmless when the
    original paths already resolve — it returns them unchanged.
    """
    def _fix(p):
        p0 = str(p)
        if os.path.exists(p0):
            return p0
        q = p0.replace("\\", "/")
        # Dataset root
        i = q.find("/Datasets/")
        if i >= 0:
            cand = os.path.join(C.DATASET_ROOT, q[i + len("/Datasets/"):])
            if os.path.exists(cand):
                return cand
        # The stored tail may or may not include the top volume folder,
        # so both spellings are tried.
        j = q.find("/volumes/")
        if j >= 0:
            rest = q[j + len("/volumes/"):]
            for cand in (os.path.join(C.TBM_ROOT, rest),
                         os.path.join(C.TBM_ROOT, "volumes", rest)):
                if os.path.exists(cand):
                    return cand
        return p0

    df = df.copy()
    df["nifti_path"] = df["nifti_path"].map(_fix)
    ok = int(sum(os.path.exists(str(p)) for p in df["nifti_path"]))
    print(f"[path] reachable images ({MODALITY}): {ok}/{len(df)}")
    if ok == 0:
        raise FileNotFoundError(
            f"No image is reachable ({MODALITY}).\n"
            f"  DATASET_ROOT = {C.DATASET_ROOT}\n"
            f"  VOLUME_ROOT  = {C.TBM_ROOT}\n"
            "YOU MUST SET YOUR OWN PATHS: see VBAI_DATASET_ROOT / "
            "VBAI_VOLUME_ROOT in config.py."
        )
    return df[df["nifti_path"].map(lambda p: os.path.exists(str(p)))].reset_index(drop=True)


@torch.no_grad()
def baseline_from_checkpoint(model, loader, device):
    """The trained mri_classifier's own score — the reference baseline."""
    preds, ys = [], []
    for batch in loader:
        if "mri" not in batch:
            continue
        out = model(mri=batch["mri"].to(device))
        preds.append(out["mri_logits"].argmax(-1).cpu())
        ys.append(batch["label"])
    preds, ys = torch.cat(preds).numpy(), torch.cat(ys).numpy()
    return {"acc": accuracy_score(ys, preds), "f1_macro": f1_score(ys, preds, average="macro")}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="Vbai-2.6AD.pt",
                    help="Vbai-2.6AD checkpoint")
    ap.add_argument("--tbm", action="store_true", help="TBM input")
    ap.add_argument("--t1", action="store_true", help="raw T1 input")
    ap.add_argument("--epochs", type=int, default=30)
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--workers", type=int, default=2)
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[device] {device}")
    print(f"[model dir] {MODEL_DIR}")

    # --- model + normalizer (from the checkpoint; NEVER re-fitted) ---
    sd = torch.load(args.ckpt, map_location=device, weights_only=False)
    mcfg = C.ModelConfig()
    for k, v in sd.get("model_cfg", {}).items():
        if hasattr(mcfg, k):
            setattr(mcfg, k, v)
    model = Vbai26ADModel(mcfg).to(device)
    # Any key mismatch is fatal on purpose: strict=False silently swallows a
    # checkpoint from a different architecture, leaving the model on random
    # weights and producing below-chance results. Older checkpoints from a
    # different architecture must fail loudly here.
    res = model.load_state_dict(sd["model"], strict=False)
    if res.missing_keys or res.unexpected_keys:
        raise RuntimeError(
            f"Checkpoint does NOT match this architecture: "
            f"{len(res.missing_keys)} missing / {len(res.unexpected_keys)} "
            f"unexpected keys.\n"
            f"  first missing    : {res.missing_keys[:3]}\n"
            f"  first unexpected : {res.unexpected_keys[:3]}\n"
            "Checkpoints from a different architecture are not compatible.\n"
            "Expected file: Vbai-2.6AD.pt"
        )
    print(f"[ckpt] {len(sd['model'])} keys matched exactly")
    if "extra" in sd and sd["extra"].get("metrics"):
        print(f"[ckpt] stored metrics: {sd['extra']['metrics']}")
    model.eval()
    for p in model.parameters():
        p.requires_grad_(False)

    norm = TabularNormalizer()
    norm.load_state_dict(sd["norm"])

    # --- subject-level split (no leakage between train and test) ---
    df = load_paired()
    df = remap_nifti_paths(df)
    train_ids, val_ids, test_ids = subject_split(df)
    tr_df = df[df["ptid"].isin(train_ids | val_ids)]     # train+val to fit the probes
    te_df = df[df["ptid"].isin(test_ids)]
    print(f"[data] probe-train {len(tr_df)} scans / test {len(te_df)} scans")

    def make_loader(d, shuffle=False):
        ds = PairedVisitDataset(d, norm, mode="mri", augment=False, mcfg=mcfg)
        return torch.utils.data.DataLoader(ds, batch_size=args.batch_size, shuffle=shuffle,
                                           collate_fn=collate_pad, num_workers=args.workers)

    tr_loader, te_loader = make_loader(tr_df), make_loader(te_df)

    # --- baseline: the checkpoint's own MRI head ---
    print("\n[1/3] Baseline (trained mri_classifier, pooled)...")
    base = baseline_from_checkpoint(model, te_loader, device)
    print(f"      acc {base['acc']:.4f} | macro-F1 {base['f1_macro']:.4f}")

    # --- token cache ---
    print("\n[2/3] Building the token cache (encoder frozen, single pass)...")
    tr_tok, tr_y, shape_info = build_token_cache(model, tr_loader, device)
    te_tok, te_y, _ = build_token_cache(model, te_loader, device)
    ch, grid = shape_info
    n_tok, dim = tr_tok.shape[1], tr_tok.shape[2]
    print(f"      ASPP output: {ch} channels @ {grid}  →  {n_tok} tokens x {dim}-d")
    print(f"      cache: train {tuple(tr_tok.shape)} / test {tuple(te_tok.shape)}")

    # --- problar ---
    print("\n[3/3] Training the probes (encoder frozen)...")
    r_mean = train_probe(MeanPoolProbe(dim), tr_tok, tr_y, te_tok, te_y,
                         device, args.epochs, args.batch_size)
    r_attn = train_probe(AttnPoolProbe(dim), tr_tok, tr_y, te_tok, te_y,
                         device, args.epochs, args.batch_size)

    # --- is the attention degenerate? ---
    a = r_attn["attn"].clamp_min(1e-9)
    ent = float((-(a * a.log()).sum(dim=1)).mean())
    max_ent = float(np.log(n_tok))

    print("\n" + "=" * 66)
    print("  TOKEN PROBE RESULT")
    print("=" * 66)
    print(f"  {'':22s} {'acc':>8s} {'macro-F1':>10s}   F1 (CN/MCI/AD)")
    print(f"  {'checkpoint (pooled)':22s} {base['acc']:8.4f} {base['f1_macro']:10.4f}")
    print(f"  {'probe: mean-pool':22s} {r_mean['acc']:8.4f} {r_mean['f1_macro']:10.4f}"
          f"   {'/'.join(f'{v:.3f}' for v in r_mean['f1_per'])}")
    print(f"  {'probe: attention':22s} {r_attn['acc']:8.4f} {r_attn['f1_macro']:10.4f}"
          f"   {'/'.join(f'{v:.3f}' for v in r_attn['f1_per'])}")
    print(f"\n  attention entropy: {ent:.3f} / {max_ent:.3f} (max = fully uniform)")

    delta = r_attn["f1_macro"] - r_mean["f1_macro"]
    print(f"  attention - mean difference (macro-F1): {delta:+.4f}")
    print("-" * 66)
    if delta > 0.02 and ent < 0.95 * max_ent:
        print("  → GO. The spatial tokens carry extra information; wiring them")
        print("    into the LLM as a sequence is justified.")
    elif ent >= 0.95 * max_ent:
        print("  → STOP. The attention is nearly uniform: the 27 positions do")
        print("    not separate. Lower the ASPP dilations to (1,2,3), or drop the")
        print("    stage-4 stride to reach a 6^3 grid, then measure again.")
    else:
        print("  → WEAK. The tokens add nothing meaningful over the pooled")
        print("    vector. Wiring them in without raising the encoder resolution")
        print("    will not help.")
    print("=" * 66)


if __name__ == "__main__":
    main()