nishparadox's picture
CodeRankEmbed int8 ONNX (dynamic QInt8 of jamie8johnson/CodeRankEmbed-onnx@151669b17375, onnxruntime 1.31.1)
79db5e4 verified
|
Raw History Blame Contribute Delete
3.63 kB
metadata
license: mit
base_model: nomic-ai/CodeRankEmbed
library_name: onnxruntime
pipeline_tag: feature-extraction
tags:
  - onnx
  - int8
  - code
  - code-retrieval
  - sentence-similarity
  - feature-extraction
language:
  - en

CodeRankEmbed, int8 ONNX

An int8-quantized ONNX build of nomic-ai/CodeRankEmbed, a 137M-parameter code embedding model (CoIR 60.1). It runs on ONNX Runtime alone: no PyTorch, no custom model code. 138 MB on disk.

Built for vibepiler, which uses it to find semantic duplicates in codebases, but it is a general-purpose code embedder.

Provenance

  • Weights: the fp32 ONNX export jamie8johnson/CodeRankEmbed-onnx at commit 151669b17375, itself a faithful conversion of the upstream PyTorch weights.
  • Quantization: onnxruntime.quantization.quantize_dynamic(weight_type=QuantType.QInt8) with onnxruntime 1.31.1. Dynamic post-training quantization of the linear layers; no fine-tuning, no calibration data.
  • tokenizer_config.json has max_length raised from a leftover 512 to 8192, matching its model_max_length and the model's n_positions; loaders that take the stricter of the two (fastembed does) otherwise truncate at 512 tokens. Everything else is a byte-for-byte copy of the export's tokenizer and config files.
  • Only the int8 file is published here. The fp32 model is available from the export above.

How to use it

Pooling is CLS (first token), then L2-normalise, matching the upstream sentence-transformers config. Context length is 8192 tokens.

For natural-language queries, prepend the upstream query prefix: Represent this query for searching relevant code: . Code (documents) gets no prefix. Code-to-code similarity needs no prefix on either side.

With fastembed:

from fastembed import TextEmbedding
from fastembed.common.model_description import ModelSource, PoolingType

TextEmbedding.add_custom_model(
    model="nishparadox/CodeRankEmbed-onnx-int8", pooling=PoolingType.CLS, normalization=True,
    sources=ModelSource(hf="nishparadox/CodeRankEmbed-onnx-int8"), dim=768, model_file="onnx/model_quantized.onnx",
)
model = TextEmbedding("nishparadox/CodeRankEmbed-onnx-int8")
vectors = list(model.embed(["def add(a, b): return a + b", "def total(xs): return sum(xs)"]))

With onnxruntime and tokenizers directly:

import numpy as np, onnxruntime as ort
from tokenizers import Tokenizer

tok = Tokenizer.from_file("tokenizer.json"); tok.enable_truncation(8192); tok.enable_padding()
sess = ort.InferenceSession("onnx/model_quantized.onnx")
enc = tok.encode_batch(texts)
ids = np.array([e.ids for e in enc], dtype=np.int64); mask = np.array([e.attention_mask for e in enc], dtype=np.int64)
tokens = sess.run(["token_embeddings"], {"input_ids": ids, "attention_mask": mask})[0]
emb = tokens[:, 0]; emb /= np.linalg.norm(emb, axis=1, keepdims=True)

What quantization cost

Measured on an Apple-silicon laptop, CPU only, embedding whole Python functions from a 1,037-function repository, batches sorted by length:

build functions / s twin-pair separation*
fp32 export 16 +0.100
this int8 build 33 +0.086

* cosine of a known pair of functions that do the same job with different code (a for loop vs a sum over a generator), minus the 99th percentile of all unrelated pairs in the same file. Both builds rank the pair as each other's nearest neighbour.

License

MIT, inherited from the upstream model.