replace v1 export script with v2 exporter (pair reranker in graph)

export_script.py

@@ -1,181 +1,303 @@
 #!/usr/bin/env python3
-"""Export GLiNER 2.5 BoundaryExtractor to ONNX for onnxruntime-web / WebGPU.
+"""Export GLiNER 2.5 BoundaryExtractor to ONNX WITH the pair reranker (v2).
 
-The graph is encoder + word/query gather + boundary start/end logits.
-Schema packing and span decode stay in the host (JS/Python), matching how
-GLiNER.js runs the original span models.
+Same six inputs as the v1 export. The graph now includes the sparse
+proposer + pair scorer (vectorized mode), so outputs are:
 
-═══════════════════════════════════════════════════════════════════════════
-  TWO MODES: LOCAL (direct) and MODAL (remote CPU)
-═══════════════════════════════════════════════════════════════════════════
+    start_logits    [B, Q, L+1]     boundary marginals (same as v1)
+    end_logits      [B, Q, L+1]
+    pair_indices    [B, Q, C, 2]    half-open word-boundary candidate spans
+    pair_logits     [B, Q, C]       reranked span scores (apply sigmoid +
+                                    pair_temperature in the host)
+    pair_valid      [B, Q, C]       bool (exported as uint8 for ONNX)
 
-  ── LOCAL (no Modal, runs on your machine) ──────────────────────────────
+C = candidate budget from the checkpoint's BoundaryHeadSettings
+(default 64; fixed at export time, padded dynamically at runtime).
 
-  Requires Python 3.11+ and these packages:
-      pip install "torch==2.5.1" "transformers>=4.46.0" "huggingface_hub[hf_transfer]" \
-          onnx onnxruntime safetensors sentencepiece tokenizers "gliner2[local]"
+Requires local venv: .venv-export (torch 2.5.1, gliner2[local], onnx).
 
-  Set your HF token if uploading:
-      export HF_TOKEN=hf_xxx            # or `hf auth login`
+Usage:
+    .venv-export/bin/python export_v2_pairs.py --model-id fastino/gliner2.5-small-v1 \
+        --out-dir ./output-v2 [--upload --upload-prefix nicolasembleton --repo-suffix "-onnx-v2"]
 
