Prince-1/jina-embeddings-v4-vllm-code

ONNX conversion of Jina AI's jina-embeddings-v4 (code task).

Jina Embeddings v4 — code → ONNX (sub-part decomposition)

Original Model | Blog | Technical Report | API

Model overview

Source checkpoint is jinaai/jina-embeddings-v4-vllm-code — a stock Qwen2.5-VL-3B with the code task LoRA merged into the base weights (no custom adapter code, single full checkpoint). It is one of three per-task variants of jina-embeddings-v4; this repo hosts the ONNX conversion of the code one.

code is the asymmetric natural-language ↔ code task: the NL query uses the Query: prefix and the code snippet uses the Passage: prefix (same convention as retrieval, and unlike text-matching, which is symmetric). Use it for code search, NL→code retrieval, and docstring↔code matching.

This is a single-vector model: the embedding is a masked mean-pool over the last hidden state, L2-normalized, 2048-d, with Matryoshka truncation to 128/256/512/1024/2048.

Decomposition

The model is split into three ONNX sub-parts (same pattern as the other recipes in this repo), so the heavy backbone is stored once and reused by both the text and image paths:

Sub-part	Input → Output	Notes
vision.onnx	pixel_values [N,1176] → image_features [N,2048]	task-agnostic (vision tower has no LoRA); grid baked at build resolution
embeddings.onnx	input_ids [B,S] (+ image_features [N,2048]) → inputs_embeds [B,S,2048]	token embeds; image features scattered into <image_pad> positions
backbone.onnx	inputs_embeds, attention_mask, position_ids [3,B,S] → last_hidden [B,S,2048]	code LoRA is merged here; MROPE position_ids host-computed

Compose at inference (all ONNX; the driver only wires sessions and pools):

text  : embeddings(ids)                     → backbone → mean-pool(attn_mask)   → L2norm
image : vision(px) → embeddings(ids, feats) → backbone → mean-pool(vision-span) → L2norm

Pooling and Matryoshka truncation happen in the driver (nothing baked), so one build serves every output dimension.

Why host-computed position_ids?

The model uses MROPE (mrope_section [16,24,24]), which onnxruntime-genai's ModelBuilder cannot emit. position_ids [3,B,S] are therefore computed on the host and fed in: cumulative positions for text, and get_rope_index(...) over the image grid for image inputs. The graph stays clean.

Prompts

The code task is asymmetric: the natural-language query uses Query: and the code snippet uses Passage:. The image prompt is the fixed template used across all tasks. manifest.json records this per build under prompts ({"query": "Query:", "document": "Passage:", "symmetric": false}).

NL query : "Query: <what the code should do>"
code     : "Passage: <source code>"
image    : "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n"

At inference pass --prefix Query for the query side and --prefix Passage for code snippets.

Files

Each build directory is self-contained (sub-parts + image_meta.npz + tokenizer/processor assets + manifest.json, which records the source hf_model id, precision, and any quantized sub-parts):

Dir	Precision	total
fp16	fp16	7.0 GB
fp32	fp32	14 GB
int8	fp16, backbone int8	4.6 GB
int4	fp16, backbone int4	3.3 GB

For the code task fp16 is the recommended production build — its quantized variants are more sensitive than the other tasks (see fidelity).

cuda_* directories, if present and empty, are placeholders. This environment's PyTorch/ORT are CPU builds, so GPU builds produce nothing there. The exported ONNX is execution-provider agnostic — the same files run on CUDAExecutionProvider via onnxruntime-gpu with no rebuild and no device flag.

Fidelity vs full PyTorch

Composed ONNX chain vs the full Qwen2_5_VLForConditionalGeneration (pooled-embedding cosine, worst of 3 text samples + 1 image; reference loaded in fp16):

Build	worst cosine	verdict
fp32	0.999973	✅
fp16	0.999972	✅
int8 (backbone)	0.998217	⚠️ just below 0.999 for this task
int4 (backbone)	0.781078	❌ not for production

