embroider

Jina v5 text + image embeddings over ONNX Runtime (Rust core, PyO3)

Contents

• What it is
• Models — registry, precisions, custom repos, licences
• Install
• Module layout
• Using embroider without Jina — session-policy recipe for other models
• Runtime: ONNX Runtime discovery
• Text embeddings
• Vision embeddings
• Running without downloads
• Session/threading policy
• Failure policy · Contract notes · Conformance · Testing
• Quick reference — copy-paste Python/Rust surface

What it is

embroider turns text — and images — into vectors: Jina v5 embeddings served through ONNX Runtime from a Rust core, with optional PyO3 bindings for Python.

One frozen contract

Task prefix → tokenize → ONNX forward → last-token pooling → L2 → Matryoshka truncate → re-normalise. Every consumer lands in the same vector space, or gets a loud error instead of a silent fork.

Two consumers

okfgraph uses the full Jina v5 text + vision contract. bobine (PDF/Office → Markdown) reuses the ONNX plumbing — session policy, provider fallback, CUDA probe — while owning its own converter-model sessions (RapidOCR, layout, tables, formulas).

One backend

No torch / transformers / optimum in the hot path, so every vector in an index stays directly comparable. If something can't be embedded exactly to contract, embroider errors loudly rather than quietly mixing vector spaces.

Models

Three blessed registry entries (src/acquire.rs), all non-commercial weights (licences tabled in docs/compat.md):

Custom models

Unknown ids fall through to the legacy path: the id itself is treated as an HF repo holding onnx/model.onnx (lookup_model → None, artifact_for default layout). Any downstream consumer can point embroider at a compatible repo with no registry change. Fully local weights skip acquisition entirely — see Running without downloads.

Install

PyPI wheels (Linux / Windows / macOS-arm64, Python 3.11–3.13) — okfgraph pulls it in automatically; standalone:

pip install embroider

From source (Rust toolchain + maturin; maturin develop needs pip, which uv venvs lack — build the wheel and install it instead):

maturin build --release
uv pip install --python <venv> target/wheels/embroider-*.whl --reinstall

Module layout

The default (pure-Rust) build is Python-free — no pyo3 in downstream trees; the extension-module Cargo feature gates the PyO3 bindings and is enabled only for wheel builds (maturin), the same pattern bobine uses.

Using embroider without Jina

