cross-backbone probe (scripts/cross_backbone_probe.py)

scripts/cross_backbone_probe.py

@@ -0,0 +1,138 @@
+"""Cross-backbone community-vector probe (LOO k-NN recall@1).
+
+For each backbone, extract the mean-pooled hidden state at a target layer for
+every val_200 sample, then compute leave-one-out 1-nearest-neighbor recall@1
+over the 35 coarse community labels. This isolates the question:
+
+    "Is the discourse-community signal that v8a's CDH lights up specific to
+    Qwen, or is it latent in other 7B causal LMs as well?"
+
+We do NOT need any adapter or training to answer this -- v8a's CDH at
+layer 4 simply mean-pools and projects, so the headline number ought to be
+recoverable from raw layer-4 features (paper's 0.484 figure is built on top
+of this same signal, just refined through 64-D community vector contrastive
+learning).
+
+Usage:
+    python scripts/cross_backbone_probe.py \
+        --val-data data/val_200.jsonl \
+        --backbones Qwen/Qwen2.5-7B mistralai/Mistral-7B-v0.3 \
+        --layers 4 8 12 \
+        --out artifacts/cross_backbone.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def mean_pool(h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    # h: (1, T, D); mask: (1, T)
+    m = mask.unsqueeze(-1).to(h.dtype)
+    return (h * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)
+
+
+def loo_knn_recall_at_1(feats: np.ndarray, labels: np.ndarray) -> float:
+    # cosine sim, leave-one-out
+    f = feats / (np.linalg.norm(feats, axis=1, keepdims=True) + 1e-9)
+    sim = f @ f.T
+    np.fill_diagonal(sim, -np.inf)
+    pred = labels[sim.argmax(axis=1)]
+    return float((pred == labels).mean())
+
+
+@torch.no_grad()
+def extract(backbone_id: str, samples: list[dict], layers: list[int],
+            max_len: int, device: str) -> dict[int, np.ndarray]:
+    print(f"\n=== {backbone_id} ===")
+    t0 = time.time()
+    tok = AutoTokenizer.from_pretrained(backbone_id)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        backbone_id, torch_dtype=torch.bfloat16, device_map=device,
+        attn_implementation="eager",
+    )
+    model.eval()
+    n_layers = model.config.num_hidden_layers
+    layers_eff = [l for l in layers if l < n_layers]
+    print(f"  loaded in {time.time()-t0:.1f}s, n_layers={n_layers}, "
+          f"using layers {layers_eff}")
+    out = {l: [] for l in layers_eff}
+    for i, ex in enumerate(samples):
+        enc = tok(ex["text"], return_tensors="pt", truncation=True,
+                  max_length=max_len).to(device)
+        h = model(**enc, output_hidden_states=True, use_cache=False).hidden_states
+        # h is tuple of length n_layers+1 (embedding + per-layer)
+        for l in layers_eff:
+            v = mean_pool(h[l + 1], enc.attention_mask)[0].float().cpu().numpy()
+            out[l].append(v)
+        if (i + 1) % 50 == 0:
+            print(f"  [{i+1}/{len(samples)}] elapsed={time.time()-t0:.1f}s")
+    out_arr = {l: np.stack(out[l]) for l in layers_eff}
+    del model, tok
+    gc.collect()
+    torch.cuda.empty_cache()
+    return out_arr
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--val-data", default="data/val_200.jsonl")
+    ap.add_argument("--backbones", nargs="+",
+                    default=["Qwen/Qwen2.5-7B",
+                             "mistralai/Mistral-7B-v0.3",
+                             "meta-llama/Meta-Llama-3-8B"])
+    ap.add_argument("--layers", type=int, nargs="+",
+                    default=[4, 8, 12, 16])
+    ap.add_argument("--max-len", type=int, default=512)
+    ap.add_argument("--out", default="artifacts/cross_backbone.json")
+    ap.add_argument("--device", default="cuda")
+    args = ap.parse_args()
+
+    samples = [json.loads(l) for l in Path(args.val_data).read_text().splitlines()]
+    labels = np.array([s["community_id"] for s in samples])
+    n_classes = len(set(labels.tolist()))
+    chance = 1.0 / n_classes
+    print(f"val_200: n={len(samples)}  classes={n_classes}  chance={chance:.3f}")
+
+    results = {
+        "n_samples": len(samples),
+        "n_classes": n_classes,
+        "chance_recall_at_1": chance,
+        "backbones": {},
+    }
+    for backbone in args.backbones:
+        try:
+            feats = extract(backbone, samples, args.layers, args.max_len, args.device)
+        except Exception as e:  # noqa: BLE001
+            print(f"  SKIP {backbone}: {e}")
+            results["backbones"][backbone] = {"error": str(e)}
+            continue
+        per_layer = {}
+        for layer, X in feats.items():
+            r1 = loo_knn_recall_at_1(X, labels)
+            per_layer[str(layer)] = {
+                "recall_at_1": r1,
+                "x_chance": r1 / chance,
+                "feature_dim": int(X.shape[1]),
+            }
+            print(f"  layer {layer:>2}  D={X.shape[1]:>5}  "
+                  f"recall@1={r1:.3f}  ({r1/chance:.1f}x chance)")
+        results["backbones"][backbone] = {"per_layer": per_layer}
+        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.out).write_text(json.dumps(results, indent=2))
+
+    print(f"\nwrote {args.out}")
+
+
+if __name__ == "__main__":
+    main()