zuliaio/coderankembed-onnx

CodeRankEmbed (ONNX)

A single-file ONNX export of CodeRankEmbed, 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 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:

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 17, 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 Represent this query for searching relevant code: . Code passages take no prefix. Maximum sequence length is 8192 tokens. Embedding dimension is 768.

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 = ["Represent this query for searching relevant code: 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=8192, 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 (Apache 2.0). See the CoRNStack project. If you use this model, please cite:

@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}
}