| --- |
| license: mit |
| base_model: nomic-ai/CodeRankEmbed |
| 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/nomic-ai/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](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 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. |
|
|
| ```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 = ["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](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} |
| } |
| ``` |
|
|
