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