Metadata-Version: 2.4
Name: okf-embed
Version: 0.2.0
Summary: Jina v5 text embeddings via ONNX Runtime (Rust core, PyO3 bindings)
Requires-Python: >=3.10
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

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

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

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

| Config | Session cold open | `encode_batch` (4 docs) | Notes |
|---|---|---|---|
| Level3, intra=16, inter=1 (**current**) | 4.7 s | **375 ms** | kept |
| Level1, intra=16, inter=1 | 5.5 s | 433 ms (+15%) | slower *and* bit-different vectors |
| Level3, intra=32, inter=1 | 4.5 s | 411 ms (+10%) | full-logical loses to phys/2 (SMT contention) |
| 4× `encode_one` vs 1× `encode_batch` | — | 389 vs 375 ms | one boundary crossing saves ~3%; sequential stays |
| Tokenizer-only cold open | 0.5 s | — | 9× cheaper than session open; budgeted reads stay cold |

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

| Level | Behaviour |
|---|---|
| Install | The wheel is a core dependency of OKFgraph; if it is missing or fails to import, the router raises a clear `RuntimeError` with the install hint — never an `ImportError` from deep inside, never a silent fallback. |
| Device | Accelerators are opportunistic: `auto`/`cuda` use CUDA when the loaded ORT registers the EP, else warn (stderr) + CPU. `used_cuda` reports the outcome. Never fatal. Unknown provider names warn and are skipped; registration failure degrades to CPU. |
| Encode | **Fail fast.** No fallback at encode time — vectors must stay bit-comparable within one index. |
| Tokenizer | No transformers in the runtime path, anywhere: internal tokenize + `count_tokens()` (== `tokenizer.encode(t, add_special_tokens=False)`) feed the context-window guard. |

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

  ```bash
  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.

