feat: production-grade bring-your-own-model (BYO) support (#733)

* fix: honor dim/MRL truncation on custom and pooled embedding paths

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix: make custom embedding models work under multiprocessing and case-insensitive

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix: skip parallel custom-model integration test on Windows spawn deadlock

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat: add CustomModelSpec one-call BYO registration helper

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat: add HF-id to ONNX export helper with lazy optional deps

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix: add NOTICE retaining fastembed (Qdrant) attribution

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat: document CustomModelSpec bring-your-own-model usage

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix: format test_export.py

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: n24q02m <n24q02m@outlook.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 3 people

NOTICE

@@ -0,0 +1,6 @@
+qwen3-embed
+Copyright 2025-2026 n24q02m
+
+This product includes software developed as part of fastembed
+(https://github.com/qdrant/fastembed), Copyright 2023 Qdrant,
+licensed under the Apache License, Version 2.0.

README.md

@@ -171,6 +171,41 @@ ONNX reranker variants are scored one sequence at a time (no padding), which kee
 RoPE positions correct regardless of batch composition. See issue
 [#725](https://github.com/n24q02m/qwen3-embed/issues/725).
 
+### Custom models (bring your own)
+
+Qwen3 is the only built-in model, but any ONNX-able embedding model can be
+registered and then loaded by id. Use `CustomModelSpec` with one of the four
+output shapes: `CLS`/`MEAN` (bert-bi), `LAST_TOKEN` (causal), or `DISABLED` (raw).
+
+```python
+from qwen3_embed import CustomModelSpec, TextEmbedding
+
+# Multilingual (incl. Vietnamese) + code, CLS-pooled, 768-dim
+CustomModelSpec(
+    model_id="onnx-community/gte-multilingual-base",
+    hf="onnx-community/gte-multilingual-base",
+    model_file="onnx/model_quantized.onnx",
+    dim=768, pooling="CLS", normalization=True,
+).register()
+
+model = TextEmbedding("onnx-community/gte-multilingual-base")
+embeddings = list(model.embed(["xin chào", "def add(a, b): return a + b"]))
+```
+
+Other verified examples: `bge-m3` (`pooling="CLS"`, `dim=1024`), `EmbeddingGemma-300m`
+(`pooling="MEAN"`, `dim=768`). MRL truncation (`embed(..., dim=256)`) works for custom
+models whose vectors are Matryoshka-trained. Custom models are scored per-row, so —
+like the built-in INT8 reranker — their scores are batch-invariant by construction.
+
+PyTorch-only models can be converted first (in a throwaway env, since the export
+deps don't co-resolve with the lean runtime pins):
+
+```python
+# pip install "optimum[exporters]" torch transformers onnx
+from qwen3_embed.export import export_to_onnx
+export_to_onnx("intfloat/multilingual-e5-base", "./e5-onnx")
+```
+
 ## Configuration
 
 ### GPU Acceleration

qwen3_embed/__init__.py

@@ -1,5 +1,6 @@
 import importlib.metadata
 
+from qwen3_embed.common.custom_model import CustomModelSpec
 from qwen3_embed.common.types import Device
 from qwen3_embed.rerank.cross_encoder import TextCrossEncoder
 from qwen3_embed.text import TextEmbedding
@@ -11,6 +12,7 @@
 
 __version__ = version
 __all__ = [
+    "CustomModelSpec",
     "Device",
     "TextEmbedding",
     "TextCrossEncoder",

qwen3_embed/common/custom_model.py

@@ -0,0 +1,62 @@
+"""One-call registration helper for bring-your-own (BYO) ONNX embedding models."""
+
+from dataclasses import dataclass, field
+
+from qwen3_embed.common.model_description import (
+    CustomDenseModelDescription,
+    ModelSource,
+    PoolingType,
+)
+
+
+@dataclass
+class CustomModelSpec:
+    """One-call registration of a BYO ONNX embedding model.
+
+    Only ONNX-able models with one of the four supported output shapes are
+    accepted: bert-bi (``CLS``/``MEAN``), causal last-token (``LAST_TOKEN``), or
+    raw 3-D output (``DISABLED``).
+
+    Usage::
+
+        from qwen3_embed import CustomModelSpec, TextEmbedding
+
+        CustomModelSpec(
+            model_id="Org/gte-multilingual-base-onnx",
+            hf="Org/gte-multilingual-base-onnx",
+            model_file="onnx/model.onnx",
+            dim=768, pooling="CLS", normalization=True,
+        ).register()
+
+        model = TextEmbedding("Org/gte-multilingual-base-onnx")
+    """
+
+    model_id: str
+    hf: str | None = None
+    url: str | None = None
+    model_file: str = "onnx/model.onnx"
+    dim: int | None = None
+    pooling: str | PoolingType = PoolingType.MEAN
+    normalization: bool = True
+    max_seq_len: int | None = None
+    additional_files: list[str] = field(default_factory=list)
+
+    def register(self) -> None:
+        """Register this model with :class:`TextEmbedding` so it can be loaded by id."""
+        from qwen3_embed import TextEmbedding
+
+        if self.dim is None:
+            raise ValueError("dim is required for an embedding model")
+
+        description = CustomDenseModelDescription(
+            model=self.model_id,
+            dim=self.dim,
+            sources=ModelSource(hf=self.hf, url=self.url),
+            model_file=self.model_file,
+            additional_files=self.additional_files,
+        )
+        TextEmbedding.add_custom_model(
+            description,
+            pooling=PoolingType(self.pooling),
+            normalization=self.normalization,
+        )

qwen3_embed/common/model_description.py

@@ -52,3 +52,9 @@ class PoolingType(StrEnum):
     MEAN = "MEAN"
     LAST_TOKEN = "LAST_TOKEN"
     DISABLED = "DISABLED"
+
+
+@dataclass(frozen=True)
+class CustomDenseModelDescription(DenseModelDescription):
+    pooling: PoolingType = PoolingType.MEAN
+    normalization: bool = True

qwen3_embed/export.py

@@ -0,0 +1,42 @@
+"""Optional HuggingFace-id to ONNX export.
+
+The runtime package stays onnxruntime-only; the heavy export dependencies
+(torch, transformers, optimum) are imported lazily so importing this module does
+not pull them in. They are NOT declared as an optional extra because the lib pins
+``tokenizers``/``huggingface-hub`` versions that current ``transformers`` releases
+cannot co-resolve (see pyproject #691) — installing them as an extra would make
+the universal lock unsatisfiable. Install them yourself in a throwaway env::
+
+    pip install "optimum[exporters]" torch transformers onnx
+"""
+
+from pathlib import Path
+
+_MISSING_EXPORT_DEPS = (
+    "ONNX export needs optimum + torch + transformers. Install them in a separate "
+    'env: pip install "optimum[exporters]" torch transformers onnx'
+)
+
+
+def export_to_onnx(model_id: str, output_dir: str, *, task: str = "feature-extraction") -> str:
+    """Export an HF model + tokenizer to ONNX under ``output_dir``.
+
+    Args:
+        model_id: HuggingFace model id (e.g. ``"intfloat/multilingual-e5-base"``).
+        output_dir: Directory to write ``onnx/model.onnx`` + tokenizer files.
+        task: optimum export task; ``"feature-extraction"`` for embeddings.
+
+    Returns:
+        The ``output_dir`` path containing the exported model.
+
+    Raises:
+        ImportError: if the ``[export]`` extra is not installed.
+    """
+    try:
+        from optimum.exporters.onnx import main_export
+    except ImportError as e:  # pragma: no cover - exercised only without the extra
+        raise ImportError(_MISSING_EXPORT_DEPS) from e
+
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    main_export(model_id, output=output_dir, task=task)
+    return output_dir

qwen3_embed/text/custom_text_embedding.py

@@ -1,11 +1,11 @@
 from collections.abc import Iterable
-from dataclasses import dataclass
 from typing import Any
 
 import numpy as np
 from numpy.typing import NDArray
 
 from qwen3_embed.common.model_description import (
+    CustomDenseModelDescription,
     DenseModelDescription,
     PoolingType,
 )
@@ -15,15 +15,11 @@
 from qwen3_embed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 
 
-@dataclass(frozen=True)
-class PostprocessingConfig:
-    pooling: PoolingType
-    normalization: bool
-
-
 class CustomTextEmbedding(OnnxTextEmbedding):
-    SUPPORTED_MODELS: list[DenseModelDescription] = []
-    POSTPROCESSING_MAPPING: dict[str, PostprocessingConfig] = {}
+    # Single source of truth: model-id (lowercased) -> description carrying
+    # pooling+normalization. The description is a frozen, picklable dataclass so it
+    # survives the deepcopy into spawned worker processes.
+    _SUPPORTED: dict[str, CustomDenseModelDescription] = {}
 
     def __init__(
         self,
@@ -34,21 +30,52 @@ def __init__(
             model_name=model_name,
             **kwargs,
         )
-        self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
-        self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization
+        desc = self._resolve_description(model_name)
+        self._pooling = desc.pooling
+        self._normalization = desc.normalization
+
+    @classmethod
+    def _register(cls, desc: CustomDenseModelDescription) -> None:
+        cls._clear_model_cache()
+        cls._SUPPORTED[desc.model.lower()] = desc
+
+    @classmethod
+    def _resolve_description(cls, model_name: str) -> CustomDenseModelDescription:
+        return cls._SUPPORTED[model_name.lower()]
+
+    @classmethod
+    def _export_registry(cls) -> list[CustomDenseModelDescription]:
+        return list(cls._SUPPORTED.values())
+
+    @classmethod
+    def _import_registry(cls, payload: list[CustomDenseModelDescription]) -> None:
+        for desc in payload:
+            cls._SUPPORTED[desc.model.lower()] = desc
 
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:
-        return cls.SUPPORTED_MODELS
+        return list(cls._SUPPORTED.values())
 
     @classmethod
     def _get_worker_class(cls) -> type["CustomTextEmbeddingWorker"]:
         return CustomTextEmbeddingWorker
 
+    def _extra_worker_params(self) -> dict[str, Any]:
+        # Propagate the runtime registry so spawned workers (fresh interpreters
+        # with an empty _SUPPORTED) can resolve + re-register this custom model.
+        return {"custom_registry": self._export_registry()}
+
     def _post_process_onnx_output(
         self, output: OnnxOutputContext, **kwargs: Any
     ) -> Iterable[NumpyArray]:
-        return self._normalize(self._pool(output.model_output, output.attention_mask))
+        embeddings = self._pool(output.model_output, output.attention_mask)
+
+        # MRL: optionally truncate to requested dimension
+        dim: int | None = kwargs.get("dim")
+        if dim is not None:
+            embeddings = embeddings[:, :dim]
+
+        return self._normalize(embeddings)
 
     def _pool(
         self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None
@@ -78,19 +105,6 @@ def _pool(
     def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
         return normalize(embeddings) if self._normalization else embeddings
 
-    @classmethod
-    def add_model(
-        cls,
-        model_description: DenseModelDescription,
-        pooling: PoolingType,
-        normalization: bool,
-    ) -> None:
-        cls._clear_model_cache()
-        cls.SUPPORTED_MODELS.append(model_description)
-        cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
-            pooling=pooling, normalization=normalization
-        )
-
 
 class CustomTextEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
@@ -99,6 +113,9 @@ def init_embedding(
         cache_dir: str,
         **kwargs: Any,
     ) -> CustomTextEmbedding:
+        registry = kwargs.pop("custom_registry", None)
+        if registry is not None:
+            CustomTextEmbedding._import_registry(registry)
         return CustomTextEmbedding(
             model_name=model_name,
             cache_dir=cache_dir,

qwen3_embed/text/onnx_embedding.py

@@ -10,7 +10,7 @@
 
 # Base class model list kept empty — Qwen3 models are registered
 # in qwen3_embedding.py. Custom models can be added at runtime
-# via CustomTextEmbedding.add_model().
+# via TextEmbedding.add_custom_model().
 supported_onnx_models: list[DenseModelDescription] = []
 
 

qwen3_embed/text/onnx_text_model.py

@@ -20,6 +20,14 @@ class OnnxTextModel(OnnxModel[T]):
     def _get_worker_class(cls) -> type["TextEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
+    def _extra_worker_params(self) -> dict[str, Any]:
+        """Extra kwargs injected into each spawned worker's init.
+
+        Subclasses override this to carry state (e.g. a runtime-registered custom
+        model registry) into worker processes that start with a fresh interpreter.
+        """
+        return {}
+
     def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
         """Post-process the ONNX model output to convert it into a usable format.
 
@@ -129,6 +137,8 @@ def _embed_documents(
             if extra_session_options is not None:
                 params.update(extra_session_options)
 
+            params.update(self._extra_worker_params())
+
             pool = ParallelWorkerPool(
                 worker=self._get_worker_class(),
                 config=PoolConfig(

qwen3_embed/text/pooled_embedding.py

@@ -11,7 +11,7 @@
 from qwen3_embed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 
 # Base class model list kept empty — mean pooling models can be added
-# at runtime via CustomTextEmbedding.add_model(pooling=PoolingType.MEAN).
+# at runtime via TextEmbedding.add_custom_model(pooling=PoolingType.MEAN).
 supported_pooled_models: list[DenseModelDescription] = []
 
 
@@ -43,7 +43,14 @@ def _post_process_onnx_output(
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return self.mean_pooling(embeddings, attn_mask)
+        pooled = self.mean_pooling(embeddings, attn_mask)
+
+        # MRL: optionally truncate to requested dimension
+        dim: int | None = kwargs.get("dim")
+        if dim is not None:
+            pooled = pooled[:, :dim]
+
+        return pooled
 
 
 class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):