The policy layer is model-agnostic — bobine uses it for converter models (RapidOCR, layout, tables, formulas) without ever touching JinaV5. The recipe, in full (reference implementation: bobine's engine::session_builder):

let builder = ort::session::Session::builder()?;
let tuned = embroider::SessionPolicy::ort_defaults().apply(builder)?;
let session = embroider::apply_providers(tuned, providers)?.commit_from_file(path)?;

plus embroider::cuda_available() for the auto-GPU decision and embroider::report() for the which-binary-serves-what diagnostic.

The boundary is deliberate:

Generic fetch + session-build helpers (a model-agnostic "N files from repo R" to complement the Jina-shaped acquisition) are a natural future addition — additive, no contract risk. Inference itself stays with the consumer by design.

Runtime: ONNX Runtime discovery

ort loads dynamically (load-dynamic, same pin as bobine: 2.0.0-rc.13). Resolution order: ORT_DYLIB_PATH first (user override always wins), else the pip-installed onnxruntime/onnxruntime-gpu build when unset. okfgraph's resolve_ort_dylib() runs before the native module is imported, so bobine and embroider share one ORT binary — no version/CUDA drift between ingest and import.

Text embeddings

Lazy sessions

JinaV5.open (model download + ONNX session build) is the single expensive step. OKFgraph therefore holds a lazy proxy: construction validates the wheel import and device string eagerly, but the session opens on the first real encode — PPR search, budgeted reads, diff, and doctor stay cold.

JinaV5.open(model_id, revision=None, cache_dir=None, truncate_dim=512, device="auto", max_length=None, precision=None, cpu_arena=False): max_length caps encodes at 1..=32768 (None = 8192 compat default; .max_length reports the effective limit), precision selects weights (auto/fp32/fp16, None = auto; .precision reports the effective choice), cpu_arena=False disables the CPU arena allocator (8x lower peak RSS for ~1.4x encode time, measured on the FP32 text model).

A failed session open is cached and re-raised — configuration errors fail fast once, not once per encode.

Weight precision (auto by default)

auto follows the resolved device (CUDA → FP16, CPU → FP32), so CUDA-requested-but-missing degrades to FP32 weights instead of stranding FP16 on CPU (FP16-on-CPU runs >40x slower than FP32-CPU — measured). fp16 maps the default model id to the published FP16 mirror repo (opticsWolf/jina-embeddings-v5-text-small-retrieval-onnx-fp16); explicit model ids (omni tower, mirrors) always win untouched. An explicit fp16-on-CPU warns loudly but is honoured. Do not mix precisions in one index — FP32 vs FP16 weights shift vectors, same as mixing tuning levels.

Memory: CPU arena off by default

cpu_arena=False registers the CPU execution provider explicitly with its arena allocator disabled — measured 8x lower peak RSS (15.3 → 1.9 GB on FP32) for ~1.4x encode time. Pass True only when peak throughput beats memory pressure. Vision slots likewise run with the CPU arena off.

Vision embeddings

Image bytes into the text-nano space (text-nano graphs only — the registry's text_partner pin refuses anything else).

Opening and encoding

JinaV5Vision.open(model_id="jina-v5-omni-nano-retrieval-vision", truncate_dim=512, device="auto", precision=None, gpu_mem_limit=None) embeds one already-resized RGB image (encode_image(rgb_bytes, h, w) where (h, w) is its own vision_target_size(h, w) — Pillow bicubic resize on the caller side; anything else is rejected rather than embedded under the wrong resolution contract). Image vectors compare against text-nano text vectors only (text_partner in the registry). The dynamic-grid graph takes six inputs; the three grid tensors (vision_pos_ids, interp_indices, interp_weights) are computed on the host, bit-identical to transformers (fixture-pinned). auto precision → fp16 on CUDA / fp32 on CPU; explicit fp16-on-CPU is an error (the graph stalls on CPU — it doesn't merely run slow, so fail-fast beats a hung batch). Sessions use the vision-slot policy (ORT defaults, never the text tuning), CUDA arena_extend_strategy=kSameAsRequested, optional gpu_mem_limit bytes, CPU arena off. Every encode re-checks seq == image_tokens + 15 so a wrong tokenizer fails loudly instead of shifting the space.

Running without downloads

Tokenizer-only counts

JinaTokenizer.open fetches only tokenizer.json for exact token counts without the session (~0.5 s cold, 9× cheaper than a session open). It never truncates, so counts report true length — JinaV5.count_tokens() instead reflects the session's max_length.

Explicit local files (air-gapped)

JinaV5.open_files(onnx_path, tokenizer_path, truncate_dim=512, device="auto", max_length=None, cpu_arena=False) and JinaTokenizer.open_files(tokenizer_path) skip every download. The sidecar (model.onnx_data-style) must sit next to the ONNX file — ORT resolves it relative to the model path, same as the HF cache layout. Precision selection does not apply here — the files are what they are (reported precision reads fp32; pass FP16 files explicitly for FP16 weights). OKFgraph's OKFRouter(model_path=..., tokenizer_path=...) uses them (both or neither; missing files raise FileNotFoundError at construction). Same bytes in → same vectors out (test-pinned against HF acquisition).

Session/threading policy (measured)

Tuning is Level3, intra = physical-cores/2, inter = 1 — kept because it measured fastest, not because it was inherited. Reference box: Windows, 32 logical cores, CPU-only ORT 1.29, warm model cache, best-of-5 reps on 4 fixed docs (short → ~400 tokens):

Two consequences:

• Do not mix tuning in one index. Level1 vs Level3 fuse the graph differently, so bits differ (hashes diverged at 1e-8 formatting). Same model + same build + same tuning, or re-embed.
• Sequential batching stays. Padded batching would waste attention on variable-length docs to save ~14 ms of boundary overhead — not worth the numerics risk.

SessionPolicy::ort_defaults() exists for consumers (bobine's converter-model sessions — RapidOCR, layout, tables, formulas) that never tuned — policy is data, never a forced default. Re-measure on new hardware/ORT before changing the policy.

Pitfall: stale onnxruntime.dll on Windows

Windows boxes can carry a stale C:\Windows\System32\onnxruntime.dll (v1.17.1 in the wild). With ORT_DYLIB_PATH unset, ort may load it and die with BadVersion { version_str: "1.17.1" }, followed by an abort at shutdown (fallout from ort's exit handler, not the root cause). Point ORT_DYLIB_PATH at a modern build — e.g. the venv's onnxruntime/capi/onnxruntime.dll. Same pitfall bobine documents in its docs/benchmarks.md.

Failure policy

Contract notes

• Session IO is discovered at load (input_ids + attention_mask required, token_type_ids fed only if declared — v5's export doesn't declare it, which is where generic runners fail). Output prefers last_hidden_state.
• truncate_dim validated (32–1024, warning off the Matryoshka ladder). max_length validated (1..=32768, None = 8192 compat default); MAX_LENGTH / MODEL_MAX_TOKENS are exposed for the window guard. Same (text, max_length) → same vector; short inputs are bit-identical at any limit, while longer inputs change once the old truncation lifts. A 32K-token forward is O(n²) memory — size the limit to the machine.
• Batch encoding is sequential by design (padded batches waste attention compute on variable-length docs). GIL is released during encode.
• Vision: one image per call (seq = patches/4 + 15, ≤1295 tokens). Inputs input_ids/attention_mask int64, pixel_values float32, vision_pos_ids/interp_indices int64, interp_weights float32; output sentence_embedding float32 [1,768], then the same L2 → truncate → re-normalise as text (same Matryoshka ladder).
• input_ids/attention_mask feed as int64; pooling takes the last attended token (mask_sum - 1, clamped ≥ 0).

Conformance

fixtures/golden_jina_v5_text_small.json pins the frozen vector space: 4 canonical texts × Query/Document prefixes × dims 64/512 (12 vectors) + exact token counts, generated with embroider 0.1.3 / onnxruntime 1.29.0 CPU / text-embed policy. Unchanged by 0.1.4 (max_length) and 0.1.5 (precision, cpu_arena) for short inputs at defaults. Consumers vendor this file and assert live vectors against it (abs=1e-6 — catches wrong model, pooling, prefix, or truncation; immune to cross-CPU noise). Regenerate only on an intentional contract change, which is a new minor version plus a re-index-everything notice. See docs/compat.md for the release matrix. Vision host math is pinned the same way: fixtures/vision/ (target_sizes.json, host_tensors.json, pixel_cases.json) generated from the spike's verified grid.py; the Rust port asserts bitwise equality, and fixtures/vision/e2e/ (5 real figures + torch goldens) backs the ignored end-to-end parity test.

Testing

• Rust unit tests (42, pure — no network, no dylib, no tokenizer file): device/precision parsing, precision-follows-device resolution, FP16 repo selection, model-id parsing, provider-matrix mapping, task-prefix idempotence, the L2 → truncate → re-normalise math, contract constants, open()/open_files() validation (dims, max_length) firing before I/O, the vision registry entry, the vision resize contract (incl. banker's rounding and aspect rejection), host tensors bit-identical to fixtures/vision/ on all 31 grids, the pixel pipeline bit-identical on 3 cases, and the vision fp16-on-CPU refusal.

  cargo test --locked
  

• Vision end-to-end (tests/vision_e2e.rs, ignored — needs ORT 1.29 dylib, the vision weight files, and ideally CUDA): 5 real figures vs torch-native goldens, cos ≥ 0.9999 fp32 / ≥ 0.999 fp16.

  ORT_DYLIB_PATH=.../onnxruntime.dll VISION_E2E_FP32=.../model.onnx \
    VISION_E2E_FP16=.../model.onnx VISION_E2E_TOK=.../tokenizer.json \
    VISION_E2E_CUDA=1 cargo test --offline --test vision_e2e -- --ignored
  

• Python parity lives with the consumers: OKFgraph's tests/test_parity.py (marked slow) pins Rust output against a numpy/transformers replication across dims × tasks × texts at ≤ 1e-5; tests/test_rust_backend.py / tests/test_rust_e2e.py cover the wheel import, the count-tokens contract, and real-model encodes.