rightkit-embed

Product-neutral local text embeddings and a durable SQLite vector store for Right Suite apps.

• Embedder runs a local ONNX embedding model. Embedder::load uses caller-supplied ONNX Runtime and model paths. It never downloads or silently switches models. It runs two models and picks one from the model directory, never from a setting: EmbeddingGemma 2 Q4 (embeddinggemma2-q4, the text-only export; selected when the directory's config.json has model_type embedding_gemma2) and the older EmbeddingGemma 300M Q4 (embeddinggemma-300m-q4, a directory with no config.json or another top-level model_type; a config.json that does not parse fails the load, src/embedder.rs:208). The suite's current model is EmbeddingGemma 2; vectors from the two models are not comparable, and pipeline_id() changes with the model. Both use the Gemma document prompt (title: none | text: ) and query prompt (task: search result | query: ), L2 normalisation, and Matryoshka prefix truncation (native 768, default 256, at most 2,048 tokens).
• VectorStore keeps little-endian f32 vectors in SQLite and ranks them by cosine similarity in Rust.

Catalog capability: embeddings. Published to crates.io (see INDEX.md). Licensed under MIT OR Apache-2.0.

When to use it

• Use it for on-device text embeddings with a model the app ships or downloads itself.
• Use VectorStore when vectors must survive restarts and be searched by cosine similarity.

Do not use it to fetch models, and do not use it to pick a model at run time. Model identity and downloads belong to the app.

Cargo features

--no-default-features --features vector builds the store without ort or tokenizers. Accelerator support comes from rightkit-onnx, which only inference enables; with inference on, it compiles CoreML on Apple targets and DirectML on Windows (rightkit-onnx/Cargo.toml:65-72).

Usage

Vector store, taken from tests/vector_e2e.rs:

use rightkit_embed::{cosine_similarity, VectorStore};

let mut store = VectorStore::open(&path, 3)?;
store.put_batch([
    ("near", "gemma-v1", &[1.0, 0.0, 0.0][..]),
    ("middle", "gemma-v1", &[0.7, 0.7, 0.0][..]),
])?;
let nearest = store.nearest("gemma-v1", &[1.0, 0.0, 0.0], 5)?;  // ranked by cosine score

Embedding, illustrative. EmbedderConfig::new(model_dir, runtime_path) and Embedder::load are the entry points in src/embedder.rs:

use rightkit_embed::{Embedder, EmbedderConfig};

let embedder = Embedder::load(EmbedderConfig::new("/abs/model-dir", "/abs/onnxruntime"))?;
let vectors = embedder.embed_batch(["first text", "second text"])?;

Batching and providers

• Batching. embed_batch and embed_query_batch tokenize all texts, order them by length, and run max_batch rows per ONNX call (EmbedderConfig::with_max_batch(n), default 16, also capped at 16384 padded tokens per call). Rows are right-padded to the chunk maximum with a zero attention mask. token_type_ids is sent only if the model declares it. Each row is pooled with its mask and normalised like a single run. Output order equals input order. Batched and single-text results agree to cosine at least 0.9999 and max element difference at most 2e-3, because Q4 kernels accumulate differently per shape. The pooling and padding math is unit-tested to 1e-6.
• Threads. EmbedderConfig::with_intra_threads(n) sets session intra-op threads. The default is 1.
• Execution provider. EmbedderConfig::with_provider(ExecutionProvider) accepts Cpu (default), Auto (DirectML on Windows, CoreML on macOS, CPU elsewhere), or an explicit DirectMl or CoreMl that errors instead of falling back. Auto registers the accelerator through rightkit-onnx, runs a probe inference, and falls back to CPU if registration, model load or the probe fails.
• Accelerators are admitted only after a vector parity check against CPU, at load. Whenever the session was built on an accelerator (Auto, DirectMl, CoreMl), Embedder::load builds a CPU session for the same model, embeds three fixed prompts (a short query, a medium document and a long document of several hundred tokens, with the crate's prefixes), each alone and then all together as one batch, through both sessions with the real inference path, and requires every pair to have cosine similarity of at least ACCELERATOR_PARITY_MIN_COSINE (0.995) on the full 768-dimension normalised vector (before Matryoshka truncation), with finite non-zero output. Under Auto a failed check falls back to the CPU session with fallback_reason such as directml -> cpu: vector parity failed: cosine 0.579 at ~20 tokens, floor 0.995. Under an explicit DirectMl or CoreMl it is EmbedError::AcceleratorParity { provider, lowest_cosine, floor, detail }, never a silent fallback and never a success. The probe still runs first under Auto; the parity check is in addition to it and cannot be turned off. Cpu does none of this (no second session, no extra inference).
• Load cost with an accelerator. One extra CPU session build and four ONNX calls on each of the two sessions (three prompts alone, then one batched call of the three; more calls only if max_batch is below 3). The CPU session is dropped before load returns when the check passes, so steady-state memory is unchanged; while loading, both sessions are resident.
• Embedder::provider_used() returns "cpu", "directml" or "coreml". Embedder::fallback_reason() says why an accelerator was dropped. Embedder::accelerator_parity() returns the parity result of the accelerator in use (AcceleratorParity { lowest_cosine, at_tokens, pairs }) so an app can show it, or None on CPU.
• Accelerated runs bucket shapes (length to 64, 128, ... 2048; rows to a power of two) to avoid per-shape recompiles.
• ExecutionProvider is defined by this crate. rightkit_onnx::ExecutionProvider values still convert.
• An accelerator needs an ONNX Runtime build that includes the provider. Pass that runtime path.
• Pipeline id. Embedder::pipeline_id() returns rkpipe1-<32 hex>, a stable fingerprint that tells the app when stored vectors need re-embedding. It covers the pipeline schema version, model id and native dimension, SHA-256 of model_q4.onnx, model_q4.onnx_data, tokenizer.json and config.json, the document and query prompt templates, pooling and normalisation descriptions, truncate dimension, max tokens, and Gemma 2 sentinel scrubbing. Provider, threads and batch size are excluded, and so is the parity check: an admitted accelerator reproduces the CPU vectors to the parity floor, so the id does not change with the provider. The first call hashes the model files; the result is cached. If a model file cannot be read, the id is rkpipe1-unavailable-… instead, which never equals a good id (src/embedder.rs:344-345).

Platform support

• CPU: all platforms.
• Accelerators (Auto, DirectMl, CoreMl): DirectML on Windows, CoreML on macOS, through rightkit-onnx. Needs a matching ONNX Runtime build, and the accelerator must pass the load-time vector parity check against CPU.

Known limits

• DirectML currently fails the parity check for EmbeddingGemma 2 Q4 on the hardware measured (Windows laptop, RTX 4070 and Intel Iris Xe, ONNX Runtime DirectML 1.24.4): cosine to the CPU vector was 0.58 at about 20 tokens, 0.59 at 40, 0.66 at 100 and 0.71-0.72 at 400, against about 0.999 between two correct runs. The output was finite, non-zero and faster than CPU, so only the parity check catches it. As a result Auto on Windows resolves to CPU with a recorded reason, and an explicit DirectMl returns EmbedError::AcceleratorParity. Other hardware or runtime versions are unmeasured and are admitted only if they pass.
• CoreML is unmeasured. It must pass the same check; nothing is known yet about the result for EmbeddingGemma 2 Q4.
• The parity prompts are subject to max_tokens: a config that truncates below the long prompt's length checks the shorter lengths only.

Tests

tests/vector_e2e.rs (vector store). tests/real_model_e2e.rs, tests/batch_real_model_e2e.rs, tests/gemma2_real_model_e2e.rs, tests/accelerator_parity_e2e.rs (real model files, accelerator parity; each skips without its inputs).

Version and changes

Published version: see INDEX.md. Release notes: CHANGELOG.md.