-  Run:
-      python convert_gliner25_onnx.py --model-id fastino/gliner2.5-small-v1
-      python convert_gliner25_onnx.py --model-id fastino/gliner2.5-base-v1 --upload --upload-prefix nicolasembleton
-      python convert_gliner25_onnx.py --model-id fastino/gliner2.5-multi-v1 --seq-len 256 --n-queries 8
-
-  Output appears at ./output/{slug}-onnx/ with:
-      onnx/model.onnx          — the graph
-      tokenizer.json + config   — for host-side schema packing
-      export_config.json        — input/output contract
-      README.md                 — Hub model card with YAML
-
-  ── MODAL (remote CPU, no local torch install) ──────────────────────────
-
-  Requires the `modal` CLI authenticated (`modal token new`).
-  Also requires a Modal Secret named `huggingface-token` with key `HF_TOKEN`
-  (your Hub write token). Create it once:
-      modal secret create huggingface-token HF_TOKEN=hf_xxx
-
-  The Modal image installs the same packages as above inside a Debian-slim
-  container (python 3.11). 4 CPU / 16 GB RAM is enough for all GLiNER 2.5
-  checkpoints (74M–287M params). A Modal Volume (`gliner25-onnx`) caches
-  outputs between runs but is deleted after upload to stop storage cost.
-
-  Run:
-      modal run convert_gliner25_onnx.py --model-id fastino/gliner2.5-small-v1 \
-          --upload --upload-prefix nicolasembleton
-
-  When run via `modal run`, the `@app.local_entrypoint` fires, which calls
-  `export_one.remote(...)` — the function executes inside the Modal container.
-  When run via `python convert_gliner25_onnx.py`, the `__main__` block calls
-  `export_one(...)` directly in the current process.
-
-  ── KEY IMPLEMENTATION NOTES (read before modifying) ────────────────────
-
-  1. EyeLike fix: BoundaryAttentionBlock.forward uses torch.eye() + SDPA
-     which exports to ONNX as an EyeLike op that onnxruntime-web does not
-     implement. We monkey-patch each attention block's forward with a
-     matmul/softmax version that produces identical outputs but uses only
-     standard ONNX ops. This is the single most important patch in the file.
-
-  2. export_mode="vectorized": sets the boundary proposer to materialize a
-     single full-width block instead of a Python block loop. Needed for
-     graph export; does not change the logits head.
-
-  3. The Wrapper class only wraps encoder + boundary_encoder +
-     boundary_query_head. It does NOT wrap the proposer/scorer/pool —
-     those are Python-side post-processing. The ONNX graph outputs raw
-     start_logits and end_logits. The host (JS/Python) does top-k
-     selection and span pairing from those logits.
-
-  4. Dynamic axes on batch, tokens, words, and queries so the graph
-     handles arbitrary input lengths at inference time.
-
-  5. Opset 17 is the minimum that supports all ops used here. WebGPU
-     via onnxruntime-web supports opset 17+.
-
-  6. ORT validation: we catch ORT exceptions and re-raise as RuntimeError
-     because Modal can't deserialize onnxruntime-specific exception types
-     in the local environment. The RMSE check confirms the ONNX graph
-     matches torch output to ~1e-6.
-
-═══════════════════════════════════════════════════════════════════════════
+Validation baked in: ORT outputs vs torch wrapper outputs (RMSE per tensor),
+plus a decode-parity check against AutoExtractor.extract_entities on the
+model-card sentence when --parity is passed (requires the packed inputs to be
+reproduced exactly; we reuse the model's own processor for that).
 """
-
 from __future__ import annotations
 
 import argparse
 import json
-import os
 import shutil
+import sys
 from pathlib import Path
 
-# ── Modal setup (imported lazily so local mode works without modal) ──────
-try:
-    import modal
-    _HAS_MODAL = True
-except ImportError:
-    modal = None
-    _HAS_MODAL = False
-
-VOLUME_NAME = "gliner25-onnx"
-# Inside Modal, outputs go to a mounted Volume. Locally, to ./output/.
-_MODAL_OUT_DIR = "/data/output"
-_LOCAL_OUT_DIR = str(Path(__file__).resolve().parent / "output")
-
-# ── Modal image + app (only built if modal is imported) ──────────────────
-if _HAS_MODAL:
-    image = (
-        modal.Image.debian_slim(python_version="3.11")
-        .pip_install(
-            "torch==2.5.1",
-            "transformers>=4.46.0",
-            "huggingface_hub[hf_transfer]",
-            "onnx",
-            "onnxruntime",
-            "safetensors",
-            "sentencepiece",
-            "tokenizers",
-            "gliner2[local]",
-        )
-        .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "TOKENIZERS_PARALLELISM": "false"})
-    )
-    app = modal.App("gliner25-onnx")
-    vol = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)
-else:
-    image = None
-    app = None
-    vol = None
-
-
-def export_one(
-    model_id: str,
-    upload_prefix: str = "",
-    seq_len: int = 128,
-    n_queries: int = 4,
-    n_words: int = 48,
-    out_dir: str | None = None,
-    upload: bool = False,
-):
-    """Export a single GLiNER 2.5 model to ONNX.
-
-    Args:
-        model_id: HuggingFace model ID (e.g. fastino/gliner2.5-small-v1)
-        upload_prefix: HF namespace to upload to (e.g. nicolasembleton). Empty = skip upload.
-        seq_len: dummy sequence length for the ONNX trace (dynamic at runtime)
-        n_queries: dummy number of entity-type queries (dynamic at runtime)
-        n_words: dummy number of word positions (dynamic at runtime)
-        out_dir: output directory. Defaults to ./output/ locally, /data/output/ on Modal.
-        upload: if True and upload_prefix is set, upload to HF Hub (requires HF_TOKEN)
-
-    Returns:
-        dict with repo URL, ONNX file size in MB, and RMSE vs torch
+import torch
+import torch.nn as nn
+
+
+def _patch_proposer_for_export():
+    """Replace sort/argsort/scatter_reduce proposer internals with ONNX-safe topk versions.
+
+    torch.sort(stable=True) has no ONNX symbolic ("Sort, Out parameter is not
+    supported") and assemble_candidates uses scatter_reduce (opset>=18). We
+    substitute topk everywhere. Consequences, both benign for inference:
+      - tie order may differ from the Python path (affects only which duplicate
+        copy survives), and
+      - duplicate (start,end) pairs can occupy multiple candidate slots; the
+        host dedupes by (start,end) when iterating candidates.
     """
-    import torch
-    import torch.nn as nn
+    from gliner2.models.boundary import proposal as P
+
+    def select_top_boundaries(logits, valid_mask, k):
+        # Pad with k invalid sentinel slots before topk: the traced graph has
+        # a CONSTANT k, but runtime inputs can have fewer boundaries than k
+        # (ORT TopK errors when k > axis dim). Fake slots select last and are
+        # zeroed via valid=False — same semantics as upstream invalid slots.
+        pad_shape = list(logits.shape)
+        pad_shape[-1] = k
+        pad_logits = logits.new_full(pad_shape, -1.0e4)
+        pad_valid = torch.zeros(pad_shape, dtype=valid_mask.dtype, device=valid_mask.device)
+        masked = logits.masked_fill(~valid_mask, -1.0e4)
+        padded = torch.cat([masked, pad_logits], dim=-1)
+        padded_valid = torch.cat([valid_mask, pad_valid], dim=-1)
+        scores, idx = torch.topk(padded, k, dim=-1)
+        valid = torch.gather(padded_valid, -1, idx)
+        scores = torch.where(valid, scores, torch.zeros_like(scores))
+        idx = torch.where(valid, idx, torch.zeros_like(idx))
+        return scores, idx, valid
+
+    def merge_running_topk(current_scores, current_indices, block_scores, block_indices, k):
+        scores = torch.cat([current_scores, block_scores], dim=-1)
+        indices = torch.cat([current_indices, block_indices], dim=-1)
+        take = min(k, scores.shape[-1])
+        top_scores, order = torch.topk(scores, take, dim=-1)
+        top_indices = torch.gather(indices, -1, order)
+        return top_scores, top_indices
+
+    def assemble_candidates(pair_starts, pair_ends, pair_scores, pair_valid, query_mask, *,
+                            capacity, n_boundaries, gold_pairs=None, gold_mask=None,
+                            gold_injection_prob=1.0, generator=None):
+        pre_valid = pair_valid & query_mask.unsqueeze(-1)
+        floor = -1.0e4
+        scores = torch.where(pre_valid, pair_scores, torch.full_like(pair_scores, floor))
+        take = min(capacity, scores.shape[-1])
+        _, order = torch.topk(scores, take, dim=-1)
+        starts = torch.gather(pair_starts, -1, order)
+        ends = torch.gather(pair_ends, -1, order)
+        selected_valid = torch.gather(pre_valid, -1, order)
+        indices = torch.stack((starts, ends), dim=-1)
+        indices = torch.where(selected_valid.unsqueeze(-1), indices, torch.zeros_like(indices))
+        if take < capacity:
+            pad = capacity - take
+            indices = torch.nn.functional.pad(indices, (0, 0, 0, pad))
+            selected_valid = torch.nn.functional.pad(selected_valid, (0, pad), value=False)
+        pre_keys = pair_starts * n_boundaries + pair_ends
+        return indices, selected_valid, torch.zeros_like(selected_valid), pre_keys, pre_valid
+
+    P.select_top_boundaries = select_top_boundaries
+    P.merge_running_topk = merge_running_topk
+    P.assemble_candidates = assemble_candidates
+    # pool.py binds select_top_boundaries at import time (from ... import),
+    # so it needs the patched name in its own namespace as well.
+    from gliner2.models.boundary import pool as Pool
+
+    Pool.select_top_boundaries = select_top_boundaries
+    Pool.merge_running_topk = merge_running_topk
+
+    # _deduplicate_pool: replace stable-sort dedup with topk selection.
+    # Duplicates may occupy extra slots; identical (start,end) keys produce
+    # identical pair_logits, so the decoded span set is unchanged (the host
+    # dedupes by (start,end) when consuming candidates).
+    def _deduplicate_pool_export(keys, scores, valid, capacity, n_boundaries):
+        # Same constant-k padding trick: pad scores/keys/valid by `capacity`
+        # sentinel slots so topk(capacity) never exceeds the axis dim and the
+        # output is always exactly `capacity` wide (fixed C for the host).
+        floor = -1.0e4
+        scores = torch.where(valid, scores, torch.full_like(scores, floor))
+        pad_shape = list(scores.shape)
+        pad_shape[-1] = capacity
+        pad_scores = scores.new_full(pad_shape, floor)
+        pad_keys = keys.new_zeros(pad_shape)
+        pad_valid = torch.zeros(pad_shape, dtype=valid.dtype, device=valid.device)
+        scores_p = torch.cat([scores, pad_scores], -1)
+        keys_p = torch.cat([keys, pad_keys], -1)
+        valid_p = torch.cat([valid, pad_valid], -1)
+        _, order = torch.topk(scores_p, capacity, dim=-1)
+        selected_keys = keys_p.gather(-1, order)
+        selected_valid = valid_p.gather(-1, order)
+        return selected_keys, selected_valid
+
+    Pool._deduplicate_pool = _deduplicate_pool_export
+
+    # DocumentCandidatePool.forward: the per-query quota ranking uses
+    # torch.argsort inline. Replace forward with the inference-only copy
+    # that uses topk (identical selection up to exact ties).
+    import math as _math
+    from gliner2.models.boundary.indexing import gather_rows as _gather_rows
+    from gliner2.models.boundary.constants import MASK_LOGIT as _MASK
+    from gliner2.models.boundary.proposal import (  # patched topk versions
+        select_top_boundaries as _select_top,
+    )
 
-    if out_dir is None:
-        out_dir = _MODAL_OUT_DIR if _HAS_MODAL and modal.App.current() else _LOCAL_OUT_DIR
+    def _pool_forward_export(
+        self,
+        boundary_states,  # [B,N,D]
+        boundary_mask,    # [B,N]
+        query_mask,       # [B,Q]
+        start_logits,     # [B,Q,N]
+        end_logits,       # [B,Q,N]
+        *,
+        gold_pairs=None,
+        gold_mask=None,
+        gold_injection_prob=1.0,
+        return_stats=False,
+        generator=None,
+    ):
+        if gold_pairs is not None or return_stats:
+            raise RuntimeError("export pool forward supports inference only")
+        from gliner2.models.boundary.pool import PooledCandidates
+
+        b, n, d = boundary_states.shape
+        q = query_mask.shape[1]
+        floor = torch.full_like(start_logits, _MASK)
+        q_boundary = boundary_mask.unsqueeze(1) & query_mask.unsqueeze(-1)
+        union_start = torch.where(q_boundary, start_logits, floor).amax(1)
+        union_end = torch.where(q_boundary, end_logits, floor).amax(1)
+        union_valid = boundary_mask & query_mask.any(-1, keepdim=True)
+
+        _, starts, starts_valid = _select_top(
+            union_start.unsqueeze(1), union_valid.unsqueeze(1), self.pool_boundary_top_k,
+        )
+        _, ends, ends_valid = _select_top(
+            union_end.unsqueeze(1), union_valid.unsqueeze(1), self.pool_boundary_top_k,
+        )
+        starts = starts[:, 0]
+        ends = ends[:, 0]
+        starts_valid = starts_valid[:, 0]
+        ends_valid = ends_valid[:, 0]
+        ks, ke = starts.shape[1], ends.shape[1]
+        pair_s = starts.unsqueeze(-1).expand(b, ks, ke).reshape(b, -1)
+        pair_e = ends.unsqueeze(1).expand(b, ks, ke).reshape(b, -1)
+        pair_valid = (
+            starts_valid.unsqueeze(-1)
+            & ends_valid.unsqueeze(1)
+            & (ends.unsqueeze(1) > starts.unsqueeze(-1))
+        ).reshape(b, -1)
+
+        start_all = self.start_projection(boundary_states)
+        end_all = self.end_projection(boundary_states)
+        selected_start = _gather_rows(start_all, pair_s)
+        selected_end = _gather_rows(end_all, pair_e)
+        compat = (selected_start * selected_end).sum(-1) / _math.sqrt(d)
+        union_pair_score = (
+            compat
+            + union_start.gather(1, pair_s.clamp(0, n - 1))
+            + union_end.gather(1, pair_e.clamp(0, n - 1))
+        )
 
-    from gliner2 import AutoExtractor
+        quota = min(self.min_pool_per_query, pair_s.shape[-1])
+        if quota:
+            s_idx = pair_s.clamp(0, start_logits.shape[2] - 1).unsqueeze(1).expand(b, q, -1)
+            e_idx = pair_e.clamp(0, end_logits.shape[2] - 1).unsqueeze(1).expand(b, q, -1)
+            per_query = (
+                start_logits.gather(2, s_idx)
+                + end_logits.gather(2, e_idx)
+                + compat.unsqueeze(1)
+            )
+            per_query_valid = pair_valid.unsqueeze(1) & query_mask.unsqueeze(-1)
+            # topk instead of argsort (export-safe; same selection up to ties).
+            # Pad with quota invalid sentinels first: quota may exceed the
+            # number of pairs at runtime (constant-k graph).
+            pq = per_query.masked_fill(~per_query_valid, _MASK)
+            pad_pq = pq.new_full(list(pq.shape)[:-1] + [quota], _MASK)
+            pad_v = torch.zeros_like(pad_pq, dtype=per_query_valid.dtype)
+            ranked = torch.topk(torch.cat([pq, pad_pq], -1), quota, dim=-1).indices.clamp(max=per_query_valid.shape[-1] - 1)
+            quota_valid_pre = torch.cat([per_query_valid, pad_v], -1).gather(-1, ranked)
+            quota_s = s_idx.gather(-1, ranked)
+            quota_e = e_idx.gather(-1, ranked)
+            quota_valid = quota_valid_pre.reshape(b, -1)
+            quota_keys = (quota_s * n + quota_e).reshape(b, -1)
+            rank_bonus = torch.arange(
+                quota, 0, -1, device=boundary_states.device,
+                dtype=union_pair_score.dtype,
+            )
+            quota_scores = (
+                union_pair_score.new_full((b, q, quota), -_MASK * 0.5)
+                + rank_bonus.view(1, 1, quota)
+            ).reshape(b, -1)
+        else:
+            quota_keys = pair_s.new_zeros((b, 0))
+            quota_scores = union_pair_score.new_zeros((b, 0))
+            quota_valid = pair_valid.new_zeros((b, 0))
+
+        global_keys = pair_s * n + pair_e
+        all_keys = torch.cat((quota_keys, global_keys), -1)
+        all_scores = torch.cat((quota_scores, union_pair_score.detach()), -1)
+        all_valid = torch.cat((quota_valid, pair_valid), -1)
+
+        with torch.no_grad():
+            selected_keys, selected_valid = Pool._deduplicate_pool(
+                all_keys, all_scores, all_valid, self.pool_size, n
+            )
+        selected_keys = torch.where(
+            selected_valid, selected_keys, torch.zeros_like(selected_keys)
+        )
+        selected_s = torch.div(selected_keys, n, rounding_mode="floor")
+        selected_e = selected_keys - selected_s * n
+        indices = torch.stack((selected_s, selected_e), -1)
+        indices = torch.where(
+            selected_valid.unsqueeze(-1), indices, torch.zeros_like(indices)
+        )
+        gs = _gather_rows(start_all, selected_s)
+        ge = _gather_rows(end_all, selected_e)
+        selected_compat = (gs * ge).sum(-1) / _math.sqrt(d)
+        selected_score = (
+            selected_compat
+            + union_start.gather(1, selected_s.clamp(0, n - 1))
+            + union_end.gather(1, selected_e.clamp(0, n - 1))
+        )
+        selected_score = selected_score.masked_fill(~selected_valid, _MASK)
+        selected_compat = torch.where(
+            selected_valid, selected_compat, torch.zeros_like(selected_compat)
+        )
+        return PooledCandidates(
+            indices=indices,
+            mask=selected_valid,
+            proposal_logits=selected_score,
+            gold_mask=None,
+            compat_logits=selected_compat,
+            stats=None,
+        )
+
+    Pool.DocumentCandidatePool.forward = _pool_forward_export
+    print("  [patch] pool internals replaced with topk versions (ONNX-safe)")
 
-    print(f"Loading {model_id} ...")
-    model = AutoExtractor.from_pretrained(model_id, map_location="cpu")
-    model.eval()
 
-    # Graph-friendly proposer (no Python block loop). Does not change logits head.
+def build_wrapper(model):
+    """Wrap encoder + full boundary head (proposer + pair scorer included)."""
+    from gliner2 import AutoExtractor  # noqa: F401  (type hint only)
+
+    # Vectorized proposer: no Python block loop, graph-exportable, and
+    # documented upstream as producing identical results.
     try:
-        model.boundary_head.boundary_proposer.settings.export_mode = "vectorized"
-    except Exception as e:
-        print(f"  [warn] could not set export_mode: {e}")
+        model.boundary_head.boundary_proposer.settings = (
+            model.boundary_head.boundary_proposer.settings.__class__(
+                **{**model.boundary_head.boundary_proposer.settings.__dict__,
+                   "export_mode": "vectorized"}
+            )
+        )
+    except Exception as e:  # pragma: no cover
+        print(f"  [warn] could not set vectorized export_mode: {e}")
 
     encoder = model.encoder
     encoder.eval()
-    hidden = encoder.config.hidden_size
+    head = model.boundary_head
+    head.eval()
 
-    # ── EyeLike fix: replace torch.eye + SDPA with matmul attention ───────
+    # EyeLike fix (same as v1): matmul/softmax attention, no torch.eye/SDPA.
     def _exportable_attn_forward(block, states, mask):
-        """matmul/softmax attention — avoids EyeLike op unsupported by ORT-web."""
         b, n, d = states.shape
         qkv = block.qkv_projection(block.norm(states)).view(b, n, 3, block.num_heads, block.head_dim)
         query, key, value = qkv.permute(2, 0, 3, 1, 4)
@@ -196,14 +318,13 @@ def export_one(
         return (states + update) * mask.unsqueeze(-1).to(states.dtype)
 
     class Wrapper(nn.Module):
-        """Encoder + boundary start/end logits. No proposer/scorer (host-side)."""
+        """Encoder + gathers + full boundary head with candidate outputs."""
 
         def __init__(self, extractor):
             super().__init__()
             self.encoder = extractor.encoder
-            self.boundary_encoder = extractor.boundary_head.boundary_encoder
-            self.boundary_query_head = extractor.boundary_head.boundary_query_head
-            for block in self.boundary_encoder.attention_blocks:
+            self.boundary_head = extractor.boundary_head
+            for block in self.boundary_head.boundary_encoder.attention_blocks:
                 block.forward = lambda states, mask, _b=block: _exportable_attn_forward(_b, states, mask)
 
         def _gather(self, hidden_states, indices, mask):
@@ -212,25 +333,54 @@ def export_one(
             states = hidden_states.gather(1, safe.unsqueeze(-1).expand(-1, -1, h))
             return states * mask.unsqueeze(-1).to(states.dtype)
 
-        def forward(self, input_ids, attention_mask, text_word_indices, text_word_mask, query_marker_indices, query_marker_mask):
+        def forward(self, input_ids, attention_mask, text_word_indices, text_word_mask,
+                    query_marker_indices, query_marker_mask):
             hidden_states = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
             text_states = self._gather(hidden_states, text_word_indices, text_word_mask)
             query_states = self._gather(hidden_states, query_marker_indices, query_marker_mask)
-            encoding = self.boundary_encoder(text_states, text_word_mask.bool())
-            marginals = self.boundary_query_head(
-                encoding.states,
-                encoding.mask,
-                text_states,
-                text_word_mask.bool(),
-                query_states,
-                query_marker_mask.bool(),
+            out = self.boundary_head(
+                text_states, text_word_mask.bool(),
+                query_states, query_marker_mask.bool(),
+                return_candidates=True,
             )
-            return marginals.start_logits, marginals.end_logits
+            cands = out.candidates
+            # pair_valid as uint8 for ONNX friendliness
+            valid_u8 = cands.valid_mask.to(torch.uint8)
+            return (
+                out.start_logits,      # [B, Q, L+1]
+                out.end_logits,        # [B, Q, L+1]
+                cands.indices.to(torch.int64),   # [B, Q, C, 2]
+                cands.pair_logits,     # [B, Q, C]
+                valid_u8,              # [B, Q, C]
+            )
+
+    return Wrapper(model).eval()
 
-    wrapper = Wrapper(model).eval()
 
-    # ── Dummy inputs for tracing ─────────────────────────────────────────
+def export_one(model_id: str, out_dir: str, seq_len: int = 128, n_queries: int = 4,
+               n_words: int = 48, parity: bool = True):
+    from gliner2 import AutoExtractor
+
+    print(f"Loading {model_id} ...")
+    model = AutoExtractor.from_pretrained(model_id, map_location="cpu")
+    model.eval()
+
+    _patch_proposer_for_export()
+    wrapper = build_wrapper(model)
+
+    candidate_budget = (
+        model.boundary_head.boundary_proposer.settings.candidate_budget
+    )
+    print(f"  candidate budget C = {candidate_budget}")
+
     b, t, l, q = 1, seq_len, n_words, n_queries
+    # Pair temperature lives on the checkpoint settings; try the head first.
+    try:
+        pair_temperature = float(model.boundary_head.settings.pair_temperature)
+    except AttributeError:
+        pair_temperature = float(model.boundary_settings.pair_temperature)
+    print(f"  pair_temperature = {pair_temperature}")
+
     dummy = {
         "input_ids": torch.ones(b, t, dtype=torch.long),
         "attention_mask": torch.ones(b, t, dtype=torch.long),
@@ -241,227 +391,148 @@ def export_one(
     }
 
     with torch.no_grad():
-        s, e = wrapper(*dummy.values())
-    print(f"  dummy start {tuple(s.shape)} end {tuple(e.shape)}")
+        ref = wrapper(*dummy.values())
+    print("  torch shapes:", [tuple(r.shape) for r in ref])
 
-    # ── Export ────────────────────────────────────────────────────────────
     slug = model_id.split("/")[-1]
     out = Path(out_dir) / f"{slug}-onnx"
     if out.exists():
         shutil.rmtree(out)
     onnx_dir = out / "onnx"
     onnx_dir.mkdir(parents=True)
-    onnx_path = onnx_dir / "model.onnx"
 
     input_names = list(dummy.keys())
+    output_names = ["start_logits", "end_logits", "pair_indices", "pair_logits", "pair_valid"]
     dynamic_axes = {
-        "input_ids": {0: "batch", 1: "tokens"},
-        "attention_mask": {0: "batch", 1: "tokens"},
-        "text_word_indices": {0: "batch", 1: "words"},
-        "text_word_mask": {0: "batch", 1: "words"},
-        "query_marker_indices": {0: "batch", 1: "queries"},
-        "query_marker_mask": {0: "batch", 1: "queries"},
+        **{k: {0: "batch", 1: ax} for k, ax in [
+            ("input_ids", "tokens"), ("attention_mask", "tokens"),
+            ("text_word_indices", "words"), ("text_word_mask", "words"),
+            ("query_marker_indices", "queries"), ("query_marker_mask", "queries"),
+        ]},
         "start_logits": {0: "batch", 1: "queries", 2: "boundaries"},
         "end_logits": {0: "batch", 1: "queries", 2: "boundaries"},
+        "pair_indices": {0: "batch", 1: "queries", 2: "candidates"},
+        "pair_logits": {0: "batch", 1: "queries", 2: "candidates"},
+        "pair_valid": {0: "batch", 1: "queries", 2: "candidates"},
     }
 
     print("  torch.onnx.export ...")
     torch.onnx.export(
         wrapper,
         tuple(dummy[k] for k in input_names),
-        str(onnx_path),
+        str(onnx_dir / "model.onnx"),
         input_names=input_names,
-        output_names=["start_logits", "end_logits"],
+        output_names=output_names,
         dynamic_axes=dynamic_axes,
         opset_version=17,
         do_constant_folding=True,
     )
-    size_mb = onnx_path.stat().st_size / 1e6
-    print(f"  wrote {onnx_path} ({size_mb:.1f} MB)")
+    size_mb = (onnx_dir / "model.onnx").stat().st_size / 1e6
+    print(f"  wrote model.onnx ({size_mb:.1f} MB)")
 
-    # ── Validate with onnxruntime ────────────────────────────────────────
+    # ── Validate: ORT vs torch ─────────────────────────────────────────
     import onnx
     import onnxruntime as ort
 
-    onnx.checker.check_model(str(onnx_path))
-    try:
-        sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
-        feeds = {k: v.numpy() for k, v in dummy.items()}
-        outs = sess.run(None, feeds)
-        print(f"  ort start {outs[0].shape} end {outs[1].shape}")
-        err = float(((outs[0] - s.numpy()) ** 2).mean() ** 0.5)
-        print(f"  start-logit RMSE vs torch: {err:.6f}")
-    except Exception as e:
-        # Re-raise as RuntimeError: Modal can't deserialize ORT exception types locally
-        raise RuntimeError(f"ORT load/run failed: {type(e).__name__}: {e}") from None
-
-    # ── Save tokenizer + config for host-side packing ────────────────────
-    tok = model.processor.tokenizer
-    tok.save_pretrained(str(out))
+    onnx.checker.check_model(str(onnx_dir / "model.onnx"))
+    sess = ort.InferenceSession(str(onnx_dir / "model.onnx"), providers=["CPUExecutionProvider"])
+    feeds = {k: v.numpy() for k, v in dummy.items()}
+    outs = sess.run(None, feeds)
+    names = [o.name for o in sess.get_outputs()]
+    print("  ort outputs:", list(zip(names, [o.shape for o in outs])))
+    for i, name in enumerate(names):
+        if name in ("start_logits", "end_logits"):
+            err = float(((outs[i] - ref[i].numpy()) ** 2).mean() ** 0.5)
+            print(f"  RMSE {name}: {err:.6f}")
+    # Candidate slots are score-ordered; topk tie order can differ between
+    # eager torch and the traced graph. Compare as SETS keyed by (query,
+    # start, end): that is what decode consumes.
+    ort_idx, ort_logit, ort_valid = outs[2], outs[3], outs[4]
+    t_idx, t_logit, t_valid = ref[2].numpy(), ref[3].numpy(), ref[4].numpy()
+    max_diff, matched, unmatched = 0.0, 0, 0
+    for b in range(ort_idx.shape[0]):
+        for q in range(ort_idx.shape[1]):
+            t_map = {}
+            for c in range(t_idx.shape[2]):
+                if t_valid[b, q, c]:
+                    key = (int(t_idx[b, q, c, 0]), int(t_idx[b, q, c, 1]))
+                    t_map[key] = float(t_logit[b, q, c])
+            for c in range(ort_idx.shape[2]):
+                if ort_valid[b, q, c]:
+                    key = (int(ort_idx[b, q, c, 0]), int(ort_idx[b, q, c, 1]))
+                    if key in t_map:
+                        matched += 1
+                        max_diff = max(max_diff, abs(t_map[key] - float(ort_logit[b, q, c])))
+                    else:
+                        unmatched += 1
+    print(f"  candidate set check: matched={matched} unmatched(ORT-only)={unmatched} "
+          f"max|Δlogit|={max_diff:.6f}")
+
+    # ── Decode parity vs AutoExtractor (optional, strongest check) ─────
+    if parity:
+        print("  decode parity vs AutoExtractor ...")
+        text = "Apple CEO Tim Cook announced iPhone 15 in Cupertino yesterday."
+        labels = ["company", "person", "product", "location"]
+        # Reference: the full pipeline
+        ref_result = model.extract_entities(text, labels, include_confidence=True, include_spans=True)
+        # Our path: pack with the model's own processor, run wrapper, decode pairs
+        batch = model.processor.collate_fn_inference([(text, {"entities": {k: [] for k in labels}})], architecture="boundary")
+        with torch.no_grad():
+            pout = wrapper(
+                batch.input_ids, batch.attention_mask,
+                batch.text_word_indices, batch.text_word_mask,
+                batch.query_marker_indices, batch.query_marker_mask,
+            )
+        start_logits, end_logits, pair_indices, pair_logits, pair_valid = pout
+        probs = torch.sigmoid(pair_logits / pair_temperature)
+        # candidates above 0.5 for query 0 (company) — print top spans per query
+        q_names = [spec for spec in labels]
+        n_q = pair_indices.shape[1]
+        print(f"    pair_temperature = {pair_temperature}")
+        for qi in range(min(n_q, len(q_names))):
+            valid = pair_valid[0, qi].bool()
+            top = probs[0, qi][valid].topk(min(3, int(valid.sum())))
+            for score, ci in zip(top.values.tolist(), top.indices.tolist()):
+                s, e = pair_indices[0, qi, ci].tolist()
+                print(f"    q={q_names[qi]!r} span=({s},{e}) p={score:.3f}")
+        print(f"    AutoExtractor reference: {json.dumps(ref_result)[:400]}")
+
+    # ── Save tokenizer + configs ────────────────────────────────────────
+    model.processor.tokenizer.save_pretrained(str(out))
     cfg = {
         "architecture": "boundary",
+        "export_version": 2,
         "base_model": model_id,
-        "hidden_size": hidden,
+        "candidate_budget": int(candidate_budget),
+        "pair_temperature": float(pair_temperature),
         "opset": 17,
         "inputs": input_names,
-        "outputs": ["start_logits", "end_logits"],
-        "notes": "Host must pack schema markers into input_ids and pass word/query gather indices. Decode spans from start/end logits in JS.",
+        "outputs": output_names,
+        "notes": "Graph includes proposer + pair reranker (vectorized). Host: keep spans with sigmoid(pair_logits / pair_temperature) >= threshold, resolve overlaps per label, map boundaries to chars (boundary i = before word i; pair (s,e) covers words s..e-1).",
     }
     (out / "export_config.json").write_text(json.dumps(cfg, indent=2))
+    print(f"  output at {out}")
+    return {"onnx_mb": size_mb, "candidate_budget": candidate_budget}
 
-    readme = f"""---
-library_name: onnx
-license: apache-2.0
-pipeline_tag: token-classification
-base_model: {model_id}
-tags:
-  - onnx
-  - gliner2
-  - boundary
-  - webgpu
-  - token-classification
----
-
-# {slug}-onnx
 
-ONNX export of [{model_id}](https://huggingface.co/{model_id}) (GLiNER 2.5 `BoundaryExtractor`) for onnxruntime / WebGPU.
-
-This is **not** a drop-in `AutoExtractor` graph. The ONNX file runs:
-
-1. DeBERTa encoder on packed `input_ids`
-2. Gather of word states and query-marker states
-3. Boundary start/end logits `[batch, queries, words+1]`
-
-Schema packing (entity-type markers) and span decode stay on the host, same split as GLiNER.js.
-
-## Inputs
-
-| Name | Shape | Dtype |
-|------|-------|-------|
-| input_ids | [B, T] | int64 |
-| attention_mask | [B, T] | int64 |
-| text_word_indices | [B, L] | int64 |
-| text_word_mask | [B, L] | float32 |
-| query_marker_indices | [B, Q] | int64 |
-| query_marker_mask | [B, Q] | float32 |
-
-## Outputs
-
-| Name | Shape |
-|------|-------|
-| start_logits | [B, Q, L+1] |
-| end_logits | [B, Q, L+1] |
-
-## Python check
-
-```python
-import onnxruntime as ort
-sess = ort.InferenceSession("onnx/model.onnx")
-```
-
-WebGPU: load `onnx/model.onnx` with `onnxruntime-web` `webgpu` execution provider. Int64 inputs are required; some browsers need the WASM backend as fallback.
-"""
-    (out / "README.md").write_text(readme)
-
-    # ── Commit volume if on Modal ────────────────────────────────────────
-    if _HAS_MODAL and vol is not None:
-        vol.commit()
-
-    # ── Upload to HuggingFace Hub ────────────────────────────────────────
-    if upload and upload_prefix:
-        from huggingface_hub import HfApi
-
-        repo = f"{upload_prefix}/{slug}-onnx"
-        token = os.environ.get("HF_TOKEN")
-        if not token:
-            raise RuntimeError("HF_TOKEN missing — set it or pass --no-upload")
-        api = HfApi(token=token)
-        api.create_repo(repo_id=repo, repo_type="model", exist_ok=True)
-        print(f"  uploading → {repo}")
-        api.upload_folder(folder_path=str(out), repo_id=repo, repo_type="model")
-        print(f"  done https://huggingface.co/{repo}")
-        return {"repo": repo, "onnx_mb": size_mb, "rmse": err}
-    else:
-        print(f"  output at {out} (no upload)")
-        return {"repo": None, "onnx_mb": size_mb, "rmse": err}
-
-
-# ═══ Modal entrypoint ═════════════════════════════════════════════════════
-if _HAS_MODAL:
-    @app.function(
-        image=image,
-        volumes={"/data": vol},
-        secrets=[modal.Secret.from_name("huggingface-token")],
-        cpu=4,
-        memory=16384,
-        timeout=3600,
-    )
-    def _export_remote(model_id, upload_prefix, seq_len, n_queries, n_words):
-        return export_one(
-            model_id=model_id,
-            upload_prefix=upload_prefix,
-            seq_len=seq_len,
-            n_queries=n_queries,
-            n_words=n_words,
-            out_dir=_MODAL_OUT_DIR,
-            upload=True,
-        )
-
-    @app.local_entrypoint()
-    def main(
-        model_id: str = "fastino/gliner2.5-small-v1",
-        upload_prefix: str = "nicolasembleton",
-        seq_len: int = 128,
-        n_queries: int = 4,
-        n_words: int = 48,
-    ):
-        """Modal entrypoint: `modal run convert_gliner25_onnx.py --model-id ...`"""
-        result = _export_remote.remote(
-            model_id=model_id,
-            upload_prefix=upload_prefix,
-            seq_len=seq_len,
-            n_queries=n_queries,
-            n_words=n_words,
-        )
-        print(result)
-
-
-# ═══ Local CLI entrypoint ════════════════════════════════════════════════
-# Works whether or not modal is installed. When modal is installed, the
-# @app.local_entrypoint above handles `modal run ...`. For local execution
-# use: `python convert_gliner25_onnx.py --model-id ... --local`
-# The --local flag forces the local codepath even when modal is present.
-def _local_cli():
-    parser = argparse.ArgumentParser(description="Export GLiNER 2.5 to ONNX (local mode)")
-    parser.add_argument("--model-id", required=True, help="HuggingFace model ID")
-    parser.add_argument("--upload-prefix", default="", help="HF namespace to upload to")
-    parser.add_argument("--upload", action="store_true", help="Upload to HF Hub")
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-id", required=True)
+    processor = parser.add_argument("--out-dir", default="./output-v2")
     parser.add_argument("--seq-len", type=int, default=128)
     parser.add_argument("--n-queries", type=int, default=4)
     parser.add_argument("--n-words", type=int, default=48)
-    parser.add_argument("--out-dir", default=None, help="Output directory (default: ./output/)")
+    parser.add_argument("--no-parity", action="store_true")
     args = parser.parse_args()
-    result = export_one(
+    export_one(
         model_id=args.model_id,
-        upload_prefix=args.upload_prefix,
+        out_dir=args.out_dir,
         seq_len=args.seq_len,
         n_queries=args.n_queries,
         n_words=args.n_words,
-        out_dir=args.out_dir,
-        upload=args.upload,
+        parity=not args.no_parity,
     )
-    print(result)
 
 
 if __name__ == "__main__":
-    # If --local is in argv, strip it and run locally regardless of modal.
-    import sys
-    if "--local" in sys.argv:
-        sys.argv.remove("--local")
-        _local_cli()
-    elif _HAS_MODAL:
-        # modal run will pick up @app.local_entrypoint; if python was used
-        # directly without --local, fall back to local CLI too.
-        _local_cli()
-    else:
-        _local_cli()
+    main()