coderankembed-onnx / export_coderankembed_onnx.py

ONNX export of nomic-ai/CodeRankEmbed

72d8b33 verified 5 days ago

11.6 kB

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

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