Merge branch 'main' into mean_pooling

fastembed/common/pooling.py

@@ -0,0 +1,38 @@
+import numpy as np
+from numpy import ufunc
+
+
+class LateInteractionPooler(object):
+    def __init__(self, agg_type="row", agg="mean"):
+        self.agg = agg
+        self.type = agg_type
+
+    def _pick_operation(self) -> ufunc:
+        if self.agg == "mean":
+            return np.mean
+        elif self.agg == "max":
+            return np.max
+        elif self.agg == "min":
+            return np.min
+        else:
+            raise NotImplementedError(
+                f"LateInteractionPooler only supports agg=mean,min,max, provided {self.agg}"
+            )
+
+    def pool(self, embeddings_batch) -> np.array:
+        if isinstance(embeddings_batch, np.ndarray) and len(embeddings_batch.shape) == 2:
+            embeddings_batch = [embeddings_batch]
+
+        if self.type == "row":
+            pooled_embedding = self.pool_row(embeddings_batch)
+        elif self.type == "col":
+            pooled_embedding = self.pool_col(embeddings_batch)
+        else:
+            raise ValueError("type must be 'row' or 'col'")
+        return pooled_embedding
+
+    def pool_row(self, embeddings_batch) -> np.array:
+        return self._pick_operation()(embeddings_batch, axis=-1)
+
+    def pool_col(self, embeddings_batch) -> np.array:
+        return self._pick_operation()(embeddings_batch, axis=-2)

fastembed/text/clip_embedding.py

@@ -22,8 +22,6 @@ supported_clip_models = [
 
 
 class CLIPOnnxEmbedding(OnnxTextEmbedding):
-    supported_models = supported_clip_models
-
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return CLIPEmbeddingWorker
@@ -35,11 +33,7 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
-        return cls.supported_models
-
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
+        return supported_clip_models
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return output.model_output

fastembed/text/multitask_embedding.py

@@ -41,7 +41,6 @@ class Task(int, Enum):
 class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
-    supported_models = supported_multitask_models
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
@@ -53,11 +52,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
-        return cls.supported_models
-
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
+        return supported_multitask_models
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, np.ndarray], **kwargs

fastembed/text/onnx_embedding.py

