qfa.adapters.embedding#

Self-hosted ONNX embedding adapter (multilingual, dense-only).

Runs a multilingual sentence-embedding model via onnxruntime, in-process, loaded once. Behind EmbeddingPort.

Two model families are supported; they differ only in how the ONNX graph’s output is turned into one dense vector per text:

  • bge-m3 — the shipped BGE-M3 build emits an already-pooled dense_vecs head (shape (batch, dim)); the adapter takes it as-is (pooling="pre_pooled").

  • e5 — multilingual-E5 ONNX exports emit token-level last_hidden_state (shape (batch, seq, hidden)); the adapter mean-pools it over the attention mask (pooling="mean") and prepends the "query: " prefix every E5 input requires.

The dimension and token cap are per-artifact, not per-family, so they are separate knobs (dense_dim / the builder’s max_tokens): both e5-base (768-d) and e5-small (384-d) use model_kind="e5".

Security posture (asserted at construction, per the design spec):

  • trust_remote_code=False — a standard-op ONNX graph cannot execute arbitrary code or perform I/O, unlike a pickle .bin checkpoint.

  • No custom-operator libraries registered — custom ops can run native code.

  • Model pinned by a revision hash and loaded from a local mirrored artifact path — never fetched from HuggingFace at runtime in prod.

The residual attack surface is onnxruntime parser CVEs (keep patched) and conversion correctness (the one-time cosine~0.999 validation against the official reference, see the e2e-marked test).

Batching & concurrency: records are embedded in sequential batches of batch_size (default 100) — one session.run() per batch — so a large corpus never materialises one giant padded-token tensor or activation map (the dominant memory cost, since padding is to the longest row in the batch). Within a batch, intra_op_num_threads saturates cores; there is no thread/process pool across batches.

Functions

build_bge_m3_embedder(*, model_path, ...[, ...])

Build the BGE-M3 (1024-d, pre-pooled) embedder — a thin family wrapper.

build_onnx_embedder(*, model_kind, ...[, ...])

Build an OnnxEmbedder for a model family from a local artifact.

Classes

OnnxEmbedder(*, model_path, revision_hash, ...)

Self-hosted ONNX dense-only embedder (explicitly inherits the port).
class qfa.adapters.embedding.OnnxEmbedder(*, model_path: str, revision_hash: str, session: Any, tokenizer: Any, pooling: str = 'pre_pooled', query_prefix: str = '', dense_dim: int = 1024, trust_remote_code: bool = False, custom_op_libraries: tuple[str, ...] = (), intra_op_num_threads: int | None = None, batch_size: int = 100, max_tokens: int = 8192)[source]#

Bases: EmbeddingPort

Self-hosted ONNX dense-only embedder (explicitly inherits the port).

Construct the embedder and assert the required security flags.

Parameters:

  • model_path (str) – Filesystem path to the mirrored ONNX artifact (never a HF URL in production).

  • revision_hash (str) – Pinned revision/content hash of the artifact. Must be non-empty.

  • session (Any) – A pre-built onnxruntime.InferenceSession (or a test fake exposing run). Injected so unit tests need no model file.

  • tokenizer (Any) – A callable tokenizer returning {"input_ids", "attention_mask"} arrays. Injected for the same reason.

  • pooling (str) – "pre_pooled" (take outputs[0] as the dense vector, BGE-M3) or "mean" (mean-pool token-level last_hidden_state over the attention mask, E5). Any other value raises.

  • query_prefix (str) – String prepended to every text before tokenizing ("query: " for E5; empty for BGE-M3).

  • dense_dim (int) – Expected output dimensionality; each batch is validated against it so a wrong artifact/config fails loud.

  • trust_remote_code (bool) – MUST be False. Any other value raises.

  • custom_op_libraries (tuple[str, ...]) – MUST be empty. Any registered library raises.

  • intra_op_num_threads (int | None) – onnxruntime thread count; None leaves the core-count default.

  • batch_size (int) – Number of records encoded per session.run call. The corpus is embedded in sequential batches of this size to bound peak memory on large inputs. Must be >= 1.

  • max_tokens (int) – The tokenizer’s truncation cap. Used only to detect and warn about silently truncated inputs (a row whose attention-mask sum reaches this cap was almost certainly cut off). Must match the enable_truncation length the builder configured so the two stay in lock-step.

Raises:

ValueError – If a security flag is violated, pooling is unknown, revision_hash is empty, or batch_size is less than 1.

embed(texts: tuple[str, ...]) → tuple[tuple[float, ...], ...][source]#

Return one dense dense_dim-d vector per input text, in input order.

Encodes the input in sequential batches of batch_size (one session.run per batch) and concatenates the results, so a large corpus never holds one giant padded-token tensor or activation map in memory at once. Each text is prefixed with query_prefix (empty for BGE-M3) before tokenizing, and the model output is reduced to one vector per row according to pooling. Empty input returns () without touching the model.

qfa.adapters.embedding.build_onnx_embedder(*, model_kind: str, model_path: str, tokenizer_path: str, revision_hash: str, dense_dim: int, max_tokens: int | None = None, intra_op_num_threads: int | None = None, batch_size: int = 100) → OnnxEmbedder[source]#

Build an OnnxEmbedder for a model family from a local artifact.

Resolves model_kind to its pooling strategy, query prefix, and natural token cap, loads the ONNX session with the standard CPU provider and the configured thread count, and loads the tokenizer from the mirrored files. Imports of onnxruntime/tokenizers are local to this function so unit tests (which inject fakes) never trigger them.

Parameters:

  • model_kind (str) – "bge-m3" or "e5" — selects pooling + query prefix + the default token cap (see _FAMILY).

  • model_path (str) – Path to the mirrored ONNX graph.

  • tokenizer_path (str) – Path to the mirrored tokenizer file.

  • revision_hash (str) – Pinned artifact hash (passed through to the constructor’s check).

  • dense_dim (int) – Expected output dimensionality, validated per batch.

  • max_tokens (int | None) – Tokenizer truncation cap; None uses the family’s natural context (8192 for bge-m3, 512 for e5).

  • intra_op_num_threads (int | None) – onnxruntime intra-op thread count; None keeps the core-count default.

  • batch_size (int) – Records encoded per session.run call (memory bound for large corpora); passed through to the constructor.

Raises:

ValueError – If model_kind is not a known family.

qfa.adapters.embedding.build_bge_m3_embedder(*, model_path: str, tokenizer_path: str, revision_hash: str, intra_op_num_threads: int | None = None, batch_size: int = 100) → OnnxEmbedder[source]#

Build the BGE-M3 (1024-d, pre-pooled) embedder — a thin family wrapper.

Kept as the named entry point for the BGE-M3 path (used by the e2e artifact-validation test); delegates to build_onnx_embedder() with model_kind="bge-m3" and the model’s 1024-d / 8192-token defaults.