bge-code-v1 β quantized ONNX
A 4-bit weight-only (MatMulNBits) ONNX export of BAAI/bge-code-v1, the code-specialized 1536-dim embedding model (Qwen2.5-Coder-1.5B backbone, 32k context, state of the art on the CoIR code-retrieval benchmark). As far as I could find there was no usable ONNX build of this model, so I made one for the local-inference tier of my code-search engine and published it here.
- File:
model_q4.onnx(~1.8 GB, 4-bit block-32 weight-only (MatMulNBits)) - Output:
last_hidden_state; apply last-token pooling over the attention mask, then L2-normalize - Query format:
<instruct>{task}\n<query>{query}β documents/passages are embedded raw - Validated for retrieval: on a query/doc test matrix the quantized model returns the same top-1 document as fp32 PyTorch for every query; embedding cosine vs fp32 is 0.952 mean / 0.929 min (full numbers in
validation_report.json) - Runs on CPU via ONNX Runtime β no GPU required, but this is a 1.5B model: expect roughly a second or more per embedding on laptop-class CPUs. For bulk indexing you want a GPU endpoint serving the original model; this build exists so the fully-local, zero-API path works.
Why this exists
Most published quantized exports tell you the method and nothing about what it cost you. The number that actually matters for an embedding model is not raw cosine similarity to fp32, it is whether retrieval still returns the same documents. This build was gated on top-1 retrieval agreement rather than a flat cosine threshold, and the first two quantization methods tried failed that gate and were thrown away.
Usage (Python, onnxruntime)
import numpy as np, onnxruntime as ort
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("disisnoturbusiness/bge-code-v1-onnx")
session = ort.InferenceSession("model_q4.onnx")
def embed(text):
enc = tok(text, return_tensors="np", truncation=True, max_length=32768)
feeds = {k: v for k, v in enc.items() if k in {i.name for i in session.get_inputs()}}
hidden = session.run(None, feeds)[0]
last = hidden[0, enc["attention_mask"][0].sum() - 1]
return last / np.linalg.norm(last)
query = embed("<instruct>Given a code search query, retrieve relevant code that answers the query.\n<query>parse json")
doc = embed("public T Deserialize<T>(string json) => JsonConvert.DeserializeObject<T>(json);")
print(float(query @ doc))
Usage (C# / .NET, Microsoft.ML.OnnxRuntime)
This is the path this export was actually built for. Tokenize with your tokenizer of choice, then:
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using var session = new InferenceSession( "model_q4.onnx" );
// inputIds / attentionMask are long[1, seqLen] from your tokenizer
var inputs = new List<NamedOnnxValue>
{
NamedOnnxValue.CreateFromTensor( "input_ids", new DenseTensor<long>( inputIds, new[] { 1, seqLen } ) ),
NamedOnnxValue.CreateFromTensor( "attention_mask", new DenseTensor<long>( attentionMask, new[] { 1, seqLen } ) )
};
using var results = session.Run( inputs );
var hidden = results.First().AsTensor<float>(); // [1, seqLen, 1536]
// last-token pooling over the attention mask, then L2 normalize
int last = seqLen - 1;
var vec = new float[1536];
for( int i = 0; i < 1536; i++ ) vec[i] = hidden[0, last, i];
double norm = Math.Sqrt( vec.Sum( v => (double)v * v ) );
for( int i = 0; i < vec.Length; i++ ) vec[i] = (float)( vec[i] / norm );
That is how it runs inside AzureDevOpsForager, a hybrid code-search engine over Azure SQL native VECTOR columns with DiskANN, RRF fusion of vector and full-text, and cross-encoder reranking (live demo).
Credits
All model weights and training by BAAI β this is just a faithful ONNX export + quantization, published under the same Apache 2.0 license.
Exported, validated and published by Dan Weaver, applied AI / RAG engineer, Michigan USA β GitHub Β· LinkedIn Β· live demo
Export: optimum (feature-extraction task) β onnxruntime quantization (4-bit block-32 weight-only (MatMulNBits)). Dynamic int8 (per-tensor and per-channel) was rejected by the parity gate β cosine vs fp32 collapsed to ~0.6 and ~0.0 respectively on this decoder. What ships is the first method that actually validated.
- Downloads last month
- 32
Model tree for disisnoturbusiness/bge-code-v1-onnx
Base model
BAAI/bge-code-v1