draft

README.md

@@ -63,23 +63,6 @@ embeddings = list(model.embed(documents))
 
 ```
 
-Dense text embedding can also be extended with models which are not in the list of supported models.
-
-```python
-from fastembed import TextEmbedding
-from fastembed.common.model_description import PoolingType, ModelSource
-
-TextEmbedding.add_custom_model(
-    model="intfloat/multilingual-e5-small",
-    pooling=PoolingType.MEAN,
-    normalization=True,
-    sources=ModelSource(hf="intfloat/multilingual-e5-small"),  # can be used with an `url` to load files from a private storage
-    dim=384,
-    model_file="onnx/model.onnx",  # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
-)
-model = TextEmbedding(model_name="intfloat/multilingual-e5-small")
-embeddings = list(model.embed(documents))
-```
 
 
 ### 🔱 Sparse text embeddings
@@ -154,27 +137,6 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
-### Late interaction multimodal models (ColPali)
-
-```python
-from fastembed import LateInteractionMultimodalEmbedding
-
-doc_images = [
-    "./path/to/qdrant_pdf_doc_1_screenshot.jpg",
-    "./path/to/colpali_pdf_doc_2_screenshot.jpg",
-]
-
-query = "What is Qdrant?"
-
-model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
-doc_images_embeddings = list(model.embed_image(doc_images))
-# shape (2, 1030, 128)
-# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
-query_embedding = model.embed_text(query)
-# shape (1, 20, 128)
-# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
-```
-
 ### 🔄 Rerankers
 ```python
 from fastembed.rerank.cross_encoder import TextCrossEncoder

fastembed/common/model_description.py

@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Optional, Any
 
 
@@ -29,7 +28,7 @@ class BaseModelDescription:
 @dataclass(frozen=True)
 class DenseModelDescription(BaseModelDescription):
     dim: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = field(default_factory=dict)
+    tasks: Optional[dict[str, Any]] = None
 
     def __post_init__(self) -> None:
         assert self.dim is not None, "dim is required for dense model description"
@@ -39,9 +38,3 @@ class DenseModelDescription(BaseModelDescription):
 class SparseModelDescription(BaseModelDescription):
     requires_idf: Optional[bool] = None
     vocab_size: Optional[int] = None
-
-
-class PoolingType(str, Enum):
-    CLS = "CLS"
-    MEAN = "MEAN"
-    DISABLED = "DISABLED"

fastembed/common/model_management.py

@@ -33,31 +33,6 @@ class ModelManagement(Generic[T]):
         """
         raise NotImplementedError()
 
-    @classmethod
-    def add_custom_model(
-        cls,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
-        """Add a custom model to the existing embedding classes based on the passed model descriptions
-
-        Model description dict should contain the fields same as in one of the model descriptions presented
-         in fastembed.common.model_description
-
-         E.g. for BaseModelDescription:
-              model: str
-              sources: ModelSource
-              model_file: str
-              description: str
-              license: str
-              size_in_GB: float
-              additional_files: list[str]
-
-        Returns:
-            None
-        """
-        raise NotImplementedError()
-
     @classmethod
     def _list_supported_models(cls) -> list[T]:
         raise NotImplementedError()

fastembed/common/utils.py

@@ -8,7 +8,6 @@ from itertools import islice
 from typing import Iterable, Optional, TypeVar
 
 import numpy as np
-from numpy.typing import NDArray
 
 from fastembed.common.types import NumpyArray
 
@@ -23,15 +22,6 @@ def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e
     return normalized_array
 
 
-def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
-    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
-    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
-    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
-    sum_mask = np.sum(input_mask_expanded, axis=1)
-    pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
-    return pooled_embeddings
-
-
 def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))

fastembed/config/dense_models.json

@@ -0,0 +1,16 @@
+{
+  "models": [
+    {
+      "model": "BAAI/bge-base-en",
+      "dim": 768,
+      "description": "Text embeddings, Unimodal (text), English...",
+      "license": "mit",
+      "size_in_GB": 0.42,
+      "sources": {
+        "hf": "Qdrant/fast-bge-base-en",
+        "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz"
+      },
+      "model_file": "model_optimized.onnx"
+    }
+  ]
+}

fastembed/late_interaction_multimodal/colpali.py

@@ -197,11 +197,12 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         Returns:
             Dict[str, NumpyArray]: ONNX input with text placeholders.
         """
