# ============================================================================ # CAPTIONBERT-8192-V2 — CAPABILITY + GEOMETRY EVAL (single standalone cell) # # Self-contained. Pulls the checkpoint from the hub, redefines the encoder # inline (no trainer import), runs the capability gauges the training loop # cannot see, and measures the baselines IN THE SAME HARNESS so the numbers # are comparable rather than cited. # # WHY THIS EXISTS # Training reports student->consensus R@1. That is MIMICRY: how well the # student reproduces its target. It says nothing about whether the space # means anything. Capability is STS/SICK against models that never saw the # consensus. Keep the two on separate lines, always. # # WHAT IT REPORTS # spearman the capability gauge (STS-B, SICK-R, STS12-16 optional) # self_cos isotropy. mean-pooled BERT sits in a narrow cone (~0.57); # a good sentence encoder is near 0 (MiniLM ~0.02) # erank participation ratio = how many directions the embedding # actually uses. THE KEY COLUMN. Measured 2026-07-31: # consensus target 28.7 / 768 # v2 in-domain 80.5 # v2 on STS-B 31.2 <- falls back to the target's rank # all-MiniLM-L6-v2 103.1 <- on the SAME sentences # Averaging teachers cannot create rank they do not share. If # v2's OOD erank stays ~30 while STS stays ~0.55, the ceiling is # the consensus construction, not the student. # # BASELINE (measured, CPU, same harness, 2026-07-31, checkpoint step 3327): # bert-base mean-pooled 109.5M STS-B .4729 self_cos .570 erank 34.3 # captionbert v2 @ 6% 58.3M STS-B .5444 self_cos .139 erank 31.2 # all-MiniLM-L6-v2 22.7M STS-B .8203 self_cos .023 erank 103.1 # (bert-base reproduced v1's published .4729073 to 7 digits -> harness valid) # # L4 (24GB) is plenty; it runs on CPU too, just slower. # ============================================================================ import subprocess, sys, json, os for _p in ("datasets", "transformers", "huggingface_hub", "scipy"): try: __import__(_p) except ImportError: subprocess.run([sys.executable, "-m", "pip", "install", "-q", _p], check=False) import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from scipy.stats import spearmanr, pearsonr from huggingface_hub import hf_hub_download from transformers import AutoTokenizer, AutoModel from datasets import load_dataset DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # ══════════════════════════════════════════════════════════════════ # CONFIG # ══════════════════════════════════════════════════════════════════ REPO = "AbstractPhil/captionbert-8192-v2" CKPT = "checkpoints/best_model.pt" # or "checkpoints/model_sNNNN.pt" TOKENIZER = "google-bert/bert-base-uncased" BASELINES = ["google-bert/bert-base-uncased", "sentence-transformers/all-MiniLM-L6-v2"] RUN_BASELINES = True # False once you have them; they do not change EXTRA_STS = False # STS12-16 as well as STS-B/SICK-R (slower) MAX_LEN = 64 BATCH = 256 GEOM_N = 1500 # sentences for the isotropy / erank probe # architecture — must match config/config.json in the repo ARCH = dict(vocab_size=30522, max_len=8192, d_model=512, n_heads=8, n_layers=12, d_ff=2048, output_dim=768, dropout=0.1, pad_token_id=0, pooling="mean") # ══════════════════════════════════════════════════════════════════ # STUDENT (inline copy — keys must match the checkpoint exactly) # ══════════════════════════════════════════════════════════════════ class CaptionEncoder(nn.Module): def __init__(self, vocab_size=30522, max_len=8192, d_model=512, n_heads=8, n_layers=12, d_ff=2048, output_dim=768, dropout=0.1, pad_token_id=0, pooling="mean"): super().__init__() self.pad_token_id, self.pooling = pad_token_id, pooling self.token_emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_token_id) self.pos_emb = nn.Embedding(max_len, d_model) self.emb_norm = nn.LayerNorm(d_model) self.emb_drop = nn.Dropout(dropout) layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=n_heads, dim_feedforward=d_ff, dropout=dropout, activation="gelu", batch_first=True, norm_first=True) self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers, enable_nested_tensor=False) self.output_proj = nn.Sequential( nn.Linear(d_model, d_model), nn.GELU(), nn.LayerNorm(d_model), nn.Linear(d_model, output_dim)) def forward(self, input_ids, attention_mask=None): L = input_ids.shape[1] pos = torch.arange(L, device=input_ids.device).unsqueeze(0) x = self.emb_drop(self.emb_norm(self.token_emb(input_ids) + self.pos_emb(pos))) kpm = (~attention_mask.bool()) if attention_mask is not None \ else (input_ids == self.pad_token_id) x = self.encoder(x, src_key_padding_mask=kpm) if self.pooling == "cls": pooled = x[:, 0] else: m = (attention_mask.unsqueeze(-1).float() if attention_mask is not None else (~kpm).unsqueeze(-1).float()) pooled = (x * m).sum(1) / m.sum(1).clamp(min=1) return F.normalize(self.output_proj(pooled), dim=-1) # ══════════════════════════════════════════════════════════════════ # GAUGES # ══════════════════════════════════════════════════════════════════ def effective_rank(x: torch.Tensor) -> float: """Participation ratio of the singular spectrum: how many directions are used.""" xc = (x - x.mean(0, keepdim=True)).double() s2 = torch.linalg.svdvals(xc) ** 2 return float((s2.sum() ** 2 / (s2 ** 2).sum()).item()) def geometry(E: torch.Tensor) -> dict: n = min(GEOM_N, E.shape[0]) X = E[:n] S = X @ X.T S.fill_diagonal_(0) return {"self_cos": float(S.sum() / (n * n - n)), "erank": effective_rank(X)} def line(t=""): print("-" * 76 if not t else f"-- {t} " + "-" * max(0, 72 - len(t))) # ══════════════════════════════════════════════════════════════════ # ENCODERS # ══════════════════════════════════════════════════════════════════ def load_student(): line("STUDENT") p = hf_hub_download(REPO, CKPT) sd = torch.load(p, weights_only=True, map_location="cpu") model = CaptionEncoder(**ARCH) model.load_state_dict(sd, strict=True) # strict: a silent mismatch is worse model.eval().to(DEVICE) n = sum(q.numel() for q in model.parameters()) print(f" {REPO}/{CKPT}") print(f" {n:,} params ({n/109_482_240:.2f}x bert-base) | strict load OK | {DEVICE}") tok = AutoTokenizer.from_pretrained(TOKENIZER) @torch.no_grad() def enc(texts): out = [] for i in range(0, len(texts), BATCH): t = tok(list(texts[i:i + BATCH]), max_length=MAX_LEN, padding=True, truncation=True, return_tensors="pt").to(DEVICE) out.append(model(t["input_ids"], t["attention_mask"]).float().cpu()) return torch.cat(out) return enc, n def load_baseline(name): tok = AutoTokenizer.from_pretrained(name) mdl = AutoModel.from_pretrained(name).eval().to(DEVICE) n = sum(q.numel() for q in mdl.parameters()) @torch.no_grad() def enc(texts): out = [] for i in range(0, len(texts), BATCH): t = tok(list(texts[i:i + BATCH]), max_length=MAX_LEN, padding=True, truncation=True, return_tensors="pt").to(DEVICE) h = mdl(**t).last_hidden_state m = t["attention_mask"].unsqueeze(-1).float() pooled = (h * m).sum(1) / m.sum(1).clamp(min=1) # mean pool, as published out.append(F.normalize(pooled, dim=-1).float().cpu()) return torch.cat(out) return enc, n, mdl # ══════════════════════════════════════════════════════════════════ # TASKS # ══════════════════════════════════════════════════════════════════ TASKS = [("STS-B", "mteb/stsbenchmark-sts", "test"), ("SICK-R", "mteb/sickr-sts", "test")] if EXTRA_STS: TASKS += [(f"STS{y}", f"mteb/sts{y}-sts", "test") for y in (12, 13, 14, 15, 16)] def load_task(path, split): ds = load_dataset(path, split=split) cols = ds.column_names a = "sentence1" if "sentence1" in cols else cols[0] b = "sentence2" if "sentence2" in cols else cols[1] s = "score" if "score" in cols else ("similarity_score" if "similarity_score" in cols else None) return list(ds[a]), list(ds[b]), np.asarray(ds[s], dtype=float) def score(enc, a, b, gold): ea, eb = enc(a), enc(b) cos = F.cosine_similarity(ea, eb, dim=-1).numpy() return (float(spearmanr(cos, gold).correlation), float(pearsonr(cos, gold)[0]), torch.cat([ea, eb])) # ══════════════════════════════════════════════════════════════════ # RUN # ══════════════════════════════════════════════════════════════════ def main(): print("=" * 76) print("CAPTIONBERT-8192-V2 - CAPABILITY + GEOMETRY") print("=" * 76) if DEVICE == "cuda": print(f"gpu={torch.cuda.get_device_name()} " f"vram={torch.cuda.get_device_properties(0).total_memory/1e9:.0f}GB") results = {} data = {} for name, path, split in TASKS: try: data[name] = load_task(path, split) print(f" {name}: {len(data[name][2])} pairs") except Exception as e: print(f" {name}: SKIPPED ({type(e).__name__}: {str(e)[:60]})") enc, n_par = load_student() line("STUDENT SCORES") results["captionbert-v2"] = {"params": n_par} for name in data: a, b, g = data[name] sp, pe, E = score(enc, a, b, g) geo = geometry(E) results["captionbert-v2"][name] = {"spearman": sp, "pearson": pe, **geo} print(f" {name:8s} spearman {sp:.4f} pearson {pe:.4f} " f"self_cos {geo['self_cos']:+.4f} erank {geo['erank']:.1f}/768") del enc if DEVICE == "cuda": torch.cuda.empty_cache() if RUN_BASELINES: for bn in BASELINES: line(f"BASELINE {bn}") benc, bn_par, mdl = load_baseline(bn) results[bn] = {"params": bn_par} for name in data: a, b, g = data[name] sp, pe, E = score(benc, a, b, g) geo = geometry(E) results[bn][name] = {"spearman": sp, "pearson": pe, **geo} print(f" {name:8s} spearman {sp:.4f} pearson {pe:.4f} " f"self_cos {geo['self_cos']:+.4f} erank {geo['erank']:.1f}") del mdl, benc if DEVICE == "cuda": torch.cuda.empty_cache() # ---- table ---- print() print("=" * 76) print("SUMMARY") print("=" * 76) tasks = list(data.keys()) hdr = f" {'model':34s}{'params':>10s}" + "".join(f"{t:>10s}" for t in tasks) \ + f"{'self_cos':>10s}{'erank':>8s}" print(hdr) for k, v in results.items(): row = f" {k[-34:]:34s}{v['params']/1e6:>9.1f}M" for t in tasks: row += f"{v[t]['spearman']:>10.4f}" if t in v else f"{'-':>10s}" ref = tasks[0] row += f"{v[ref]['self_cos']:>+10.4f}{v[ref]['erank']:>8.1f}" if ref in v else "" print(row) # ---- the read ---- print() line("READ") cb = results.get("captionbert-v2", {}) ref = tasks[0] if tasks else None if ref and ref in cb: er = cb[ref]["erank"] print(f" erank on {ref} = {er:.1f}. Consensus target measured 28.7/768;") print(f" v2 in-domain (CC12M val) measured 80.5. On out-of-domain text the") print(f" student falls back toward its target's intrinsic rank.") mini = results.get("sentence-transformers/all-MiniLM-L6-v2") if mini and ref in mini: print(f" all-MiniLM uses {mini[ref]['erank']:.1f} directions on the SAME " f"sentences at {mini['params']/1e6:.1f}M params.") print(f" Averaging teachers cannot create rank they do not share -- if this") print(f" gap holds, the ceiling is the CONSENSUS, not the student, and the") print(f" fix is heterogeneous teachers rather than a bigger model.") print(" Training's student->consensus R@1 is MIMICRY. This table is capability.") with open("v2_capability.json", "w") as f: json.dump(results, f, indent=2) print("\n wrote v2_capability.json") return results RESULTS = main()