fix: Pass providers and cuda to multimodal models

* fix: fix onnx text embedding list supported models, do not add already registered models, add tests

* fix: autouse fixture in custom model tests

* Refactor custom models (#482)

* refactor: refactor custom models

* fix: fix types

* remove commented out code

.github/workflows/python-tests.yml

@@ -1,4 +1,5 @@
 name: Tests
+run-name: Tests (gpu)
 
 on:
   push:
@@ -21,8 +22,6 @@ jobs:
           - '3.13.x'
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
 
     runs-on: ${{ matrix.os }}
 

README.md

@@ -63,6 +63,23 @@ 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
@@ -137,6 +154,27 @@ 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,4 +1,5 @@
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import Optional, Any
 
 
@@ -28,7 +29,7 @@ class BaseModelDescription:
 @dataclass(frozen=True)
 class DenseModelDescription(BaseModelDescription):
     dim: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = None
+    tasks: Optional[dict[str, Any]] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         assert self.dim is not None, "dim is required for dense model description"
@@ -38,3 +39,9 @@ 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,6 +33,31 @@ 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/onnx_model.py

@@ -68,7 +68,15 @@ class OnnxModel(Generic[T]):
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
-                onnx_providers = [("CUDAExecutionProvider", {"device_id": device_id})]
+                # kSameAsRequested: Allocates only the requested memory, avoiding over-allocation.
+                # more precise than 'kNextPowerOfTwo', which grows memory aggressively.
+                # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
+                onnx_providers = [
+                    (
+                        "CUDAExecutionProvider",
+                        {"device_id": device_id, "arena_extend_strategy": "kSameAsRequested"},
+                    )
+                ]
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
@@ -132,5 +140,7 @@ class EmbeddingWorker(Worker, Generic[T]):
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError("Subclasses must implement this method")

fastembed/common/utils.py

@@ -5,11 +5,12 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar
+from typing import Iterable, Optional, TypeVar, Sequence
 
 import numpy as np
+from numpy.typing import NDArray
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, OnnxProvider
 
 T = TypeVar("T")
 
@@ -22,6 +23,15 @@ 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))
@@ -57,3 +67,18 @@ def get_all_punctuation() -> set[str]:
 
 def remove_non_alphanumeric(text: str) -> str:
     return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)
+
+
+def is_cuda_enabled(cuda: bool, providers: Optional[Sequence[OnnxProvider]]) -> bool:
+    """
+    Check if CUDA is enabled based on the `cuda` and `providers` parameters
+    """
+    if cuda:
+        return True
+    if not providers:
+        return False
+    if isinstance(providers, str):
+        return "CUDAExecutionProvider" in providers
+    return isinstance(providers, (list, tuple)) and any(
+        isinstance(p, str) and "CUDAExecutionProvider" in p for p in providers
+    )

fastembed/image/onnx_image_model.py

@@ -6,13 +6,14 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from PIL import Image
+import onnxruntime as ort
 
 from fastembed.image.transform.operators import Compose
 from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
-from fastembed.common.utils import iter_batch
+from fastembed.common.utils import iter_batch, is_cuda_enabled
 from fastembed.parallel_processor import ParallelWorkerPool
 
 # Holds type of the embedding result
@@ -74,7 +75,21 @@ class OnnxImageModel(OnnxModel[T]):
             encoded = np.array(self.processor(image_files))
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
-        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+
+        run_options = ort.RunOptions()
+        providers = kwargs.get("providers", None)
+        cuda = kwargs.get("cuda", False)
+        if is_cuda_enabled(cuda, providers):
+            device_id = kwargs.get("device_id", None)
+            device_id = str(device_id if isinstance(device_id, int) else 0)
+            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
+            # helps prevent excessive memory retention, especially for dynamic workloads.
+            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
+            run_options.add_run_config_entry(
+                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
+            )
+
+        model_output = self.model.run(None, onnx_input, run_options)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -104,7 +119,9 @@ class OnnxImageModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+                yield from self._post_process_onnx_output(
+                    self.onnx_embed(batch, cuda=cuda, providers=providers)
+                )
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -129,7 +146,9 @@ class OnnxImageModel(OnnxModel[T]):
 
 
 class ImageEmbeddingWorker(EmbeddingWorker[T]):
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed(batch)
+            embeddings = self.model.onnx_embed(batch, **kwargs)
             yield idx, embeddings

fastembed/late_interaction_multimodal/colpali.py

@@ -197,12 +197,11 @@ 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["input_ids"]]
+            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["pixel_values"]]
         )
         onnx_input["attention_mask"] = np.array(
-            [self.EVEN_ATTENTION_MASK for _ in onnx_input["input_ids"]]
+            [self.EVEN_ATTENTION_MASK for _ in onnx_input["pixel_values"]]
         )
         return onnx_input
 

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -6,13 +6,14 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from PIL import Image
+import onnxruntime as ort
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common import OnnxProvider, ImageInput
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer, load_preprocessor
 from fastembed.common.types import NumpyArray