+
         onnx_input["input_ids"] = np.array(
-            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["pixel_values"]]
+            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["input_ids"]]
         )
         onnx_input["attention_mask"] = np.array(
-            [self.EVEN_ATTENTION_MASK for _ in onnx_input["pixel_values"]]
+            [self.EVEN_ATTENTION_MASK for _ in onnx_input["input_ids"]]
         )
         return onnx_input
 

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -73,8 +73,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
             cuda=cuda,
             device_id=device_id,
         )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         self.processor = load_preprocessor(model_dir=model_dir)
 
     def load_onnx_model(self) -> None:
@@ -159,6 +159,10 @@ class OnnxMultimodalModel(OnnxModel[T]):
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 yield from self._post_process_onnx_text_output(batch)  # type: ignore
 
+    def _build_onnx_image_input(self, encoded: NumpyArray) -> dict[str, NumpyArray]:
+        input_name = self.model.get_inputs()[0].name  # type: ignore[union-attr]
+        return {input_name: encoded}
+
     def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
         with contextlib.ExitStack():
             image_files = [
@@ -167,7 +171,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             ]
             assert self.processor is not None, "Processor is not initialized"
             encoded = np.array(self.processor(image_files))
-        onnx_input = {"pixel_values": encoded}
+        onnx_input = self._build_onnx_image_input(encoded)
         onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
         model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)

fastembed/model_loader.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+import json
+from typing import Dict, List
+
+
+class ModelLoader:
+    def __init__(self):
+        self.config_dir = Path(__file__).parent / "configs"
+        self._models: Dict[str, List[Dict]] = {}
+
+    def load_models(self, model_type: str) -> List[Dict]:
+        if model_type not in self._models:
+            config_path = self.config_dir / f"{model_type}_models.json"
+            with open(config_path) as f:
+                self._models[model_type] = json.load(f)["models"]
+        return self._models[model_type]

fastembed/text/custom_text_embedding.py

@@ -1,91 +0,0 @@
-from typing import Optional, Sequence, Any, Iterable
-
-from dataclasses import dataclass
-
-import numpy as np
-from numpy.typing import NDArray
-
-from fastembed.common import OnnxProvider
-from fastembed.common.model_description import (
-    PoolingType,
-    DenseModelDescription,
-)
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray
-from fastembed.common.utils import normalize, mean_pooling
-from fastembed.text.onnx_embedding import OnnxTextEmbedding
-
-
-@dataclass(frozen=True)
-class PostprocessingConfig:
-    pooling: PoolingType
-    normalization: bool
-
-
-class CustomTextEmbedding(OnnxTextEmbedding):
-    SUPPORTED_MODELS: list[DenseModelDescription] = []
-    POSTPROCESSING_MAPPING: dict[str, PostprocessingConfig] = {}
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
-        lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
-    ):
-        super().__init__(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_ids=device_ids,
-            lazy_load=lazy_load,
-            device_id=device_id,
-            specific_model_path=specific_model_path,
-            **kwargs,
-        )
-        self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
-        self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        return cls.SUPPORTED_MODELS
-
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
-        return self._normalize(self._pool(output.model_output, output.attention_mask))
-
-    def _pool(
-        self, embeddings: NumpyArray, attention_mask: Optional[NDArray[np.int64]] = None
-    ) -> NumpyArray:
-        if self._pooling == PoolingType.CLS:
-            return embeddings[:, 0]
-
-        if self._pooling == PoolingType.MEAN:
-            if attention_mask is None:
-                raise ValueError("attention_mask must be provided for mean pooling")
-            return mean_pooling(embeddings, attention_mask)
-
-        if self._pooling == PoolingType.DISABLED:
-            return embeddings
-
-    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.SUPPORTED_MODELS.append(model_description)
-        cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
-            pooling=pooling, normalization=normalization
-        )

fastembed/text/onnx_embedding.py

@@ -93,6 +93,18 @@ supported_onnx_models: list[DenseModelDescription] = [
         ),
         model_file="model_optimized.onnx",
     ),