@@ -173,8 +173,6 @@ supported_onnx_models = [
 class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
     """Implementation of the Flag Embedding model."""
 
-    supported_models = supported_onnx_models
-
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
         """
@@ -183,11 +181,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
-        return cls.supported_models
-
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
+        return supported_onnx_models
 
     def __init__(
         self,

fastembed/text/pooled_embedding.py

@@ -79,8 +79,6 @@ supported_pooled_models = [
 
 
 class PooledEmbedding(OnnxTextEmbedding):
-    supported_models = supported_pooled_models
-
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return PooledEmbeddingWorker
@@ -103,11 +101,7 @@ class PooledEmbedding(OnnxTextEmbedding):
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
-        return cls.supported_models
-
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
+        return supported_pooled_models
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         if output.attention_mask is None:

fastembed/text/pooled_normalized_embedding.py

@@ -88,8 +88,6 @@ supported_pooled_normalized_models = [
 
 
 class PooledNormalizedEmbedding(PooledEmbedding):
-    supported_models = supported_pooled_normalized_models
-
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return PooledNormalizedEmbeddingWorker
@@ -101,11 +99,7 @@ class PooledNormalizedEmbedding(PooledEmbedding):
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
-        return cls.supported_models
-
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
+        return supported_pooled_normalized_models
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         if output.attention_mask is None:

fastembed/text/text_embedding.py

@@ -50,38 +50,6 @@ class TextEmbedding(TextEmbeddingBase):
             result.extend(embedding.list_supported_models())
         return result
 
-    @classmethod
-    def add_custom_model(
-        cls, model_info: dict[str, Any], mean_pooling: bool = True, normalization: bool = False
-    ) -> None:
-        """
-        Register a custom model so that TextEmbedding(...) can find it later.
-
-        Args:
-            model_info: Dictionary describing the model, e.g.:
-                {
-                    "model": "alibaba/blablabla",
-                    "dim": 512,
-                    "description": "...",
-                    "license": "apache-2.0",
-                    "size_in_GB": 1.23,
-                    "sources": { ... }  # optional
-                }
-            mean_pooling: apply mean_pooling or not.
-            normalization: apply normalization or not.
-
-        Returns:
-            None
-        """
-        if mean_pooling and not normalization:
-            PooledEmbedding.add_custom_model(model_info)
-        elif mean_pooling and normalization:
-            PooledNormalizedEmbedding.add_custom_model(model_info)
-        elif "clip" in model_info["model"].lower():
-            CLIPOnnxEmbedding.add_custom_model(model_info)
-        else:
-            OnnxTextEmbedding.add_custom_model(model_info)
-
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",

pyproject.toml

@@ -28,7 +28,7 @@ requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
-pillow = ">=10.3.0,<12.0.0"
+pillow = "^10.3.0"
 mmh3 = "^4.1.0"
 py-rust-stemmers = "^0.1.0"
 
@@ -44,7 +44,7 @@ onnx = ">=1.15.0"
 [tool.poetry.group.docs.dependencies]
 mkdocs-material = "^9.5.10"
 mkdocstrings = "^0.24.0"
-pillow = ">=10.3.0,<12.0.0"
+pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 

tests/test_add_custom_model.py

@@ -1,115 +0,0 @@
-import os
-import numpy as np
-import pytest
-
-from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
-
-canonical_vectors = [
-    {
-        "model": "intfloat/multilingual-e5-small",
-        "mean_pooling": True,
-        "normalization": True,
-        "canonical_vector": [3.1317e-02, 3.0939e-02, -3.5117e-02, -6.7274e-02, 8.5084e-02],
-    },
-    {
-        "model": "intfloat/multilingual-e5-small",
-        "mean_pooling": True,
-        "normalization": False,
-        "canonical_vector": [1.4604e-01, 1.4428e-01, -1.6376e-01, -3.1372e-01, 3.9677e-01],
-    },
-    {
-        "model": "mixedbread-ai/mxbai-embed-xsmall-v1",
-        "mean_pooling": False,
-        "normalization": False,
-        "canonical_vector": [
-            2.49407589e-02,
-            1.00189969e-02,
-            1.07807154e-02,
-            3.63860987e-02,
-            -2.27128249e-02,
-        ],
-    },
-]
-
-DIMENSIONS = {
-    "intfloat/multilingual-e5-small": 384,
-    "mixedbread-ai/mxbai-embed-xsmall-v1": 384,
-}
-
-SOURCES = {
-    "intfloat/multilingual-e5-small": "intfloat/multilingual-e5-small",
-    "mixedbread-ai/mxbai-embed-xsmall-v1": "mixedbread-ai/mxbai-embed-xsmall-v1",
-}
-
-
-@pytest.mark.parametrize("scenario", canonical_vectors)
-def test_add_custom_model_variations(scenario):
-    is_ci = bool(os.getenv("CI", False))
-
-    base_model_name = scenario["model"]
-    mean_pooling = scenario["mean_pooling"]
-    normalization = scenario["normalization"]
-    cv = np.array(scenario["canonical_vector"], dtype=np.float32)
-
-    backup_supported_models = {}
-    for embedding_cls in TextEmbedding.EMBEDDINGS_REGISTRY:
-        backup_supported_models[embedding_cls] = embedding_cls.list_supported_models().copy()
-
-    suffixes = []
-    suffixes.append("mean" if mean_pooling else "no-mean")
-    suffixes.append("norm" if normalization else "no-norm")
-    suffix_str = "-".join(suffixes)
-
-    custom_model_name = f"{base_model_name}-{suffix_str}"
-
-    dim = DIMENSIONS[base_model_name]
-    hf_source = SOURCES[base_model_name]
-
-    model_info = {
-        "model": custom_model_name,
-        "dim": dim,
-        "description": f"{base_model_name} with {suffix_str}",
-        "license": "mit",
-        "size_in_GB": 0.13,
-        "sources": {
-            "hf": hf_source,
-        },
-        "model_file": "onnx/model.onnx",
-        "additional_files": [],
-    }
-
-    if is_ci and model_info["size_in_GB"] > 1.0:
-        pytest.skip(
-            f"Skipping {custom_model_name} on CI due to size_in_GB={model_info['size_in_GB']}"
-        )
-
-    try:
-        TextEmbedding.add_custom_model(
-            model_info=model_info, mean_pooling=mean_pooling, normalization=normalization
-        )
-
-        model = TextEmbedding(model_name=custom_model_name)
-
-        docs = ["hello world", "flag embedding"]
-        embeddings = list(model.embed(docs))
-        embeddings = np.stack(embeddings, axis=0)
-
-        assert embeddings.shape == (
-            2,
-            dim,
-        ), f"Expected shape (2, {dim}) for {custom_model_name}, but got {embeddings.shape}"
-
-        num_compare_dims = cv.shape[0]
-        assert np.allclose(
-            embeddings[0, :num_compare_dims], cv, atol=1e-3
-        ), f"Embedding mismatch for {custom_model_name} (first {num_compare_dims} dims)."
-
-        assert not np.allclose(embeddings[0, :], 0.0), "Embedding should not be all zeros."
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-    finally:
-        for embedding_cls, old_list in backup_supported_models.items():
-            embedding_cls.supported_models = old_list

tests/test_pooling.py

@@ -0,0 +1,34 @@
+import os
+
+import numpy as np
+
+from fastembed.late_interaction.late_interaction_text_embedding import (
+    LateInteractionTextEmbedding,
+)
+from fastembed.common.pooling import LateInteractionPooler
+
+from tests.utils import delete_model_cache
+
+CANONICAL_COLUMN_VALUES = {
+    "colbert-ir/colbertv2.0": np.array(
+        [4.0727495e-03, -2.4026826e-03, -6.8204990e-04, -7.1383954e-05, 4.4963313e-03]
+    ),
+}
+
+docs = ["Hello World"]
+
+
+def test_batch_embedding():
+    is_ci = os.getenv("CI")
+    docs_to_embed = docs * 10
+
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionTextEmbedding(model_name=model_name)
+        pooler = LateInteractionPooler()
+        result = list(model.embed(docs_to_embed, batch_size=6))
+        pooled_result = pooler.pool(result)
+        assert np.allclose(pooled_result[0], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)