-from fastembed.common.utils import iter_batch
+from fastembed.common.utils import iter_batch, is_cuda_enabled
 from fastembed.image.transform.operators import Compose
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -73,8 +74,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
             cuda=cuda,
             device_id=device_id,
         )
-        assert self.tokenizer is not None
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        assert self.tokenizer is not None
         self.processor = load_preprocessor(model_dir=model_dir)
 
     def load_onnx_model(self) -> None:
@@ -103,7 +104,21 @@ class OnnxMultimodalModel(OnnxModel[T]):
             )
 
         onnx_input = self._preprocess_onnx_text_input(onnx_input, **kwargs)
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+
+        run_options = ort.RunOptions()
+        providers = kwargs.get("providers", None)
+        cuda = kwargs.get("cuda", False)
+        if is_cuda_enabled(cuda, providers):
+            device_id = kwargs.get("device_id", None)
+            device_id = str(device_id if isinstance(device_id, int) else 0)
+            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
+            # helps prevent excessive memory retention, especially for dynamic workloads.
+            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
+            run_options.add_run_config_entry(
+                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
+            )
+
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -136,7 +151,9 @@ class OnnxMultimodalModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_text_output(self.onnx_embed_text(batch))
+                yield from self._post_process_onnx_text_output(
+                    self.onnx_embed_text(batch, cuda=cuda, providers=providers)
+                )
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -159,10 +176,6 @@ 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 = [
@@ -171,9 +184,23 @@ class OnnxMultimodalModel(OnnxModel[T]):
             ]
             assert self.processor is not None, "Processor is not initialized"
             encoded = np.array(self.processor(image_files))
-        onnx_input = self._build_onnx_image_input(encoded)
+        onnx_input = {"pixel_values": encoded}
         onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
-        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+
+        run_options = ort.RunOptions()
+        providers = kwargs.get("providers", None)
+        cuda = kwargs.get("cuda", False)
+        if is_cuda_enabled(cuda, providers):
+            device_id = kwargs.get("device_id", None)
+            device_id = str(device_id if isinstance(device_id, int) else 0)
+            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
+            # helps prevent excessive memory retention, especially for dynamic workloads.
+            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
+            run_options.add_run_config_entry(
+                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
+            )
+
+        model_output = self.model.run(None, onnx_input, run_options)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -203,7 +230,9 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_image_output(self.onnx_embed_image(batch))
+                yield from self._post_process_onnx_image_output(
+                    self.onnx_embed_image(batch, cuda=cuda, providers=providers)
+                )
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -245,9 +274,11 @@ class TextEmbeddingWorker(EmbeddingWorker[T]):
     ) -> OnnxMultimodalModel:
         raise NotImplementedError()
 
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed_text(batch)
+            onnx_output = self.model.onnx_embed_text(batch, **kwargs)
             yield idx, onnx_output
 
 
@@ -269,7 +300,9 @@ class ImageEmbeddingWorker(EmbeddingWorker[T]):
     ) -> OnnxMultimodalModel:
         raise NotImplementedError()
 
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed_image(batch)
+            embeddings = self.model.onnx_embed_image(batch, **kwargs)
             yield idx, embeddings

fastembed/parallel_processor.py

@@ -28,7 +28,9 @@ class Worker:
     def start(cls, *args: Any, **kwargs: Any) -> "Worker":
         raise NotImplementedError()
 
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError()
 
 
@@ -63,7 +65,7 @@ def _worker(
                     break
                 yield item
 
-        for processed_item in worker.process(input_queue_iterable()):
+        for processed_item in worker.process(input_queue_iterable(), **kwargs):
             output_queue.put(processed_item)
     except Exception as e:  # pylint: disable=broad-except
         logging.exception(e)

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
+import onnxruntime as ort
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -14,7 +15,7 @@ from fastembed.common.onnx_model import (
 )
 from fastembed.common.types import NumpyArray
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.utils import iter_batch
+from fastembed.common.utils import iter_batch, is_cuda_enabled
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
@@ -71,7 +72,21 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         tokenized_input = self.tokenize(pairs, **kwargs)
         inputs = self._build_onnx_input(tokenized_input)
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
-        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+
+        run_options = ort.RunOptions()
+        providers = kwargs.get("providers", None)
+        cuda = kwargs.get("cuda", False)
+        if is_cuda_enabled(cuda, providers):
+            device_id = kwargs.get("device_id", None)
+            device_id = str(device_id if isinstance(device_id, int) else 0)
+            # Enables memory arena shrinkage, freeing unused memory after each Run() cycle.
+            # Helps prevent excessive memory retention, especially for dynamic workloads.
+            # Source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
+            run_options.add_run_config_entry(
+                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
+            )
+
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
         relevant_output = outputs[0]
         scores: NumpyArray = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
@@ -110,7 +125,9 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(pairs, batch_size):
-                yield from self._post_process_onnx_output(self.onnx_embed_pairs(batch, **kwargs))
+                yield from self._post_process_onnx_output(
+                    self.onnx_embed_pairs(batch, cuda=cuda, providers=providers, **kwargs)
+                )
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -163,7 +180,9 @@ class TextRerankerWorker(EmbeddingWorker[float]):
     ) -> OnnxCrossEncoderModel:
         raise NotImplementedError()
 
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed_pairs(batch)
+            onnx_output = self.model.onnx_embed_pairs(batch, **kwargs)
             yield idx, onnx_output

