ONNX export of nomic-ai/CodeRankEmbed

README.md

@@ -0,0 +1,80 @@
+---
+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}
+}
+```

config.json

@@ -0,0 +1,108 @@
+{
+  "prenorm": false,
+  "parallel_block": false,
+  "parallel_block_tied_norm": false,
+  "rotary_emb_fraction": 1.0,
+  "tie_word_embeddings": true,
+  "fused_dropout_add_ln": true,
+  "fused_bias_fc": true,
+  "use_flash_attn": true,
+  "use_xentropy": true,
+  "qkv_proj_bias": false,
+  "rotary_emb_base": 1000,
+  "rotary_emb_scale_base": null,
+  "rotary_emb_interleaved": false,
+  "mlp_fc1_bias": false,
+  "mlp_fc2_bias": false,
+  "use_rms_norm": false,
+  "causal": false,
+  "type_vocab_size": 2,
+  "dense_seq_output": true,
+  "pad_vocab_size_multiple": 64,
+  "rotary_scaling_factor": null,
+  "max_trained_positions": 2048,
+  "vocab_size": 30528,
+  "n_positions": 8192,
+  "n_embd": 768,
+  "n_layer": 12,
+  "n_head": 12,
+  "n_inner": 3072,
+  "activation_function": "swiglu",
+  "resid_pdrop": 0.1,
+  "embd_pdrop": 0.1,
+  "attn_pdrop": 0.0,
+  "layer_norm_epsilon": 1e-12,
+  "initializer_range": 0.02,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "scale_attn_weights": true,
+  "use_cache": true,
+  "scale_attn_by_inverse_layer_idx": false,
+  "reorder_and_upcast_attn": false,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "float32",
+  "pruned_heads": {},
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "architectures": [
+    "NomicBertModel"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "task_specific_params": null,
+  "problem_type": null,
+  "tokenizer_class": null,
+  "prefix": null,
+  "pad_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "_name_or_path": "nomic-ai/CodeRankEmbed",
+  "transformers_version": "4.57.6",
+  "model_type": "nomic_bert",
+  "tf_legacy_loss": false,
+  "use_bfloat16": false,
+  "output_attentions": false
+}

export_coderankembed_onnx.py

@@ -0,0 +1,274 @@
+#!/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())

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12ba359756210927940f8a9ec569cc7d6d6f091e7fcfae6cbe09bba6aab8830d
+size 547637263

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,64 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "max_length": 512,
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}