# /// 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:])