ONNX conversion of Jina AI's jina-embeddings-v4 (code task).
# Jina Embeddings v4 — code → ONNX (sub-part decomposition)
[Original Model](https://huggingface.co/jinaai/jina-embeddings-v4) | [Blog](https://jina.ai/news/jina-embeddings-v4-universal-embeddings-for-multimodal-multilingual-retrieval) | [Technical Report](https://arxiv.org/abs/2506.18902) | [API](https://jina.ai/embeddings)
## 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 `` 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: "
code : "Passage: "
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:
```bash
# 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)
```python
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()))
```