+    DenseModelDescription(
+        model="thenlper/gte-large",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=1.20,
+        sources=ModelSource(hf="qdrant/gte-large-onnx"),
+        model_file="model.onnx",
+    ),
     DenseModelDescription(
         model="mixedbread-ai/mxbai-embed-large-v1",
         dim=1024,
@@ -302,7 +314,6 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         embeddings = output.model_output
-
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
             processed_embeddings = embeddings[:, 0]
         elif embeddings.ndim == 2:  # (batch_size, embedding_dim)

fastembed/text/pooled_embedding.py

@@ -1,11 +1,9 @@
 from typing import Any, Iterable, Type
 
 import numpy as np
-from numpy.typing import NDArray
 
 from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import mean_pooling
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 from fastembed.common.model_description import DenseModelDescription, ModelSource
 
@@ -95,10 +93,16 @@ class PooledEmbedding(OnnxTextEmbedding):
         return PooledEmbeddingWorker
 
     @classmethod
-    def mean_pooling(
-        cls, model_output: NumpyArray, attention_mask: NDArray[np.int64]
-    ) -> NumpyArray:
-        return mean_pooling(model_output, attention_mask)
+    def mean_pooling(cls, model_output: NumpyArray, attention_mask: NumpyArray) -> NumpyArray:
+        token_embeddings = model_output.astype(np.float32)
+        attention_mask = attention_mask.astype(np.float32)
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
+        input_mask_expanded = input_mask_expanded.astype(np.float32)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
 
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:

fastembed/text/pooled_normalized_embedding.py

@@ -109,18 +109,6 @@ supported_pooled_normalized_models: list[DenseModelDescription] = [
         sources=ModelSource(hf="thenlper/gte-base"),
         model_file="onnx/model.onnx",
     ),
-    DenseModelDescription(
-        model="thenlper/gte-large",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=1.20,
-        sources=ModelSource(hf="qdrant/gte-large-onnx"),
-        model_file="model.onnx",
-    ),
 ]
 
 

fastembed/text/text_embedding.py

@@ -4,13 +4,12 @@ from dataclasses import asdict
 
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
-from fastembed.text.custom_text_embedding import CustomTextEmbedding
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.multitask_embedding import JinaEmbeddingV3
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
-from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
+from fastembed.common.model_description import DenseModelDescription
 
 
 class TextEmbedding(TextEmbeddingBase):
@@ -20,7 +19,6 @@ class TextEmbedding(TextEmbeddingBase):
         PooledNormalizedEmbedding,
         PooledEmbedding,
         JinaEmbeddingV3,
-        CustomTextEmbedding,
     ]
 
     @classmethod
@@ -39,43 +37,6 @@ class TextEmbedding(TextEmbeddingBase):
             result.extend(embedding._list_supported_models())
         return result
 
-    @classmethod
-    def add_custom_model(
-        cls,
-        model: str,
-        pooling: PoolingType,
-        normalization: bool,
-        sources: ModelSource,
-        dim: int,
-        model_file: str = "onnx/model.onnx",
-        description: str = "",
-        license: str = "",
-        size_in_gb: float = 0.0,
-        additional_files: Optional[list[str]] = None,
-    ) -> None:
-        registered_models = cls._list_supported_models()
-        for registered_model in registered_models:
-            if model == registered_model.model:
-                raise ValueError(
-                    f"Model {model} is already registered in TextEmbedding, if you still want to add this model, "
-                    f"please use another model name"
-                )
-
-        CustomTextEmbedding.add_model(
-            DenseModelDescription(
-                model=model,
-                sources=sources,
-                dim=dim,
-                model_file=model_file,
-                description=description,
-                license=license,
-                size_in_GB=size_in_gb,
-                additional_files=additional_files or [],
-            ),
-            pooling=pooling,
-            normalization=normalization,
-        )
-
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
@@ -90,24 +51,29 @@ class TextEmbedding(TextEmbeddingBase):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
             warnings.warn(
-                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
-                "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
+                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. "
+                "Please review the latest documentation and release notes to ensure compatibility with your workflow. ",
+                UserWarning,
+                stacklevel=2,
+            )
+        if model_name == "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":
+            warnings.warn(
+                "The model 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2' has been updated to "
+                "include a mean pooling layer. Please ensure your usage aligns with the new functionality. "
+                "Support for the previous version without mean pooling will be removed as of version 0.5.2.",
                 UserWarning,
                 stacklevel=2,
             )
         if model_name in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-            "thenlper/gte-large",
-            "intfloat/multilingual-e5-large",
             "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+            "intfloat/multilingual-e5-large",
         }:
             warnings.warn(
-                f"The model {model_name} now uses mean pooling instead of CLS embedding. "
-                f"In order to preserve the previous behaviour, consider either pinning fastembed version to 0.5.1 or "
-                "using `add_custom_model` functionality.",
+                f"{model_name} has been updated as of fastembed 0.5.2, outputs are now average pooled.",
                 UserWarning,
                 stacklevel=2,
             )
