Upload folder using huggingface_hub

.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
+sample_page.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,24 @@
 ---
-title: V Splade Document Retrieval
-emoji: 🏆
-colorFrom: yellow
-colorTo: green
+title: V-SPLADE Quality Document Retrieval
+emoji: 🔍
+colorFrom: pink
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.19.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Visual document retrieval via sparse lexical vectors
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# V-SPLADE Quality — Visual Document Retrieval
+
+This Space demonstrates **V-SPLADE Quality** ([naver/v-splade-quality](https://huggingface.co/naver/v-splade-quality)),
+a 0.25B inference-free sparse retriever for visual document retrieval.
+
+Upload a document page image and enter text queries to see:
+- The **top activated vocabulary tokens** — the sparse lexical representation of the image
+- **Similarity scores** between each query and the document, with the top contributing tokens
+
+V-SPLADE encodes document pages directly into sparse vocabulary vectors without OCR or captioning,
+enabling retrieval 20× faster than caption-based pipelines.

app.py

@@ -0,0 +1,190 @@
+"""V-SPLADE Quality — Visual Document Retrieval Demo.
+
+Upload a document page image and enter text queries to see:
+  1. The top activated vocabulary tokens (sparse representation)
+  2. Similarity scores between each query and the document image
+"""
+
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST be first (before torch)
+import torch
+import gradio as gr
+from PIL import Image
+from transformers import AutoProcessor
+
+# Make the v-splade train/models package importable
+import sys
+from pathlib import Path
+
+_TRAIN_DIR = Path(__file__).resolve().parent / "train"
+if str(_TRAIN_DIR) not in sys.path:
+    sys.path.insert(0, str(_TRAIN_DIR))
+
+from models import build_model
+
+MODEL_ID = "naver/v-splade-quality"
+
+
+# ── Module-scope model load (eager, as required by ZeroGPU) ──────────────────
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = build_model(MODEL_ID, mode="inference_only", dtype=torch.bfloat16).to("cuda")
+# Pre-build the Li-LSR vocab lookup table so the first query is fast.
+model.query_encoder.to("cuda", dtype=torch.bfloat16)
+model.query_encoder.build_lookup_table()
+tokenizer = processor.tokenizer
+
+
+@spaces.GPU(duration=60)
+def retrieve(image: Image.Image, queries: str, topk: int = 15) -> tuple:
+    """Encode a document image and score it against text queries.
+
+    Args:
+        image: A document page image (PNG/JPEG).
+        queries: Newline-separated text queries to score against the image.
+        topk: Number of top activated vocabulary tokens to display.
+
+    Returns:
+        A tuple of (top tokens HTML, query scores HTML, sparse stats text).
+    """
+    if image is None:
+        return "<p style='color:red'>Please upload a document image.</p>", "", ""
+    if not queries.strip():
+        return "", "<p style='color:red'>Please enter at least one query.</p>", ""
+
+    # ── Encode image → sparse embedding ──────────────────────────────────
+    chat = [{"role": "user",
+             "content": [{"type": "image"}, {"type": "text", "text": ""}]}]
+    prompt = processor.apply_chat_template(chat, add_generation_prompt=True)
+    inputs = processor(images=[image.convert("RGB")], text=prompt,
+                       return_tensors="pt")
+    inputs = {k: v.to("cuda") if torch.is_tensor(v) else v
+              for k, v in inputs.items()}
+    if "pixel_values" in inputs:
+        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
+    doc_vec = model.encode_passage(**inputs)[0].cpu()
+
+    nnz = int((doc_vec > 0).sum())
+    max_val = float(doc_vec.max())
+
+    # ── Top-k activated vocabulary tokens ────────────────────────────────
+    top_w, top_ids = torch.topk(doc_vec.float(), k=min(topk, doc_vec.shape[0]))
+    tokens_html = "<table style='width:100%; border-collapse:collapse;'>"
+    tokens_html += "<tr style='background:#f0f0f0;'><th style='padding:6px; text-align:left;'>Rank</th><th style='padding:6px; text-align:left;'>Token</th><th style='padding:6px; text-align:right;'>Weight</th></tr>"
+    for rank, (idx, w) in enumerate(zip(top_ids, top_w), 1):
+        tok_str = tokenizer.decode([int(idx)]).strip() or f"<id:{int(idx)}>"
+        bar_width = max(2, int(float(w) / max_val * 200))
+        tokens_html += (
+            f"<tr><td style='padding:4px;'>{rank}</td>"
+            f"<td style='padding:4px; font-family:monospace;'>{tok_str}</td>"
+            f"<td style='padding:4px; text-align:right;'>"
+            f"<div style='display:flex; align-items:center; gap:8px;'>"
+            f"<div style='background:linear-gradient(90deg, #ff6b9d, #8b5cf6); width:{bar_width}px; height:16px; border-radius:3px;'></div>"
+            f"<span>{float(w):.4f}</span></div></td></tr>"
+        )
+    tokens_html += "</table>"
+
+    # ── Query similarity scores ──────────────────────────────────────────
+    query_list = [q.strip() for q in queries.strip().split("\n") if q.strip()]
+    scores_html = "<table style='width:100%; border-collapse:collapse;'>"
+    scores_html += "<tr style='background:#f0f0f0;'><th style='padding:6px; text-align:left;'>Query</th><th style='padding:6px; text-align:right;'>Score</th><th style='padding:6px;'>Top matching tokens</th></tr>"
+
+    max_score = 0.0
+    results = []
+    for q in query_list:
+        tok = tokenizer(q, return_tensors="pt", add_special_tokens=False)
+        q_vec = model.encode_query(
+            tok["input_ids"].to("cuda"),
+            tok["attention_mask"].to("cuda"),
+        )[0].cpu()
+        score = float((q_vec.float() * doc_vec.float()).sum())
+        max_score = max(max_score, abs(score))
+        # Top contributing tokens
+        contrib = (q_vec.float() * doc_vec.float()).cpu()
+        top_cw, top_cids = torch.topk(contrib, k=5)
+        contribs = ", ".join(
+            f"{tokenizer.decode([int(i)]).strip()}({float(w):.3f})"
+            for i, w in zip(top_cids, top_cw) if float(w) > 0
+        )
+        results.append((q, score, contribs))
+
+    for q, score, contribs in results:
+        bar_width = max(2, int(abs(score) / max(max_score, 1e-6) * 200))
+        color = "#22c55e" if score > 0 else "#ef4444"
+        scores_html += (
+            f"<tr><td style='padding:4px;'>{q}</td>"
+            f"<td style='padding:4px; text-align:right;'><b>{score:.4f}</b></td>"
+            f"<td style='padding:4px; font-size:0.9em; font-family:monospace;'>{contribs or '—'}</td></tr>"
+        )
+    scores_html += "</table>"
+
+    stats = (
+        f"Sparse vector: vocab_size={doc_vec.shape[0]}, nnz={nnz} "
+        f"({nnz/doc_vec.shape[0]*100:.1f}% active), max={max_val:.4f}"
+    )
+
+    return tokens_html, scores_html, stats
+
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown(
+        "# 🔍 V-SPLADE Quality — Visual Document Retrieval\n"
+        "Upload a document page image and enter text queries to see the "
+        "sparse lexical representation and similarity scores in real-time. "
+        "[Model card](https://huggingface.co/naver/v-splade-quality) · "
+        "[Paper](https://arxiv.org/abs/2605.30917) · "
+        "[Code](https://github.com/naver/v-splade)"
+    )
+
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(
+                    label="Document page", type="pil", height=400
+                )
+                queries_input = gr.Textbox(
+                    label="Queries (one per line)",
+                    value="financial summary\nmarket share\ncarbon emissions",
+                    lines=5,
+                    placeholder="Enter one query per line…",
+                )
+                run_btn = gr.Button("Retrieve", variant="primary")
+            with gr.Column(scale=1):
+                stats_output = gr.Textbox(label="Sparse vector stats", interactive=False)
+                tokens_output = gr.HTML(label="Top activated vocabulary tokens")
+                scores_output = gr.HTML(label="Query similarity scores")
+
+        with gr.Accordion("Advanced settings", open=False):
+            topk_slider = gr.Slider(
+                minimum=5, maximum=50, value=15, step=1,
+                label="Top-k tokens to display",
+            )
+
+        gr.Examples(
+            examples=[
+                ["sample_page.png", "financial summary\nmarket share\ncarbon emissions", 15],
+                ["sample_page.png", "annual revenue 2024\nboard of directors", 15],
+                ["sample_page.png", "risk factors\nstrategic priorities", 20],
+            ],
+            inputs=[image_input, queries_input, topk_slider],
+            outputs=[tokens_output, scores_output, stats_output],
+            fn=retrieve,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    run_btn.click(
+        fn=retrieve,
+        inputs=[image_input, queries_input, topk_slider],
+        outputs=[tokens_output, scores_output, stats_output],
+        api_name="retrieve",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.8.0
+torchvision
+transformers==4.57.6
+colpali-engine==0.3.13
+peft
+accelerate
+safetensors
+Pillow
+numpy
+scipy
+einops
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt28-cu128-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl

train/models/__init__.py

@@ -0,0 +1,142 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+V_SPLADE modular components.
+
+A V_SPLADE retriever is composed of:
+    Encoder + Pooling + SparseHead + (BOW | Li-LSR) QueryEncoder + Losses
+"""
+
+from pathlib import Path
+
+import torch
+
+from models.model import UnifiedRetriever, RetrievalOutput, compute_logits
+from models.encoder import EncoderType, build_encoder
+from models.pooling import PoolingType, Pooling
+from models.head import HeadType, build_head, SparseHead
+from models.query_encoder import (
+    QueryEncoderType,
+    build_query_encoder,
+    BOWQueryEncoder,
+    InferenceFreeQueryEncoder,
+)
+from models.losses import FLOPSLoss, NCELoss, CaptionPushUpLoss
+
+
+DEFAULT_POOLING = {
+    "vbert": "max",
+}
+
+
+def build_model(
+    path: str = None,
+    mode: str = "inference_only",
+    *,
+    encoder_type: str = "vbert",
+    head_type: str = "sparse",
+    query_encoder_type: str = "li_lsr",
+    pooling_type: str = None,
+    query_lsr_lora_r: int = 0,
+    query_lsr_activation: str = "softplus",
+    dtype: torch.dtype = torch.bfloat16,
+    **kwargs,
+) -> UnifiedRetriever:
+    """Factory: build a V-SPLADE retriever in one of two modes.
+
+    ``mode='inference_only'`` (default):
+        ``path`` is a V-SPLADE HF export directory containing
+        ``model.safetensors`` + ``config.json``. The retriever is constructed
+        as an empty shell and every weight (backbone + SPLADE head + Li-LSR
+        query head) is dispatched from the export in a single pass. No base
+        model download, no LoRA wrapping.
+
+    ``mode='from_scratch'``:
+        ``path`` is the base BiModernVBert backbone directory (e.g. the
+        canonical ``ModernVBERT/modernvbert`` checkpoint). The retriever is
+        built for training — encoder/LM-head LoRA, fresh query head,
+        loss/regularizer hooks. Extra ``**kwargs`` are forwarded to
+        :class:`UnifiedRetriever`.
+    """
+    if mode == "inference_only":
+        if path is None:
+            raise ValueError("inference_only mode requires path= to the HF export dir")
+        model = UnifiedRetriever.from_hf_export(
+            path,
+            query_lsr_activation=query_lsr_activation,
+            dtype=dtype,
+        )
+        load_hf_export(model, path, dtype=dtype)
+        return model
+
+    if mode == "from_scratch":
+        if pooling_type is None:
+            pooling_type = DEFAULT_POOLING.get(encoder_type, "max")
+        if path is not None:
+            kwargs.setdefault("model_name", path)
+        return UnifiedRetriever(
+            encoder_type=encoder_type,
+            pooling_type=pooling_type,
+            head_type=head_type,
+            query_encoder_type=query_encoder_type,
+            query_lsr_lora_r=query_lsr_lora_r,
+            query_lsr_activation=query_lsr_activation,
+            **kwargs,
+        )
+
+    raise ValueError(f"Unknown mode: {mode!r}. Choose 'inference_only' or 'from_scratch'.")
+
+
+def _resolve_export_file(hf_dir: str, filename: str) -> str:
+    """Resolve a file from a V-SPLADE HF export given as a local dir or Hub id.
+
+    Returns a local filesystem path: the file under ``hf_dir`` if it exists,
+    otherwise the result of downloading ``filename`` from the Hub repo
+    ``hf_dir`` (cached by huggingface_hub).
+    """
+    local = Path(hf_dir) / filename
+    if local.is_file():
+        return str(local)
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(repo_id=str(hf_dir), filename=filename)
+
+
+def load_hf_export(model: UnifiedRetriever, hf_dir: str,
+                   dtype: torch.dtype = torch.bfloat16) -> None:
+    """Load a V-SPLADE HF export ``model.safetensors`` into ``model``.
+
+    Dispatches the export's three logical slices to the right sub-modules
+    by stripping the training-wrapper prefix:
+
+        encoder.encoder.model.*   -> model.encoder.encoder.model.*
+        encoder.mlm_head.*        -> model.encoder.mlm_head.*
+        query_encoder.*           -> model.query_encoder.*
+
+    ``hf_dir`` may be a local directory or a HuggingFace Hub repo id; in the
+    latter case ``model.safetensors`` is downloaded automatically.
+
+    Raises if any safetensors tensor is not consumed by these three slices.
+    """
+    from safetensors.torch import load_file
+
+    full_sd = load_file(_resolve_export_file(hf_dir, "model.safetensors"))
+
+    dispatch = [
+        (model.encoder,       "encoder."),
+        (model.query_encoder, "query_encoder."),
+    ]
+    consumed = set()
+    for module, prefix in dispatch:
+        slice_sd = {k[len(prefix):]: v.to(dtype)
+                    for k, v in full_sd.items() if k.startswith(prefix)}
+        module.load_state_dict(slice_sd, strict=False)
+        consumed.update(prefix + k for k in slice_sd)
+
+    leftover = set(full_sd) - consumed
+    if leftover:
+        raise RuntimeError(
+            f"{len(leftover)} tensor(s) in {hf_dir}/model.safetensors were not "
+            f"dispatched to any sub-module. First few: {sorted(leftover)[:3]}"
+        )

train/models/convert.py

@@ -0,0 +1,214 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+Backbone conversion: upstream ModernVBERT -> V-SPLADE-compatible layout.
+
+The public ``ModernVBERT/modernvbert`` checkpoint stores its embeddings/LM-head
+in a layout the decoupled-embedding model can't load (combined 50408 embedding,
+plain LM head, double-nested vision keys, no decoupled-vocab config). This
+module re-packages it.
+
+It is used two ways:
+  * ``ensure_compatible_backbone(ref)`` — called automatically by the training
+    loader; converts on the fly and caches, so ``from_scratch`` "just works"
+    with the upstream Hub id.
+  * ``scripts/convert_modernvbert_backbone.py`` — a thin CLI over the same code.
+
+Verified transform (every tensor reproduced bit-identically from upstream):
+
+  tok_embeddings.weight                            <- upstream[:50368]
+  tok_embeddings.additional_embedding.weight       <- upstream[50368:50408]
+  lm_head.decoder.weight / .bias                   <- upstream lm_head.weight/bias[:50368]
+  additional_fc.weight                             <- upstream lm_head.weight[50368:50408]
+  lm_head.head.dense.weight                        <- upstream projection_head.dense.weight
+  lm_head.head.norm.weight                         <- upstream projection_head.norm.weight
+  model.vision_model.X                             <- upstream model.vision_model.vision_model.X
+  model.connector.modality_projection.proj.weight  <- upstream ...modality_projection.weight
+  (all other model.text_model.* keys copied unchanged)
+"""
+
+import json
+import os
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+MAIN_VOCAB = 50368
+FULL_VOCAB = 50408
+ADDITIONAL_VOCAB = FULL_VOCAB - MAIN_VOCAB   # 40
+
+UPSTREAM_REPO = "ModernVBERT/modernvbert"
+TOKENIZER_REPO = "jhu-clsp/ettin-encoder-150m"
+
+PREPROCESSOR_CONFIG = {
+    "do_convert_rgb": True, "do_image_splitting": True, "do_normalize": True,
+    "do_pad": True, "do_rescale": True, "do_resize": True,
+    "image_mean": [0.5, 0.5, 0.5], "image_std": [0.5, 0.5, 0.5],
+    "image_processor_type": "Idefics3ImageProcessor",
+    "processor_class": "Idefics3Processor",
+    "max_image_size": {"longest_edge": 512}, "resample": 1,
+    "rescale_factor": 0.00392156862745098, "size": {"longest_edge": 2048},
+}
+PROCESSOR_CONFIG = {"image_seq_len": 64, "processor_class": "Idefics3Processor"}
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  Core transform
+# ──────────────────────────────────────────────────────────────────────────────
+def transform_state_dict(u: dict) -> dict:
+    """Apply the verified upstream -> V-SPLADE backbone weight recipe."""
+    out = {}
+    tok = u["model.text_model.embeddings.tok_embeddings.weight"]
+    assert tok.shape[0] == FULL_VOCAB, f"unexpected embedding rows: {tuple(tok.shape)}"
+    out["model.text_model.embeddings.tok_embeddings.weight"] = tok[:MAIN_VOCAB].clone()
+    out["model.text_model.embeddings.tok_embeddings.additional_embedding.weight"] = \
+        tok[MAIN_VOCAB:FULL_VOCAB].clone()
+
+    lw, lb = u["lm_head.weight"], u["lm_head.bias"]
+    out["lm_head.decoder.weight"] = lw[:MAIN_VOCAB].clone()
+    out["lm_head.decoder.bias"]   = lb[:MAIN_VOCAB].clone()
+    out["additional_fc.weight"]   = lw[MAIN_VOCAB:FULL_VOCAB].clone()
+
+    out["lm_head.head.dense.weight"] = u["projection_head.dense.weight"].clone()
+    out["lm_head.head.norm.weight"]  = u["projection_head.norm.weight"].clone()
+
+    out["model.connector.modality_projection.proj.weight"] = \
+        u["model.connector.modality_projection.weight"].clone()
+
+    vp = "model.vision_model.vision_model."
+    for k, v in u.items():
+        if k.startswith(vp):
+            out["model.vision_model." + k[len(vp):]] = v.clone()
+
+    for k, v in u.items():
+        if k.startswith("model.text_model.") and "tok_embeddings" not in k:
+            out[k] = v.clone()
+    return out
+
+
+def build_config(u_cfg: dict) -> dict:
+    """Patch the upstream config into the decoupled-embedding V-SPLADE layout."""
+    cfg = json.loads(json.dumps(u_cfg))
+    cfg["vocab_size"] = MAIN_VOCAB
+    cfg["additional_vocab_size"] = ADDITIONAL_VOCAB
+    cfg["freeze_config"] = {
+        "freeze_lm_head": True, "freeze_text_layers": True, "freeze_vision_layers": True,
+    }
+    cfg.setdefault("architectures", ["BiModernVBert"])
+    if "text_config" in cfg:
+        cfg["text_config"]["vocab_size"] = MAIN_VOCAB
+    return cfg
+
+
+def is_compatible_config(cfg: dict) -> bool:
+    """True if the config already uses the decoupled-embedding V-SPLADE layout."""
+    return cfg.get("freeze_config") is not None and cfg.get("additional_vocab_size") is not None
+
+
+# ─��────────────────────────────────────────────────────────────────────────────
+#  Conversion + auto-ensure
+# ──────────────────────────────────────────────────────────────────────────────
+def _load_config(ref: str) -> dict:
+    if os.path.isdir(ref):
+        return json.load(open(os.path.join(ref, "config.json")))
+    from huggingface_hub import hf_hub_download
+    return json.load(open(hf_hub_download(ref, "config.json")))
+
+
+def convert_backbone(ref: str, out_dir, tokenizer_repo: str = TOKENIZER_REPO,
+                     with_tokenizer: bool = True) -> str:
+    """Convert backbone ``ref`` (Hub id or local dir) into ``out_dir``. Returns out_dir."""
+    out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
+
+    if os.path.isdir(ref):
+        sd_path = os.path.join(ref, "model.safetensors")
+        cfg = json.load(open(os.path.join(ref, "config.json")))
+    else:
+        from huggingface_hub import hf_hub_download
+        sd_path = hf_hub_download(ref, "model.safetensors")
+        cfg = json.load(open(hf_hub_download(ref, "config.json")))
+
+    save_file(transform_state_dict(load_file(sd_path)),
+              str(out / "model.safetensors"), metadata={"format": "pt"})
+    json.dump(build_config(cfg), open(out / "config.json", "w"), indent=2)
+
+    if with_tokenizer:
+        from huggingface_hub import hf_hub_download
+        for fn in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
+            try:
+                shutil.copy2(hf_hub_download(tokenizer_repo, fn), out / fn)
+            except Exception:
+                pass
+        json.dump(PREPROCESSOR_CONFIG, open(out / "preprocessor_config.json", "w"), indent=2)
+        json.dump(PROCESSOR_CONFIG, open(out / "processor_config.json", "w"), indent=2)
+    return str(out)
+
+
+def _cache_root() -> Path:
+    root = os.environ.get("VSPLADE_BACKBONE_CACHE")
+    if root:
+        return Path(root)
+    return Path.home() / ".cache" / "v-splade" / "backbones"
+
+
+def ensure_compatible_backbone(ref: str, tokenizer_repo: str = TOKENIZER_REPO,
+                               verbose: bool = True) -> str:
+    """Return a local path to a V-SPLADE-compatible backbone for ``ref``.
+
+    If ``ref`` already uses the decoupled layout, it is returned unchanged
+    (``from_pretrained`` will download a Hub id as usual). Otherwise the upstream
+    checkpoint is converted once into a cache directory and that path is returned,
+    so ``from_scratch`` training works directly from the raw upstream Hub id.
+    """
+    try:
+        cfg = _load_config(ref)
+    except Exception:
+        return ref   # can't introspect (offline/unknown) — let from_pretrained handle it
+    if is_compatible_config(cfg):
+        return ref
+
+    out = _cache_root() / ref.replace("/", "__")
+    if (out / "model.safetensors").is_file() and (out / "config.json").is_file():
+        if verbose:
+            print(f"[convert] using cached converted backbone: {out}")
+        return str(out)
+    if verbose:
+        print(f"[convert] '{ref}' is an upstream-layout backbone; "
+              f"converting once -> {out}")
+    return convert_backbone(ref, out, tokenizer_repo=tokenizer_repo)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  Double-check (used by the CLI)
+# ──────────────────────────────────────────────────────────────────────────────
+def double_check(out_dir, ref_dir) -> bool:
+    out_dir, ref_dir = Path(out_dir), Path(ref_dir)
+    print(f"\n[verify] comparing {out_dir}  vs  reference {ref_dir}")
+    o = load_file(out_dir / "model.safetensors")
+    r = load_file(ref_dir / "model.safetensors")
+    ok = True
+    if set(o) != set(r):
+        ok = False
+        print(f"  [FAIL] key sets differ "
+              f"(only_out={sorted(set(o)-set(r))[:3]} only_ref={sorted(set(r)-set(o))[:3]})")
+    else:
+        print(f"  [ok] key sets match ({len(o)} tensors)")
+    mismatched = [k for k in set(o) & set(r)
+                  if o[k].shape != r[k].shape or not torch.equal(o[k].float(), r[k].float())]
+    if mismatched:
+        ok = False
+        print(f"  [FAIL] {len(mismatched)} tensor(s) differ, e.g. {mismatched[:5]}")
+    else:
+        print(f"  [ok] all {len(set(o) & set(r))} shared tensors bit-identical")
+    oc, rc = json.load(open(out_dir / "config.json")), json.load(open(ref_dir / "config.json"))
+    for f in ["vocab_size", "additional_vocab_size", "freeze_config"]:
+        if oc.get(f) != rc.get(f):
+            ok = False; print(f"  [FAIL] config.{f}: {oc.get(f)} != {rc.get(f)}")
+        else:
+            print(f"  [ok] config.{f} == {oc.get(f)}")
+    print(f"\n[verify] {'PASSED' if ok else 'FAILED'}.")
+    return ok

train/models/encoder.py

@@ -0,0 +1,294 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+Encoder Layer — backbone model that produces hidden states.
+
+V_SPLADE uses the VBert (BiModernVBERT) encoder as its sole backbone.
+The encoder exposes a unified API:
+    encode_passage(inputs) -> (hidden_states, attention_mask)
+    encode_text(inputs)    -> (hidden_states, attention_mask)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from enum import Enum
+from typing import Optional, Tuple
+
+
+class EncoderType(Enum):
+    VBERT = "vbert"
+
+
+# --------------------------------------------------------------
+# Abstract Base
+# --------------------------------------------------------------
+
+class BaseEncoder(nn.Module):
+    """Abstract encoder base.
+
+    Unified API:
+        encode_passage(inputs) -> (hidden_states, attention_mask)
+        encode_text(inputs)    -> (hidden_states, attention_mask)
+    """
+
+    vocab_size: int = 0
+    hidden_size: int = 0
+
+    def encode_passage(self, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError
+
+    def encode_text(self, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        pass
+
+    def get_text_embeddings(self) -> Optional[nn.Embedding]:
+        """Return the text embedding layer (for query-encoder initialization)."""
+        return None
+
+
+# --------------------------------------------------------------
+# MLM head used by VBert
+# --------------------------------------------------------------
+
+class ModernVBertMLMHead(nn.Module):
+    """MLM head: dense(768->768) -> GELU -> LayerNorm(768) -> decoder(768->50368)."""
+
+    def __init__(self, hidden_size: int = 768, vocab_size: int = 50368):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.norm = nn.LayerNorm(hidden_size)
+        self.decoder = nn.Linear(hidden_size, vocab_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        h = self.dense(hidden_states)
+        h = F.gelu(h)
+        h = self.norm(h)
+        h = self.decoder(h)
+        return h
+
+    @classmethod
+    def from_safetensors(cls, safetensors_path: str, **kwargs):
+        from safetensors import safe_open
+        head = cls(**kwargs)
+        with safe_open(safetensors_path, framework="pt") as f:
+            head.dense.weight.data.copy_(f.get_tensor("lm_head.head.dense.weight"))
+            head.norm.weight.data.copy_(f.get_tensor("lm_head.head.norm.weight"))
+            head.decoder.weight.data.copy_(f.get_tensor("lm_head.decoder.weight"))
+            head.decoder.bias.data.copy_(f.get_tensor("lm_head.decoder.bias"))
+        return head
+
+
+# --------------------------------------------------------------
+# VBert Encoder (BiModernVBERT)
+# --------------------------------------------------------------
+
+class VBertEncoder(BaseEncoder):
+    """BiModernVBERT encoder + external MLM head, with optional LoRA."""
+
+    def __init__(
+        self,
+        model_name: str = "ModernVBERT/bimodernvbert",
+        lm_head_model: str = "ModernVBERT/ModernVBERT",
+        lm_head_lora_r: int = 32,
+        encoder_lora_r: int = 32,
+        lm_head_full: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        from peft import LoraConfig, get_peft_model
+        from models.convert import ensure_compatible_backbone
+
+        # 0. Auto-convert the backbone if it uses the upstream ModernVBERT layout.
+        #    Compatible backbones (local or Hub) pass through unchanged; the raw
+        #    upstream checkpoint is downloaded + converted once (cached) so that
+        #    from_scratch training works directly from the Hub id.
+        model_name = ensure_compatible_backbone(model_name)
+        lm_head_model = ensure_compatible_backbone(lm_head_model) if lm_head_model else model_name
+
+        # 1. Load encoder backbone.
+        model_cls = self._resolve_model_cls(model_name)
+        self.encoder = model_cls.from_pretrained(model_name, dtype=torch.bfloat16)
+
+        # Disable compiled_mlp - FX tracing in gradient_checkpointing traces
+        # both branches of the if/else in ModernBertEncoderLayer.forward(),
+        # hitting compiled_mlp even when reference_compile is None/False.
+        def _set_reference_compile_false(module):
+            if hasattr(module, "config") and hasattr(module.config, "reference_compile"):
+                module.config.reference_compile = False
+        for m in self.encoder.modules():
+            _set_reference_compile_false(m)
+
+        # 2. Merge any existing LoRA adapters into base weights.
+        has_lora = any("lora" in k for k in self.encoder.state_dict().keys())
+        if has_lora:
+            from peft.tuners.lora.layer import Linear as LoraLinear
+            for _, mod in self.encoder.named_modules():
+                if isinstance(mod, LoraLinear) and hasattr(mod, "merge"):
+                    mod.merge()
+
+        # 3. Apply a fresh full LoRA on encoder (all layers: attn + mlp).
+        if encoder_lora_r > 0:
+            self.encoder.model.text_model = get_peft_model(
+                self.encoder.model.text_model,
+                LoraConfig(
+                    r=encoder_lora_r, lora_alpha=encoder_lora_r,
+                    target_modules=["Wqkv", "Wo", "Wi"],
+                    bias="none",
+                ),
+            )
+
+        # 4. Load MLM head - from same model dir or separate model.
+        import os as _os
+        encoder_sf = _os.path.join(model_name, "model.safetensors")
+        has_lm_head_in_encoder = False
+        if _os.path.isfile(encoder_sf):
+            from safetensors import safe_open as _safe_open
+            with _safe_open(encoder_sf, framework="pt") as _f:
+                has_lm_head_in_encoder = any("lm_head" in k for k in _f.keys())
+        if has_lm_head_in_encoder:
+            self.mlm_head = ModernVBertMLMHead.from_safetensors(
+                encoder_sf, hidden_size=768, vocab_size=50368,
+            ).to(torch.bfloat16)
+        else:
+            safetensors_path = self._find_safetensors(lm_head_model)
+            self.mlm_head = ModernVBertMLMHead.from_safetensors(
+                safetensors_path, hidden_size=768, vocab_size=50368,
+            ).to(torch.bfloat16)
+
+        # 5. Apply LoRA to MLM head (dense + decoder).
+        if lm_head_lora_r > 0 and not lm_head_full:
+            self.mlm_head = get_peft_model(self.mlm_head, LoraConfig(
+                r=lm_head_lora_r, lora_alpha=lm_head_lora_r,
+                target_modules=["dense", "decoder"], bias="none",
+            ))
+
+        self.vocab_size = 50368
+        self.hidden_size = 768
+
+        # Freeze base weights, keep LoRA trainable.
+        for name, param in self.named_parameters():
+            if "lora" in name.lower():
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+
+        # Optional full-parameter tuning for the MLM head (no LoRA).
+        if lm_head_full:
+            for param in self.mlm_head.parameters():
+                param.requires_grad = True
+
+    @classmethod
+    def from_hf_export(cls, hf_dir: str, dtype: torch.dtype = torch.bfloat16) -> "VBertEncoder":
+        """Build an empty VBertEncoder shell from a V-SPLADE HF export.
+
+        Constructs `BiModernVBert(config)` + `ModernVBertMLMHead(...)` with
+        randomly-initialized weights — the caller is expected to populate them
+        via :func:`models.load_hf_export`. Used by `build_model(mode='inference_only')`.
+        """
+        from colpali_engine.models import BiModernVBert
+        instance = cls.__new__(cls)
+        nn.Module.__init__(instance)
+        config = BiModernVBert.config_class.from_pretrained(hf_dir)
+        instance.encoder  = BiModernVBert(config).to(dtype=dtype)
+        instance.mlm_head = ModernVBertMLMHead(
+            hidden_size=config.hidden_size, vocab_size=50368,
+        ).to(dtype=dtype)
+        instance.vocab_size = 50368
+        instance.hidden_size = config.hidden_size
+        # No grad needed at inference; trainer-side flags are not touched.
+        for p in instance.parameters():
+            p.requires_grad = False
+        return instance
+
+    @staticmethod
+    def _resolve_model_cls(model_name: str):
+        import json, os
+        from colpali_engine.models import BiModernVBert
+
+        config_path = os.path.join(model_name, "config.json")
+        adapter_config_path = os.path.join(model_name, "adapter_config.json")
+
+        if os.path.isfile(adapter_config_path):
+            with open(adapter_config_path) as f:
+                adapter_cfg = json.load(f)
+            base_path = adapter_cfg.get("base_model_name_or_path", "")
+            base_config = os.path.join(base_path, "config.json")
+            if os.path.isfile(base_config):
+                config_path = base_config
+
+        if os.path.isfile(config_path):
+            with open(config_path) as f:
+                cfg = json.load(f)
+            archs = cfg.get("architectures", [])
+            # V_SPLADE only uses the bidirectional encoder variant.
+            if "BiModernVBert" in archs:
+                return BiModernVBert
+
+        return BiModernVBert
+
+    @staticmethod
+    def _find_safetensors(model_name: str) -> str:
+        import os
+        local = os.path.join(model_name, "model.safetensors")
+        if os.path.isfile(local):
+            return local
+        from huggingface_hub import hf_hub_download
+        return hf_hub_download(model_name, "model.safetensors")
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        kwargs = gradient_checkpointing_kwargs or {"use_reentrant": False}
+        text_model = self.encoder.model.text_model
+        if hasattr(text_model, "gradient_checkpointing_enable"):
+            text_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=kwargs)
+
+    def _get_hidden_states(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        kw = dict(input_ids=input_ids, attention_mask=attention_mask)
+        if pixel_values is not None:
+            kw["pixel_values"] = pixel_values
+        if pixel_attention_mask is not None:
+            kw["pixel_attention_mask"] = pixel_attention_mask
+        outputs = self.encoder.model(**kw)
+        return outputs[0]
+
+    def encode_passage(self, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        hidden = self._get_hidden_states(
+            kwargs["input_ids"], kwargs["attention_mask"],
+            kwargs.get("pixel_values"), kwargs.get("pixel_attention_mask"),
+        )
+        return hidden, kwargs["attention_mask"]
+
+    def encode_text(self, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        hidden = self._get_hidden_states(kwargs["input_ids"], kwargs["attention_mask"])
+        return hidden, kwargs["attention_mask"]
+
+    def get_lm_head(self):
+        return self.mlm_head
+
+    def get_text_embeddings(self) -> Optional[nn.Module]:
+        return self.encoder.model.text_model.get_input_embeddings()
+
+    @property
+    def image_token_id(self) -> int:
+        return 50407  # BiModernVBERT <image> token
+
+
+# --------------------------------------------------------------
+# Factory
+# --------------------------------------------------------------
+
+def build_encoder(encoder_type: str, **kwargs) -> BaseEncoder:
+    """Build encoder by type string. V_SPLADE only ships the vbert backbone."""
+    if encoder_type == "vbert":
+        return VBertEncoder(**kwargs)
+    raise ValueError(f"Unknown encoder_type: {encoder_type}. Choose: vbert")

train/models/head.py

@@ -0,0 +1,61 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+Output head — projects encoder representations into the retrieval space.
+
+V_SPLADE uses SparseHead: the encoder's LM head is reused to project
+hidden states to vocab-dim sparse weights via log1p(relu(.)).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from enum import Enum
+from typing import Tuple
+
+
+class HeadType(Enum):
+    DENSE = "dense"
+    SPARSE = "sparse"
+
+
+class DenseHead(nn.Module):
+    """L2-normalized dense embedding."""
+
+    def forward(self, pooled: torch.Tensor) -> torch.Tensor:
+        return F.normalize(pooled, dim=-1)
+
+
+class SparseHead(nn.Module):
+    """Hidden states → LM head → log1p(relu) → sparse vocab-dim weights.
+
+    Returns (h_raw, w_sparse).  The LM head is held as a *reference* (not a
+    registered sub-module) to avoid duplicated state-dict keys, since the
+    encoder already owns the same module.
+    """
+
+    def __init__(self, lm_head: nn.Module, hidden_size: int):
+        super().__init__()
+        object.__setattr__(self, "_lm_head_ref", lm_head)
+        self.scale = hidden_size ** -0.25
+
+    @property
+    def lm_head(self):
+        return self._lm_head_ref
+
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        h = self._lm_head_ref(hidden_states) * self.scale
+        w = torch.log1p(torch.relu(h))
+        return h, w
+
+
+def build_head(head_type: str, lm_head=None, hidden_size: int = 768, **kwargs) -> nn.Module:
+    if head_type == "dense":
+        return DenseHead()
+    if head_type == "sparse":
+        if lm_head is None:
+            raise ValueError("sparse head requires lm_head")
+        return SparseHead(lm_head, hidden_size)
+    raise ValueError(f"Unknown head_type: {head_type}. Choose: dense | sparse")

train/models/losses.py

@@ -0,0 +1,183 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+V_SPLADE loss functions.
+
+- FLOPSLoss:           Expected retrieval cost regularization (vocab-wise L2).
+- NCELoss:             InfoNCE contrastive loss with in-batch + hard negatives.
+- CaptionPushUpLoss:   Caption-gated token supervision loss — biases passage
+                       sparse activations toward the vocab positions activated
+                       by the paired caption.
+- ZipfianPushUpLoss:   Same as ``CaptionPushUpLoss`` with a temperature
+                       applied to the overlap distribution
+                       (``push_focus_tau``): τ<1 hard-mines a few
+                       strongly-weighted tokens; τ=1 recovers
+                       ``CaptionPushUpLoss``.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FLOPSLoss(nn.Module):
+    """Expected retrieval cost: sum_v mean(|w|_v)^2."""
+
+    def forward(self, reps: torch.Tensor) -> torch.Tensor:
+        return torch.sum(torch.mean(torch.abs(reps), dim=0) ** 2)
+
+
+class NCELoss(nn.Module):
+    """InfoNCE contrastive loss with in-batch negatives + optional hard negatives."""
+
+    def __init__(self, temperature: float = 1.0):
+        super().__init__()
+        self.temperature = temperature
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def forward(
+        self,
+        q_reps: torch.Tensor,
+        p_reps: torch.Tensor,
+        hn_reps: torch.Tensor = None,
+    ) -> torch.Tensor:
+        B = q_reps.size(0)
+        device = q_reps.device
+        all_p = torch.cat([p_reps, hn_reps], dim=0) if hn_reps is not None else p_reps
+        logits = (q_reps @ all_p.t()) / self.temperature
+        labels = torch.arange(B, device=device)
+        return self.loss_fn(logits, labels)
+
+
+class CaptionPushUpLoss(nn.Module):
+    """Caption-gated token supervision loss on passage sparse activations.
+
+    Given the passage's sparse weights p_w (post log1p(relu)) and the
+    caption's sparse weights cap_w (detached), an overlap distribution is
+    formed from (p_w + alpha * mean(p_w)) * cap_w and renormalized.  The
+    loss then maximizes the expected log-prob of activation under that
+    distribution, biasing the passage to put mass on the vocab positions
+    activated by the paired caption while remaining sparse elsewhere.
+    """
+
+    def __init__(self, cap_loss_mode: str = "logsigmoid_h", p_mean_alpha: float = 1.0):
+        super().__init__()
+        assert cap_loss_mode in ("logsigmoid_h", "raw_h", "raw_reps")
+        self.cap_loss_mode = cap_loss_mode
+        self.p_mean_alpha = p_mean_alpha
+
+    def forward(
+        self,
+        p_h: torch.Tensor,
+        p_reps: torch.Tensor,
+        cap_sparse: torch.Tensor,
+        overlap_type: str = "passage_mean",
+    ) -> torch.Tensor:
+        cap_w = cap_sparse.detach()
+        p_w = p_reps.detach()
+
+        if overlap_type == "passage_mean":
+            p_mean = p_w.mean(dim=-1, keepdim=True)
+            overlap = (p_w + self.p_mean_alpha * p_mean) * cap_w
+        else:
+            overlap = p_w * cap_w
+
+        overlap = overlap / (overlap.sum(dim=-1, keepdim=True) + 1e-8)
+
+        if self.cap_loss_mode == "raw_h":
+            push_up_target_rep = p_h
+        elif self.cap_loss_mode == "raw_reps":
+            push_up_target_rep = p_reps
+        else:  # logsigmoid_h (default) — BCE-style log P(active)
+            push_up_target_rep = F.logsigmoid(p_h)
+
+        per_sample = -(overlap * push_up_target_rep).sum(dim=-1)
+        return per_sample.mean()
+
+
+class ZipfianPushUpLoss(nn.Module):
+    """Caption-gated token supervision — identical objective to
+    ``CaptionPushUpLoss`` with a temperature on the overlap distribution.
+    The name is a dev-time convenience for distinguishing the tempered
+    variant in experiment configs.
+
+    Applies a per-sample log-space softmax with ``push_focus_tau`` to the
+    overlap distribution:
+
+        o_sharp[v] = overlap[v]^(1/τ) / Σ_v' overlap[v']^(1/τ)
+
+    - ``push_focus_tau < 1``: hard-mines a few strongly-weighted tokens.
+    - ``push_focus_tau = 1``: identical to ``CaptionPushUpLoss``.
+
+    The rest of the push-up pipeline matches ``CaptionPushUpLoss``.
+    """
+
+    def __init__(
+        self,
+        push_focus_tau: float = 1.0,
+        p_mean_alpha: float = 1.0,
+        cap_loss_mode: str = "logsigmoid_h",
+    ):
+        super().__init__()
+        assert cap_loss_mode in ("logsigmoid_h", "raw_h", "raw_reps")
+        self.push_focus_tau = push_focus_tau
+        self.p_mean_alpha = p_mean_alpha
+        self.cap_loss_mode = cap_loss_mode
+
+    def forward(
+        self,
+        p_h: torch.Tensor,
+        p_reps: torch.Tensor,
+        cap_sparse: torch.Tensor,
+        overlap_type: str = "passage_mean",
+    ) -> torch.Tensor:
+        cap_w = cap_sparse.detach()
+        p_w = p_reps.detach()
+
+        if overlap_type == "passage_mean":
+            p_mean = p_w.mean(dim=-1, keepdim=True)
+            overlap = (p_w + self.p_mean_alpha * p_mean) * cap_w
+        else:  # "passage"
+            overlap = p_w * cap_w
+        # overlap is non-negative and zero wherever cap_w == 0.
+
+        # Target sharpening via push_focus_tau (log-space, numerically stable).
+        mask = (overlap > 0)
+        neg_inf = torch.full_like(overlap, float("-inf"))
+        log_overlap = torch.where(mask, torch.log(overlap.clamp(min=1e-30)), neg_inf)
+        log_o_scaled = (1.0 / self.push_focus_tau) * log_overlap
+        lse = torch.logsumexp(log_o_scaled, dim=-1, keepdim=True)
+        o_sharp = torch.exp(log_o_scaled - lse)
+        # Rows with no caption-active tokens → all -inf → exp(nan) → 0.
+        o_sharp = torch.nan_to_num(o_sharp, nan=0.0, posinf=0.0, neginf=0.0)
+
+        if self.cap_loss_mode == "raw_h":
+            push_up_target_rep = p_h
+        elif self.cap_loss_mode == "raw_reps":
+            push_up_target_rep = p_reps
+        else:  # logsigmoid_h (default)
+            push_up_target_rep = F.logsigmoid(p_h)
+
+        per_sample = -(o_sharp * push_up_target_rep).sum(dim=-1)
+        return per_sample.mean()
+
+
+def all_gather_2d_with_grad(local_tensor: torch.Tensor) -> torch.Tensor:
+    """All-gather (B, D) across ranks, keeping gradient on the local slot.
+
+    Used to give each query access to all ranks' passages as in-batch
+    negatives during contrastive training.  Returns the input unchanged
+    when distributed is not initialized or world_size <= 1.
+    """
+    import torch.distributed as dist
+
+    if not dist.is_initialized() or dist.get_world_size() <= 1:
+        return local_tensor
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    gathered = [torch.zeros_like(local_tensor) for _ in range(world_size)]
+    dist.all_gather(gathered, local_tensor)
+    gathered[rank] = local_tensor  # preserve grad on local slot
+    return torch.cat(gathered, dim=0)

train/models/model.py

@@ -0,0 +1,572 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+UnifiedRetriever -- V_SPLADE core model.
+
+Combines:
+    VBert encoder + Sparse head (SPLADE) + BOW query encoder + Losses
+
+Training objective (per forward call):
+    L_total = NCE(q, p, hn) + lambda_p * FLOPS(p) + lambda_cap * CaptionPushUp
+
+The model expects a passage (image) and an optional caption; the query is
+encoded as a Bag-of-Words vocabulary indicator. The sparse passage
+representation is supervised by:
+  - NCE in-batch ranking against the query (and hard negatives, if any)
+  - FLOPS regularization to encourage sparsity
+  - Caption-gated token supervision: pull passage activations toward overlapping
+    caption-vocabulary tokens.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from typing import Optional
+
+from models.encoder import build_encoder
+from models.pooling import Pooling
+from models.head import SparseHead
+from models.query_encoder import (
+    build_query_encoder,
+    BOWQueryEncoder,
+    InferenceFreeQueryEncoder,
+)
+from models.losses import FLOPSLoss, NCELoss, CaptionPushUpLoss, ZipfianPushUpLoss  # noqa: F401
+
+
+def compute_logits(q, p, hn, temperature):
+    """Compute in-batch negative logits and labels.
+
+    Args:
+        q:  (B, V) query sparse reps
+        p:  (B, V) passage sparse reps
+        hn: (B*num_hn, V) hard-negative passage reps, or None
+        temperature: scaling factor for the inner product
+
+    Returns:
+        logits: (B, B [+ B*num_hn]) score matrix
+        labels: (B,) target index for cross-entropy
+    """
+    B = q.size(0)
+    device = q.device
+    all_p = torch.cat([p, hn], dim=0) if hn is not None else p
+    logits = (q @ all_p.t()) / temperature
+    logits = logits.clamp(min=-100, max=100)  # prevent exp() overflow in cross_entropy
+    labels = torch.arange(B, device=device)
+    return logits, labels
+
+
+# Default pooling for V_SPLADE (single backbone): max over sequence positions.
+DEFAULT_POOLING = "max"
+
+
+@dataclass
+class RetrievalOutput:
+    """Container for forward() outputs."""
+    loss: Optional[torch.Tensor] = None
+    rank_loss: Optional[torch.Tensor] = None
+    reg_loss: Optional[torch.Tensor] = None        # total reg (passage FLOPS + caption FLOPS)
+    reg_loss_p: Optional[torch.Tensor] = None      # passage FLOPS only
+    reg_loss_cap: Optional[torch.Tensor] = None    # caption FLOPS only
+    cap_loss: Optional[torch.Tensor] = None        # caption-gated token supervision loss
+    cap_sparse_rank_loss: Optional[torch.Tensor] = None  # caption-gated ranking loss
+    # Sparsity / retrieval-cost diagnostics (no_grad, mean over batch).
+    train_p_nnz: Optional[torch.Tensor] = None     # avg nonzero per passage
+    train_p_max: Optional[torch.Tensor] = None     # avg max|p| per passage
+    train_flops_qd: Optional[torch.Tensor] = None  # sum_v mean|q|_v * mean|p|_v
+    train_cap_bow_killed_pct: Optional[torch.Tensor] = None  # % bow-active tokens zeroed by model
+    query_reps: Optional[torch.Tensor] = None
+    passage_reps: Optional[torch.Tensor] = None
+    hard_neg_reps: Optional[torch.Tensor] = None
+    debug_tokens: Optional[dict] = None
+
+
+class UnifiedRetriever(nn.Module):
+    """V_SPLADE retriever: VBert encoder + SPLADE sparse head + (BOW | Li-LSR) query encoder."""
+
+    def __init__(
+        self,
+        encoder_type: str = "vbert",
+        pooling_type: str = None,
+        head_type: str = "sparse",
+        query_encoder_type: str = "bow",
+        # Model paths
+        model_name: str = "",
+        lm_head_model: str = None,
+        # Training config
+        temperature: float = 1.0,
+        # Caption-gated token supervision
+        cap_weight: float = 0.0,
+        cap_sparse_rank_weight: float = 0.0,
+        cap_loss_mode: str = "logsigmoid_h",
+        overlap_type: str = "passage_mean",
+        p_mean_alpha: float = 1.0,
+        use_zipfian_pushup: bool = False,
+        push_focus_tau: float = 1.0,
+        # Regularization weights
+        reg_weight_p: float = 0.0,
+        reg_weight_cap: float = 0.0,
+        # SPLADE pooling
+        splade_pooling: str = "max",
+        # Encoder-specific kwargs
+        lm_head_lora_r: int = 32,
+        encoder_lora_r: int = 32,
+        lm_head_full: bool = False,
+        # Li-LSR query encoder kwargs
+        query_lsr_lora_r: int = 0,
+        query_lsr_activation: str = "relu",
+        # Misc
+        **kwargs,
+    ):
+        super().__init__()
+
+        # V_SPLADE only supports the sparse head with a single (vbert) backbone.
+        assert encoder_type == "vbert", "V_SPLADE only supports the vbert encoder."
+        assert head_type == "sparse", "V_SPLADE only supports the sparse head."
+
+        if pooling_type is None:
+            pooling_type = DEFAULT_POOLING
+
+        self.encoder_type = encoder_type
+        self.head_type = head_type
+        self.query_encoder_type = query_encoder_type
+        self.splade_pooling = splade_pooling
+
+        # Build encoder.
+        encoder_kwargs = dict(
+            model_name=model_name,
+            lm_head_model=lm_head_model or model_name,
+            lm_head_lora_r=lm_head_lora_r,
+            encoder_lora_r=encoder_lora_r,
+            lm_head_full=lm_head_full,
+            **kwargs,
+        )
+        self.encoder = build_encoder(encoder_type, **encoder_kwargs)
+
+        # Build pooling and sparse head.
+        self.pooling = Pooling(pooling_type)
+        lm_head = self.encoder.get_lm_head()
+        self.head = SparseHead(lm_head, self.encoder.hidden_size)
+
+        self.vocab_size = self.encoder.vocab_size
+        self.hidden_size = self.encoder.hidden_size
+
+        # Build query encoder: BOW (cheapest) or Li-LSR (inference-free learned).
+        assert query_encoder_type in ("bow", "li_lsr"), (
+            "V_SPLADE only supports the BOW or Li-LSR query encoders."
+        )
+        if query_encoder_type == "li_lsr":
+            embed_layer = self.encoder.get_text_embeddings()
+            if embed_layer is not None:
+                embed_weight = embed_layer.weight
+            else:
+                # Fallback: use lm_head weights (tied with embed_tokens in causal LMs).
+                _lm_head = self.encoder.get_lm_head()
+                embed_weight = _lm_head.weight if _lm_head is not None else None
+            self.query_encoder = build_query_encoder(
+                "li_lsr",
+                vocab_size=self.vocab_size,
+                hidden_size=self.hidden_size,
+                embed_weight=embed_weight,
+                lora_r=query_lsr_lora_r,
+                activation=query_lsr_activation,
+            )
+        else:
+            self.query_encoder = build_query_encoder("bow", vocab_size=self.vocab_size)
+
+        # Loss / regularization config.
+        self.temperature = temperature
+        self.cap_weight = cap_weight
+        self.cap_sparse_rank_weight = cap_sparse_rank_weight
+        self.cap_loss_mode = cap_loss_mode
+        self.overlap_type = overlap_type
+        self.reg_weight_p = reg_weight_p
+        self.reg_weight_cap = reg_weight_cap
+        self.loss_fn = nn.CrossEntropyLoss()
+        self.reg_fn = FLOPSLoss()
+        if cap_weight > 0:
+            if use_zipfian_pushup:
+                self.cap_push_up = ZipfianPushUpLoss(
+                    push_focus_tau=push_focus_tau,
+                    p_mean_alpha=p_mean_alpha,
+                    cap_loss_mode=cap_loss_mode,
+                )
+            else:
+                self.cap_push_up = CaptionPushUpLoss(
+                    cap_loss_mode, p_mean_alpha=p_mean_alpha,
+                )
+        else:
+            self.cap_push_up = None
+
+        # Special-token mask for the VBert vocabulary. Prevents tokens such as
+        # [CLS]/[SEP]/[PAD]/[MASK] (and out-of-MLM image tokens) from activating
+        # in the sparse representation.
+        _special_ids = [
+            50280,  # [UNK]
+            50281,  # [CLS]
+            50282,  # [SEP]
+            50283,  # [PAD]
+            50284,  # [MASK]
+            # Additional image/layout tokens beyond MLM head vocab (50368);
+            # kept here for safety in case downstream code changes vocab_size.
+            *range(50368, 50408),
+        ]
+        mask = torch.ones(self.vocab_size)
+        for sid in _special_ids:
+            if sid < self.vocab_size:
+                mask[sid] = 0.0
+        self.register_buffer("special_token_mask", mask, persistent=False)
+
+    @classmethod
+    def from_hf_export(cls, hf_dir: str,
+                       query_lsr_activation: str = "softplus",
+                       dtype: torch.dtype = torch.bfloat16) -> "UnifiedRetriever":
+        """Build an empty V-SPLADE retriever shell from a V-SPLADE HF export.
+
+        Only the module *structure* is created (random weights); call
+        :func:`models.load_hf_export` afterwards to populate every tensor
+        from the export's `model.safetensors` in one pass.
+        """
+        from models.encoder import VBertEncoder
+
+        instance = cls.__new__(cls)
+        nn.Module.__init__(instance)
+        instance.encoder_type      = "vbert"
+        instance.head_type         = "sparse"
+        instance.query_encoder_type = "li_lsr"
+        instance.splade_pooling    = "max"
+        instance.temperature       = 1.0
+        instance.cap_weight        = 0.0
+        instance.cap_sparse_rank_weight = 0.0
+        instance.cap_loss_mode     = "logsigmoid_h"
+        instance.overlap_type      = "passage_mean"
+        instance.reg_weight_p      = 0.0
+        instance.reg_weight_cap    = 0.0
+        instance.cap_push_up       = None
+        instance.loss_fn           = nn.CrossEntropyLoss()
+        instance.reg_fn            = FLOPSLoss()
+
+        # Encoder + sparse head (shares the encoder's LM head via SparseHead).
+        instance.encoder      = VBertEncoder.from_hf_export(hf_dir, dtype=dtype)
+        instance.vocab_size   = instance.encoder.vocab_size
+        instance.hidden_size  = instance.encoder.hidden_size
+        instance.pooling      = Pooling("max")
+        instance.head         = SparseHead(instance.encoder.get_lm_head(),
+                                           instance.encoder.hidden_size)
+
+        # Li-LSR inference-free query encoder (softplus activation, LoRA off).
+        embed_layer = instance.encoder.get_text_embeddings()
+        embed_weight = (embed_layer.weight if embed_layer is not None
+                        else instance.encoder.get_lm_head().weight)
+        instance.query_encoder = build_query_encoder(
+            "li_lsr",
+            vocab_size=instance.vocab_size,
+            hidden_size=instance.hidden_size,
+            embed_weight=embed_weight,
+            lora_r=0,
+            activation=query_lsr_activation,
+        )
+
+        # Special-token mask (same as training).
+        _special_ids = [50280, 50281, 50282, 50283, 50284, *range(50368, 50408)]
+        mask = torch.ones(instance.vocab_size)
+        for sid in _special_ids:
+            if sid < instance.vocab_size:
+                mask[sid] = 0.0
+        instance.register_buffer("special_token_mask", mask, persistent=False)
+
+        for p in instance.parameters():
+            p.requires_grad = False
+        return instance.eval()
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs)
+
+    # ---- Encoding helpers ----------------------------------------------------
+
+    def _encode_passage_sparse(self, **kwargs):
+        """Encode passage -> (h_raw, w_sparse) via encoder + sparse head."""
+        hidden, mask = self.encoder.encode_passage(**kwargs)
+        return self._apply_sparse_head(hidden, mask)
+
+    def _encode_text_sparse(self, **kwargs):
+        """Encode text (e.g. caption) -> (h_raw, w_sparse)."""
+        hidden, mask = self.encoder.encode_text(**kwargs)
+        return self._apply_sparse_head(hidden, mask)
+
+    def _apply_sparse_head(self, hidden_states, attention_mask):
+        """SPLADE sparse head: LM head (per position) -> max/mean pool over sequence."""
+        scale = self.hidden_size ** -0.25
+
+        # max/mean pooling: apply LM head over all positions, then pool.
+        h = self.head.lm_head(hidden_states) * scale  # (B, seq, vocab)
+        w = torch.log1p(torch.relu(h))
+        mask = attention_mask.unsqueeze(-1).to(w.dtype)
+
+        if self.splade_pooling == "max":
+            w_out = (w * mask).max(dim=1).values
+            h_masked = h * mask + (~mask.bool()) * (-1e9)
+            h_out = h_masked.max(dim=1).values
+        else:  # mean
+            seq_len = mask.sum(dim=1, keepdim=True).clamp(min=1)
+            w_out = (w * mask).sum(dim=1) / seq_len.squeeze(1)
+            h_out = (h * mask).sum(dim=1) / seq_len.squeeze(1)
+
+        # Mask special tokens.
+        w_out = w_out * self.special_token_mask.to(w_out.dtype)
+        return h_out, w_out
+
+    def _encode_query(self, input_ids, attention_mask):
+        """Encode query via BOW or Li-LSR (training: forward, eval: lookup)."""
+        if isinstance(self.query_encoder, InferenceFreeQueryEncoder):
+            if self.training:
+                q = self.query_encoder(input_ids, attention_mask)
+            else:
+                q = self.query_encoder.encode_with_lookup(input_ids, attention_mask)
+        else:
+            q = self.query_encoder(input_ids, attention_mask)
+        # Apply special-token mask to the query output.
+        if q.dim() == 2 and q.size(-1) == self.special_token_mask.size(0):
+            q = q * self.special_token_mask.to(q.dtype).to(q.device)
+        return q
+
+    def encode_query(self, input_ids, attention_mask) -> torch.Tensor:
+        """Public API for query encoding."""
+        return self._encode_query(input_ids, attention_mask)
+
+    def encode_passage(
+        self, input_ids=None, attention_mask=None,
+        pixel_values=None, pixel_attention_mask=None,
+    ) -> torch.Tensor:
+        """Public API for passage encoding."""
+        kwargs = {}
+        if input_ids is not None:
+            kwargs["input_ids"] = input_ids
+        if attention_mask is not None:
+            kwargs["attention_mask"] = attention_mask
+        if pixel_values is not None:
+            kwargs["pixel_values"] = pixel_values
+        if pixel_attention_mask is not None:
+            kwargs["pixel_attention_mask"] = pixel_attention_mask
+        _, w = self._encode_passage_sparse(**kwargs)
+        return w
+
+    # ---- Forward -------------------------------------------------------------
+
+    def forward(
+        self,
+        query_input_ids, query_attention_mask,
+        passage_input_ids, passage_attention_mask,
+        passage_pixel_values=None, passage_pixel_attention_mask=None,
+        caption_input_ids=None, caption_attention_mask=None,
+        hard_neg_passage_input_ids=None, hard_neg_passage_attention_mask=None,
+        hard_neg_passage_pixel_values=None, hard_neg_passage_pixel_attention_mask=None,
+        **kwargs,
+    ) -> RetrievalOutput:
+        """V_SPLADE forward pass.
+
+        Computes NCE rank loss + FLOPS regularization + (optional) caption-gated token
+        loss. Multi-vector / self-distillation / dense / multi-backbone paths
+        are out of scope for this code release.
+        """
+        has_caption = caption_input_ids is not None
+        use_cap = (
+            self.cap_sparse_rank_weight > 0
+            or self.cap_weight > 0
+            or self.reg_weight_cap > 0
+        ) and has_caption
+
+        # ---- Query (BOW) ----
+        q_reps = self._encode_query(query_input_ids, query_attention_mask)
+
+        # ---- Positive passage ----
+        p_kwargs = dict(input_ids=passage_input_ids, attention_mask=passage_attention_mask)
+        if passage_pixel_values is not None:
+            p_kwargs["pixel_values"] = passage_pixel_values
+        if passage_pixel_attention_mask is not None:
+            p_kwargs["pixel_attention_mask"] = passage_pixel_attention_mask
+        p_h, p_reps = self._encode_passage_sparse(**p_kwargs)
+
+        # ---- Hard negatives (optional) ----
+        hn_reps = None
+        if hard_neg_passage_input_ids is not None:
+            if hard_neg_passage_pixel_values is not None:
+                # Image hard negatives.
+                hn_kwargs = dict(
+                    input_ids=hard_neg_passage_input_ids,
+                    attention_mask=hard_neg_passage_attention_mask,
+                    pixel_values=hard_neg_passage_pixel_values,
+                    pixel_attention_mask=hard_neg_passage_pixel_attention_mask,
+                )
+                _, hn_reps = self._encode_passage_sparse(**hn_kwargs)
+            else:
+                # Text hard negatives.
+                _, hn_reps = self._encode_text_sparse(
+                    input_ids=hard_neg_passage_input_ids,
+                    attention_mask=hard_neg_passage_attention_mask,
+                )
+
+        # ---- Caption encoding ----
+        cap_sparse = None
+        cap_sparse_raw = None  # kept for diagnostics
+        cap_bow = None
+        if use_cap:
+            _, cap_sparse_raw = self._encode_text_sparse(
+                input_ids=caption_input_ids, attention_mask=caption_attention_mask,
+            )
+            # Use plain caption tokens for BOW if available (so that any instruction
+            # tokens added by a caption prompt are not counted in the BOW mask).
+            bow_ids = kwargs.get("caption_bow_input_ids", caption_input_ids)
+            bow_mask = kwargs.get("caption_bow_attention_mask", caption_attention_mask)
+            cap_bow = BOWQueryEncoder(self.vocab_size).to(q_reps.device)(bow_ids, bow_mask)
+            cap_bow = cap_bow * self.special_token_mask.to(cap_bow.dtype).to(cap_bow.device)
+            cap_sparse = cap_sparse_raw * cap_bow.to(cap_sparse_raw.dtype)
+
+        # ---- dtype alignment (BOW query is float32, passage is bf16) ----
+        if q_reps.dtype != p_reps.dtype:
+            q_reps = q_reps.to(p_reps.dtype)
+
+        # ---- Ranking loss (NCE) — cross-GPU gather for in-batch negatives ----
+        # Symmetric gather across DDP ranks: every gathered query sees
+        # ``world_size × B`` passages as negatives. The local rank's slice
+        # keeps gradient via ``all_gather_2d_with_grad``; other ranks are
+        # detached. Asymmetric gathers (only p) under-train the passage
+        # encoder because each p sees gradients from only ``B`` queries.
+        import torch.distributed as _dist_nce
+        if _dist_nce.is_initialized() and _dist_nce.get_world_size() > 1:
+            from models.losses import all_gather_2d_with_grad as _gather_nce
+            q_reps_g = _gather_nce(q_reps)
+            p_reps_g = _gather_nce(p_reps)
+            hn_reps_g = _gather_nce(hn_reps) if hn_reps is not None else None
+            logits, labels = compute_logits(
+                q_reps_g, p_reps_g, hn_reps_g, self.temperature,
+            )
+        else:
+            logits, labels = compute_logits(
+                q_reps, p_reps, hn_reps, self.temperature,
+            )
+        rank_loss = self.loss_fn(logits, labels)
+
+        # ---- FLOPS regularization ----
+        reg_loss_p = self.reg_weight_p * self.reg_fn(p_reps)
+        if hn_reps is not None:
+            reg_loss_p = reg_loss_p + self.reg_weight_p * self.reg_fn(hn_reps)
+        reg_loss_cap = q_reps.new_tensor(0.0)
+        if cap_sparse is not None:
+            reg_loss_cap = self.reg_weight_cap * self.reg_fn(cap_sparse)
+        reg_loss = reg_loss_p + reg_loss_cap
+
+        # ---- Caption loss ----
+        cap_loss = q_reps.new_tensor(0.0)
+        cap_sparse_rank_loss = q_reps.new_tensor(0.0)
+
+        if use_cap and self.cap_sparse_rank_weight > 0 and cap_sparse is not None:
+            # Cross-GPU gather for caption-side ranking too (parallel to the
+            # passage-side NCE above): caption push-up still uses LOCAL p_reps.
+            import torch.distributed as _dist_capr
+            if _dist_capr.is_initialized() and _dist_capr.get_world_size() > 1:
+                from models.losses import all_gather_2d_with_grad as _gather_capr
+                q_reps_capr = _gather_capr(q_reps)
+                cap_sparse_capr = _gather_capr(cap_sparse)
+            else:
+                q_reps_capr = q_reps
+                cap_sparse_capr = cap_sparse
+            cap_s_logits, cap_s_labels = compute_logits(
+                q_reps_capr, cap_sparse_capr, None, self.temperature,
+            )
+            cap_sparse_rank_loss = self.loss_fn(cap_s_logits, cap_s_labels) * self.cap_sparse_rank_weight
+
+        if use_cap and self.cap_weight > 0 and cap_sparse is not None and self.cap_push_up is not None:
+            cap_loss = self.cap_push_up(p_h, p_reps, cap_sparse, self.overlap_type) * self.cap_weight
+
+        # ---- Total loss ----
+        loss = rank_loss + reg_loss + cap_loss + cap_sparse_rank_loss
+
+        # ---- Diagnostics (no_grad) ----
+        debug_tokens = None
+        train_p_nnz = train_p_max = train_flops_qd = None
+        train_cap_bow_killed_pct = None
+        with torch.no_grad():
+            def _topk_lowk(vec, k=5):
+                """Return (top-k values, indices, low-k nonzero values, indices)."""
+                nz_mask = vec > 0
+                nz_n = int(nz_mask.sum().item())
+                if nz_n == 0:
+                    return (torch.empty(0), torch.empty(0, dtype=torch.long),
+                            torch.empty(0), torch.empty(0, dtype=torch.long))
+                t_k = min(k, vec.size(0))
+                l_k = min(k, nz_n)
+                top_v, top_i = vec.topk(t_k)
+                masked = torch.where(nz_mask, vec, torch.full_like(vec, float('inf')))
+                low_vals_asc, low_idx_asc = masked.topk(l_k, largest=False)
+                return top_v.cpu(), top_i.cpu(), low_vals_asc.cpu(), low_idx_asc.cpu()
+
+            q0, p0 = q_reps[0], p_reps[0]
+            q_top_v, q_top_i, q_low_v, q_low_i = _topk_lowk(q0)
+            p_top_v, p_top_i, p_low_v, p_low_i = _topk_lowk(p0)
+            debug_tokens = {
+                "q_indices": q_top_i, "q_values": q_top_v,
+                "q_low_indices": q_low_i, "q_low_values": q_low_v,
+                "p_indices": p_top_i, "p_values": p_top_v,
+                "p_low_indices": p_low_i, "p_low_values": p_low_v,
+                "q_input_ids": query_input_ids[0].cpu(),
+                "q_attention_mask": query_attention_mask[0].cpu(),
+                "q_nonzero": (q0 > 0).sum().item(),
+                "p_nonzero": (p0 > 0).sum().item(),
+            }
+            # Batch sparsity / FLOPs diagnostics.
+            train_p_nnz = (p_reps > 0).float().sum(-1).mean().detach()
+            train_p_max = p_reps.max(dim=-1).values.mean().detach()
+            train_flops_qd = (q_reps.abs().mean(dim=0) * p_reps.abs().mean(dim=0)).sum().detach()
+
+            if cap_sparse is not None:
+                cap0 = cap_sparse[0]
+                c_top_v, c_top_i, c_low_v, c_low_i = _topk_lowk(cap0)
+                debug_tokens["cap_indices"] = c_top_i
+                debug_tokens["cap_values"] = c_top_v
+                debug_tokens["cap_low_indices"] = c_low_i
+                debug_tokens["cap_low_values"] = c_low_v
+                debug_tokens["cap_nonzero"] = (cap0 > 0).sum().item()
+
+                # Push-up overlap top-5 / low-5.
+                p_w_d = p0.detach()
+                p_mean_d = p_w_d.mean(dim=-1, keepdim=True)
+                alpha = self.cap_push_up.p_mean_alpha if self.cap_push_up is not None else 1.0
+                overlap = (p_w_d + alpha * p_mean_d) * cap0.detach()
+                overlap_norm = overlap / (overlap.sum() + 1e-8)
+                ov_top_v, ov_top_i, ov_low_v, ov_low_i = _topk_lowk(overlap_norm)
+                debug_tokens["pushup_indices"] = ov_top_i
+                debug_tokens["pushup_values"] = ov_top_v
+                debug_tokens["pushup_low_indices"] = ov_low_i
+                debug_tokens["pushup_low_values"] = ov_low_v
+
+                # cap_bow_killed_pct: among bow-active caption tokens, what fraction
+                # did the SPLADE head zero out? (Higher = model treats more cap
+                # tokens as noise.)
+                if cap_bow is not None and cap_sparse_raw is not None:
+                    bow_active = (cap_bow > 0).float()
+                    raw_active = (cap_sparse_raw > 0).float()
+                    killed = bow_active * (1.0 - raw_active)
+                    denom = bow_active.sum(-1).clamp(min=1.0)
+                    train_cap_bow_killed_pct = (killed.sum(-1) / denom).mean().detach()
+
+        return RetrievalOutput(
+            loss=loss,
+            rank_loss=rank_loss,
+            reg_loss=reg_loss,
+            reg_loss_p=reg_loss_p,
+            reg_loss_cap=reg_loss_cap,
+            cap_loss=cap_loss,
+            cap_sparse_rank_loss=cap_sparse_rank_loss,
+            train_p_nnz=train_p_nnz,
+            train_p_max=train_p_max,
+            train_flops_qd=train_flops_qd,
+            train_cap_bow_killed_pct=train_cap_bow_killed_pct,
+            query_reps=q_reps,
+            passage_reps=p_reps,
+            hard_neg_reps=hn_reps,
+            debug_tokens=debug_tokens,
+        )

train/models/pooling.py

@@ -0,0 +1,62 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+Pooling layer: (B, L, D) → (B, D).
+
+V_SPLADE uses MAX pooling (SPLADE-style position-wise max over the
+sequence). Other strategies are provided for completeness.
+"""
+
+import torch
+import torch.nn as nn
+from enum import Enum
+
+
+class PoolingType(Enum):
+    MAX = "max"
+    MEAN = "mean"
+    EOS = "eos"
+    CLS = "cls"
+
+
+class Pooling(nn.Module):
+    """Pool a sequence of hidden states to a single vector."""
+
+    def __init__(self, pooling_type: str):
+        super().__init__()
+        self.pooling_type = PoolingType(pooling_type)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.pooling_type == PoolingType.MAX:
+            return self._max_pool(hidden_states, attention_mask)
+        if self.pooling_type == PoolingType.MEAN:
+            return self._mean_pool(hidden_states, attention_mask)
+        if self.pooling_type == PoolingType.EOS:
+            return self._eos_pool(hidden_states, attention_mask)
+        if self.pooling_type == PoolingType.CLS:
+            return hidden_states[:, 0]
+        raise ValueError(f"Unknown pooling type: {self.pooling_type}")
+
+    @staticmethod
+    def _max_pool(hidden_states, attention_mask):
+        mask = attention_mask.unsqueeze(-1).to(hidden_states.dtype)
+        h_masked = hidden_states * mask + (~mask.bool()) * (-1e9)
+        return h_masked.max(dim=1).values
+
+    @staticmethod
+    def _mean_pool(hidden_states, attention_mask):
+        mask = attention_mask.unsqueeze(-1).to(hidden_states.dtype)
+        seq_len = mask.sum(dim=1, keepdim=True).clamp(min=1)
+        return (hidden_states * mask).sum(dim=1) / seq_len.squeeze(1)
+
+    @staticmethod
+    def _eos_pool(hidden_states, attention_mask):
+        seq_lengths = (attention_mask.sum(dim=1) - 1).clamp(min=0)
+        batch_idx = torch.arange(hidden_states.size(0), device=hidden_states.device)
+        return hidden_states[batch_idx, seq_lengths]

train/models/query_encoder.py

@@ -0,0 +1,159 @@
+# V-SPLADE
+# Copyright (c) 2026-present NAVER Corp.
+# Apache-2.0
+
+"""
+Query encoder strategies.
+
+V_SPLADE quality variant uses an inference-free *Learned Sparse Retriever*
+query encoder (``li_lsr``): a frozen copy of the backbone word embeddings,
+optionally adapted with a low-rank LoRA, followed by a 1-dim projection
+and an activation. At inference time the projection becomes a
+``vocab_size``-entry lookup table, so encoding a query is just an
+``Embedding(input_ids)`` — no transformer forward.
+
+The simpler ``bow`` encoder turns the query into a binary token-presence
+indicator over the vocab. Useful for smoke tests and the cheapest
+inference-time setting.
+"""
+
+import torch
+import torch.nn as nn
+from enum import Enum
+from typing import Optional
+
+
+class QueryEncoderType(Enum):
+    BOW = "bow"
+    LI_LSR = "li_lsr"
+
+
+class InferenceFreeQueryEncoder(nn.Module):
+    """Li-LSR style inference-free sparse query encoder.
+
+    Training:  token_id → embedding (frozen + LoRA) → projection(D→1) →
+               activation → scatter into a vocab-dim vector.
+    Inference: token_id → precomputed lookup table (no forward pass).
+    """
+
+    def __init__(
+        self,
+        embed_weight: torch.Tensor,
+        vocab_size: int,
+        hidden_size: int,
+        lora_r: int = 0,
+        activation: str = "relu",
+    ):
+        super().__init__()
+        assert activation in ("relu", "softplus"), f"Unknown activation: {activation}"
+        self.activation = activation
+
+        # Frozen copy of backbone word embeddings.
+        self.embeddings = nn.Embedding(vocab_size, hidden_size)
+        self.embeddings.weight.data.copy_(embed_weight.detach())
+        self.embeddings.weight.requires_grad = False
+
+        # Optional LoRA on the embedding lookup.
+        self.use_lora = lora_r > 0
+        if self.use_lora:
+            self.lora_A = nn.Parameter(torch.randn(vocab_size, lora_r) * 0.01)
+            self.lora_B = nn.Parameter(torch.zeros(lora_r, hidden_size))
+
+        # Projection: hidden_size → 1 (fully tuned).
+        self.projection = nn.Linear(hidden_size, 1, bias=True)
+
+        self.vocab_size = vocab_size
+        self._lookup_table: Optional[torch.Tensor] = None
+
+    def _activate(self, x: torch.Tensor) -> torch.Tensor:
+        if self.activation == "relu":
+            return torch.log1p(torch.relu(x))
+        return torch.nn.functional.softplus(x)
+
+    def _valid_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return input_ids < self.vocab_size
+
+    def _get_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        safe_ids = input_ids.clamp(max=self.vocab_size - 1)
+        e = self.embeddings(safe_ids)
+        if self.use_lora:
+            e = e + self.lora_A[safe_ids] @ self.lora_B
+        return e
+
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Training forward. Returns (B, vocab_size) sparse vector."""
+        valid = self._valid_mask(input_ids)
+        mask = attention_mask.float() * valid.float()
+        scores = self.projection(self._get_embed(input_ids)).squeeze(-1)
+        scores = self._activate(scores) * mask
+        B = input_ids.size(0)
+        safe_ids = input_ids.clamp(max=self.vocab_size - 1)
+        out = torch.zeros(B, self.vocab_size, device=input_ids.device, dtype=scores.dtype)
+        out.scatter_add_(1, safe_ids, scores)
+        return out
+
+    @torch.no_grad()
+    def build_lookup_table(self) -> torch.Tensor:
+        device = self.embeddings.weight.device
+        all_ids = torch.arange(self.vocab_size, device=device)
+        e = self._get_embed(all_ids.unsqueeze(0))
+        s = self.projection(e).squeeze(-1).squeeze(0)
+        self._lookup_table = self._activate(s)
+        return self._lookup_table
+
+    def encode_with_lookup(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Inference: pure lookup, no embedding/projection forward."""
+        if self._lookup_table is None:
+            self.build_lookup_table()
+        valid = self._valid_mask(input_ids)
+        mask = attention_mask.float() * valid.float()
+        safe_ids = input_ids.clamp(max=self.vocab_size - 1)
+        scores = self._lookup_table.to(input_ids.device)[safe_ids] * mask
+        B = input_ids.size(0)
+        out = torch.zeros(B, self.vocab_size, device=input_ids.device, dtype=scores.dtype)
+        out.scatter_add_(1, safe_ids, scores)
+        return out
+
+
+class BOWQueryEncoder(nn.Module):
+    """Bag-of-Words: binary token indicator over the vocab."""
+
+    def __init__(self, vocab_size: int):
+        super().__init__()
+        self.vocab_size = vocab_size
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        B, _ = input_ids.shape
+        device = input_ids.device
+        bow = torch.zeros(B, self.vocab_size, device=device, dtype=torch.bfloat16)
+        safe_ids = input_ids.clamp(0, self.vocab_size - 1) * attention_mask.long()
+        bow.scatter_(1, safe_ids, 1.0)
+        bow[:, 0] = 0.0  # clear padding token (id=0) accumulation
+        return bow
+
+
+def build_query_encoder(
+    query_encoder_type: str,
+    vocab_size: int,
+    hidden_size: int = 768,
+    embed_weight: Optional[torch.Tensor] = None,
+    lora_r: int = 0,
+    activation: str = "relu",
+    **kwargs,
+) -> nn.Module:
+    if query_encoder_type == "bow":
+        return BOWQueryEncoder(vocab_size)
+    if query_encoder_type == "li_lsr":
+        if embed_weight is None:
+            raise ValueError("li_lsr requires embed_weight from the encoder")
+        return InferenceFreeQueryEncoder(
+            embed_weight, vocab_size, hidden_size,
+            lora_r=lora_r, activation=activation,
+        )
+    raise ValueError(
+        f"Unknown query_encoder_type: {query_encoder_type}. Choose: bow | li_lsr"
+    )