v0.1: fix 2_Dense use_residual

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+logo_reason_mxai.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+training/
+.DS_Store
+.idea/
+.vscode/
+__pycache__/
+*.pyc

1_Dense/config.json

@@ -0,0 +1 @@
+{"in_features": 384, "out_features": 768, "bias": false, "activation_function": "torch.nn.modules.linear.Identity", "use_residual": false}

1_Dense/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd2aa41a2974566861ff9690be4f4c2b9a917ee0f97e338911bd66fa06dafddd
+size 589912

2_Dense/config.json

@@ -0,0 +1 @@
+{"in_features": 768, "out_features": 768, "bias": false, "activation_function": "torch.nn.modules.linear.Identity", "use_residual": true}

2_Dense/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec4fecf007916948b029915b76dbf37efd5241a21a497dfb4f2abe622410b55a
+size 1179736

3_Dense/config.json

@@ -0,0 +1 @@
+{"in_features": 768, "out_features": 128, "bias": false, "activation_function": "torch.nn.modules.linear.Identity", "use_residual": false}

3_Dense/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c99e270c45b0db9b5ffbf5e0de10ba729680aa778807b2329189b682d686203d
+size 196696

README.md

@@ -0,0 +1,133 @@
+---
+language:
+- en
+tags:
+- ColBERT
+- PyLate
+- sentence-transformers
+- sentence-similarity
+- feature-extraction
+- late-interaction
+- reasoning-retrieval
+- edge
+- generated_from_trainer
+- loss:CachedContrastive
+base_model: mixedbread-ai/mxbai-edge-colbert-v0-32m
+datasets:
+- hanhainebula/bge-reasoner-data
+- reasonir/reasonir-data
+pipeline_tag: sentence-similarity
+library_name: PyLate
+license: cc-by-nc-4.0
+---
+
+<img src="./logo_reason_mxai.png" width="500" height="auto">
+
+# Reason-mxbai-colbert-v0.1-32m
+
+**v0.1** of the Reason-mxbai-colbert series — same edge-scale late-interaction retriever as v0, **retrained with the correct projection-head architecture** (`use_residual: true` in the 2_Dense layer). Mean BRIGHT nDCG@10 improves from **19.00 → 19.61** (+0.61), with the largest gains on the natural-language and formal-math splits that v0 struggled with.
+
+## What changed vs v0
+
+The v0 release inherited `use_residual: false` in the 2_Dense (768→768) layer from the upstream `mixedbread-ai/mxbai-edge-colbert-v0-32m` config. That turned out to be a shipped-config bug — the base weights were trained with a residual connection on 2_Dense, but the config flag said otherwise. PyLate respects the flag, so every downstream fine-tune was working with a silently-broken architecture.
+
+v0.1 fixes this by:
+1. Copying the base model and flipping `2_Dense/config.json: "use_residual": true`.
+2. Re-widening the 3_Dense projection from 64 → 128 dims on top of the corrected base (small-random init for the new channels).
+3. Retraining the full curriculum (Stage 1 VL warmup → Stage 2 BGE-reasoner + HQ-hn polish) on the corrected architecture.
+
+Under the corrected architecture, stage-2 training loss drops from **38.73 → 31.71** (−7.02) — real learning. The v0 equivalent run under the broken architecture was stuck at loss ~45-51 and never moved the new-dim weights.
+
+## BRIGHT results
+
+Evaluated with MTEB `BrightRetrieval` (brute-force MaxSim, `query_length=256`, `document_length=2048`).
+
+| Split | v0 | **v0.1** | Δ v0 | Reason-MCB (150M, paper) |
+|---|---:|---:|---:|---:|
+| aops | 5.05 | 4.89 | −0.16 | 9.17 |
+| biology | 32.71 | **33.16** | +0.45 | 33.25 |
+| earth_science | 43.88 | **44.28** | +0.40 | 41.02 |
+| economics | 18.70 | **20.25** | **+1.54** | 24.93 |
+| leetcode | 17.67 | 17.40 | −0.28 | 31.07 |
+| pony | 20.73 | **22.77** | **+2.03** | 8.51 |
+| psychology | 22.62 | **24.91** | **+2.29** | 30.73 |
+| robotics | 18.43 | **18.65** | +0.23 | 21.12 |
+| stackoverflow | 16.78 | 16.66 | −0.12 | 20.62 |
+| sustainable_living | 20.77 | 20.11 | −0.66 | 20.31 |
+| theoremqa_questions | 8.38 | **9.04** | **+0.66** | 19.51 |
+| theoremqa_theorems | 2.25 | **3.19** | **+0.94** | 11.24 |
+| **Full mean** | **19.00** | **19.61** | **+0.61** | **22.62** |
+
+**8 wins / 4 losses.** Highlights:
+- **Natural-language splits** (psychology +2.29, economics +1.54, pony +2.03) — the correct architecture lets the model learn meaningfully richer representations.
+- **Formal-math splits** (theoremqa_theorems +0.94, theoremqa_questions +0.66) — previously stuck near the floor, now moving.
+- **Small code-split regressions** (leetcode −0.28, stackoverflow −0.12) are within noise.
+
+## Model Details
+
+- **Model Type:** PyLate ColBERT (late-interaction, multi-vector)
+- **Base model:** [mixedbread-ai/mxbai-edge-colbert-v0-32m](https://huggingface.co/mixedbread-ai/mxbai-edge-colbert-v0-32m) (with `2_Dense/use_residual: true` patched to match trained weights)
+- **Parameters:** ~32M backbone + widened 128-dim projection
+- **Document Length (training):** 2048 tokens
+- **Query Length (training):** 256 tokens
+- **Output Dimensionality:** 128 per token (widened from base's 64)
+- **Similarity Function:** MaxSim
+- **Training Data:** `hanhainebula/bge-reasoner-data` + `reasonir/reasonir-data` (VL split for warmup, HQ with hard negatives for polish)
+- **Language:** en
+- **License:** CC-BY-NC-4.0
+
+## Model Architecture
+
+```
+ColBERT(
+  (0): Transformer({'max_seq_length': 127, 'do_lower_case': True}) with ModernBertModel
+      hidden_size=384, num_hidden_layers=10, num_attention_heads=6,
+      position_embedding_type='sans_pos', max_position_embeddings=7999
+  (1): Dense(384 → 768, bias=False, use_residual=False)
+  (2): Dense(768 → 768, bias=False, use_residual=True)     # ← THE FIX (was False in v0)
+  (3): Dense(768 → 128, bias=False, use_residual=False)    # widened from 64 → 128 on correct base
+)
+```
+
+## Usage
+
+```python
+from pylate import indexes, models, retrieve
+
+model = models.ColBERT(model_name_or_path="<hf_user>/Reason-mxbai-colbert-v0.1-32m")
+
+# Retrieval (see https://lightonai.github.io/pylate/ for full API)
+index = indexes.Voyager(index_folder="pylate-index", index_name="index", override=True)
+docs = ["document 1 text", "document 2 text"]
+doc_embs = model.encode(docs, is_query=False, batch_size=32, show_progress_bar=True)
+index.add_documents(documents_ids=["1","2"], documents_embeddings=doc_embs)
+
+retriever = retrieve.ColBERT(index=index)
+q_embs = model.encode(
+    ["Given a Psychology post, retrieve relevant passages that help answer the post.\nQuery: why do I procrastinate?"],
+    is_query=True,
+)
+scores = retriever.retrieve(queries_embeddings=q_embs, k=10)
+```
+
+## Training
+
+See `training/` (not uploaded to the HF repo — see `training/README.md` for the full reproduction recipe).
+
+Pipeline (2 nodes × 4 H100):
+
+1. **Widen** `mixedbread-ai/mxbai-edge-colbert-v0-32m` (patched `2_Dense/use_residual: true`) from 64 → 128-dim projection, small-random init.
+2. **Stage 1 (VL warmup):** 1 epoch on `reasonir-vl`, lr=1e-5, global batch 2048, ~30 min.
+3. **Stage 2 (BGE + HQ-hn polish):** 1 epoch on merged `bge-reasoner + reasonir-hq-hn` (~2.7M triples), lr=5e-6, ~2h.
+
+Key training-infra fixes already baked into this release:
+- `max_grad_norm=10000` (default `1.0` over-clips cached-contrastive's summed gradients; original v0 training silently froze).
+- `use_residual: true` in 2_Dense (matches base weights).
+
+## Acknowledgements
+
+Thanks to Antoine Chaffin (LightOn / Reason-ModernColBERT author) for flagging the upstream 2_Dense config bug in `mxbai-edge-colbert-v0-32m` — without that heads-up, v0.1 wouldn't exist.
+
+## License
+
+CC-BY-NC-4.0 (inherited from ReasonIR / BGE-reasoner data).

config.json

@@ -0,0 +1,50 @@
+{
+  "_name_or_path": "/data/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output_edge32m_d128_fixed/stage-vl/final",
+  "architectures": [
+    "ModernBertModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 50281,
+  "causal_mask": false,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "mean",
+  "cls_token_id": 50281,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 50282,
+  "global_attn_every_n_layers": 3,
+  "global_rope_theta": 160000.0,
+  "gradient_checkpointing": false,
+  "hidden_activation": "gelu",
+  "hidden_size": 384,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 576,
+  "is_causal": false,
+  "layer_norm_eps": 1e-05,
+  "local_attention": 128,
+  "local_rope_theta": 160000.0,
+  "max_position_embeddings": 7999,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 6,
+  "num_hidden_layers": 10,
+  "pad_token_id": 50283,
+  "position_embedding_type": "sans_pos",
+  "reference_compile": false,
+  "repad_logits_with_grad": false,
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.2",
+  "vocab_size": 50370
+}

config_sentence_transformers.json

@@ -0,0 +1,21 @@
+{
+  "model_type": "ColBERT",
+  "__version__": {
+    "sentence_transformers": "4.0.2",
+    "transformers": "4.48.2",
+    "pytorch": "2.5.1"
+  },
+  "prompts": {
+    "query": "",
+    "document": ""
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "MaxSim",
+  "query_prefix": "[Q] ",
+  "document_prefix": "[D] ",
+  "query_length": 256,
+  "document_length": 2048,
+  "attend_to_expansion_tokens": false,
+  "skiplist_words": [],
+  "do_query_expansion": false
+}

evaluation/evaluate_bright.py

@@ -0,0 +1,470 @@
+"""Official BRIGHT evaluation for LaRLI / ColBERT models.
+
+Uses the same protocol as the Reason-ModernColBERT authors:
+- MTEB BrightRetrieval task for data loading
+- Brute-force MaxSim scoring (exact, no approximation)
+- Filters excluded_ids from scores
+- query_length=256 (except pony=32)
+
+Output layout
+-------------
+By default each run writes into a unique directory:
+    BRIGHT_scores_<label>[__<run_name>]__<YYYYMMDD-HHMMSS>/
+where <label> is derived from the model path (HF ids become "org__name";
+local paths use the last meaningful components so ".../final" directories
+don't collide). This avoids the old behaviour of silently reloading cached
+per-split scores from a stale dir.
+
+Key flags:
+- --run_name LABEL   : extra label appended to the auto dir name
+- --no_timestamp     : drop the timestamp suffix (old stable-path behaviour)
+- --overwrite        : recompute a split even if its JSON already exists
+- --output_dir PATH  : override the auto dir entirely
+- --output_root DIR  : parent dir for the auto dir (default: CWD)
+
+Usage:
+    # Default: auto dir with timestamp (will never overwrite a prior run)
+    python LaRLI/evaluate_bright.py \
+        --model_path $PWD/repro_reason_moderncolbert/output_edge17m_v2a/pos4_D/stage-reason-embed/final \
+        --model_version baseline \
+        --run_name v2a
+
+    # Pin a stable directory (re-runs reuse per-split cache unless --overwrite)
+    python LaRLI/evaluate_bright.py --model_path lightonai/Reason-ModernColBERT \
+        --model_version baseline --run_name ref --no_timestamp
+
+    # PLAID index mode
+    python LaRLI/evaluate_bright.py --model_path lightonai/Reason-ModernColBERT \
+        --model_version baseline --use_plaid
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime as _dt
+import gc
+import os
+import re
+import sys
+
+import numpy as np
+import torch
+
+_project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, _project_root)
+sys.path.insert(0, os.path.join(_project_root, "pylate"))
+
+import mteb
+import srsly
+
+from pylate import evaluation, models
+
+
+_UNINFORMATIVE_TAIL = {"final", "checkpoint-last", ""}
+_SANITIZE_RE = re.compile(r"[^A-Za-z0-9._-]+")
+
+
+def derive_model_label(model_path: str) -> str:
+    """Derive a human-readable label from a model path.
+
+    Goals:
+      - Huggingface refs like "org/model" -> "org__model"
+      - Local paths like ".../output_edge17m_v2a/pos4_D/stage-reason-embed/final"
+        -> "pos4_D__stage-reason-embed__final" so "final" never stands alone.
+      - Always filesystem-safe (no slashes, no weird chars).
+    """
+    p = str(model_path).rstrip("/")
+    if not p:
+        return "unknown"
+
+    # HF-style "org/name" (not an absolute local path)
+    if not p.startswith(("/", ".")) and p.count("/") == 1:
+        label = p.replace("/", "__")
+        return _SANITIZE_RE.sub("-", label)
+
+    parts = [seg for seg in p.split("/") if seg]
+    # Walk from the tail; collect segments until we've seen at least two
+    # *informative* ones (i.e. not "final" / "checkpoint-*"). This keeps the
+    # run-identifying parent in the label so
+    #   .../output_edge17m_v2a/pos4_D/stage-reason-embed/final
+    # becomes 'pos4_D__stage-reason-embed__final', not just '…__final'.
+    collected: list[str] = []
+    informative_seen = 0
+    for seg in reversed(parts):
+        collected.append(seg)
+        is_uninformative = seg in _UNINFORMATIVE_TAIL or seg.startswith("checkpoint-")
+        if not is_uninformative:
+            informative_seen += 1
+            if informative_seen >= 2:
+                break
+    label = "__".join(reversed(collected)) if collected else parts[-1]
+    return _SANITIZE_RE.sub("-", label)
+
+
+def build_default_output_dir(model_path: str, run_name: str | None,
+                             with_timestamp: bool, root: str) -> str:
+    label = derive_model_label(model_path)
+    if run_name:
+        label = f"{label}__{_SANITIZE_RE.sub('-', run_name)}"
+    if with_timestamp:
+        stamp = _dt.datetime.now().strftime("%Y%m%d-%H%M%S")
+        label = f"{label}__{stamp}"
+    return os.path.join(root, f"BRIGHT_scores_{label}")
+
+
+def load_model(model_path: str, model_version: str, query_length: int, document_length: int | None = None):
+    """Load a ColBERT or LaRLI model."""
+    if model_version == "baseline":
+        kwargs = {"model_name_or_path": model_path, "query_length": query_length}
+        if document_length is not None:
+            kwargs["document_length"] = document_length
+        return models.ColBERT(**kwargs)
+    elif model_version == "v7":
+        from LaRLI.models import LaRLIColBERTv7
+        return LaRLIColBERTv7(
+            model_name_or_path=model_path,
+            num_heads=4,
+            query_length=query_length,
+        )
+    elif model_version == "v6":
+        from LaRLI.models import LaRLIColBERTv6
+        return LaRLIColBERTv6(
+            model_name_or_path=model_path,
+            num_reasoning_tokens=8,
+            query_length=query_length,
+        )
+    elif model_version == "v5":
+        from LaRLI.models import LaRLIColBERTv5
+        return LaRLIColBERTv5(
+            model_name_or_path=model_path,
+            num_reasoning_tokens=32,
+            query_length=query_length,
+        )
+    elif model_version == "v4":
+        from LaRLI.models import LaRLIColBERTv4
+        return LaRLIColBERTv4(
+            model_name_or_path=model_path,
+            num_reasoning_tokens=32,
+            query_length=query_length,
+        )
+    elif model_version == "v3":
+        from LaRLI.models import LaRLIColBERTv3
+        return LaRLIColBERTv3(
+            model_name_or_path=model_path,
+            num_reasoning_tokens=32,
+            query_length=query_length,
+        )
+    else:
+        raise ValueError(f"Unknown model version: {model_version}")
+
+
+def bruteforce_retrieve(
+    query_embeddings: list,
+    doc_embeddings: list,
+    doc_ids: list[str],
+    k: int = 100,
+    doc_chunk_size: int = 512,
+) -> list[list[dict]]:
+    """Brute-force MaxSim retrieval. One query at a time, docs in chunks."""
+    all_scores = []
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_docs = len(doc_embeddings)
+    dim = doc_embeddings[0].shape[1]
+
+    # Pre-pad documents and build tensor on CPU
+    max_doc_len = max(d.shape[0] for d in doc_embeddings)
+    doc_tensor = torch.zeros(n_docs, max_doc_len, dim, dtype=torch.float16)
+    doc_mask = torch.zeros(n_docs, max_doc_len, dtype=torch.bool)
+    for i, d in enumerate(doc_embeddings):
+        doc_len = d.shape[0]
+        doc_tensor[i, :doc_len] = torch.as_tensor(d, dtype=torch.float16)
+        doc_mask[i, :doc_len] = True
+
+    for qi in range(len(query_embeddings)):
+        if qi % 10 == 0:
+            print(f"  Scoring query {qi+1}/{len(query_embeddings)}...", end="\r")
+
+        q = torch.as_tensor(query_embeddings[qi], dtype=torch.float16).to(device)  # [Q, D]
+        query_scores = torch.zeros(n_docs, dtype=torch.float32)
+
+        for d_start in range(0, n_docs, doc_chunk_size):
+            d_end = min(d_start + doc_chunk_size, n_docs)
+            d_gpu = doc_tensor[d_start:d_end].to(device)  # [C, L, D]
+            m_gpu = doc_mask[d_start:d_end].to(device)     # [C, L]
+
+            # [Q, D] x [C, L, D]^T -> [Q, C, L] via matmul
+            # q: [Q, D], d_gpu: [C, L, D] -> d_gpu.permute(0,2,1): [C, D, L]
+            # sim[c] = q @ d_gpu[c].T -> [Q, L]  for each doc c
+            sim = torch.einsum("qd,cld->qcl", q, d_gpu)  # [Q, C, L]
+            sim.masked_fill_(~m_gpu.unsqueeze(0), -1e4)   # mask padding
+            max_sim = sim.max(dim=-1).values               # [Q, C]
+            chunk_scores = max_sim.sum(dim=0)              # [C]
+            query_scores[d_start:d_end] = chunk_scores.float().cpu()
+
+            del d_gpu, m_gpu, sim, max_sim, chunk_scores
+
+        # Top-k
+        topk_vals, topk_idxs = query_scores.topk(min(k, n_docs))
+        query_results = [
+            {"id": doc_ids[idx], "score": val}
+            for val, idx in zip(topk_vals.tolist(), topk_idxs.tolist())
+        ]
+        all_scores.append(query_results)
+
+    print()
+    return all_scores
+
+
+def main():
+    parser = argparse.ArgumentParser(description="BRIGHT evaluation (official protocol)")
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--model_version", type=str, default="baseline",
+                        choices=["baseline", "v3", "v4", "v5", "v6", "v7"])
+    parser.add_argument("--query_length", type=int, default=256,
+                        help="Default query length (256 for all splits except pony=32)")
+    parser.add_argument("--pony_query_length", type=int, default=32,
+                        help="Query length for the pony split")
+    parser.add_argument("--document_length", type=int, default=None,
+                        help="Cap document token length during encoding. Default: use model config "
+                             "(usually 8192). Set lower (e.g. 2048) to avoid CPU-RAM OOM on "
+                             "large-corpus splits when the model has a wide projection dim.")
+    parser.add_argument("--doc_batch_size", type=int, default=100)
+    parser.add_argument("--query_batch_size", type=int, default=32)
+    parser.add_argument("--output_dir", type=str, default=None,
+                        help="Output directory for scores. If not set, a dir "
+                             "'BRIGHT_scores_<label>[__<run_name>]__<timestamp>' "
+                             "is auto-generated under --output_root so each run "
+                             "is isolated. Pass an explicit path to override.")
+    parser.add_argument("--output_root", type=str, default=".",
+                        help="Parent dir for auto-generated --output_dir. Default: CWD.")
+    parser.add_argument("--run_name", type=str, default=None,
+                        help="Optional extra label appended to the auto dir name "
+                             "(e.g. 'v2a' or 'polish-only').")
+    parser.add_argument("--no_timestamp", action="store_true",
+                        help="Disable the timestamp suffix in the auto dir name. "
+                             "With --run_name and --no_timestamp you'll get a "
+                             "stable path across runs (old behaviour).")
+    parser.add_argument("--overwrite", action="store_true",
+                        help="Re-compute per-split results even if the JSON "
+                             "already exists in --output_dir (default is to "
+                             "load-and-skip).")
+    parser.add_argument("--splits", nargs="+", default=None,
+                        help="Specific splits to evaluate (default: all)")
+    parser.add_argument("--use_plaid", action="store_true",
+                        help="Use PLAID index instead of brute-force (needs more GPU memory)")
+    parser.add_argument("--nbits", type=int, default=4,
+                        help="PLAID index nbits (only with --use_plaid)")
+    parser.add_argument("--retrieve_k", type=int, default=100,
+                        help="Number of documents to retrieve per query")
+    args = parser.parse_args()
+
+    # Load BRIGHT task via MTEB
+    print("Loading BrightRetrieval task...")
+    tasks = mteb.get_tasks(tasks=["BrightRetrieval"])
+    tasks[0].load_data()
+
+    model_name = args.model_path.rstrip("/")
+    short_name = derive_model_label(model_name)
+
+    if args.output_dir is None:
+        args.output_dir = build_default_output_dir(
+            model_path=args.model_path,
+            run_name=args.run_name,
+            with_timestamp=not args.no_timestamp,
+            root=args.output_root,
+        )
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Write a small run-manifest so you can later remember what a dir contains.
+    srsly.write_json(
+        os.path.join(args.output_dir, "run_meta.json"),
+        {
+            "model_path": args.model_path,
+            "model_version": args.model_version,
+            "query_length": args.query_length,
+            "pony_query_length": args.pony_query_length,
+            "use_plaid": args.use_plaid,
+            "nbits": args.nbits if args.use_plaid else None,
+            "retrieve_k": args.retrieve_k,
+            "overwrite": args.overwrite,
+            "run_name": args.run_name,
+            "label": short_name,
+            "started_at": _dt.datetime.now().isoformat(timespec="seconds"),
+            "splits_requested": args.splits,
+        },
+    )
+
+    all_splits = list(tasks[0].queries.keys())
+    eval_splits = args.splits if args.splits else all_splits
+    print(f"Evaluating {len(eval_splits)} splits: {eval_splits}")
+    print(f"Mode: {'PLAID index' if args.use_plaid else 'brute-force MaxSim'}")
+    print(f"Output dir: {args.output_dir}")
+    if args.overwrite:
+        print("Overwrite: ON — existing per-split JSONs will be recomputed.")
+
+    all_results = {}
+
+    for eval_set in eval_splits:
+        # Use query_length=32 for pony, 256 for everything else
+        if "pony" in eval_set:
+            query_length = args.pony_query_length
+        else:
+            query_length = args.query_length
+
+        print(f"\n{'='*60}")
+        print(f"Split: {eval_set}  (query_length={query_length})")
+        print(f"{'='*60}")
+
+        # Check if results already exist
+        out_file = os.path.join(
+            args.output_dir,
+            f"BrightRetrieval_{eval_set.replace('/', '_')}_evaluation_scores_qlen{query_length}.json",
+        )
+        if os.path.exists(out_file) and not args.overwrite:
+            print(f"Results already exist at {out_file}. Loading "
+                  f"(pass --overwrite to recompute).")
+            existing = srsly.read_json(out_file)
+            all_results[eval_set] = existing
+            ndcg10 = existing.get("ndcg@10", "N/A")
+            print(f"  nDCG@10 = {ndcg10}")
+            continue
+        if os.path.exists(out_file) and args.overwrite:
+            print(f"--overwrite set; recomputing {eval_set} (old file at {out_file}).")
+
+        # Load model with appropriate query_length (and optional document_length cap)
+        model = load_model(args.model_path, args.model_version, query_length,
+                           document_length=args.document_length)
+
+        corpus = tasks[0].corpus[eval_set]["standard"]
+        queries = tasks[0].queries[eval_set]["standard"]
+        doc_ids = list(corpus.keys())
+        doc_texts = list(corpus.values())
+        query_ids = list(queries.keys())
+        query_texts = list(queries.values())
+
+        # Encode documents
+        print(f"Encoding {len(doc_texts)} documents...")
+        documents_embeddings = model.encode(
+            sentences=doc_texts,
+            batch_size=args.doc_batch_size,
+            is_query=False,
+            show_progress_bar=True,
+        )
+
+        # Encode queries
+        print(f"Encoding {len(query_texts)} queries...")
+        queries_embeddings = model.encode(
+            sentences=query_texts,
+            is_query=True,
+            show_progress_bar=True,
+            batch_size=args.query_batch_size,
+        )
+
+        if args.use_plaid:
+            from pylate import indexes, retrieve
+
+            index = indexes.PLAID(
+                override=True,
+                nbits=args.nbits,
+                index_name=f"BrightRetrieval_{eval_set}_{short_name}_{query_length}_{args.nbits}bits",
+            )
+            retriever = retrieve.ColBERT(index=index)
+
+            index.add_documents(
+                documents_ids=doc_ids,
+                documents_embeddings=documents_embeddings,
+            )
+            scores = retriever.retrieve(
+                queries_embeddings=queries_embeddings, k=args.retrieve_k
+            )
+        else:
+            # Brute-force MaxSim (exact, memory-efficient)
+            print("Computing brute-force MaxSim scores...")
+            scores = bruteforce_retrieve(
+                query_embeddings=queries_embeddings,
+                doc_embeddings=documents_embeddings,
+                doc_ids=doc_ids,
+                k=args.retrieve_k,
+            )
+
+        # Filter excluded_ids (official protocol)
+        rel_docs = tasks[0].relevant_docs[eval_set]
+        has_excluded = "excluded" in rel_docs
+        qrels = rel_docs["standard"] if "standard" in rel_docs else rel_docs
+
+        if has_excluded:
+            excluded_map = rel_docs["excluded"]
+            filtered_scores = []
+            for query_scores, excluded_ids in zip(
+                scores, excluded_map.values()
+            ):
+                if excluded_ids == "N/A" or not excluded_ids:
+                    filtered_scores.append(query_scores)
+                    continue
+                filtered_query_scores = [
+                    qs for qs in query_scores if qs["id"] not in excluded_ids
+                ]
+                filtered_scores.append(filtered_query_scores)
+        else:
+            filtered_scores = scores
+            print("  (no excluded_ids found, skipping filtering)")
+
+        # Evaluate
+        evaluation_scores = evaluation.evaluate(
+            scores=filtered_scores,
+            qrels=qrels,
+            queries=query_ids,
+            metrics=["map", "ndcg@1", "ndcg@10", "ndcg@100", "recall@10", "recall@100"],
+        )
+
+        print(evaluation_scores)
+        all_results[eval_set] = evaluation_scores
+
+        # Save per-split results
+        os.makedirs(args.output_dir, exist_ok=True)
+        srsly.write_json(out_file, evaluation_scores)
+
+        # Free memory
+        del model, documents_embeddings, queries_embeddings
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    # Print summary
+    print(f"\n{'='*60}")
+    print("SUMMARY")
+    print(f"{'='*60}")
+    ndcg10_values = []
+    for split, scores_dict in sorted(all_results.items()):
+        ndcg10 = scores_dict.get("ndcg@10", 0)
+        ndcg10_values.append(ndcg10)
+        print(f"  {split:30s}  nDCG@10 = {ndcg10:.4f}")
+    if ndcg10_values:
+        print(f"  {'AVERAGE':30s}  nDCG@10 = {sum(ndcg10_values)/len(ndcg10_values):.4f}")
+
+    # Save summary
+    os.makedirs(args.output_dir, exist_ok=True)
+    srsly.write_json(os.path.join(args.output_dir, "summary.json"), all_results)
+    print(f"\nResults saved to {args.output_dir}/")
+
+
+if __name__ == "__main__":
+    main()
+
+
+ # Your distilled model
+"""
+ python evaluation/evaluate_bright.py \
+     --model_path output/r1colbert_distilled/R1-ColBERT-distill-alpha0.3/checkpoint-1500 \
+     --model_version baseline
+# Baseline comparison
+ python evaluation/evaluate_bright.py \
+     --model_path lightonai/Reason-ModernColBERT \
+     --model_version baseline
+
+
+python LaRLI/evaluate_bright.py \
+    --model_path /scratch/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output/baseline-ColBERT-ReasonIR/final \
+    --model_version baseline \
+    --splits biology
+ """

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d86cdb2c133be712ea760ea2c053d0eed2dc712d99981ecef4d83eb1827d95b
+size 63773752

modules.json

@@ -0,0 +1,26 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Dense",
+    "type": "pylate.models.Dense.Dense"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Dense",
+    "type": "pylate.models.Dense.Dense"
+  },
+  {
+    "idx": 3,
+    "name": "3",
+    "path": "3_Dense",
+    "type": "pylate.models.Dense.Dense"
+  }
+]

results/BrightRetrieval_aops_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.0328788796,
+  "ndcg@1":0.009009009,
+  "ndcg@10":0.0488934387,
+  "ndcg@100":0.0961218916,
+  "recall@10":0.0662269412,
+  "recall@100":0.2151973402
+}

results/BrightRetrieval_biology_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.2627349551,
+  "ndcg@1":0.3398058252,
+  "ndcg@10":0.3316177299,
+  "ndcg@100":0.4150093687,
+  "recall@10":0.3784616664,
+  "recall@100":0.6726602104
+}

results/BrightRetrieval_earth_science_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.3771350277,
+  "ndcg@1":0.4568965517,
+  "ndcg@10":0.4427963898,
+  "ndcg@100":0.5147986924,
+  "recall@10":0.4756853467,
+  "recall@100":0.7088055504
+}

results/BrightRetrieval_economics_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.1544231851,
+  "ndcg@1":0.1941747573,
+  "ndcg@10":0.2024816779,
+  "ndcg@100":0.2653182706,
+  "recall@10":0.1939382382,
+  "recall@100":0.48742489
+}

results/BrightRetrieval_leetcode_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.1243283683,
+  "ndcg@1":0.0,
+  "ndcg@10":0.1739560791,
+  "ndcg@100":0.2252391566,
+  "recall@10":0.2919014085,
+  "recall@100":0.510915493
+}

results/BrightRetrieval_pony_evaluation_scores_qlen32.json

@@ -0,0 +1,8 @@
+{
+  "map":0.1189579454,
+  "ndcg@1":0.2767857143,
+  "ndcg@10":0.2276578569,
+  "ndcg@100":0.3506526988,
+  "recall@10":0.1008905099,
+  "recall@100":0.4784767044
+}

results/BrightRetrieval_psychology_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.2074368567,
+  "ndcg@1":0.198019802,
+  "ndcg@10":0.2490934222,
+  "ndcg@100":0.3159161008,
+  "recall@10":0.2984747676,
+  "recall@100":0.5615963777
+}

results/BrightRetrieval_robotics_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.1487403514,
+  "ndcg@1":0.1683168317,
+  "ndcg@10":0.1865310541,
+  "ndcg@100":0.2415133828,
+  "recall@10":0.216579428,
+  "recall@100":0.4233349665
+}

results/BrightRetrieval_stackoverflow_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.1215601626,
+  "ndcg@1":0.1623931624,
+  "ndcg@10":0.1666014935,
+  "ndcg@100":0.2295905877,
+  "recall@10":0.175730543,
+  "recall@100":0.3979622293
+}

results/BrightRetrieval_sustainable_living_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.1692097115,
+  "ndcg@1":0.1574074074,
+  "ndcg@10":0.2010654629,
+  "ndcg@100":0.2989909999,
+  "recall@10":0.2462456855,
+  "recall@100":0.5804351005
+}

results/BrightRetrieval_theoremqa_questions_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.0720669895,
+  "ndcg@1":0.0243902439,
+  "ndcg@10":0.0904299064,
+  "ndcg@100":0.1150213787,
+  "recall@10":0.1296747967,
+  "recall@100":0.2348432056
+}

results/BrightRetrieval_theoremqa_theorems_evaluation_scores_qlen256.json

@@ -0,0 +1,8 @@
+{
+  "map":0.0268134408,
+  "ndcg@1":0.0153846154,
+  "ndcg@10":0.0318862999,
+  "ndcg@100":0.0559448533,
+  "recall@10":0.0512820513,
+  "recall@100":0.1697435897
+}

results/run_meta.json

@@ -0,0 +1,16 @@
+{
+  "model_path":"/data/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output_edge32m_d128_fixed/stage-bge-hn/final",
+  "model_version":"baseline",
+  "query_length":256,
+  "pony_query_length":32,
+  "use_plaid":false,
+  "nbits":null,
+  "retrieve_k":100,
+  "overwrite":false,
+  "run_name":"edge32m_d128_fixed",
+  "label":"output_edge32m_d128_fixed__stage-bge-hn__final",
+  "started_at":"2026-04-24T21:47:24",
+  "splits_requested":[
+    "stackoverflow"
+  ]
+}

results/summary.json

@@ -0,0 +1,18 @@
+{
+  "leetcode":{
+    "map":0.1243283683,
+    "ndcg@1":0.0,
+    "ndcg@10":0.1739560791,
+    "ndcg@100":0.2252391566,
+    "recall@10":0.2919014085,
+    "recall@100":0.510915493
+  },
+  "aops":{
+    "map":0.0328788796,
+    "ndcg@1":0.009009009,
+    "ndcg@10":0.0488934387,
+    "ndcg@100":0.0961218916,
+    "recall@10":0.0662269412,
+    "recall@100":0.2151973402
+  }
+}

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 255,
+  "do_lower_case": true
+}

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "[MASK]",
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,968 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "1": {
+      "content": "<|padding|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50254": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50255": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50256": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50257": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50280": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50281": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50282": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50283": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50284": {
+      "content": "[MASK]",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50285": {
+      "content": "[unused0]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "[unused1]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "[unused2]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "[unused3]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "[unused4]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "[unused5]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "[unused6]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "[unused7]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "[unused8]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "[unused9]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50295": {
+      "content": "[unused10]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50296": {
+      "content": "[unused11]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50297": {
+      "content": "[unused12]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50298": {
+      "content": "[unused13]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50299": {
+      "content": "[unused14]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50300": {
+      "content": "[unused15]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50301": {
+      "content": "[unused16]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50302": {
+      "content": "[unused17]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50303": {
+      "content": "[unused18]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50304": {
+      "content": "[unused19]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50305": {
+      "content": "[unused20]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50306": {
+      "content": "[unused21]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50307": {
+      "content": "[unused22]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50308": {
+      "content": "[unused23]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50309": {
+      "content": "[unused24]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50310": {
+      "content": "[unused25]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50311": {
+      "content": "[unused26]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50312": {
+      "content": "[unused27]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50313": {
+      "content": "[unused28]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50314": {
+      "content": "[unused29]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50315": {
+      "content": "[unused30]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50316": {
+      "content": "[unused31]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50317": {
+      "content": "[unused32]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50318": {
+      "content": "[unused33]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50319": {
+      "content": "[unused34]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50320": {
+      "content": "[unused35]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50321": {
+      "content": "[unused36]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50322": {
+      "content": "[unused37]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50323": {
+      "content": "[unused38]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50324": {
+      "content": "[unused39]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50325": {
+      "content": "[unused40]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50326": {
+      "content": "[unused41]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50327": {
+      "content": "[unused42]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50328": {
+      "content": "[unused43]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50329": {
+      "content": "[unused44]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50330": {
+      "content": "[unused45]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50331": {
+      "content": "[unused46]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50332": {
+      "content": "[unused47]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50333": {
+      "content": "[unused48]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50334": {
+      "content": "[unused49]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50335": {
+      "content": "[unused50]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50336": {
+      "content": "[unused51]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50337": {
+      "content": "[unused52]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50338": {
+      "content": "[unused53]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50339": {
+      "content": "[unused54]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50340": {
+      "content": "[unused55]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50341": {
+      "content": "[unused56]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50342": {
+      "content": "[unused57]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50343": {
+      "content": "[unused58]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50344": {
+      "content": "[unused59]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50345": {
+      "content": "[unused60]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50346": {
+      "content": "[unused61]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50347": {
+      "content": "[unused62]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50348": {
+      "content": "[unused63]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50349": {
+      "content": "[unused64]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50350": {
+      "content": "[unused65]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50351": {
+      "content": "[unused66]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50352": {
+      "content": "[unused67]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50353": {
+      "content": "[unused68]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50354": {
+      "content": "[unused69]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50355": {
+      "content": "[unused70]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50356": {
+      "content": "[unused71]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50357": {
+      "content": "[unused72]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50358": {
+      "content": "[unused73]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50359": {
+      "content": "[unused74]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50360": {
+      "content": "[unused75]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50361": {
+      "content": "[unused76]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50362": {
+      "content": "[unused77]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50363": {
+      "content": "[unused78]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50364": {
+      "content": "[unused79]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50365": {
+      "content": "[unused80]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50366": {
+      "content": "[unused81]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50367": {
+      "content": "[unused82]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50368": {
+      "content": "[Q] ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50369": {
+      "content": "[D] ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "max_length": 255,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 255,
+  "pad_to_multiple_of": null,
+  "pad_token": "[MASK]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}

training/launch_train_multinode.py

@@ -0,0 +1,126 @@
+"""Launcher that forwards to LaRLI/train_baseline.py, with a multi-node fix.
+
+Background
+----------
+`sentence_transformers.util.get_device_name()` contains a latent bug that
+affects multi-node torchrun runs:
+
+    if torch.distributed.is_initialized():
+        local_rank = torch.distributed.get_rank()  # actually the GLOBAL rank
+    else:
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    return f"cuda:{local_rank}"
+
+`torch.distributed.get_rank()` returns the **global** rank (0..world_size-1),
+not the local rank within the node. On node 1 of a 2-node 4-GPU run, ranks
+4-7 end up trying to bind to `cuda:4..cuda:7`, but only `cuda:0..cuda:3` are
+visible per node → `RuntimeError: CUDA error: invalid device ordinal`.
+
+This launcher:
+  1. Pins the current process to `cuda:LOCAL_RANK` from torchrun.
+  2. Monkey-patches `sentence_transformers.util.get_device_name` (and the
+     local alias in `SentenceTransformer`) to return the correct local rank
+     instead of the global one.
+  3. Hands off to `LaRLI.train_baseline.main()` unchanged.
+
+Nothing in LaRLI/ or pylate/ or sentence_transformers/ is modified.
+
+Usage (from your sbatch, inside srun + bash -c):
+    torchrun --nproc_per_node=4 --nnodes=2 --node_rank=$SLURM_PROCID ... \
+        widen_projection/launch_train_multinode.py \
+        --base_model ...  --train_data_path ...  (same args as train_baseline.py)
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+# ── 1. Resolve local rank from torchrun env ────────────────────────────
+_LOCAL_RANK = int(os.environ.get("LOCAL_RANK", 0))
+_GLOBAL_RANK = int(os.environ.get("RANK", 0))
+_WORLD = int(os.environ.get("WORLD_SIZE", 1))
+print(f"[launch] RANK={_GLOBAL_RANK} LOCAL_RANK={_LOCAL_RANK} WORLD={_WORLD}", flush=True)
+
+import torch  # noqa: E402
+
+if torch.cuda.is_available():
+    print(f"[launch] cuda.device_count={torch.cuda.device_count()} (before pin)", flush=True)
+    if _LOCAL_RANK < torch.cuda.device_count():
+        torch.cuda.set_device(_LOCAL_RANK)
+        print(f"[launch] torch.cuda.set_device({_LOCAL_RANK}); current={torch.cuda.current_device()}", flush=True)
+    else:
+        print(f"[launch] WARNING: LOCAL_RANK={_LOCAL_RANK} >= device_count={torch.cuda.device_count()}", flush=True)
+
+# ── 2. Monkey-patch ST's get_device_name before anything imports ST ───
+import sentence_transformers.util as _st_util  # noqa: E402
+
+
+def _patched_get_device_name() -> str:
+    """Return `cuda:LOCAL_RANK` under torchrun, else fall back to ST's own logic."""
+    if torch.cuda.is_available():
+        dev = f"cuda:{_LOCAL_RANK}"
+        print(f"[launch] get_device_name() -> {dev} (rank {_GLOBAL_RANK})", flush=True)
+        return dev
+    return _ORIG_GET_DEVICE_NAME()
+
+
+_ORIG_GET_DEVICE_NAME = _st_util.get_device_name
+_st_util.get_device_name = _patched_get_device_name
+print(f"[launch] patched sentence_transformers.util.get_device_name", flush=True)
+
+# SentenceTransformer imports `get_device_name` by name at module load time,
+# so we also need to patch the already-bound reference inside that module.
+# IMPORTANT: `sentence_transformers/__init__.py` does
+#   `from .SentenceTransformer import SentenceTransformer`
+# which shadows the submodule name with the class. `import
+# sentence_transformers.SentenceTransformer` therefore returns the CLASS, not
+# the module. The real module is in `sys.modules`.
+_ST_MODULE = sys.modules.get("sentence_transformers.SentenceTransformer")
+if _ST_MODULE is None:
+    # Force-import the submodule so it lands in sys.modules.
+    import importlib
+    _ST_MODULE = importlib.import_module("sentence_transformers.SentenceTransformer")
+
+if hasattr(_ST_MODULE, "get_device_name"):
+    _ST_MODULE.get_device_name = _patched_get_device_name
+    print(f"[launch] patched sentence_transformers.SentenceTransformer (module).get_device_name",
+          flush=True)
+else:
+    print(f"[launch] WARNING: module has no get_device_name; attrs with 'device': "
+          f"{[a for a in dir(_ST_MODULE) if 'device' in a.lower()]}", flush=True)
+
+# ── 3. Override max_grad_norm if requested via env var ────────────────
+# Widened projection heads start with zero-init weights in the new dims.
+# Those zeros produce pathological gradients at bootstrap (norm ~30-50k),
+# which HF Trainer's default max_grad_norm=1.0 clips away to near-nothing.
+# Set MAX_GRAD_NORM=100 (or higher / "inf") in the sbatch to let the new
+# dims actually learn; the original 64-dim weights are already trained so
+# they won't blow up.
+_MAX_GRAD_NORM_OVERRIDE = os.environ.get("MAX_GRAD_NORM")
+if _MAX_GRAD_NORM_OVERRIDE:
+    import sentence_transformers.training_args as _st_args
+
+    _orig_ta_init = _st_args.SentenceTransformerTrainingArguments.__post_init__
+
+    def _patched_post_init(self):
+        try:
+            val = float(_MAX_GRAD_NORM_OVERRIDE)
+        except ValueError:
+            val = float("inf")
+        self.max_grad_norm = val
+        print(f"[launch] forced max_grad_norm = {val} (was default 1.0)", flush=True)
+        _orig_ta_init(self)
+
+    _st_args.SentenceTransformerTrainingArguments.__post_init__ = _patched_post_init
+
+# ── 4. Make the project importable and hand off ───────────────────────
+_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+for p in (_PROJECT, os.path.join(_PROJECT, "pylate")):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+from LaRLI.train_baseline import main  # noqa: E402
+
+if __name__ == "__main__":
+    main()

training/materialize_bge_reasoner.py

@@ -0,0 +1,149 @@
+"""Materialize hanhainebula/bge-reasoner-data to {query, positive, negative}.
+
+Upstream dataset: https://huggingface.co/datasets/hanhainebula/bge-reasoner-data
+Files: 12 per-BRIGHT-domain JSONL files under ``bge-reasoner-data-0904/``.
+Schema per row: ``{prompt, query, pos:list[str], neg:list[str]}``.
+
+We build BRIGHT-aligned training triples by:
+  - Prepending the ``prompt`` to the query in the format used by ReasonEmbed/D:
+        ``f"{prompt}\nQuery: {query}"``
+    so the model sees the same instruction-style prefix it would see in the
+    vm2825 polish stage.
+  - Exploding each row into up to ``max_pos × max_neg`` (query, pos, neg) triples.
+  - Concatenating all 12 domains (or a single named domain) and shuffling.
+
+Usage
+-----
+  python LaRLI/data/materialize_bge_reasoner.py \
+      --output-dir data/bge_reasoner_processed \
+      --max-pos-per-query 2 --max-neg-per-query 4
+
+  # single domain (e.g. aops only) — useful to target a weak split:
+  python LaRLI/data/materialize_bge_reasoner.py \
+      --split aops --output-dir data/bge_reasoner_aops_processed
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import shutil
+import sys
+from pathlib import Path
+
+_PROJECT_ROOT = Path(__file__).resolve().parents[2]
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from datasets import Dataset, concatenate_datasets  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+REPO_ID = "hanhainebula/bge-reasoner-data"
+DATA_SUBDIR = "bge-reasoner-data-0904"
+BRIGHT_DOMAINS = [
+    "biology", "earth_science", "economics", "psychology",
+    "robotics", "stackoverflow", "sustainable_living",
+    "leetcode", "pony",
+    "aops", "theoremqa_questions", "theoremqa_theorems",
+]
+
+
+def _load_domain_jsonl(domain: str, cache_dir: str | None) -> list[dict]:
+    path = hf_hub_download(
+        repo_id=REPO_ID,
+        filename=f"{DATA_SUBDIR}/{domain}.jsonl",
+        repo_type="dataset",
+        cache_dir=cache_dir,
+    )
+    with open(path) as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def _format_query(prompt: str, query: str) -> str:
+    prompt = (prompt or "").strip()
+    query = (query or "").strip()
+    if prompt and query:
+        return f"{prompt}\nQuery: {query}"
+    return prompt or query
+
+
+def _explode_row(row: dict, max_pos: int, max_neg: int) -> list[dict]:
+    q = _format_query(row.get("prompt", ""), row.get("query", ""))
+    pos_list = row.get("pos") or []
+    neg_list = row.get("neg") or []
+    if not q or not pos_list or not neg_list:
+        return []
+    out: list[dict] = []
+    for pos in pos_list[:max_pos]:
+        pos = (pos or "").strip()
+        if len(pos) < 10:
+            continue
+        for neg in neg_list[:max_neg]:
+            neg = (neg or "").strip()
+            if len(neg) < 10:
+                continue
+            out.append({"query": q, "positive": pos, "negative": neg})
+    return out
+
+
+def _save_atomically(dataset: Dataset, output_path: Path) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = output_path.parent / f".{output_path.name}.tmp"
+    if tmp.exists():
+        shutil.rmtree(tmp)
+    tmp.mkdir(parents=True, exist_ok=True)
+    dataset.save_to_disk(str(tmp))
+    if output_path.exists():
+        shutil.rmtree(output_path)
+    tmp.replace(output_path)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Materialize BGE-Reasoner data to triples.")
+    p.add_argument("--split", default="all",
+                   help=f"'all' or one of {BRIGHT_DOMAINS}. Default: all.")
+    p.add_argument("--output-dir", required=True)
+    p.add_argument("--max-pos-per-query", type=int, default=2)
+    p.add_argument("--max-neg-per-query", type=int, default=4)
+    p.add_argument("--cache-dir", default=None,
+                   help="HF cache dir. Defaults to HF_HOME.")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    out = Path(args.output_dir)
+    domains = BRIGHT_DOMAINS if args.split == "all" else [args.split]
+    if args.split != "all" and args.split not in BRIGHT_DOMAINS:
+        raise ValueError(f"Unknown split {args.split!r}. Must be 'all' or one of {BRIGHT_DOMAINS}.")
+
+    per_domain: list[Dataset] = []
+    grand_raw = grand_trip = 0
+    for dom in domains:
+        logger.info("Downloading %s/%s ...", DATA_SUBDIR, dom)
+        raw = _load_domain_jsonl(dom, args.cache_dir)
+        rows: list[dict] = []
+        for r in raw:
+            rows.extend(_explode_row(r, args.max_pos_per_query, args.max_neg_per_query))
+        logger.info("  %s: %d raw rows → %d triples", dom, len(raw), len(rows))
+        grand_raw += len(raw)
+        grand_trip += len(rows)
+        if rows:
+            per_domain.append(Dataset.from_list(rows))
+
+    if not per_domain:
+        raise RuntimeError("No rows materialized.")
+    combined = concatenate_datasets(per_domain).shuffle(seed=42)
+    logger.info("Combined bge-reasoner (%d domains): %d raw rows → %d triples",
+                len(domains), grand_raw, grand_trip)
+
+    _save_atomically(combined, out)
+    logger.info("Saved to %s (%d triples)", out, len(combined))
+
+
+if __name__ == "__main__":
+    main()

training/materialize_reasonir.py

@@ -0,0 +1,108 @@
+"""Materialize ReasonIR HQ and/or VL to disk as flat {query, positive, negative} datasets.
+
+Mirrors the existing `data/reasonir_hq_processed/` layout so both splits are loaded
+the same way at train time.
+
+Usage
+-----
+    python LaRLI/data/materialize_reasonir.py --config hq --output-dir data/reasonir_hq_processed
+    python LaRLI/data/materialize_reasonir.py --config vl --output-dir data/reasonir_vl_processed
+    python LaRLI/data/materialize_reasonir.py --config both   # materializes both
+
+The script is idempotent: if the target directory already contains a saved dataset,
+it exits early.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import shutil
+import sys
+from pathlib import Path
+
+_PROJECT_ROOT = Path(__file__).resolve().parents[2]
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from LaRLI.data.prepare_data import (
+    _is_saved_dataset,
+    load_reasonir_hq_dataset,
+    load_reasonir_vl_dataset,
+)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def _save_atomically(dataset, output_path: Path) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = output_path.parent / f".{output_path.name}.tmp"
+    if tmp_path.exists():
+        shutil.rmtree(tmp_path)
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    dataset.save_to_disk(str(tmp_path))
+    if output_path.exists():
+        shutil.rmtree(output_path)
+    tmp_path.replace(output_path)
+
+
+def materialize(config: str, output_dir: Path, cache_dir: str | None) -> Path:
+    if _is_saved_dataset(output_dir):
+        logger.info("%s already materialized at %s, skipping.", config, output_dir)
+        return output_dir
+
+    if config == "hq":
+        ds = load_reasonir_hq_dataset(cache_dir=cache_dir)
+    elif config == "vl":
+        ds = load_reasonir_vl_dataset(cache_dir=cache_dir)
+    else:
+        raise ValueError(f"Unknown config: {config}")
+
+    ds = ds.select_columns(["query", "positive", "negative"])
+    logger.info("Saving %s (%d rows) to %s", config, len(ds), output_dir)
+    _save_atomically(ds, output_dir)
+    return output_dir
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Materialize ReasonIR splits to disk.")
+    parser.add_argument("--config", choices=["hq", "vl", "both"], required=True)
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Target directory (required when --config is hq or vl).",
+    )
+    parser.add_argument(
+        "--hq-output-dir",
+        type=str,
+        default=str(_PROJECT_ROOT / "data" / "reasonir_hq_processed"),
+        help="Used only when --config=both.",
+    )
+    parser.add_argument(
+        "--vl-output-dir",
+        type=str,
+        default=str(_PROJECT_ROOT / "data" / "reasonir_vl_processed"),
+        help="Used only when --config=both.",
+    )
+    parser.add_argument("--cache-dir", type=str, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.config in ("hq", "vl"):
+        if not args.output_dir:
+            raise SystemExit("--output-dir is required when --config is hq or vl")
+        out = materialize(args.config, Path(args.output_dir), args.cache_dir)
+        print(out)
+    else:
+        hq_out = materialize("hq", Path(args.hq_output_dir), args.cache_dir)
+        vl_out = materialize("vl", Path(args.vl_output_dir), args.cache_dir)
+        print(hq_out)
+        print(vl_out)
+
+
+if __name__ == "__main__":
+    main()

training/merge_datasets.py

@@ -0,0 +1,168 @@
+"""Concatenate several saved-to-disk {query, positive, negative} datasets into one.
+
+Useful for building a single Stage-B data dir from the three RaDeR splits, or a
+single Stage-D data dir from multiple reason-embed variants.
+
+Usage
+-----
+  python LaRLI/data/merge_datasets.py \
+      --inputs data/rader_numina_all_processed \
+               data/rader_math_qas_lex_processed \
+               data/rader_math_cot_lex_processed \
+      --output-dir data/rader_merged_processed
+
+Preserves only the three columns {query, positive, negative}, shuffles the
+result with a fixed seed so the training loop sees interleaved sources.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import shutil
+from pathlib import Path
+
+from datasets import concatenate_datasets, load_from_disk
+
+
+def _free_bytes(path: Path) -> int:
+    """Available bytes on the filesystem hosting `path`."""
+    p = path if path.exists() else path.parent
+    while not p.exists():
+        p = p.parent
+    stats = shutil.disk_usage(str(p))
+    return stats.free
+
+
+def _safe_rmtree(path: Path) -> None:
+    """Remove `path` even if partial writes left stray files behind."""
+    if not path.exists():
+        return
+    shutil.rmtree(path)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+
+REQUIRED_COLS = ("query", "positive", "negative")
+
+
+def _load_one(path: Path):
+    if not (path / "dataset_info.json").exists():
+        raise FileNotFoundError(f"{path} is not a saved HF dataset (missing dataset_info.json)")
+    ds = load_from_disk(str(path))
+    missing = set(REQUIRED_COLS) - set(ds.column_names)
+    if missing:
+        raise ValueError(f"{path} missing columns: {missing}")
+    return ds.select_columns(list(REQUIRED_COLS))
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Merge saved-to-disk triplet datasets.")
+    p.add_argument("--inputs", nargs="+", required=True,
+                   help="Paths to input dataset directories.")
+    p.add_argument("--output-dir", required=True)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--num-shards", type=int, default=None,
+                   help="Force shard count (default: datasets auto). Use 1 for a "
+                        "single .arrow file — slower but safer on flaky FS.")
+    p.add_argument("--num-proc", type=int, default=1,
+                   help="Parallel shard writers. Default=1 (serial) because HF "
+                        "save_to_disk has been observed to race on NFS/lustre. "
+                        "Bump to 4 only on a healthy local FS.")
+    p.add_argument("--min-free-gb", type=int, default=20,
+                   help="Abort early if the filesystem has less than this many "
+                        "GB free (guards against mid-write disk-full failures).")
+    args = p.parse_args()
+
+    parts = []
+    total = 0
+    for raw in args.inputs:
+        path = Path(raw)
+        ds = _load_one(path)
+        logger.info("  + %s (%d rows)", path, len(ds))
+        parts.append(ds)
+        total += len(ds)
+
+    combined = concatenate_datasets(parts).shuffle(seed=args.seed)
+    logger.info("Combined: %d rows across %d inputs", len(combined), len(parts))
+    assert len(combined) == total
+
+    out = Path(args.output_dir)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    free_gb = _free_bytes(out.parent) / (1 << 30)
+    logger.info("Free space on target FS: %.1f GB", free_gb)
+    if free_gb < args.min_free_gb:
+        raise SystemExit(
+            f"Only {free_gb:.1f} GB free on {out.parent}, below --min-free-gb "
+            f"{args.min_free_gb}. Aborting to avoid a mid-write failure."
+        )
+
+    tmp = out.parent / f".{out.name}.tmp"
+    # Defensive: blow away ANY leftover from a prior crash before starting.
+    _safe_rmtree(tmp)
+    tmp.mkdir(parents=True, exist_ok=True)
+
+    save_kwargs = {}
+    if args.num_shards is not None:
+        save_kwargs["num_shards"] = args.num_shards
+    if args.num_proc and args.num_proc > 1:
+        save_kwargs["num_proc"] = args.num_proc
+    # If num_proc is 1 (default), DO NOT pass it — HF's default is serial and
+    # we avoid the multiprocessing path entirely. Passing num_proc=1 explicitly
+    # still triggers the multiproc code path in some datasets versions.
+
+    try:
+        combined.save_to_disk(str(tmp), **save_kwargs)
+    except Exception:
+        logger.exception("save_to_disk failed; leaving %s for inspection.", tmp)
+        raise
+
+    # Verify every shard claimed in state.json actually exists AND is
+    # internally readable. On networked filesystems we've seen shards that
+    # have the expected byte length but are structurally malformed — the
+    # training process hits `OSError: Expected to be able to read N bytes for
+    # message body, got M` only at load time. Fail loudly HERE instead.
+    state_file = tmp / "state.json"
+    if not state_file.exists():
+        raise SystemExit(f"state.json missing from {tmp}; save_to_disk incomplete.")
+    import json
+    state = json.loads(state_file.read_text())
+    expected_shards = state.get("_data_files") or []
+    import pyarrow as pa
+    import pyarrow.ipc as ipc
+    for shard in expected_shards:
+        path = tmp / shard["filename"]
+        if not path.exists():
+            raise SystemExit(
+                f"Shard missing from {tmp}: {shard['filename']}. "
+                "Re-run the merge — the filesystem dropped shards mid-write."
+            )
+        try:
+            with pa.memory_map(str(path), "r") as src:
+                reader = ipc.open_stream(src)
+                _ = reader.read_all()
+        except Exception as e:
+            raise SystemExit(
+                f"Shard {path} exists ({path.stat().st_size} bytes) but is "
+                f"structurally malformed: {type(e).__name__}: {e}. "
+                "Re-run the merge (preferably with fewer --num-shards)."
+            )
+    logger.info("Verified %d shards in %s (existence + readability).",
+                len(expected_shards), tmp)
+
+    # Atomic-ish swap.
+    if out.exists():
+        backup = out.parent / f".{out.name}.old"
+        _safe_rmtree(backup)
+        out.replace(backup)
+    tmp.replace(out)
+    if (out.parent / f".{out.name}.old").exists():
+        _safe_rmtree(out.parent / f".{out.name}.old")
+    logger.info("Saved to %s", out)
+    print(out)
+
+
+if __name__ == "__main__":
+    main()

training/train_fixed_stage1.sbatch

@@ -0,0 +1,133 @@
+#!/bin/bash
+#ASC --vanilla
+
+#SBATCH --output=logs/%j.out
+#SBATCH --error=logs/%j.err
+#SBATCH --job-name=edge32m-fixed-s1
+#SBATCH --partition=zen4_0768_h100x4
+#SBATCH --qos=zen4_0768_h100x4
+#SBATCH --time=08:00:00
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:4
+#SBATCH --gpus-per-task=4
+#SBATCH --cpus-per-task=8
+#SBATCH --threads-per-core=1
+#
+# Stage 1 of the widened-32M curriculum: train the widened base on VL data.
+# Base must already exist at output_edge32m_d128_fixed/base/ (with random-init new
+# dims; rebuild with widen_colbert_projection.py if missing).
+#
+# Submit:
+#   mkdir -p logs
+#   sbatch widen_projection/train_widened_stage1.sbatch
+#   # then after it finishes, stage 2:
+#   sbatch widen_projection/train_widened_stage2.sbatch
+
+set -euo pipefail
+cd /data/fs201059/aa17626/modernTransformer
+
+eval "$('/data/fs201059/aa17626/miniforge3/bin/conda' shell.bash hook)"
+conda activate reason-mcb-repro
+
+# Resolve master node hostname to an IP on the cluster ethernet (avoids
+# NCCL picking a link-local/virbr bridge)
+HEAD=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_ADDR=$(srun --nodes=1 --ntasks=1 -w "$HEAD" \
+    bash -c "ip -4 -o addr show eno2np0 | awk '{print \$4}' | cut -d/ -f1")
+# Unique rendezvous port per job — avoids collisions with other running jobs
+export MASTER_PORT=$(( 20000 + (SLURM_JOB_ID % 40000) ))
+export GPUS_PER_NODE=4
+
+# HF caches on shared FS
+export HF_HOME=/data/fs201059/aa17626/.cache/huggingface
+export HF_HUB_CACHE="$HF_HOME/hub"
+export HF_DATASETS_CACHE="$HF_HOME/datasets"
+
+# Project paths + pylate on PYTHONPATH
+export PYTHONNOUSERSITE=1
+export PYTHONPATH="/data/fs201059/aa17626/modernTransformer:/data/fs201059/aa17626/modernTransformer/pylate${PYTHONPATH:+:${PYTHONPATH}}"
+
+# NCCL: pin to cluster ethernet for rendezvous + InfiniBand for the collective.
+# Link-local / virbr / docker / tun interfaces on some nodes were tripping up
+# auto-detection (169.254.x.x unroutable between nodes).
+export NCCL_SOCKET_IFNAME=eno2np0     # 10.97.x.x cluster ethernet
+export NCCL_IB_HCA=mlx5               # auto-match all InfiniBand HCAs on the node
+export NCCL_IB_DISABLE=0              # use IB for collectives (fast)
+export NCCL_IB_TIMEOUT=22
+export NCCL_ASYNC_ERROR_HANDLING=1
+export NCCL_DEBUG=WARN
+export NCCL_TIMEOUT=3600
+
+# ── Training config ─────────────────────────────────────────────────────
+BASE_MODEL=/data/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output_edge32m_d128_fixed/base
+DATA_PATH=/data/fs201059/aa17626/modernTransformer/data/reasonir_vl_processed
+OUTPUT_DIR=/data/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output_edge32m_d128_fixed
+RUN_NAME=stage-vl
+
+EPOCHS=1
+LR=1e-5                          # VL warmup lr (matches v2b's Stage A)
+BATCH_SIZE=32                   # per device → 256 × 4 × 2 = 2048 global
+MINI_BATCH_SIZE=32
+GRADIENT_ACCUMULATION_STEPS=4
+# Let zero-init-adjacent dims (new 64-dim rows) learn without clip saturating.
+# Grad norms spike high at bootstrap while the small-random new rows find
+# sensible values; 100 lets them flow.
+export MAX_GRAD_NORM=10000
+QUERY_LENGTH=256                 # matches BRIGHT eval qlen; avoids truncation on long VL queries
+DOCUMENT_LENGTH=2048             # covers p99 of all data; 8192 was wasted padding
+ATTN=flash_attention_2
+DTYPE=bfloat16
+
+if [ ! -d "$BASE_MODEL" ]; then
+  echo "ERROR: widened base missing at $BASE_MODEL"
+  echo "Run: python widen_projection/widen_colbert_projection.py \\"
+  echo "       --src <HF snapshot of mxbai-edge-colbert-v0-32m> \\"
+  echo "       --dst $BASE_MODEL --new-dim 128 --verify"
+  exit 1
+fi
+if [ ! -f "$DATA_PATH/dataset_info.json" ]; then
+  echo "ERROR: VL data missing at $DATA_PATH"
+  echo "Run: python LaRLI/data/materialize_reasonir.py --config vl --output-dir $DATA_PATH"
+  exit 1
+fi
+
+echo "============================================================"
+echo "Stage 1 (widened 32M d128): VL warmup"
+echo "nodes=$SLURM_JOB_NUM_NODES  gpus/node=$GPUS_PER_NODE  global_batch=$((BATCH_SIZE * GPUS_PER_NODE * SLURM_JOB_NUM_NODES))"
+echo "  base   : $BASE_MODEL"
+echo "  data   : $DATA_PATH"
+echo "  output : $OUTPUT_DIR/$RUN_NAME/final"
+echo "  lr/ep  : $LR / $EPOCHS"
+echo "  qlen   : $QUERY_LENGTH   doclen: $DOCUMENT_LENGTH   max_grad_norm: $MAX_GRAD_NORM"
+echo "  master : $MASTER_ADDR:$MASTER_PORT   (from $HEAD via eno2np0)"
+echo "  nccl   : IFNAME=$NCCL_SOCKET_IFNAME  IB_HCA=$NCCL_IB_HCA"
+echo "============================================================"
+
+srun --ntasks=$SLURM_JOB_NUM_NODES --ntasks-per-node=1 \
+     --gpus-per-task=$GPUS_PER_NODE \
+     --gpu-bind=none \
+     bash -c '
+       torchrun \
+         --nproc_per_node='"$GPUS_PER_NODE"' \
+         --nnodes='"$SLURM_JOB_NUM_NODES"' \
+         --node_rank="$SLURM_PROCID" \
+         --master_addr="'"$MASTER_ADDR"'" \
+         --master_port="'"$MASTER_PORT"'" \
+         widen_projection/launch_train_multinode.py \
+         --base_model '"$BASE_MODEL"' \
+         --train_data_path '"$DATA_PATH"' \
+         --epochs '"$EPOCHS"' --lr '"$LR"' \
+         --batch_size '"$BATCH_SIZE"' \
+         --mini_batch_size '"$MINI_BATCH_SIZE"' \
+         --gradient_accumulation_steps '"$GRADIENT_ACCUMULATION_STEPS"' \
+         --query_length '"$QUERY_LENGTH"' \
+         --document_length '"$DOCUMENT_LENGTH"' \
+         --attn_implementation '"$ATTN"' \
+         --torch_dtype '"$DTYPE"' \
+         --output_dir '"$OUTPUT_DIR"' \
+         --run_name '"$RUN_NAME"'
+     '
+# Stage 1 output: $OUTPUT_DIR/$RUN_NAME/final
+# Then submit stage 2:
+#   sbatch widen_projection/train_widened_stage2.sbatch

training/train_fixed_stage2.sbatch

@@ -0,0 +1,128 @@
+#!/bin/bash
+#ASC --vanilla
+
+#SBATCH --output=logs/%j.out
+#SBATCH --error=logs/%j.err
+#SBATCH --job-name=edge32m-fixed-s2
+#SBATCH --partition=zen4_0768_h100x4
+#SBATCH --qos=zen4_0768_h100x4
+#SBATCH --time=12:00:00
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:4
+#SBATCH --gpus-per-task=4
+#SBATCH --cpus-per-task=8
+#SBATCH --threads-per-core=1
+#
+# Stage 2 (widened 32M: VL → BGE+HQ-hn polish).
+# Request: 2 nodes × 4 H100 = 8 GPUs (matches Reason-MCB's 8-GPU training).
+# per_device_batch_size = 256  → global / in-batch negatives = 2048,
+# very close to the paper's 8×258 = 2064.
+#
+# Stage VL must already be at
+#   repro_reason_moderncolbert/output_edge32m_d128_fixed/stage-vl/final
+# and the merged dataset at
+#   data/bge_hn_merged_processed
+# before submission.
+#
+# Submit with:
+#   mkdir -p logs
+#   sbatch widen_projection/train_widened_stage2.sbatch
+
+set -euo pipefail
+cd /data/fs201059/aa17626/modernTransformer
+
+eval "$('/data/fs201059/aa17626/miniforge3/bin/conda' shell.bash hook)"
+conda activate reason-mcb-repro
+
+# Resolve master node hostname to an IP on the cluster ethernet (avoids
+# NCCL picking a link-local/virbr bridge).
+HEAD=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_ADDR=$(srun --nodes=1 --ntasks=1 -w "$HEAD" \
+    bash -c "ip -4 -o addr show eno2np0 | awk '{print \$4}' | cut -d/ -f1")
+export MASTER_PORT=$(( 20000 + (SLURM_JOB_ID % 40000) ))   # unique per job
+export GPUS_PER_NODE=4
+
+# NCCL: pin to cluster ethernet for rendezvous + InfiniBand for collectives.
+export NCCL_SOCKET_IFNAME=eno2np0     # 10.97.x.x cluster ethernet
+export NCCL_IB_HCA=mlx5               # match all InfiniBand HCAs on the node
+export NCCL_IB_DISABLE=0
+export NCCL_IB_TIMEOUT=22
+export NCCL_ASYNC_ERROR_HANDLING=1
+export NCCL_DEBUG=WARN
+export NCCL_TIMEOUT=3600
+
+# HF caches on shared FS
+export HF_HOME=/data/fs201059/aa17626/.cache/huggingface
+export HF_HUB_CACHE="$HF_HOME/hub"
+export HF_DATASETS_CACHE="$HF_HOME/datasets"
+
+# Project paths + pylate on PYTHONPATH
+export PYTHONNOUSERSITE=1
+export PYTHONPATH="/data/fs201059/aa17626/modernTransformer:/data/fs201059/aa17626/modernTransformer/pylate${PYTHONPATH:+:${PYTHONPATH}}"
+
+# ── Training config ─────────────────────────────────────────────────────
+BASE_MODEL=/data/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output_edge32m_d128_fixed/stage-vl/final
+DATA_PATH=/data/fs201059/aa17626/modernTransformer/data/bge_hn_merged_processed
+OUTPUT_DIR=/data/fs201059/aa17626/modernTransformer/repro_reason_moderncolbert/output_edge32m_d128_fixed
+RUN_NAME=stage-bge-hn
+
+EPOCHS=1
+LR=5e-6
+BATCH_SIZE=8                   # per device → 256 × 4 × 2 = 2048 global (≈ paper 2064)
+MINI_BATCH_SIZE=8
+# Let the zero-init widened projection dims actually learn. With default
+# max_grad_norm=1.0, initial grad norms ~35k get clipped to 1/35000 of lr and the
+# model is effectively frozen. 100 lets the new dims find sensible weights.
+export MAX_GRAD_NORM=10000
+GRADIENT_ACCUMULATION_STEPS=2
+QUERY_LENGTH=256           # BRIGHT eval uses 256; training at 128 truncated 20% of our queries (p90=444 tokens)
+DOCUMENT_LENGTH=2048       # covers p99 of every BRIGHT split and training data; 8192 was 4x wasted compute
+ATTN=flash_attention_2
+DTYPE=bfloat16
+
+# Sanity
+if [ ! -d "$BASE_MODEL" ]; then
+  echo "ERROR: stage-vl final missing at $BASE_MODEL"; exit 1
+fi
+if [ ! -f "$DATA_PATH/dataset_info.json" ]; then
+  echo "ERROR: merged dataset missing at $DATA_PATH"; exit 1
+fi
+
+echo "============================================================"
+echo "Stage 2: widened 32M + VL → BGE+HQ-hn polish"
+echo "nodes=$SLURM_JOB_NUM_NODES  gpus/node=$GPUS_PER_NODE  global_batch=$((BATCH_SIZE * GPUS_PER_NODE * SLURM_JOB_NUM_NODES))"
+echo "  base   : $BASE_MODEL"
+echo "  data   : $DATA_PATH"
+echo "  output : $OUTPUT_DIR/$RUN_NAME/final"
+echo "  lr/ep  : $LR / $EPOCHS"
+echo "  master : $MASTER_ADDR:$MASTER_PORT"
+echo "============================================================"
+
+# One srun task per node; each task launches a local torchrun process group
+# that binds to the 4 GPUs on its node. All nodes meet at $MASTER_ADDR.
+srun --ntasks=$SLURM_JOB_NUM_NODES --ntasks-per-node=1 \
+     --gpus-per-task=$GPUS_PER_NODE \
+     --gpu-bind=none \
+     bash -c '
+       torchrun \
+         --nproc_per_node='"$GPUS_PER_NODE"' \
+         --nnodes='"$SLURM_JOB_NUM_NODES"' \
+         --node_rank="$SLURM_PROCID" \
+         --master_addr="'"$MASTER_ADDR"'" \
+         --master_port="'"$MASTER_PORT"'" \
+         widen_projection/launch_train_multinode.py \
+         --base_model '"$BASE_MODEL"' \
+         --train_data_path '"$DATA_PATH"' \
+         --epochs '"$EPOCHS"' --lr '"$LR"' \
+         --batch_size '"$BATCH_SIZE"' \
+         --mini_batch_size '"$MINI_BATCH_SIZE"' \
+         --gradient_accumulation_steps '"$GRADIENT_ACCUMULATION_STEPS"' \
+         --query_length '"$QUERY_LENGTH"' \
+         --document_length '"$DOCUMENT_LENGTH"' \
+         --attn_implementation '"$ATTN"' \
+         --torch_dtype '"$DTYPE"' \
+         --output_dir '"$OUTPUT_DIR"' \
+         --run_name '"$RUN_NAME"'
+     '
+# Final model: $OUTPUT_DIR/$RUN_NAME/final

training/widen_colbert_projection.py

@@ -0,0 +1,262 @@
+"""Widen the final projection head of a PyLate ColBERT checkpoint from N→M dims.
+
+The edge-ColBERT 32M model ships a 3-layer projection head ending at 64 dims
+(see `3_Dense/config.json`). That 64-dim MaxSim channel is a bottleneck on
+symbol-dense BRIGHT splits (leetcode, aops, theoremqa) where many tokens look
+alike and need more channels to distinguish.
+
+This script takes an EXISTING saved ColBERT checkpoint dir and writes a NEW
+checkpoint dir whose final Dense layer has a widened `out_features`, without
+modifying anything else (backbone weights, tokenizer, query/doc prefixes,
+skiplist, etc. are all hard-linked from the source dir).
+
+Zero-padding strategy (safe at init):
+    old linear.weight  shape (old_dim, in_features)
+    new linear.weight  shape (new_dim, in_features)
+        first old_dim rows = old rows
+        rows old_dim .. new_dim-1 = 0
+After L2 normalisation (which the encoder applies before MaxSim), a vector
+`[x; 0]` normalises to `[x/||x||; 0]`, so the initial cosine similarities and
+MaxSim scores are IDENTICAL to the source model. Training from this init only
+learns what to put in the extra channels — the original behaviour is
+preserved unless the model finds something better.
+
+Usage
+-----
+  python widen_projection/widen_colbert_projection.py \
+      --src  repro_reason_moderncolbert/output_edge32m_v2b/pos1_A/stage-vl/final \
+      --dst  repro_reason_moderncolbert/output_edge32m_vl_d128/initial \
+      --new-dim 128
+
+  # Sanity-check the widened checkpoint loads and produces identical outputs
+  # on the first 64 dims:
+  python widen_projection/widen_colbert_projection.py \
+      --src  repro_reason_moderncolbert/output_edge32m_v2b/pos1_A/stage-vl/final \
+      --dst  repro_reason_moderncolbert/output_edge32m_vl_d128/initial \
+      --new-dim 128 --verify
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+import shutil
+import sys
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file as load_safetensors
+from safetensors.torch import save_file as save_safetensors
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# The pylate project root (so we can import models.ColBERT for --verify).
+_PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _hardlink_tree(src: Path, dst: Path, skip: set[str]) -> None:
+    """Mirror `src` to `dst` with hardlinks, skipping paths in `skip`.
+
+    HF Hub snapshot dirs contain symlinks pointing at `../../blobs/<sha>`;
+    hardlinking a symlink preserves the relative link text, which breaks in
+    the new location. We resolve symlinks to their real path before hardlinking
+    so the destination files have the actual content reachable.
+    """
+    for root, dirs, files in os.walk(src, followlinks=False):
+        rel_root = Path(root).relative_to(src)
+        if any(str(rel_root).startswith(s) for s in skip):
+            continue
+        (dst / rel_root).mkdir(parents=True, exist_ok=True)
+        for f in files:
+            rel_file = rel_root / f
+            if str(rel_file) in skip:
+                continue
+            src_f = Path(root) / f
+            dst_f = dst / rel_file
+            if dst_f.exists() or dst_f.is_symlink():
+                dst_f.unlink()
+            # Resolve if symlink so we hardlink the actual blob, not the
+            # symlink text (which would break when relative paths change).
+            real = src_f.resolve() if src_f.is_symlink() else src_f
+            try:
+                os.link(real, dst_f)
+            except OSError:
+                shutil.copy2(real, dst_f)
+
+
+def _widen_dense(src_dense: Path, dst_dense: Path, new_out: int) -> tuple[int, int]:
+    """Rewrite one Dense/ subdirectory with widened `out_features`.
+
+    Returns (old_out, in_features).
+    """
+    cfg_path = src_dense / "config.json"
+    w_path = src_dense / "model.safetensors"
+    if not cfg_path.exists() or not w_path.exists():
+        raise FileNotFoundError(f"{src_dense} missing config.json or model.safetensors")
+
+    with open(cfg_path) as f:
+        cfg = json.load(f)
+    old_out = int(cfg["out_features"])
+    in_features = int(cfg["in_features"])
+
+    if old_out == new_out:
+        logger.info("%s already %d-dim; copying unchanged.", src_dense.name, new_out)
+        dst_dense.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(cfg_path, dst_dense / "config.json")
+        shutil.copy2(w_path, dst_dense / "model.safetensors")
+        return old_out, in_features
+    if new_out < old_out:
+        raise ValueError(
+            f"Refusing to NARROW projection {src_dense.name} ({old_out} → {new_out}); "
+            "this tool only widens."
+        )
+
+    # Load old weight. Init new rows with SMALL RANDOM values (not zeros!).
+    #
+    # Why not zeros: after L2-normalisation, a zero-padded row produces 0 in
+    # the new dims on BOTH the query and document sides. The MaxSim score
+    # becomes identical to the old 64-dim model, and the gradient w.r.t. the
+    # new weights is exactly zero (dead-neuron deadlock). They never learn.
+    #
+    # Small Gaussian init (std matches the old rows' per-row stats) unlocks
+    # gradient flow while only changing MaxSim by O(std^2) at init — small
+    # enough that it doesn't undo the VL-stage training.
+    state = load_safetensors(str(w_path))
+    if "linear.weight" not in state:
+        raise KeyError(f"{w_path} missing key 'linear.weight' (got {list(state)})")
+    w_old = state["linear.weight"]
+    if w_old.shape != (old_out, in_features):
+        raise ValueError(
+            f"Unexpected weight shape {tuple(w_old.shape)}; "
+            f"expected ({old_out}, {in_features}) from config.json"
+        )
+
+    # Use the std of the existing rows so the new rows are scaled consistently.
+    old_std = float(w_old.float().std().item())
+    init_std = max(old_std * 0.1, 1e-3)  # 10% of old magnitude; never below 1e-3
+    gen = torch.Generator().manual_seed(0)
+    pad = torch.randn((new_out - old_out, in_features), generator=gen,
+                      dtype=torch.float32) * init_std
+    pad = pad.to(w_old.dtype)
+    w_new = torch.cat([w_old, pad], dim=0).contiguous()
+    logger.info(
+        "new rows: std=%.3e (10%% of old std=%.3e)  max|abs|=%.3e",
+        init_std, old_std, pad.abs().max().item(),
+    )
+
+    new_state = {"linear.weight": w_new}
+    if "linear.bias" in state:
+        b_old = state["linear.bias"]
+        b_pad = torch.zeros((new_out - old_out,), dtype=b_old.dtype)
+        new_state["linear.bias"] = torch.cat([b_old, b_pad], dim=0).contiguous()
+
+    dst_dense.mkdir(parents=True, exist_ok=True)
+    save_safetensors(new_state, str(dst_dense / "model.safetensors"))
+
+    cfg_new = dict(cfg)
+    cfg_new["out_features"] = new_out
+    with open(dst_dense / "config.json", "w") as f:
+        json.dump(cfg_new, f, indent=2)
+    logger.info(
+        "widened %s: out_features %d → %d  (weight shape %s → %s)",
+        src_dense.name, old_out, new_out, tuple(w_old.shape), tuple(w_new.shape),
+    )
+    return old_out, in_features
+
+
+def _find_final_dense(src: Path) -> Path:
+    """Find the last Dense subdir referenced in modules.json."""
+    mod_path = src / "modules.json"
+    if not mod_path.exists():
+        raise FileNotFoundError(f"{src} missing modules.json — not a PyLate/ST model?")
+    with open(mod_path) as f:
+        modules = json.load(f)
+    dense_paths = [m["path"] for m in modules if m.get("path")
+                   and (src / m["path"]).is_dir()
+                   and m.get("type", "").endswith("Dense")]
+    if not dense_paths:
+        raise RuntimeError(f"{src}/modules.json has no Dense entries.")
+    final = dense_paths[-1]
+    return src / final
+
+
+def verify(src: Path, dst: Path) -> None:
+    """Encode a tiny batch with both models and check the first-N dims match."""
+    sys.path.insert(0, str(_PROJECT_ROOT))
+    sys.path.insert(0, str(_PROJECT_ROOT / "pylate"))
+    from pylate import models  # noqa: E402
+
+    sentences = ["hello world", "a quick brown fox"]
+    logger.info("loading source model: %s", src)
+    m_old = models.ColBERT(str(src))
+    logger.info("loading widened model: %s", dst)
+    m_new = models.ColBERT(str(dst))
+
+    emb_old = m_old.encode(sentences, is_query=False, convert_to_numpy=True, show_progress_bar=False)
+    emb_new = m_new.encode(sentences, is_query=False, convert_to_numpy=True, show_progress_bar=False)
+
+    import numpy as np  # local import so widening path doesn't need numpy
+    ok = True
+    for i, (e_o, e_n) in enumerate(zip(emb_old, emb_new)):
+        old_dim = e_o.shape[-1]
+        # With small-random init the first old_dim coordinates DRIFT SLIGHTLY
+        # from the source (L2-normalisation couples them to the new dims via
+        # the norm denominator). Check the tail is non-zero (gradient will
+        # flow) and that the drift on the old dims is within a tolerance.
+        diff = np.abs(e_n[:, :old_dim] - e_o).max()
+        tail = np.abs(e_n[:, old_dim:]).max() if e_n.shape[-1] > old_dim else 0.0
+        logger.info("sentence %d: max |old - new[:, :%d]| = %.2e   "
+                    "max |new[:, %d:]| = %.2e  (should be > 0 → grad flows)",
+                    i, old_dim, diff, old_dim, tail)
+        if abs(float(diff)) > 5e-2 or abs(float(tail)) < 1e-4:
+            ok = False
+    if ok:
+        logger.info("VERIFY OK — widened model reproduces source behaviour on first %d dims.",
+                    old_dim)
+    else:
+        logger.warning("VERIFY MISMATCH — differences above float tolerance; check script.")
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Widen ColBERT final projection head.")
+    p.add_argument("--src", required=True, type=Path,
+                   help="Source ColBERT checkpoint directory.")
+    p.add_argument("--dst", required=True, type=Path,
+                   help="Destination directory (created if missing).")
+    p.add_argument("--new-dim", required=True, type=int,
+                   help="New out_features for the final Dense layer.")
+    p.add_argument("--verify", action="store_true",
+                   help="Load both models and check the first old_dim dims match.")
+    args = p.parse_args()
+
+    if not args.src.is_dir():
+        raise FileNotFoundError(f"--src {args.src} is not a directory")
+    args.dst.mkdir(parents=True, exist_ok=True)
+
+    final_dense_src = _find_final_dense(args.src)
+    final_dense_rel = final_dense_src.relative_to(args.src)
+    logger.info("Final Dense in source: %s (will widen)", final_dense_rel)
+
+    # Hardlink everything except the final Dense dir.
+    skip = {str(final_dense_rel)}
+    _hardlink_tree(args.src, args.dst, skip=skip)
+
+    # Widen the final Dense into the destination.
+    final_dense_dst = args.dst / final_dense_rel
+    if final_dense_dst.exists():
+        shutil.rmtree(final_dense_dst)
+    _widen_dense(final_dense_src, final_dense_dst, args.new_dim)
+
+    logger.info("Wrote widened checkpoint to %s", args.dst)
+    logger.info("Backbone + tokenizer + other Dense layers are hardlinks to source "
+                "(zero extra disk).")
+
+    if args.verify:
+        verify(args.src, args.dst)
+
+
+if __name__ == "__main__":
+    main()