+
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
             if any(model_name.lower() == model.model.lower() for model in supported_models):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.6.0"
+version = "0.5.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -29,7 +29,7 @@ tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
 pillow = ">=10.3.0,<12.0.0"
-mmh3 = ">=4.1.0,<6.0.0"
+mmh3 = "^4.1.0"
 py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.test.dependencies]

tests/test_custom_models.py

@@ -1,162 +0,0 @@
-import itertools
-import os
-import numpy as np
-import pytest
-
-from fastembed.common.model_description import PoolingType, ModelSource, DenseModelDescription
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, mean_pooling
-from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
-from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
-
-
-@pytest.fixture(autouse=True)
-def restore_custom_models_fixture():
-    CustomTextEmbedding.SUPPORTED_MODELS = []
-    yield
-    CustomTextEmbedding.SUPPORTED_MODELS = []
-
-
-def test_text_custom_model():
-    is_ci = os.getenv("CI")
-    custom_model_name = "intfloat/multilingual-e5-small"
-    canonical_vector = np.array(
-        [3.1317e-02, 3.0939e-02, -3.5117e-02, -6.7274e-02, 8.5084e-02], dtype=np.float32
-    )
-    pooling = PoolingType.MEAN
-    normalization = True
-    dim = 384
-    size_in_gb = 0.47
-    source = ModelSource(hf=custom_model_name)
-
-    TextEmbedding.add_custom_model(
-        custom_model_name,
-        pooling=pooling,
-        normalization=normalization,
-        sources=source,
-        dim=dim,
-        size_in_gb=size_in_gb,
-    )
-
-    assert CustomTextEmbedding.SUPPORTED_MODELS[0] == DenseModelDescription(
-        model=custom_model_name,
-        sources=source,
-        model_file="onnx/model.onnx",
-        description="",
-        license="",
-        size_in_GB=size_in_gb,
-        additional_files=[],
-        dim=dim,
-        tasks={},
-    )
-    assert CustomTextEmbedding.POSTPROCESSING_MAPPING[custom_model_name] == PostprocessingConfig(
-        pooling=pooling, normalization=normalization
-    )
-
-    model = TextEmbedding(custom_model_name)
-    docs = ["hello world", "flag embedding"]
-    embeddings = list(model.embed(docs))
-    embeddings = np.stack(embeddings, axis=0)
-    assert embeddings.shape == (2, dim)
-
-    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
-def test_mock_add_custom_models():
-    dim = 5
-    size_in_gb = 0.1
-    source = ModelSource(hf="artificial")
-
-    num_tokens = 10
-    dummy_pooled_embedding = np.random.random((1, dim)).astype(np.float32)
-    dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(np.float32)
-    dummy_attention_mask = np.ones((1, num_tokens)).astype(np.int64)
-
-    dummy_token_output = OnnxOutputContext(
-        model_output=dummy_token_embedding, attention_mask=dummy_attention_mask
-    )
-    dummy_pooled_output = OnnxOutputContext(model_output=dummy_pooled_embedding)
-    input_data = {
-        f"{PoolingType.MEAN.lower()}-normalized": dummy_token_output,
-        f"{PoolingType.MEAN.lower()}": dummy_token_output,
-        f"{PoolingType.CLS.lower()}-normalized": dummy_token_output,
-        f"{PoolingType.CLS.lower()}": dummy_token_output,
-        f"{PoolingType.DISABLED.lower()}-normalized": dummy_pooled_output,
-        f"{PoolingType.DISABLED.lower()}": dummy_pooled_output,
-    }
-
-    expected_output = {
-        f"{PoolingType.MEAN.lower()}-normalized": normalize(
-            mean_pooling(dummy_token_embedding, dummy_attention_mask)
-        ).astype(np.float32),
-        f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
-        f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]).astype(
-            np.float32
-        ),
-        f"{PoolingType.CLS.lower()}": dummy_token_embedding[:, 0],
-        f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding).astype(
-            np.float32
-        ),
-        f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
-    }
-
-    for pooling, normalization in itertools.product(
-        (PoolingType.MEAN, PoolingType.CLS, PoolingType.DISABLED), (True, False)
-    ):
-        model_name = f"{pooling.name.lower()}{'-normalized' if normalization else ''}"
-        TextEmbedding.add_custom_model(
-            model_name,
-            pooling=pooling,
-            normalization=normalization,
-            sources=source,
-            dim=dim,
-            size_in_gb=size_in_gb,
-        )
-
-        custom_text_embedding = CustomTextEmbedding(
-            model_name,
-            lazy_load=True,
-            specific_model_path="./",  # disable model downloading and loading
-        )
-
-        post_processed_output = next(
-            iter(custom_text_embedding._post_process_onnx_output(input_data[model_name]))
-        )
-        assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3)
-
-
-def test_do_not_add_existing_model():
-    existing_base_model = "sentence-transformers/all-MiniLM-L6-v2"
-    custom_model_name = "intfloat/multilingual-e5-small"
-
-    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
-        TextEmbedding.add_custom_model(
-            existing_base_model,
-            pooling=PoolingType.MEAN,
-            normalization=True,
-            sources=ModelSource(hf=existing_base_model),
-            dim=384,
-            size_in_gb=0.47,
-        )
-
-    TextEmbedding.add_custom_model(
-        custom_model_name,
-        pooling=PoolingType.MEAN,
-        normalization=False,
-        sources=ModelSource(hf=existing_base_model),
-        dim=384,
-        size_in_gb=0.47,
-    )
-
-    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
-        TextEmbedding.add_custom_model(
-            custom_model_name,
-            pooling=PoolingType.MEAN,
-            normalization=True,
-            sources=ModelSource(hf=custom_model_name),
-            dim=384,
-            size_in_gb=0.47,
-        )

