Re-export model.onnx and model_int8.onnx from model.safetensors; fix onnx_config and tokenizer_config

export_onnx.py

@@ -0,0 +1,386 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#   "pylate==1.6.0",
+#   "torch==2.11.0",
+#   "transformers==5.3.0",
+#   "sentence-transformers==5.3.0",
+#   "onnx==1.19.1",
+#   "onnxscript==0.7.2",
+#   "onnxruntime==1.23.2",
+#   "numpy==2.5.3",
+# ]
+# ///
+"""Export a PyLate ColBERT checkpoint to ONNX, quantize it to INT8, and validate both.
+
+    uv run export_onnx.py export   --model-id lightonai/mLateOn --out-dir mlateon-onnx
+    uv run export_onnx.py quantize --model-id lightonai/mLateOn --out-dir mlateon-onnx
+    uv run export_onnx.py check    --model-id lightonai/mLateOn --onnx mlateon-onnx/model.onnx
+
+    inputs : input_ids, attention_mask  int64 [batch, sequence], [Q]/[D] prefix already inserted
+    output : output                     float32 [batch, sequence, dim], L2-normalized per token.
+             Padding positions are NOT zeroed; mask them with attention_mask.
+
+Validation compares ONNX Runtime against PyLate's encode() on the same tokens and checks
+three things, because per-token cosine alone does not catch a wrong checkpoint:
+    1. per-token embeddings (max |diff| and cosine) for queries and documents,
+    2. MaxSim scores on a query x document matrix, plus top-1 ranking agreement,
+    3. sequence lengths from a single unpadded sentence up to the model's maximum, so a
+       sliding-window mask traced as a constant would be caught.
+For the INT8 model the fp32 ONNX is also used as a reference, which separates export
+error from quantization error.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+# Short multilingual texts used for the embedding and MaxSim checks. The last one
+# repeats a sentence 30 times and is numerically sensitive: float32 runs (PyLate's
+# included) sit about 1e-5 away from float64 on it, versus about 1e-7 elsewhere.
+DOCUMENTS = [
+    "Paris is the capital and most populous city of France.",
+    "Ein HNSW-Index baut einen hierarchischen Graphen für die schnelle Nachbarsuche auf.",
+    "Le Louvre est le musée le plus visité au monde.",
+    "El Volga es el río más largo de Europa.",
+    "Roma è la capitale d'Italia.",
+    "Lisboa é a capital de Portugal.",
+    "القاهرة هي عاصمة جمهورية مصر العربية.",
+    "def add(a, b):\n    return a + b",
+    "東京は日本の首都です。",
+    "Late interaction models keep one vector per token and score with MaxSim. " * 30,
+]
+
+# One query per document, in matching order, so the expected top-1 is the diagonal.
+QUERIES = [
+    "What is the capital of France?",
+    "Wie funktioniert ein HNSW-Index?",
+    "Quel est le musée le plus visité du monde ?",
+    "¿Cuál es el río más largo de Europa?",
+    "Qual è la capitale d'Italia?",
+    "capital de Portugal",
+    "ما هي عاصمة مصر؟",
+    "python function that adds two numbers",
+    "日本の首都はどこですか",
+    "how do late interaction models score documents",
+]
+
+# Word counts for the length sweep. Each is encoded alone (batch of one, no padding).
+LENGTH_SWEEP_WORDS = [8, 100, 400, 1500, 3000, 6000]
+
+
+# --------------------------------------------------------------------------- helpers
+
+
+def log(msg: str = "") -> None:
+    print(msg, flush=True)
+
+
+def load_reference(model_id: str, **model_kwargs):
+    from pylate import models
+
+    return models.ColBERT(model_id, device="cpu", model_kwargs=model_kwargs or None)
+
+
+def load_session(path: Path):
+    import onnxruntime as ort
+
+    return ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+
+
+def run_onnx(session, features) -> list[np.ndarray]:
+    """Run the graph on PyLate-tokenized features and strip padding, as encode() does."""
+    output = session.run(
+        None,
+        {
+            "input_ids": features["input_ids"].numpy(),
+            "attention_mask": features["attention_mask"].numpy(),
+        },
+    )[0]
+    mask = features["attention_mask"].numpy().astype(bool)
+    return [output[i][mask[i]] for i in range(output.shape[0])]
+
+
+def maxsim(queries: list[np.ndarray], documents: list[np.ndarray]) -> np.ndarray:
+    """ColBERT score matrix: sum over query tokens of the max dot product over doc tokens."""
+    scores = np.zeros((len(queries), len(documents)), dtype=np.float64)
+    for i, q in enumerate(queries):
+        for j, d in enumerate(documents):
+            scores[i, j] = (q.astype(np.float64) @ d.astype(np.float64).T).max(axis=1).sum()
+    return scores
+
+
+def long_text(words: int) -> str:
+    base = (
+        "Late interaction retrieval keeps one embedding per token and scores a query "
+        "against a document with the MaxSim operator, which is cheap and robust. "
+    ).split()
+    return " ".join(base[i % len(base)] for i in range(words))
+
+
+# --------------------------------------------------------------------------- checks
+
+
+def compare_embeddings(expected, actual, label: str, tol: float | None):
+    """Per-text max |diff| and cosine. Returns (max_diff, mean_cos) over all tokens."""
+    all_cos, worst = [], 0.0
+    log(f"\n{label}:")
+    for i, (e, a) in enumerate(zip(expected, actual)):
+        if e.shape != a.shape:
+            raise SystemExit(f"shape mismatch on text {i}: {e.shape} vs {a.shape}")
+        max_diff = float(np.abs(e - a).max())
+        cos = (e * a).sum(-1)  # both are unit vectors
+        all_cos.append(cos)
+        worst = max(worst, max_diff)
+        log(f"  text {i}: tokens={len(e):5d}  max|diff|={max_diff:.2e}  cos min={cos.min():.7f} mean={cos.mean():.7f}")
+    all_cos = np.concatenate(all_cos)
+    log(f"  all {len(all_cos)} tokens: max|diff|={worst:.2e}  cos min={all_cos.min():.7f} mean={all_cos.mean():.7f}")
+    if tol is not None and worst >= tol:
+        raise SystemExit(f"FAIL {label}: max|diff| {worst:.2e} >= {tol:.0e}")
+    return worst, float(all_cos.mean())
+
+
+def compare_maxsim(ref_q, ref_d, act_q, act_d, label: str, score_tol: float | None):
+    ref = maxsim(ref_q, ref_d)
+    act = maxsim(act_q, act_d)
+    diff = np.abs(ref - act)
+    top1_ref, top1_act = ref.argmax(1), act.argmax(1)
+    order_ref, order_act = ref.argsort(1), act.argsort(1)
+    log(f"\n{label}: MaxSim {ref.shape[0]}x{ref.shape[1]}")
+    log(
+        f"  score max|diff|={diff.max():.2e}  mean|diff|={diff.mean():.2e}  score range=[{ref.min():.2f}, {ref.max():.2f}]"
+    )
+    log(
+        f"  top-1 agreement: {(top1_ref == top1_act).mean():.0%}   full-order agreement: {(order_ref == order_act).all(1).mean():.0%}"
+    )
+    log(f"  reference top-1 hits diagonal: {(top1_ref == np.arange(len(top1_ref))).mean():.0%}")
+    if score_tol is not None and diff.max() >= score_tol:
+        raise SystemExit(f"FAIL {label}: MaxSim max|diff| {diff.max():.2e} >= {score_tol:.0e}")
+    if (top1_ref != top1_act).any():
+        raise SystemExit(f"FAIL {label}: top-1 ranking differs from reference")
+    return float(diff.max())
+
+
+def length_sweep(reference, session, cos_tol: float | None, max_words: int):
+    """Encode single long documents, one at a time, so there is no padding at all.
+
+    Gated on cosine rather than absolute difference: past a few thousand tokens the
+    float32 reference itself drifts. On mLateOn at 6753 tokens, PyLate float32 is 1.7e-4
+    from a float64 run while the ONNX graph is 6e-6 from it, so a 1e-4 absolute gate
+    would fail the better of the two.
+    """
+    log("\nlength sweep (batch of one, no padding, is_query=False):")
+    worst = 0.0
+    for words in [w for w in LENGTH_SWEEP_WORDS if w <= max_words]:
+        text = long_text(words)
+        features = reference.tokenize([text], is_query=False)
+        n = int(features["attention_mask"].sum())
+        t0 = time.time()
+        expected = reference.encode([text], is_query=False, show_progress_bar=False)[0]
+        actual = run_onnx(session, features)[0]
+        max_diff = float(np.abs(expected - actual).max())
+        cos_min = float((expected * actual).sum(-1).min())
+        worst = max(worst, 1.0 - cos_min)
+        log(f"  tokens={n:5d}  max|diff|={max_diff:.2e}  cos min={cos_min:.7f}  ({time.time() - t0:.0f}s)")
+    if cos_tol is not None and worst >= cos_tol:
+        raise SystemExit(f"FAIL length sweep: 1 - cos min = {worst:.2e} >= {cos_tol:.0e}")
+    return worst
+
+
+def check(
+    model_id: str,
+    onnx_path: Path,
+    tol: float | None,
+    score_tol: float | None,
+    max_words: int,
+    fp32_path: Path | None = None,
+    cos_tol: float | None = None,
+) -> dict:
+    """Validate onnx_path against PyLate (and optionally against a fp32 ONNX)."""
+    reference = load_reference(model_id)
+    session = load_session(onnx_path)
+    results = {}
+
+    enc = {}
+    for is_query, texts in ((False, DOCUMENTS), (True, QUERIES)):
+        kind = "queries" if is_query else "documents"
+        expected = reference.encode(texts, is_query=is_query, show_progress_bar=False)
+        actual = run_onnx(session, reference.tokenize(texts, is_query=is_query))
+        results[f"{kind}_max_diff"], results[f"{kind}_mean_cos"] = compare_embeddings(
+            expected, actual, f"{kind} vs PyLate", tol
+        )
+        enc[kind] = (expected, actual)
+
+    results["maxsim_max_diff"] = compare_maxsim(
+        enc["queries"][0],
+        enc["documents"][0],
+        enc["queries"][1],
+        enc["documents"][1],
+        "ONNX vs PyLate",
+        score_tol,
+    )
+    results["length_sweep_one_minus_cos"] = length_sweep(reference, session, cos_tol, max_words)
+
+    if fp32_path is not None:
+        fp32 = load_session(fp32_path)
+        for is_query, texts in ((False, DOCUMENTS), (True, QUERIES)):
+            kind = "queries" if is_query else "documents"
+            features = reference.tokenize(texts, is_query=is_query)
+            compare_embeddings(
+                run_onnx(fp32, features),
+                run_onnx(session, features),
+                f"{kind}: {onnx_path.name} vs {fp32_path.name}",
+                None,
+            )
+    return results
+
+
+# --------------------------------------------------------------------------- export
+
+
+def export(model_id: str, out: Path, opset: int, ir_version: int) -> None:
+    import onnx
+    import torch
+    from pylate import models
+
+    class ColBERTForONNX(torch.nn.Module):
+        """Transformer -> PyLate Dense layers -> L2 normalization, as in PyLate's encode()."""
+
+        def __init__(self, colbert: models.ColBERT):
+            super().__init__()
+            transformer, *dense_layers = colbert
+            self.transformer = transformer.auto_model
+            # Call the whole Dense modules, not their .linear: Dense.forward adds the
+            # residual projection when use_residual is set (layers 1 and 2 of mLateOn).
+            self.dense_layers = torch.nn.ModuleList(dense_layers)
+
+        def forward(self, input_ids, attention_mask):
+            hidden = self.transformer(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+            features = {"token_embeddings": hidden}
+            for dense in self.dense_layers:
+                features = dense(features)
+            return torch.nn.functional.normalize(features["token_embeddings"], p=2, dim=-1)
+
+    # Eager attention exports to plain MatMul/Softmax ops.
+    colbert = load_reference(model_id, attn_implementation="eager")
+    local_window = getattr(colbert[0].auto_model.config, "local_attention", 0)
+
+    # Trace with a padded batch longer than the local attention window, so both the
+    # padding mask and the sliding-window mask are present in the graph.
+    example = colbert.tokenize([long_text(max(local_window, 128) * 2), "A short document."], is_query=False)
+    log(f"tracing with input_ids {tuple(example['input_ids'].shape)} (local window {local_window})")
+
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with torch.no_grad():
+        torch.onnx.export(
+            ColBERTForONNX(colbert).eval(),
+            (example["input_ids"], example["attention_mask"]),
+            str(out),
+            input_names=["input_ids", "attention_mask"],
+            output_names=["output"],
+            dynamic_shapes=(
+                {0: "batch_size", 1: "sequence_length"},
+                {0: "batch_size", 1: "sequence_length"},
+            ),
+            opset_version=opset,
+            dynamo=True,
+            external_data=False,  # single file like the previous model.onnx (protobuf limit is 2 GB)
+            optimize=True,
+        )
+
+    # The dynamo exporter stores Python stack traces (with local paths) in node metadata.
+    # Strip them, then lower the IR version so older onnxruntime releases can load the
+    # file. The full checker fails if anything else in the graph needs a newer IR.
+    model = onnx.load(str(out))
+    for node in model.graph.node:
+        del node.metadata_props[:]
+    del model.graph.metadata_props[:]
+    model.ir_version = ir_version
+    model.producer_name = "pylate-onnx-export"
+    onnx.save(model, str(out))
+    onnx.checker.check_model(str(out), full_check=True)
+    load_session(out)  # must load with the pinned onnxruntime
+    log(f"saved {out} ({out.stat().st_size / 1e6:.0f} MB, opset {opset}, IR {ir_version})")
+
+
+def quantize(fp32: Path, int8: Path, per_channel: bool, reduce_range: bool) -> None:
+    from onnxruntime.quantization import QuantType, quantize_dynamic
+
+    # Dynamic quantization: int8 weights, activations quantized per batch at runtime.
+    # Only MatMuls with constant weights are converted; attention score MatMuls stay fp32.
+    quantize_dynamic(
+        str(fp32),
+        str(int8),
+        weight_type=QuantType.QInt8,
+        per_channel=per_channel,
+        reduce_range=reduce_range,
+    )
+    log(f"saved {int8} ({int8.stat().st_size / 1e6:.0f} MB, per_channel={per_channel}, reduce_range={reduce_range})")
+
+
+# --------------------------------------------------------------------------- cli
+
+
+def main(argv=None) -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    common = argparse.ArgumentParser(add_help=False)
+    common.add_argument("--model-id", default="lightonai/mLateOn")
+    common.add_argument("--out-dir", type=Path, default=Path("mlateon-onnx"))
+    common.add_argument(
+        "--max-words",
+        type=int,
+        default=6000,
+        help="cap for the length sweep (6000 words is about 8k tokens; lower it for a quick run)",
+    )
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    e = sub.add_parser(
+        "export", parents=[common], help="export model.onnx from the safetensors checkpoint and validate it"
+    )
+    e.add_argument("--opset", type=int, default=18)
+    e.add_argument("--ir-version", type=int, default=9)
+    e.add_argument("--tol", type=float, default=1e-4, help="max |diff| per token embedding vs PyLate")
+    e.add_argument("--score-tol", type=float, default=1e-3, help="max |diff| on MaxSim scores vs PyLate")
+    e.add_argument("--cos-tol", type=float, default=1e-5, help="max (1 - cosine) per token in the length sweep")
+    e.add_argument("--no-check", action="store_true")
+
+    q = sub.add_parser("quantize", parents=[common], help="quantize model.onnx to model_int8.onnx and validate it")
+    q.add_argument("--per-tensor", action="store_true", help="per-tensor weight scales instead of per-channel")
+    q.add_argument(
+        "--reduce-range",
+        action="store_true",
+        help="7-bit weights; avoids saturation on x86 CPUs without VNNI at a small accuracy cost",
+    )
+    q.add_argument("--no-check", action="store_true")
+
+    c = sub.add_parser("check", parents=[common], help="validate an existing ONNX file against PyLate")
+    c.add_argument("--onnx", type=Path, required=True)
+    c.add_argument("--fp32", type=Path, default=None, help="also compare against this fp32 ONNX")
+    c.add_argument("--tol", type=float, default=None)
+    c.add_argument("--score-tol", type=float, default=None)
+    c.add_argument("--cos-tol", type=float, default=None)
+
+    a = p.parse_args(argv)
+    fp32, int8 = a.out_dir / "model.onnx", a.out_dir / "model_int8.onnx"
+
+    if a.cmd == "export":
+        export(a.model_id, fp32, a.opset, a.ir_version)
+        if not a.no_check:
+            check(a.model_id, fp32, a.tol, a.score_tol, a.max_words, cos_tol=a.cos_tol)
+    elif a.cmd == "quantize":
+        quantize(fp32, int8, per_channel=not a.per_tensor, reduce_range=a.reduce_range)
+        if not a.no_check:
+            # INT8 cannot match fp32 to 1e-4; the MaxSim top-1 check is the hard gate.
+            check(a.model_id, int8, None, None, a.max_words, fp32_path=fp32)
+    else:
+        check(a.model_id, a.onnx, a.tol, a.score_tol, a.max_words, fp32_path=a.fp32, cos_tol=a.cos_tol)
+    log("\nOK")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c3c2be391b1b101273bab96c5f241b0b93922e39947d5ac6d3883b05cbcd2e2e
-size 1247529629
+oid sha256:8ef72047e438bdb17993fd75fe68944bb91d6fd9208881c4ca638f0c337a4ca5
+size 1247204767

model_int8.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:02260bc68c4d942a3a65b1b0624c11e69b71eeb0be1a0c280ee834b39d3d5b27
-size 312704983
+oid sha256:470c17cd1a6bc2fe2553126f23694ba4b9f76c8c811b863c6fbfa65445a70ac3
+size 312794488

onnx_config.json

@@ -1,6 +1,6 @@
 {
   "model_type": "ColBERT",
-  "model_name": "lightonai/LateOn-multilingual",
+  "model_name": "lightonai/mLateOn",
   "model_class": "ModernBertModel",
   "uses_token_type_ids": false,
   "query_prefix": "[Q] ",
@@ -16,4 +16,4 @@
   "query_prefix_id": 256000,
   "document_prefix_id": 256001,
   "do_lower_case": false
-}
+}

tokenizer_config.json

@@ -6,12 +6,12 @@
   "eos_token": "<eos>",
   "is_local": true,
   "mask_token": "<mask>",
-  "max_length": 299,
+  "max_length": 8191,
   "model_input_names": [
     "input_ids",
     "attention_mask"
   ],
-  "model_max_length": 299,
+  "model_max_length": 8191,
   "pad_to_multiple_of": null,
   "pad_token": "<mask>",
   "pad_token_type_id": 0,