embroider — Jina v5 text + image embeddings (Rust core, PyO3)

embroider turns text into vectors — Jina v5 embeddings served through ONNX Runtime from a Rust core, with optional PyO3 bindings for Python. One frozen contract covers the whole path: task prefix → tokenize → ONNX forward → last-token pooling → L2 → Matryoshka truncate → re-normalise, so every consumer lands in the same vector space.

One engine, two consumers: bobine (PDF/Office → Markdown) reuses the ONNX plumbing, okfgraph uses the full Jina v5 text-embedding contract. It began as a clean move out of OKFgraph's rust/okf-embed — an exact port of EmbeddingEngine._encode — and stays pinned against a numpy/transformers replication by OKFgraph's parity harness (tests/test_parity.py, max abs diff ≤ 1e-5).

Single backend, kept comparable: one lean runtime with 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.

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.

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.

Lifecycle: lazy session, cheap tokenizer

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

JinaTokenizer.open fetches only tokenizer.json for exact token counts without the session. It never truncates, so counts report true length — JinaV5.count_tokens() instead reflects the session's max_length. 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.

Vision: image embeddings sharing the text-nano space

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.

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.

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 vision sessions) 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 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.