The code task is more quantization-sensitive than the others. int8 clears 0.999 for retrieval / text-matching, but on code it lands at ~0.998 (worst case on the image + code samples), so build.py fails the int8 sanity gate here — use fp16 for code, or accept ~0.998 if that's tolerable for your recall target. int4 collapses to ~0.78 (vs ~0.91 for text-matching); do not use it for code.

Reproducing / using

CPU only; runs in the repo's uv project env (transformers 5.x, torchvision for the image processor). The pipeline is three task-agnostic scripts sharing common.py — point --model at the code source:

# build sub-parts — one --precision flag: fp16 (default) | fp32 | int8 | int4
uv run build.py --model vllm-text-code --output onnx/fp16                  # fp16 (recommended)
uv run build.py --model vllm-text-code --output onnx/fp32 --precision fp32
uv run build.py --model vllm-text-code --output onnx/int8 --precision int8 # ~0.998 → fails the gate
uv run build.py --model vllm-text-code --output onnx/int4 --precision int4 # lossy (see above)

# eval — accepts multiple build dirs (positional), dedupes repeats, auto-detects each one's precision
uv run eval.py --model vllm-text-code onnx/fp16 onnx/fp32 onnx/int8 onnx/int4

# inference (no PyTorch load) — code task: Query: for the NL query, Passage: for code
uv run inference.py --onnx-dir onnx/fp16 --text "read a file line by line in python" --prefix Query
uv run inference.py --onnx-dir onnx/fp16 --text "def add(a,b): return a+b"           --prefix Passage
uv run inference.py --onnx-dir onnx/fp16 --image doc.png

int8/int4 build the fp16 graph then weight-quantize the backbone in place (block-wise MatMulNBits; vision/embeddings stay fp16). build.py runs a composed self-sanity check: fp16 / fp32 / int8 must hit cosine ≥ 0.999 or the build fails (int8 fails for code, as above), while int4 only warns. The vision sub-part is identical across all tasks (no LoRA), so a single vision.onnx can be shared to save disk.

Same scripts serve the other tasks — --model vllm-retrieval (asymmetric Query: / Passage:) or --model vllm-text-matching (symmetric Query: / Query:).

Minimal ONNX Runtime example (NL query → code, cosine)

import json, numpy as np, onnxruntime as ort
from pathlib import Path
from transformers import AutoTokenizer

d = Path("onnx/fp16")
man = json.loads((d / "manifest.json").read_text())
npdt = np.float16 if man["precision"] == "fp16" else np.float32
tok = AutoTokenizer.from_pretrained(str(d))

def sess(name):  # log level raised to silence the harmless constant-fold notice
    so = ort.SessionOptions(); so.log_severity_level = 3
    return ort.InferenceSession(str(d / name), so, providers=["CPUExecutionProvider"])

emb_s, back_s = sess("embeddings.onnx"), sess("backbone.onnx")

def embed(text):
    enc = tok([text], return_tensors="np", padding="longest")
    ids, am = enc["input_ids"], enc["attention_mask"]
    pos = np.clip(np.cumsum(am, -1) - 1, 0, None)[None].repeat(3, 0)     # MROPE (text)
    e = emb_s.run(None, {"input_ids": ids, "image_features": np.zeros((0, 2048), npdt)})[0]
    h = back_s.run(None, {"inputs_embeds": e, "attention_mask": am, "position_ids": pos})[0]
    p = (h * am[..., None]).sum(1) / am.sum(1, keepdims=True)            # masked mean-pool
    return p / np.linalg.norm(p, axis=-1, keepdims=True)                # L2-norm → [1, 2048]

# asymmetric: query with Query:, code with Passage:
q = embed("Query: add two numbers in python")
c = embed("Passage: def add(a, b): return a + b")
print("cosine(query, code) =", float((q * c).sum()))