jina-embeddings-v4-vllm-code

<br><br>

<p align="center"> <img src="https://huggingface.co/datasets/jinaai/documentation-images/resolve/main/logo.webp" alt="Jina AI: Your Search Foundation, Supercharged!" width="150px"> </p> <p align="center"> <b>ONNX conversion of <a href="https://jina.ai/"><b>Jina AI</b></a>'s jina-embeddings-v4 (code task).</b> </p>

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:

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):

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):

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()))