export_coderankembed_onnx.py

#!/usr/bin/env python3
"""Export CodeRankEmbed to a single ONNX graph for Zulia.

nomic-ai/CodeRankEmbed is a NomicBERT encoder (custom modeling code, rotary positions, SwiGLU) fine-tuned for
natural language to code retrieval. Its Hugging Face repository ships PyTorch weights only. This script traces the
encoder into one ONNX file that DJL can load through KnownEmbeddingModel with CLS pooling.

Requirements (CPU is fine):
    pip install torch --index-url https://download.pytorch.org/whl/cpu
    pip install "transformers>=4.51,<5" einops onnx onnxruntime huggingface_hub

Usage:
    python export_coderankembed_onnx.py --out ./coderankembed-onnx
    python export_coderankembed_onnx.py --out ./coderankembed-onnx --push zuliaio/coderankembed-onnx

This script is published in the model repo it produces, so every revision of the ONNX file carries the exact
exporter that made it. The upstream model is MIT licensed (see the generated model card).
"""

import argparse
import json
import shutil
import sys
import textwrap
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer

DEFAULT_MODEL = "nomic-ai/CodeRankEmbed"
OPSET = 17
QUERY_PREFIX = "Represent this query for searching relevant code: "
MAX_TOKENS = 8192

TRACE_SAMPLES = [
    QUERY_PREFIX + "Calculate the n-th factorial",
    "def fact(n):\n if n < 0:\n  raise ValueError\n return 1 if n == 0 else n * fact(n - 1)",
]

# Validation covers short inputs, a padded batch of mixed lengths, and inputs far longer than the trace,
# which is where a rotary table baked in at the trace length would show up
LONG_CODE = "public int add(int a, int b) {\n\treturn a + b;\n}\n" * 250
VALIDATION_BATCHES = [
    [QUERY_PREFIX + "Calculate the n-th factorial",
     "def fact(n):\n if n < 0:\n  raise ValueError\n return 1 if n == 0 else n * fact(n - 1)",
     "The history of Renaissance painting in Florence."],
    [QUERY_PREFIX + "Add two integers", LONG_CODE[:400], LONG_CODE[:2000]],
    [LONG_CODE],
]


class LastHiddenState(torch.nn.Module):
    """Wrap the encoder so the ONNX graph has one output named last_hidden_state."""

    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_ids, attention_mask, token_type_ids):
        return self.model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids).last_hidden_state


def reset_rotary_caches(model):
    """The rotary modules cache cos/sin tables from earlier forward passes. Clearing them makes the trace compute the
    tables from the input length instead of capturing a cached tensor as a fixed size constant."""
    for module in model.modules():
        if hasattr(module, "_seq_len_cached"):
            module._seq_len_cached = 0
            module._cos_cached = module._sin_cached = None
            module._cos_k_cached = module._sin_k_cached = None


def encode(tokenizer, texts, tensors):
    return tokenizer(texts, padding=True, truncation=True, max_length=MAX_TOKENS, return_tensors=tensors)


def reference_outputs(model, tokenizer):
    outputs = []
    with torch.no_grad():
        for batch in VALIDATION_BATCHES:
            outputs.append(model(**encode(tokenizer, batch, "pt")).last_hidden_state.numpy())
    return outputs


def export(model, tokenizer, onnx_path: Path):
    reset_rotary_caches(model)
    enc = encode(tokenizer, TRACE_SAMPLES, "pt")
    dynamic = {0: "batch", 1: "sequence"}
    with torch.no_grad():
        torch.onnx.export(
            LastHiddenState(model),
            (enc["input_ids"], enc["attention_mask"], enc["token_type_ids"]),
            str(onnx_path),
            input_names=["input_ids", "attention_mask", "token_type_ids"],
            output_names=["last_hidden_state"],
            dynamic_axes={
                "input_ids": dynamic,
                "attention_mask": dynamic,
                "token_type_ids": dynamic,
                "last_hidden_state": dynamic,
            },
            opset_version=OPSET,
            do_constant_folding=True,
            dynamo=False,
        )


def cls_normalized(hidden: np.ndarray) -> np.ndarray:
    cls = hidden[:, 0, :]
    return cls / np.linalg.norm(cls, axis=1, keepdims=True)


def validate(references, tokenizer, onnx_path: Path):
    import onnx
    import onnxruntime as ort

    onnx.checker.check_model(str(onnx_path), full_check=False)
    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    print("ONNX inputs:", [(i.name, i.shape) for i in session.get_inputs()])
    print("ONNX outputs:", [(o.name, o.shape) for o in session.get_outputs()])

    first = None
    for batch, torch_hidden in zip(VALIDATION_BATCHES, references):
        enc = encode(tokenizer, batch, "np")
        ort_hidden = session.run(None, {k: v for k, v in enc.items()})[0]
        mask = enc["attention_mask"][..., None].astype(bool)
        # padded positions are never pooled, so only real tokens are compared
        max_diff = float(np.abs(np.where(mask, torch_hidden - ort_hidden, 0)).max())
        cls_cos = float((cls_normalized(torch_hidden) * cls_normalized(ort_hidden)).sum(1).min())
        print(f"sequence length {enc['input_ids'].shape[1]}: max |torch - onnx| {max_diff:.3e}, min CLS cosine {cls_cos:.6f}")
        if max_diff > 1e-3 or cls_cos < 0.9999:
            raise SystemExit(f"ONNX output diverges from PyTorch at sequence length {enc['input_ids'].shape[1]}")
        if first is None:
            first = cls_normalized(ort_hidden)

    near = float(first[0] @ first[1])
    far = float(first[0] @ first[2])
    print(f"cosine near={near:.4f} far={far:.4f}")
    if near <= far:
        raise SystemExit("sanity check failed: the factorial query should score its code above the unrelated text")


