okf-embed — Jina v5 text embeddings (Rust core, PyO3)

Exact port of EmbeddingEngine._encode: task prefix → tokenize (8192) → ONNX forward → last-token pooling → L2 → Matryoshka truncate → re-normalise. Pinned against a numpy/transformers replication of that pipeline by tests/test_parity.py (max abs diff ≤ 1e-5, cosine ≥ 0.999999).

The only embedding backend. There is no Python fallback stack, no embedding_backend selector, and no optimum/transformers in the runtime path — a mid-run stack switch would silently mix vector spaces in one index, so the design is fail-fast instead.

Install

Published to PyPI — okfgraph pulls it in automatically (platform wheels for Linux / Windows / macOS-arm64, Python 3.11–3.13):

uv sync          # editable path source, builds via maturin

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

cd rust/okf-embed
maturin build --release
uv pip install --python <venv> target/wheels/okf_embed-*.whl --reinstall

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: Windows uses capi/onnxruntime.dll, macOS uses capi/libonnxruntime.dylib, and Linux prefers versioned capi/libonnxruntime.so.* with capi/libonnxruntime.so as fallback. GPU DLL warming and Windows DLL-directory setup are best-effort and never fatal. OKFgraph's resolve_ort_dylib() runs before the native module is imported and exposes the choice as OKFRouter.ort_dylib, so both bobine and okf-embed 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. OKFRouter 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.

JinaTokenizer.open fetches only tokenizer.json for exact token counts without the session. The truncation policy is shared, so counts are identical to the session path (verified). A failed session open is cached and re-raised — configuration errors fail fast once, not once per encode.

Explicit local files (air-gapped)

JinaV5.open_files(onnx_path, tokenizer_path) 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. 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, Phase 6)

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.

Methodology: scratch script + temporary env-knob patch (both reverted; only this table committed). Re-measure on new hardware/ORT before changing the policy.

Pitfall: stale onnxruntime.dll on Windows

This dev machine carries C:\Windows\System32\onnxruntime.dll (v1.17.1). With ORT_DYLIB_PATH unset, ort loads it and dies with BadVersion { version_str: "1.17.1" }, followed by an abort at shutdown (STATUS_STACK_BUFFER_OVERRUN from ort's exit handler — fallout, not the root cause). resolve_ort_dylib() exists precisely so normal entry points never hit this; bare JinaV5.open without it does. Same pitfall as bobine (see 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 like the router (32–1024, warning off the Matryoshka ladder). MAX_LENGTH (8192) is exposed for the window guard.
• Batch encoding is sequential by design (padded batches waste attention compute on variable-length docs). GIL is released during encode.
• input_ids/attention_mask feed as int64; pooling takes the last attended token (mask_sum - 1, clamped ≥ 0).

Testing

• Rust unit tests (18, pure — no network, no dylib, no tokenizer file): device parsing, model-id parsing, provider-matrix mapping, task-prefix idempotence, the L2 → truncate → re-normalise math, contract constants, and open() validation firing before I/O.

  cd rust/okf-embed && cargo test --locked
  

  Runs in CI (Ubuntu, --locked) alongside OKFgraph's pytest jobs.

• Python parity (tests/test_parity.py, marked slow): Rust output vs a numpy/transformers replication across dims × tasks × texts, ≤ 1e-5. Needs the omni extra (transformers rides in via sentence-transformers).

• Python e2e (tests/test_rust_backend.py, tests/test_rust_e2e.py): wheel import, count_tokens contract, encode against the real model.