fastembed/sparse/bm25.py

@@ -344,7 +344,7 @@ class Bm25Worker(Worker):
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(
-        self, items: Iterable[tuple[int, Any]]
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
     ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
         for idx, batch in items:
             onnx_output = self.model.raw_embed(batch)

fastembed/text/custom_text_embedding.py

@@ -0,0 +1,91 @@
+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,18 +93,6 @@ 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,
@@ -314,6 +302,7 @@ 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/onnx_text_model.py

@@ -4,13 +4,14 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
+import onnxruntime as ort
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.utils import iter_batch
+from fastembed.common.utils import iter_batch, is_cuda_enabled
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
@@ -82,7 +83,21 @@ class OnnxTextModel(OnnxModel[T]):
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        run_options = ort.RunOptions()
+        providers = kwargs.get("providers", None)
+        cuda = kwargs.get("cuda", False)
+
+        if is_cuda_enabled(cuda, providers):
+            device_id = kwargs.get("device_id", None)
+            device_id = str(device_id if isinstance(device_id, int) else 0)
+            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
+            # helps prevent excessive memory retention, especially for dynamic workloads.
+            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
+            run_options.add_run_config_entry(
+                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
+            )
+
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -115,7 +130,9 @@ class OnnxTextModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+                yield from self._post_process_onnx_output(
+                    self.onnx_embed(batch, cuda=cuda, providers=providers)
+                )
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -140,7 +157,9 @@ class OnnxTextModel(OnnxModel[T]):
 
 
 class TextEmbeddingWorker(EmbeddingWorker[T]):
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+    ) -> Iterable[tuple[int, OnnxOutputContext]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed(batch)
+            onnx_output = self.model.onnx_embed(batch, **kwargs)
             yield idx, onnx_output

fastembed/text/pooled_embedding.py

@@ -1,9 +1,11 @@
 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
 
@@ -93,16 +95,10 @@ class PooledEmbedding(OnnxTextEmbedding):
         return PooledEmbeddingWorker
 
     @classmethod
-    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
+    def mean_pooling(
+        cls, model_output: NumpyArray, attention_mask: NDArray[np.int64]
+    ) -> NumpyArray:
+        return mean_pooling(model_output, attention_mask)
 
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:

fastembed/text/pooled_normalized_embedding.py

@@ -109,6 +109,18 @@ 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,12 +4,13 @@ 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
+from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
 
 
 class TextEmbedding(TextEmbeddingBase):
@@ -19,6 +20,7 @@ class TextEmbedding(TextEmbeddingBase):
         PooledNormalizedEmbedding,
         PooledEmbedding,
         JinaEmbeddingV3,
+        CustomTextEmbedding,
     ]
 
     @classmethod
@@ -37,6 +39,43 @@ 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",
@@ -51,29 +90,24 @@ 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 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.",
+                "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. ",
                 UserWarning,
                 stacklevel=2,
             )
         if model_name in {
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+            "thenlper/gte-large",
             "intfloat/multilingual-e5-large",
+            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
         }:
             warnings.warn(
-                f"{model_name} has been updated as of fastembed 0.5.2, outputs are now average pooled.",
+                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.",
                 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.5.1"
+name = "fastembed-gpu"
+version = "0.6.0"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -18,7 +18,7 @@ numpy = [
     { version = ">=2.1.0", python = ">=3.13" },
     { version = ">=1.21,<2.1.0", python = "<3.10" },
 ]
-onnxruntime = [
+onnxruntime-gpu = [
     { version = ">1.20.0", python = ">=3.13" },
     { version = ">=1.17.0,<1.20.0", python = "<3.10" },
     { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
@@ -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"
+mmh3 = ">=4.1.0,<6.0.0"
 py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.test.dependencies]

tests/test_custom_models.py

@@ -0,0 +1,162 @@
+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,15 +11,13 @@ 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],
         ]
     ),
 }
@@ -56,8 +54,8 @@ def test_batch_embedding():
             result = list(model.embed_image(images, batch_size=2))
 
             for value in result:
-                batch_size, token_num, abridged_dim = expected_result.shape
-                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=1e-3)
+                token_num, abridged_dim = expected_result.shape
+                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_single_embedding():
@@ -67,7 +65,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)))
-            batch_size, token_num, abridged_dim = expected_result.shape
+            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.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
+        [-0.00986551, -0.00018734, 0.00605892, -0.03289612, -0.0387564],
     ),
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]