tests/test_late_interaction_multimodal.py

@@ -11,13 +11,15 @@ from tests.config import TEST_MISC_DIR
 CANONICAL_IMAGE_VALUES = {
     "Qdrant/colpali-v1.3-fp16": np.array(
         [
-            [-0.0345, -0.022, 0.0567, -0.0518, -0.0782, 0.1714, -0.1738],
-            [-0.1181, -0.099, 0.0268, 0.0774, 0.0228, 0.0563, -0.1021],
-            [-0.117, -0.0683, 0.0371, 0.0921, 0.0107, 0.0659, -0.0666],
-            [-0.1393, -0.0948, 0.037, 0.0951, -0.0126, 0.0678, -0.087],
-            [-0.0957, -0.081, 0.0404, 0.052, 0.0409, 0.0335, -0.064],
-            [-0.0626, -0.0445, 0.056, 0.0592, -0.0229, 0.0409, -0.0301],
-            [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
+            [
+                [-0.0345, -0.022, 0.0567, -0.0518, -0.0782, 0.1714, -0.1738],
+                [-0.1181, -0.099, 0.0268, 0.0774, 0.0228, 0.0563, -0.1021],
+                [-0.117, -0.0683, 0.0371, 0.0921, 0.0107, 0.0659, -0.0666],
+                [-0.1393, -0.0948, 0.037, 0.0951, -0.0126, 0.0678, -0.087],
+                [-0.0957, -0.081, 0.0404, 0.052, 0.0409, 0.0335, -0.064],
+                [-0.0626, -0.0445, 0.056, 0.0592, -0.0229, 0.0409, -0.0301],
+                [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
+            ]
         ]
     ),
 }
@@ -54,8 +56,8 @@ def test_batch_embedding():
             result = list(model.embed_image(images, batch_size=2))
 
             for value in result:
-                token_num, abridged_dim = expected_result.shape
-                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+                batch_size, token_num, abridged_dim = expected_result.shape
+                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=1e-3)
 
 
 def test_single_embedding():
@@ -65,7 +67,7 @@ def test_single_embedding():
             print("evaluating", model_name)
             model = LateInteractionMultimodalEmbedding(model_name=model_name)
             result = next(iter(model.embed_image(images, batch_size=6)))
-            token_num, abridged_dim = expected_result.shape
+            batch_size, token_num, abridged_dim = expected_result.shape
             assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 

tests/test_text_onnx_embeddings.py

@@ -52,7 +52,7 @@ CANONICAL_VECTOR_VALUES = {
         [0.0802303, 0.3700881, -4.3053818, 0.4431803, -0.271572]
     ),
     "thenlper/gte-large": np.array(
-        [-0.00986551, -0.00018734, 0.00605892, -0.03289612, -0.0387564],
+        [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
     ),
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]