def write_model_card(out: Path, model_id: str):
    card = textwrap.dedent(f"""\
        ---
        license: mit
        base_model: {model_id}
        pipeline_tag: feature-extraction
        library_name: onnxruntime
        tags:
          - onnx
          - code
          - code-retrieval
          - sentence-similarity
          - feature-extraction
        language:
          - en
        ---

        # CodeRankEmbed (ONNX)

        A single-file ONNX export of [CodeRankEmbed](https://huggingface.co/{model_id}), a 137M parameter bi-encoder
        for natural language to code retrieval (Suresh et al., 2025). The upstream repository ships PyTorch weights
        with custom modeling code. This export runs on ONNX Runtime without that code.

        It is used by [Zulia](https://github.com/zuliaio/zuliasearch) as `KnownEmbeddingModel.CODE_RANK_EMBED`.

        ## Reproducing

        `export_coderankembed_onnx.py` in this repository is the exact script that produced `model.onnx`. It loads the
        upstream model, traces the encoder, checks the ONNX output against PyTorch at several sequence lengths and
        writes this folder:

        ```bash
        pip install torch --index-url https://download.pytorch.org/whl/cpu
        pip install "transformers>=4.51,<5" einops onnx onnxruntime huggingface_hub
        python export_coderankembed_onnx.py --out ./coderankembed-onnx
        ```

        ## Modifications from the upstream model

        - The encoder was traced to ONNX (opset {OPSET}, fp32). No weights were changed. Outputs match the PyTorch model
          to within floating point tolerance, including inputs of several thousand tokens.
        - The graph takes `input_ids`, `attention_mask` and `token_type_ids` and returns `last_hidden_state`.
          Pooling is not part of the graph. Apply CLS pooling (take position 0) and L2 normalize, as the upstream
          sentence-transformers configuration does.

        ## Usage

        Queries must start with the task instruction `{QUERY_PREFIX}`. Code passages take no prefix. Maximum sequence
        length is {MAX_TOKENS} tokens. Embedding dimension is 768.

        ```python
        import numpy as np
        import onnxruntime as ort
        from transformers import AutoTokenizer

        tok = AutoTokenizer.from_pretrained("zuliaio/coderankembed-onnx")
        sess = ort.InferenceSession("model.onnx")
        texts = ["{QUERY_PREFIX}Calculate the n-th factorial", "def fact(n): return 1 if n == 0 else n * fact(n - 1)"]
        enc = tok(texts, padding=True, truncation=True, max_length={MAX_TOKENS}, return_tensors="np")
        hidden = sess.run(None, dict(enc))[0]
        emb = hidden[:, 0, :]
        emb = emb / np.linalg.norm(emb, axis=1, keepdims=True)
        ```

        ## License and attribution

        MIT, unchanged from the upstream model released by the CoRNStack authors and Nomic AI. CodeRankEmbed was
        initialized from [Snowflake Arctic Embed M Long](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-long)
        (Apache 2.0). See the [CoRNStack project](https://gangiswag.github.io/cornstack/). If you use this model,
        please cite:

        ```bibtex
        @misc{{suresh2025cornstackhighqualitycontrastivedata,
          title={{CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking}},
          author={{Tarun Suresh and Revanth Gangi Reddy and Yifei Xu and Zach Nussbaum and Andriy Mulyar and Brandon Duderstadt and Heng Ji}},
          year={{2025}},
          eprint={{2412.01007}},
          archivePrefix={{arXiv}},
          primaryClass={{cs.CL}},
          url={{https://arxiv.org/abs/2412.01007}}
        }}
        ```
        """)
    (out / "README.md").write_text(card)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--revision", default=None, help="upstream commit to export, defaults to the latest")
    parser.add_argument("--out", required=True, help="output directory for model.onnx, tokenizer and model card")
    parser.add_argument("--push", metavar="REPO_ID", help="upload the output directory to this Hugging Face repo")
    args = parser.parse_args()

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    print(f"loading {args.model}")
    tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision)
    model = AutoModel.from_pretrained(args.model, revision=args.revision, trust_remote_code=True,
                                      add_pooling_layer=False, safe_serialization=True)
    model.eval()

    references = reference_outputs(model, tokenizer)

    onnx_path = out / "model.onnx"
    print(f"exporting to {onnx_path}")
    export(model, tokenizer, onnx_path)
    validate(references, tokenizer, onnx_path)

    tokenizer.save_pretrained(out)
    config = model.config.to_dict()
    config.pop("auto_map", None)  # the custom modeling code is not shipped, the ONNX graph replaces it
    (out / "config.json").write_text(json.dumps(config, indent=2))
    write_model_card(out, args.model)

    shutil.copy(Path(__file__).resolve(), out / Path(__file__).name)
    print("wrote", sorted(p.name for p in out.iterdir()))

    if args.push:
        from huggingface_hub import HfApi

        api = HfApi()
        api.create_repo(args.push, repo_type="model", exist_ok=True)
        api.upload_folder(folder_path=str(out), repo_id=args.push, repo_type="model",
                          commit_message=f"ONNX export of {args.model}")
        print(f"pushed to https://huggingface.co/{args.push}")


if __name__ == "__main__":
    sys.exit(main())