#!/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())