new: gpu package (#224)

When HF_HUB_OFFLINE is set to a truthy value (1, true, yes, on),
download_model() should treat local_files_only=True to avoid any
network calls. Currently, even with the local-cache-first pass (which
may fail due to missing metadata), the retry loop still calls
download_files_from_huggingface() without local_files_only, which
triggers model_info() — a network API call that immediately fails in
offline mode. This causes an unnecessary fallback to GCS download from
storage.googleapis.com.

By setting local_files_only=True when HF_HUB_OFFLINE is enabled:

1. The HF local cache pass works if the model is cached
2. The retry loop skips the network-dependent HF path entirely
3. retrieve_model_gcs() only checks for local fast-* directories
4. No network calls are attempted at all

The truthy value check aligns with huggingface_hub's own parsing of
HF_HUB_OFFLINE, which accepts "1", "true", "yes", "on" (case-insensitive).

This is critical for air-gapped / restricted environments where both
HuggingFace and Google Cloud Storage are unreachable.

Made-with: Cursor

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -39,7 +39,7 @@ body:
     attributes:
       label: FastEmbed version
       description: What version of FastEmbed are you running? python -c "import fastembed; print(fastembed.__version__)". If you're not on the latest, please upgrade and see if the problem persists.
-      placeholder: v0.5.1
+      placeholder: v0.7.4
     validations:
       required: true
   - type: dropdown

.github/workflows/python-publish.yml

@@ -25,7 +25,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: '3.9.x'
+        python-version: '3.10.x'
     - name: Install dependencies
       run: |
         python -m pip install poetry

.github/workflows/python-tests.yml

@@ -1,10 +1,11 @@
 name: Tests
+run-name: Tests (gpu)
 
 on:
   pull_request:
     branches: [ master, main, gpu ]
   workflow_dispatch:
-  
+
 
 env:
   CARGO_TERM_COLOR: always
@@ -15,15 +16,12 @@ jobs:
     strategy:
       matrix:
         python-version:
-          - '3.9.x'
           - '3.10.x'
           - '3.11.x'
           - '3.12.x'
           - '3.13.x'
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
 
     runs-on: ${{ matrix.os }}
 

.github/workflows/type-checkers.yml

@@ -8,7 +8,7 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
         os: [ubuntu-latest]
 
     name: Python ${{ matrix.python-version }} test

fastembed/common/model_description.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Optional, Any
+from typing import Any
 
 
 @dataclass(frozen=True)
 class ModelSource:
-    hf: Optional[str] = None
-    url: Optional[str] = None
+    hf: str | None = None
+    url: str | None = None
     _deprecated_tar_struct: bool = False
 
     @property
@@ -33,8 +33,8 @@ class BaseModelDescription:
 
 @dataclass(frozen=True)
 class DenseModelDescription(BaseModelDescription):
-    dim: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = field(default_factory=dict)
+    dim: int | None = None
+    tasks: dict[str, Any] | None = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         assert self.dim is not None, "dim is required for dense model description"
@@ -42,8 +42,8 @@ class DenseModelDescription(BaseModelDescription):
 
 @dataclass(frozen=True)
 class SparseModelDescription(BaseModelDescription):
-    requires_idf: Optional[bool] = None
-    vocab_size: Optional[int] = None
+    requires_idf: bool | None = None
+    vocab_size: int | None = None
 
 
 class PoolingType(str, Enum):

fastembed/common/model_management.py

@@ -3,8 +3,9 @@ import time
 import json
 import shutil
 import tarfile
+from copy import deepcopy
 from pathlib import Path
-from typing import Any, Optional, Union, TypeVar, Generic
+from typing import Any, TypeVar, Generic
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -179,8 +180,8 @@ class ModelManagement(Generic[T]):
 
         def _collect_file_metadata(
             model_dir: Path, repo_files: list[RepoFile]
-        ) -> dict[str, dict[str, Union[int, str]]]:
-            meta: dict[str, dict[str, Union[int, str]]] = {}
+        ) -> dict[str, dict[str, int | str]]:
+            meta: dict[str, dict[str, int | str]] = {}
             file_info_map = {f.path: f for f in repo_files}
             for file_path in model_dir.rglob("*"):
                 if file_path.is_file() and file_path.name != cls.METADATA_FILE:
@@ -192,9 +193,7 @@ class ModelManagement(Generic[T]):
                         }
             return meta
 
-        def _save_file_metadata(
-            model_dir: Path, meta: dict[str, dict[str, Union[int, str]]]
-        ) -> None:
+        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) -> None:
             try:
                 if not model_dir.exists():
                     model_dir.mkdir(parents=True, exist_ok=True)
@@ -224,11 +223,6 @@ class ModelManagement(Generic[T]):
                     logger.warning(
                         "Local file sizes do not match the metadata."
                     )  # do not raise, still make an attempt to load the model
-            else:
-                logger.warning(
-                    "Metadata file not found. Proceeding without checking local files."
-                )  # if users have downloaded models from hf manually, or they're updating from previous versions of
-                # fastembed
             result = snapshot_download(
                 repo_id=hf_source_repo,
                 allow_patterns=allow_patterns,
@@ -401,21 +395,43 @@ class ModelManagement(Generic[T]):
             Path: The path to the downloaded model directory.
         """
         local_files_only = kwargs.get("local_files_only", False)
-        specific_model_path: Optional[str] = kwargs.pop("specific_model_path", None)
+        hf_offline = os.environ.get("HF_HUB_OFFLINE", "").strip().upper()
+        if not local_files_only and hf_offline in {"1", "TRUE", "YES", "ON"}:
+            local_files_only = True
+            kwargs["local_files_only"] = True
+        specific_model_path: str | None = kwargs.pop("specific_model_path", None)
         if specific_model_path:
             return Path(specific_model_path)
         retries = 1 if local_files_only else retries
         hf_source = model.sources.hf
         url_source = model.sources.url
 
+        extra_patterns = [model.model_file]
+        extra_patterns.extend(model.additional_files)
+
+        if hf_source:
+            try:
+                cache_kwargs = deepcopy(kwargs)
+                cache_kwargs["local_files_only"] = True
+                return Path(
+                    cls.download_files_from_huggingface(
+                        hf_source,
+                        cache_dir=cache_dir,
+                        extra_patterns=extra_patterns,
+                        **cache_kwargs,
+                    )
+                )
+            except Exception:
+                pass
+            finally:
+                enable_progress_bars()
+
         sleep = 3.0
         while retries > 0:
             retries -= 1
 
-            if hf_source:
-                extra_patterns = [model.model_file]
-                extra_patterns.extend(model.additional_files)
-
+            if hf_source and not local_files_only:
+                # we have already tried loading with `local_files_only=True` via hf and we failed
                 try:
                     return Path(
                         cls.download_files_from_huggingface(
@@ -448,11 +464,12 @@ class ModelManagement(Generic[T]):
 
             if local_files_only:
                 logger.error("Could not find model in cache_dir")
+                break
             else:
                 logger.error(
                     f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
                 )
-            time.sleep(sleep)
-            sleep *= 3
+                time.sleep(sleep)
+                sleep *= 3
 
         raise ValueError(f"Could not load model {model.model} from any source.")

fastembed/common/onnx_model.py

@@ -1,7 +1,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
+from typing import Any, Generic, Iterable, Sequence, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort
@@ -9,7 +9,7 @@ import onnxruntime as ort
 from numpy.typing import NDArray
 from tokenizers import Tokenizer
 
-from fastembed.common.types import OnnxProvider, NumpyArray
+from fastembed.common.types import OnnxProvider, NumpyArray, Device
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -19,11 +19,14 @@ T = TypeVar("T")
 @dataclass
 class OnnxOutputContext:
     model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
+    attention_mask: NDArray[np.int64] | None = None
+    input_ids: NDArray[np.int64] | None = None
+    metadata: dict[str, Any] | None = None
 
 
 class OnnxModel(Generic[T]):
+    EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)
+
     @classmethod
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
@@ -41,8 +44,8 @@ class OnnxModel(Generic[T]):
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
-        self.model: Optional[ort.InferenceSession] = None
-        self.tokenizer: Optional[Tokenizer] = None
+        self.model: ort.InferenceSession | None = None
+        self.tokenizer: Tokenizer | None = None
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -56,24 +59,30 @@ class OnnxModel(Generic[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        available_providers = ort.get_available_providers()
+        cuda_available = "CUDAExecutionProvider" in available_providers
+        explicit_cuda = cuda is True or cuda == Device.CUDA
 
-        if cuda and providers is not None:
+        if explicit_cuda and providers is not None:
             warnings.warn(
-                f"`cuda` and `providers` are mutually exclusive parameters, cuda: {cuda}, providers: {providers}",
+                f"`cuda` and `providers` are mutually exclusive parameters, "
+                f"cuda: {cuda}, providers: {providers}. If you'd like to use providers, cuda should be one of "
+                f"[False, Device.CPU, Device.AUTO].",
                 category=UserWarning,
                 stacklevel=6,
             )
 
         if providers is not None:
             onnx_providers = list(providers)
-        elif cuda:
+        elif explicit_cuda or (cuda == Device.AUTO and cuda_available):
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
@@ -81,7 +90,6 @@ class OnnxModel(Generic[T]):
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
-        available_providers = ort.get_available_providers()
         requested_provider_names: list[str] = []
         for provider in onnx_providers:
             # check providers available
@@ -99,6 +107,9 @@ class OnnxModel(Generic[T]):
             so.intra_op_num_threads = threads
             so.inter_op_num_threads = threads
 
+        if extra_session_options is not None:
+            self.add_extra_session_options(so, extra_session_options)
+
         self.model = ort.InferenceSession(
             str(model_path), providers=onnx_providers, sess_options=so
         )
@@ -113,6 +124,38 @@ class OnnxModel(Generic[T]):
                     RuntimeWarning,
                 )
 
+    @classmethod
+    def _select_exposed_session_options(cls, model_kwargs: dict[str, Any]) -> dict[str, Any]:
+        """A convenience method to select the exposed session options in models
+
+        Args:
+            model_kwargs (dict[str, Any]): The model kwargs.
+
+        Returns:
+            dict[str, Any]: a dict with filtered exposed session options.
+        """
+        return {k: v for k, v in model_kwargs.items() if k in cls.EXPOSED_SESSION_OPTIONS}
+
+    @classmethod
+    def add_extra_session_options(
+        cls, session_options: ort.SessionOptions, extra_options: dict[str, Any]
+    ) -> None:
+        """Add extra session options to the existing options object in-place
+
+        Args:
+            session_options (ort.SessionOptions): The existing session options object.
+            extra_options (dict[str, Any]): The extra session options available in cls.EXPOSED_SESSION_OPTIONS.
+
+        Returns:
+            None
+        """
+        for option in extra_options:
+            assert (
+                option in cls.EXPOSED_SESSION_OPTIONS
+            ), f"{option} is unknown or not exposed (exposed options: {cls.EXPOSED_SESSION_OPTIONS})"
+        if "enable_cpu_mem_arena" in extra_options:
+            session_options.enable_cpu_mem_arena = extra_options["enable_cpu_mem_arena"]
+
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 

fastembed/common/preprocessor_utils.py

@@ -50,9 +50,10 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=max_context)
-    tokenizer.enable_padding(
-        pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
-    )
+    if not tokenizer.padding:
+        tokenizer.enable_padding(
+            pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
+        )
 
     for token in tokens_map.values():
         if isinstance(token, str):

fastembed/common/types.py

@@ -1,25 +1,27 @@
+from enum import Enum
 from pathlib import Path
-import sys
-from PIL import Image
-from typing import Any, Union
+from typing import Any, TypeAlias
+
 import numpy as np
 from numpy.typing import NDArray
-
-if sys.version_info >= (3, 10):
-    from typing import TypeAlias
-else:
-    from typing_extensions import TypeAlias
+from PIL import Image
 
 
-PathInput: TypeAlias = Union[str, Path]
-ImageInput: TypeAlias = Union[PathInput, Image.Image]
+class Device(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+    AUTO = "auto"
 
-OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]
-NumpyArray = Union[
-    NDArray[np.float64],
-    NDArray[np.float32],
-    NDArray[np.float16],
-    NDArray[np.int8],
-    NDArray[np.int64],
-    NDArray[np.int32],
-]
+
+PathInput: TypeAlias = str | Path
+ImageInput: TypeAlias = PathInput | Image.Image
+
+OnnxProvider: TypeAlias = str | tuple[str, dict[Any, Any]]
+NumpyArray: TypeAlias = (
+    NDArray[np.float64]
+    | NDArray[np.float32]
+    | NDArray[np.float16]
+    | NDArray[np.int8]
+    | NDArray[np.int64]
+    | NDArray[np.int32]
+)

fastembed/common/utils.py

@@ -5,7 +5,7 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar
+from typing import Iterable, TypeVar
 
 import numpy as np
 from numpy.typing import NDArray
@@ -45,7 +45,7 @@ def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
         yield b
 
 
-def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
+def define_cache_dir(cache_dir: str | None = None) -> Path:
     """
     Define the cache directory for fastembed
     """

fastembed/embedding.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any
+from typing import Any
 
 from loguru import logger
 
@@ -17,8 +17,8 @@ class JinaEmbedding(TextEmbedding):
     def __init__(
         self,
         model_name: str = "jinaai/jina-embeddings-v2-base-en",
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)

fastembed/image/image_embedding.py

@@ -1,7 +1,7 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
@@ -48,11 +48,11 @@ class ImageEmbedding(ImageEmbeddingBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -98,7 +98,7 @@ class ImageEmbedding(ImageEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -113,9 +113,9 @@ class ImageEmbedding(ImageEmbeddingBase):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

fastembed/image/image_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Any, Union
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -10,21 +10,21 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

fastembed/image/onnx_embedding.py

@@ -1,7 +1,7 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
@@ -63,14 +63,15 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
     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,
+
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -82,10 +83,11 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -98,13 +100,14 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -134,6 +137,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     @classmethod
@@ -148,9 +152,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -180,6 +184,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/image/onnx_image_model.py

@@ -2,13 +2,13 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from PIL import Image
 
 from fastembed.image.transform.operators import Compose
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
@@ -37,7 +37,7 @@ class OnnxImageModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.processor: Optional[Compose] = None
+        self.processor: Compose | None = None
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -51,10 +51,11 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -63,6 +64,7 @@ class OnnxImageModel(OnnxModel[T]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
@@ -74,9 +76,11 @@ class OnnxImageModel(OnnxModel[T]):
         return {input_name: encoded}
 
     def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack():
+        with contextlib.ExitStack() as stack:
             image_files = [
-                Image.open(image) if not isinstance(image, Image.Image) else image
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
                 for image in images
             ]
             assert self.processor is not None, "Processor is not initialized"
@@ -91,14 +95,15 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -130,6 +135,9 @@ class OnnxImageModel(OnnxModel[T]):
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),

fastembed/image/transform/functional.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import numpy as np
 from PIL import Image
 
@@ -15,7 +13,7 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Union[Image.Image, NumpyArray],
+    image: Image.Image | NumpyArray,
     size: tuple[int, int],
 ) -> NumpyArray:
     if isinstance(image, np.ndarray):
@@ -64,8 +62,8 @@ def center_crop(
 
 def normalize(
     image: NumpyArray,
-    mean: Union[float, list[float]],
-    std: Union[float, list[float]],
+    mean: float | list[float],
+    std: float | list[float],
 ) -> NumpyArray:
     num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
 
@@ -96,8 +94,8 @@ def normalize(
 
 def resize(
     image: Image.Image,
-    size: Union[int, tuple[int, int]],
-    resample: Union[int, Image.Resampling] = Image.Resampling.BILINEAR,
+    size: int | tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.BILINEAR,
 ) -> Image.Image:
     if isinstance(size, tuple):
         return image.resize(size, resample)
@@ -117,7 +115,7 @@ def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyA
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
+def pil2ndarray(image: Image.Image | NumpyArray) -> NumpyArray:
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image
@@ -126,7 +124,7 @@ def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
 def pad2square(
     image: Image.Image,
     size: int,
-    fill_color: Union[str, int, tuple[int, ...]] = 0,
+    fill_color: str | int | tuple[int, ...] = 0,
 ) -> Image.Image:
     height, width = image.height, image.width
 
@@ -147,3 +145,77 @@ def pad2square(
     new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
     new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
     return new_image
+
+
+def resize_longest_edge(
+    image: Image.Image,
+    max_size: int,
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+) -> Image.Image:
+    height, width = image.height, image.width
+    aspect_ratio = width / height
+
+    if width >= height:
+        # Width is longer
+        new_width = max_size
+        new_height = int(new_width / aspect_ratio)
+    else:
+        # Height is longer
+        new_height = max_size
+        new_width = int(new_height * aspect_ratio)
+
+    # Ensure even dimensions
+    if new_height % 2 != 0:
+        new_height += 1
+    if new_width % 2 != 0:
+        new_width += 1
+
+    return image.resize((new_width, new_height), resample)
+
+
+def crop_ndarray(
+    image: NumpyArray,
+    x1: int,
+    y1: int,
+    x2: int,
+    y2: int,
+    channel_first: bool = True,
+) -> NumpyArray:
+    if channel_first:
+        # (C, H, W) format
+        return image[:, y1:y2, x1:x2]
+    else:
+        # (H, W, C) format
+        return image[y1:y2, x1:x2, :]
+
+
+def resize_ndarray(
+    image: NumpyArray,
+    size: tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+    channel_first: bool = True,
+) -> NumpyArray:
+    # Convert to PIL-friendly format (H, W, C)
+    if channel_first:
+        img_hwc = image.transpose((1, 2, 0))
+    else:
+        img_hwc = image
+
+    # Handle different dtypes
+    if img_hwc.dtype == np.float32 or img_hwc.dtype == np.float64:
+        # Assume normalized, scale to 0-255 for PIL
+        img_hwc_scaled = (img_hwc * 255).astype(np.uint8)
+        pil_img = Image.fromarray(img_hwc_scaled, mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized).astype(np.float32) / 255.0
+    else:
+        # uint8 or similar
+        pil_img = Image.fromarray(img_hwc.astype(np.uint8), mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized)
+
+    # Convert back to original format
+    if channel_first:
+        result = result.transpose((2, 0, 1))
+
+    return result

fastembed/image/transform/operators.py

@@ -1,4 +1,5 @@
-from typing import Any, Union, Optional
+from typing import Any
+import math
 
 from PIL import Image
 
@@ -6,16 +7,19 @@ from fastembed.common.types import NumpyArray
 from fastembed.image.transform.functional import (
     center_crop,
     convert_to_rgb,
+    crop_ndarray,
     normalize,
     pil2ndarray,
     rescale,
     resize,
+    resize_longest_edge,
+    resize_ndarray,
     pad2square,
 )
 
 
 class Transform:
-    def __call__(self, images: list[Any]) -> Union[list[Image.Image], list[NumpyArray]]:
+    def __call__(self, images: list[Any]) -> list[Image.Image] | list[NumpyArray]:
         raise NotImplementedError("Subclasses must implement this method")
 
 
@@ -33,18 +37,28 @@ class CenterCrop(Transform):
 
 
 class Normalize(Transform):
-    def __init__(self, mean: Union[float, list[float]], std: Union[float, list[float]]):
+    def __init__(self, mean: float | list[float], std: float | list[float]):
         self.mean = mean
         self.std = std
 
-    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
-        return [normalize(image, mean=self.mean, std=self.std) for image in images]
+    def __call__(  # type: ignore[override]
+        self, images: list[NumpyArray] | list[list[NumpyArray]]
+    ) -> list[NumpyArray] | list[list[NumpyArray]]:
+        if images and isinstance(images[0], list):
+            # Nested structure from ImageSplitter
+            return [
+                [normalize(image, mean=self.mean, std=self.std) for image in img_patches]  # type: ignore[arg-type]
+                for img_patches in images
+            ]
+        else:
+            # Flat structure (backward compatibility)
+            return [normalize(image, mean=self.mean, std=self.std) for image in images]  # type: ignore[arg-type]
 
 
 class Resize(Transform):
     def __init__(
         self,
-        size: Union[int, tuple[int, int]],
+        size: int | tuple[int, int],
         resample: Image.Resampling = Image.Resampling.BICUBIC,
     ):
         self.size = size
@@ -58,12 +72,22 @@ class Rescale(Transform):
     def __init__(self, scale: float = 1 / 255):
         self.scale = scale
 
-    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
-        return [rescale(image, scale=self.scale) for image in images]
+    def __call__(  # type: ignore[override]
+        self, images: list[NumpyArray] | list[list[NumpyArray]]
+    ) -> list[NumpyArray] | list[list[NumpyArray]]:
+        if images and isinstance(images[0], list):
+            # Nested structure from ImageSplitter
+            return [
+                [rescale(image, scale=self.scale) for image in img_patches]  # type: ignore[arg-type]
+                for img_patches in images
+            ]
+        else:
+            # Flat structure (backward compatibility)
+            return [rescale(image, scale=self.scale) for image in images]  # type: ignore[arg-type]
 
 
 class PILtoNDarray(Transform):
-    def __call__(self, images: list[Union[Image.Image, NumpyArray]]) -> list[NumpyArray]:
+    def __call__(self, images: list[Image.Image | NumpyArray]) -> list[NumpyArray]:
         return [pil2ndarray(image) for image in images]
 
 
@@ -71,7 +95,7 @@ class PadtoSquare(Transform):
     def __init__(
         self,
         size: int,
-        fill_color: Union[str, int, tuple[int, ...]],
+        fill_color: str | int | tuple[int, ...],
     ):
         self.size = size
         self.fill_color = fill_color
@@ -82,13 +106,174 @@ class PadtoSquare(Transform):
         ]
 
 
+class ResizeLongestEdge(Transform):
+    """Resize images so the longest edge equals target size, preserving aspect ratio."""
+
+    def __init__(
+        self,
+        size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [resize_longest_edge(image, self.size, self.resample) for image in images]
+
+
+class ResizeForVisionEncoder(Transform):
+    """
+    Resize both dimensions to be multiples of vision_encoder_max_size.
+    Preserves aspect ratio approximately.
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        max_size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.max_size = max_size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+        result = []
+        for image in images:
+            # Assume (C, H, W) format
+            _, height, width = image.shape
+
+            aspect_ratio = width / height
+
+            if width >= height:
+                # Calculate new width as multiple of max_size
+                new_width = math.ceil(width / self.max_size) * self.max_size
+                new_height = int(new_width / aspect_ratio)
+                new_height = math.ceil(new_height / self.max_size) * self.max_size
+            else:
+                # Calculate new height as multiple of max_size
+                new_height = math.ceil(height / self.max_size) * self.max_size
+                new_width = int(new_height * aspect_ratio)
+                new_width = math.ceil(new_width / self.max_size) * self.max_size
+
+            # Resize using the ndarray resize function
+            resized = resize_ndarray(
+                image,
+                size=(new_width, new_height),  # PIL expects (width, height)
+                resample=self.resample,
+                channel_first=True,
+            )
+            result.append(resized)
+
+        return result
+
+
+class ImageSplitter(Transform):
+    """
+    Split images into grid of patches plus a global view.
+
+    If image dimensions exceed max_size:
+    - Divide into ceil(H/max_size) x ceil(W/max_size) patches
+    - Each patch is cropped from the image
+    - Add a global view (original resized to max_size x max_size)
+
+    If image is smaller than max_size:
+    - Return single image unchanged
+
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        max_size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.max_size = max_size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        result = []
+
+        for image in images:
+            # Assume (C, H, W) format
+            _, height, width = image.shape
+            max_height = max_width = self.max_size
+
+            frames = []
+
+            if height > max_height or width > max_width:
+                # Calculate the number of splits needed
+                num_splits_h = math.ceil(height / max_height)
+                num_splits_w = math.ceil(width / max_width)
+
+                # Calculate optimal patch dimensions
+                optimal_height = math.ceil(height / num_splits_h)
+                optimal_width = math.ceil(width / num_splits_w)
+
+                # Generate patches in grid order (row by row)
+                for r in range(num_splits_h):
+                    for c in range(num_splits_w):
+                        # Calculate crop coordinates
+                        start_x = c * optimal_width
+                        start_y = r * optimal_height
+                        end_x = min(start_x + optimal_width, width)
+                        end_y = min(start_y + optimal_height, height)
+
+                        # Crop the patch
+                        cropped = crop_ndarray(
+                            image, x1=start_x, y1=start_y, x2=end_x, y2=end_y, channel_first=True
+                        )
+                        frames.append(cropped)
+
+                # Add global view (resized to max_size x max_size)
+                global_view = resize_ndarray(
+                    image,
+                    size=(max_width, max_height),  # PIL expects (width, height)
+                    resample=self.resample,
+                    channel_first=True,
+                )
+                frames.append(global_view)
+            else:
+                # Image is small enough, no splitting needed
+                frames.append(image)
+
+            # Append (not extend) to preserve per-image grouping
+            result.append(frames)
+
+        return result
+
+
+class SquareResize(Transform):
+    """
+    Resize images to square dimensions (max_size x max_size).
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        return [
+            [
+                resize_ndarray(
+                    image, size=(self.size, self.size), resample=self.resample, channel_first=True
+                )
+            ]
+            for image in images
+        ]
+
+
 class Compose:
     def __init__(self, transforms: list[Transform]):
         self.transforms = transforms
 
     def __call__(
-        self, images: Union[list[Image.Image], list[NumpyArray]]
-    ) -> Union[list[NumpyArray], list[Image.Image]]:
+        self, images: list[Image.Image] | list[NumpyArray]
+    ) -> list[NumpyArray] | list[Image.Image]:
         for transform in self.transforms:
             images = transform(images)
         return images
@@ -118,6 +303,7 @@ class Compose:
                 Valid size keys (nested):
                     - {"height", "width"}
                     - {"shortest_edge"}
+                    - {"longest_edge"}
 
         Returns:
             Compose: Image processor.
@@ -128,6 +314,7 @@ class Compose:
         cls._get_pad2square(transforms, config)
         cls._get_center_crop(transforms, config)
         cls._get_pil2ndarray(transforms, config)
+        cls._get_image_splitting(transforms, config)
         cls._get_rescale(transforms, config)
         cls._get_normalize(transforms, config)
         return cls(transforms=transforms)
@@ -196,6 +383,25 @@ class Compose:
                             resample=resample,
                         )
                     )
+        elif mode == "Idefics3ImageProcessor":
+            if config.get("do_resize", False):
+                size = config.get("size", {})
+                if "longest_edge" not in size:
+                    raise ValueError(
+                        "Size dictionary must contain 'longest_edge' key for Idefics3ImageProcessor"
+                    )
+
+                # Handle resample parameter - can be int enum or PIL.Image.Resampling
+                resample = config.get("resample", Image.Resampling.LANCZOS)
+                if isinstance(resample, int):
+                    resample = Image.Resampling(resample)
+
+                transforms.append(
+                    ResizeLongestEdge(
+                        size=size["longest_edge"],
+                        resample=resample,
+                    )
+                )
         else:
             raise ValueError(f"Preprocessor {mode} is not supported")
 
@@ -217,6 +423,8 @@ class Compose:
             pass
         elif mode == "JinaCLIPImageProcessor":
             pass
+        elif mode == "Idefics3ImageProcessor":
+            pass
         else:
             raise ValueError(f"Preprocessor {mode} is not supported")
 
@@ -224,6 +432,28 @@ class Compose:
     def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
         transforms.append(PILtoNDarray())
 
+    @classmethod
+    def _get_image_splitting(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
+        """
+        Add image splitting transforms for Idefics3.
+        Handles conditional logic: splitting vs square resize.
+        Must be called AFTER PILtoNDarray.
+        """
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+
+        if mode == "Idefics3ImageProcessor":
+            do_splitting = config.get("do_image_splitting", False)
+            max_size = config.get("max_image_size", {}).get("longest_edge", 512)
+            resample = config.get("resample", Image.Resampling.LANCZOS)
+            if isinstance(resample, int):
+                resample = Image.Resampling(resample)
+
+            if do_splitting:
+                transforms.append(ResizeForVisionEncoder(max_size, resample))
+                transforms.append(ImageSplitter(max_size, resample))
+            else:
+                transforms.append(SquareResize(max_size, resample))
+
     @staticmethod
     def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
         if config.get("do_rescale", True):
@@ -253,7 +483,7 @@ class Compose:
             )
 
     @staticmethod
-    def _interpolation_resolver(resample: Optional[str] = None) -> Image.Resampling:
+    def _interpolation_resolver(resample: str | None = None) -> Image.Resampling:
         interpolation_map = {
             "nearest": Image.Resampling.NEAREST,
             "lanczos": Image.Resampling.LANCZOS,

fastembed/late_interaction/colbert.py

@@ -1,14 +1,14 @@
 import string
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import define_cache_dir
+from fastembed.common.utils import define_cache_dir, iter_batch
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
@@ -19,7 +19,7 @@ supported_colbert_models: list[DenseModelDescription] = [
     DenseModelDescription(
         model="colbert-ir/colbertv2.0",
         dim=128,
-        description="Late interaction model",
+        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2023 year",
         license="mit",
         size_in_GB=0.44,
         sources=ModelSource(hf="colbert-ir/colbertv2.0"),
@@ -28,7 +28,7 @@ supported_colbert_models: list[DenseModelDescription] = [
     DenseModelDescription(
         model="answerdotai/answerai-colbert-small-v1",
         dim=96,
-        description="Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2024 year",
         license="apache-2.0",
         size_in_GB=0.13,
         sources=ModelSource(hf="answerdotai/answerai-colbert-small-v1"),
@@ -96,6 +96,38 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         encoded = self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
         return encoded
 
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        is_doc: bool = True,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        tokenizer = self.tokenizer if is_doc else self.query_tokenizer
+        assert tokenizer is not None
+        for batch in iter_batch(texts, batch_size):
+            for tokens in tokenizer.encode_batch(batch):
+                if is_doc:
+                    token_num += sum(tokens.attention_mask)
+                else:
+                    attend_count = sum(tokens.attention_mask)
+                    if include_extension:
+                        token_num += max(attend_count, self.MIN_QUERY_LENGTH)
+
+                    else:
+                        token_num += attend_count
+            if include_extension:
+                token_num += len(
+                    batch
+                )  # add 1 for each cls.DOC_MARKER_TOKEN_ID or cls.QUERY_MARKER_TOKEN_ID
+
+        return token_num
+
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
@@ -108,14 +140,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -127,10 +159,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -143,13 +176,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -165,11 +199,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
         )
-        self.mask_token_id: Optional[int] = None
-        self.pad_token_id: Optional[int] = None
+        self.mask_token_id: int | None = None
+        self.pad_token_id: int | None = None
         self.skip_list: set[int] = set()
 
-        self.query_tokenizer: Optional[Tokenizer] = None
+        self.query_tokenizer: Tokenizer | None = None
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -182,6 +216,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
         self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir)
 
@@ -204,9 +239,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -235,10 +270,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         if isinstance(query, str):
             query = [query]
 

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -9,21 +9,21 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
@@ -43,7 +43,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -69,3 +69,12 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def embedding_size(self) -> int:
         """Returns embedding size for the current model"""
         raise NotImplementedError("Subclasses must implement this method")
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,8 +1,8 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -51,11 +51,11 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -101,7 +101,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -116,9 +116,9 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -138,7 +138,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -151,3 +151,30 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.model.query_embed(query, **kwargs)
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        is_doc: bool = True,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+            is_doc (bool): Whether the texts are documents (disable embedding a query with include_mask=True).
+            include_extension (bool): Turn on to count DOC / QUERY marker tokens, and [MASK] token in query mode.
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(
+            texts,
+            batch_size=batch_size,
+            is_doc=is_doc,
+            include_extension=include_extension,
+            **kwargs,
+        )

fastembed/late_interaction/token_embeddings.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict
-from typing import Union, Iterable, Optional, Any, Type
+from typing import Iterable, Any, Type
 
 from fastembed.common.model_description import DenseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -63,9 +63,9 @@ class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         yield from super().embed(documents, batch_size=batch_size, parallel=parallel, **kwargs)

fastembed/late_interaction_multimodal/colmodernvbert.py

@@ -0,0 +1,532 @@
+import contextlib
+from typing import Any, Iterable, Type, Optional, Sequence
+import json
+
+import numpy as np
+from tokenizers import Encoding
+from PIL import Image
+
+from fastembed.common import ImageInput
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+from fastembed.late_interaction_multimodal.onnx_multimodal_model import (
+    OnnxMultimodalModel,
+    TextEmbeddingWorker,
+    ImageEmbeddingWorker,
+)
+
+supported_colmodernvbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/colmodernvbert",
+        dim=128,
+        description="The late-interaction version of ModernVBERT, CPU friendly, English, 2025.",
+        license="mit",
+        size_in_GB=1.0,
+        sources=ModelSource(hf="Qdrant/colmodernvbert"),
+        additional_files=["processor_config.json"],
+        model_file="model.onnx",
+    ),
+]
+
+
+class ColModernVBERT(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyArray]):
+    """
+    The ModernVBERT/colmodernvbert model implementation. This model uses
+    bidirectional attention, which proves to work better for retrieval.
+
+    See: https://huggingface.co/ModernVBERT/colmodernvbert
+    """
+
+    VISUAL_PROMPT_PREFIX = (
+        "<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:"
+    )
+    QUERY_AUGMENTATION_TOKEN = "<end_of_utterance>"
+
+    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,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: Optional[int] = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.image_seq_len: Optional[int] = None
+        self.max_image_size: Optional[int] = None
+        self.image_size: Optional[int] = None
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_colmodernvbert_models
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+        # Load image processing configuration
+        processor_config_path = self._model_dir / "processor_config.json"
+        with open(processor_config_path) as f:
+            processor_config = json.load(f)
+            self.image_seq_len = processor_config.get("image_seq_len", 64)
+
+        preprocessor_config_path = self._model_dir / "preprocessor_config.json"
+        with open(preprocessor_config_path) as f:
+            preprocessor_config = json.load(f)
+            self.max_image_size = preprocessor_config.get("max_image_size", {}).get(
+                "longest_edge", 512
+            )
+
+        # Load model configuration
+        config_path = self._model_dir / "config.json"
+        with open(config_path) as f:
+            model_config = json.load(f)
+            vision_config = model_config.get("vision_config", {})
+            self.image_size = vision_config.get("image_size", 512)
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        batch_size, seq_length = onnx_input["input_ids"].shape
+        empty_image_placeholder: NumpyArray = np.zeros(
+            (batch_size, seq_length, 3, self.image_size, self.image_size),
+            dtype=np.float32,  # type: ignore[type-var,arg-type,assignment]
+        )
+        onnx_input["pixel_values"] = empty_image_placeholder
+        return onnx_input
+
+    def _post_process_onnx_text_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        return output.model_output
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        # Add query augmentation tokens (matching process_queries logic from colpali-engine)
+        augmented_queries = [doc + self.QUERY_AUGMENTATION_TOKEN * 10 for doc in documents]
+        encoded = self.tokenizer.encode_batch(augmented_queries)  # type: ignore[union-attr]
+        return encoded
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        assert self.tokenizer is not None
+        tokenize_func = self.tokenize if include_extension else self.tokenizer.encode_batch
+        for batch in iter_batch(texts, batch_size):
+            token_num += sum([sum(encoding.attention_mask) for encoding in tokenize_func(batch)])
+        return token_num
+
+    def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+        with contextlib.ExitStack() as stack:
+            image_files = [
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
+                for image in images
+            ]
+            assert self.processor is not None, "Processor is not initialized"
+            processed = self.processor(image_files)
+            encoded, attention_mask, metadata = self._process_nested_patches(processed)  # type: ignore[arg-type]
+
+        onnx_input = {"pixel_values": encoded, "attention_mask": attention_mask}
+        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+
+        return OnnxOutputContext(
+            model_output=model_output[0],
+            attention_mask=attention_mask,  # type: ignore[arg-type]
+            metadata=metadata,
+        )
+
+    @staticmethod
+    def _process_nested_patches(
+        processed: list[list[NumpyArray]],
+    ) -> tuple[NumpyArray, NumpyArray, dict[str, Any]]:
+        """
+        Process nested image patches (from ImageSplitter).
+
+        Args:
+            processed: List of patch lists, one per image [[img1_patches], [img2_patches], ...]
+
+        Returns:
+            tuple: (encoded array, attention_mask, metadata)
+                - encoded: (batch_size, max_patches, C, H, W)
+                - attention_mask: (batch_size, max_patches) with 1 for real patches, 0 for padding
+                - metadata: Dict with 'patch_counts' key
+        """
+        patch_counts = [len(patches) for patches in processed]
+        max_patches = max(patch_counts)
+
+        # Get dimensions from first patch
+        channels, height, width = processed[0][0].shape
+        batch_size = len(processed)
+
+        # Create padded array
+        encoded = np.zeros(
+            (batch_size, max_patches, channels, height, width), dtype=processed[0][0].dtype
+        )
+
+        # Create attention mask (1 for real patches, 0 for padding)
+        attention_mask = np.zeros((batch_size, max_patches), dtype=np.int64)
+
+        # Fill in patches and attention mask
+        for i, patches in enumerate(processed):
+            for j, patch in enumerate(patches):
+                encoded[i, j] = patch
+                attention_mask[i, j] = 1
+
+        metadata = {"patch_counts": patch_counts}
+        return encoded, attention_mask, metadata  # type: ignore[return-value]
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Add text input placeholders for image data, following Idefics3 processing logic.
+
+        Constructs input_ids dynamically based on the actual number of image patches,
+        using the same token expansion logic as Idefics3Processor.
+
+        Args:
+            onnx_input: Dict with 'pixel_values' (batch, num_patches, C, H, W)
+                        and 'attention_mask' (batch, num_patches) indicating real patches
+            **kwargs: Additional arguments
+
+        Returns:
+            Updated onnx_input with 'input_ids' and updated 'attention_mask' for token sequence
+        """
+        # The attention_mask in onnx_input has a shape of (batch_size, num_patches),
+        # and should be used to create an attention mask matching the input_ids shape.
+        patch_attention_mask = onnx_input["attention_mask"]
+        pixel_values = onnx_input["pixel_values"]
+
+        batch_size = pixel_values.shape[0]
+        batch_input_ids = []
+
+        # Build input_ids for each image based on its actual patch count
+        for i in range(batch_size):
+            # Count real patches (non-padded) from attention mask
+            patch_count = int(np.sum(patch_attention_mask[i]))
+
+            # Compute rows/cols from patch count
+            rows, cols = self._compute_rows_cols_from_patches(patch_count)
+
+            # Build input_ids for this image
+            input_ids = self._build_input_ids_for_image(rows, cols)
+            batch_input_ids.append(input_ids)
+
+        # Pad sequences to max length in batch
+        max_len = max(len(ids) for ids in batch_input_ids)
+
+        # Get padding config from tokenizer
+        padding_direction = self.tokenizer.padding["direction"]  # type: ignore[index,union-attr]
+        pad_token_id = self.tokenizer.padding["pad_id"]  # type: ignore[index,union-attr]
+
+        # Initialize with pad token
+        padded_input_ids = np.full((batch_size, max_len), pad_token_id, dtype=np.int64)
+        attention_mask = np.zeros((batch_size, max_len), dtype=np.int64)
+
+        for i, input_ids in enumerate(batch_input_ids):
+            seq_len = len(input_ids)
+            if padding_direction == "left":
+                # Left padding: place tokens at the END of the array
+                start_idx = max_len - seq_len
+                padded_input_ids[i, start_idx:] = input_ids
+                attention_mask[i, start_idx:] = 1
+            else:
+                # Right padding: place tokens at the START of the array
+                padded_input_ids[i, :seq_len] = input_ids
+                attention_mask[i, :seq_len] = 1
+
+        onnx_input["input_ids"] = padded_input_ids
+        # Update attention_mask with token-level data
+        onnx_input["attention_mask"] = attention_mask
+        return onnx_input
+
+    @staticmethod
+    def _compute_rows_cols_from_patches(patch_count: int) -> tuple[int, int]:
+        if patch_count <= 1:
+            return 0, 0
+
+        # Subtract 1 for the global image
+        grid_patches = patch_count - 1
+
+        # Find rows and cols (assume square or near-square grid)
+        rows = int(grid_patches**0.5)
+        cols = grid_patches // rows
+
+        # Verify the calculation
+        if rows * cols + 1 != patch_count:
+            # Handle non-square grids
+            for r in range(1, grid_patches + 1):
+                if grid_patches % r == 0:
+                    c = grid_patches // r
+                    if r * c + 1 == patch_count:
+                        return r, c
+            # Fallback: treat as unsplit
+            return 0, 0
+
+        return rows, cols
+
+    def _create_single_image_prompt_string(self) -> str:
+        return (
+            "<fake_token_around_image>"
+            + "<global-img>"
+            + "<image>" * self.image_seq_len  # type: ignore[operator]
+            + "<fake_token_around_image>"
+        )
+
+    def _create_split_image_prompt_string(self, rows: int, cols: int) -> str:
+        text_split_images = ""
+
+        # Add tokens for each patch in the grid
+        for n_h in range(rows):
+            for n_w in range(cols):
+                text_split_images += (
+                    "<fake_token_around_image>"
+                    + f"<row_{n_h + 1}_col_{n_w + 1}>"
+                    + "<image>" * self.image_seq_len  # type: ignore[operator]
+                )
+            text_split_images += "\n"
+
+        # Add global image at the end
+        text_split_images += (
+            "\n<fake_token_around_image>"
+            + "<global-img>"
+            + "<image>" * self.image_seq_len  # type: ignore[operator]
+            + "<fake_token_around_image>"
+        )
+
+        return text_split_images
+
+    def _build_input_ids_for_image(self, rows: int, cols: int) -> np.ndarray:
+        # Create the appropriate image prompt string
+        if rows == 0 and cols == 0:
+            image_prompt_tokens = self._create_single_image_prompt_string()
+        else:
+            image_prompt_tokens = self._create_split_image_prompt_string(rows, cols)
+
+        # Replace <image> in visual prompt with expanded tokens
+        # The visual prompt is: "<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:"
+        expanded_prompt = self.VISUAL_PROMPT_PREFIX.replace("<image>", image_prompt_tokens)
+
+        # Tokenize the complete prompt
+        encoded = self.tokenizer.encode(expanded_prompt)  # type: ignore[union-attr]
+
+        # Convert to numpy array
+        return np.array(encoded.ids, dtype=np.int64)
+
+    def _post_process_onnx_image_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        assert self.model_description.dim is not None, "Model dim is not defined"
+        return output.model_output.reshape(
+            output.model_output.shape[0], -1, self.model_description.dim
+        )
+
+    def embed_text(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def embed_image(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+
+        Args:
+            images: Iterator of image paths or single image path to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_text_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return ColModernVBERTTextEmbeddingWorker
+
+    @classmethod
+    def _get_image_worker_class(cls) -> Type[ImageEmbeddingWorker[NumpyArray]]:
+        return ColModernVBERTImageEmbeddingWorker
+
+
+class ColModernVBERTTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColModernVBERT:
+        return ColModernVBERT(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+
+
+class ColModernVBERTImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColModernVBERT:
+        return ColModernVBERT(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction_multimodal/colpali.py

@@ -1,12 +1,12 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from tokenizers import Encoding
 
 from fastembed.common import OnnxProvider, ImageInput
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray
-from fastembed.common.utils import define_cache_dir
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.utils import define_cache_dir, iter_batch
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
 )
@@ -46,14 +46,14 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -65,10 +65,11 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -80,13 +81,14 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -125,6 +127,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     def _post_process_onnx_image_output(
@@ -170,6 +173,23 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         encoded = self.tokenizer.encode_batch(texts_query)  # type: ignore[union-attr]
         return encoded
 
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        assert self.tokenizer is not None
+        tokenize_func = self.tokenize if include_extension else self.tokenizer.encode_batch
+        for batch in iter_batch(texts, batch_size):
+            token_num += sum([sum(encoding.attention_mask) for encoding in tokenize_func(batch)])
+        return token_num
+
     def _preprocess_onnx_text_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
     ) -> dict[str, NumpyArray]:
@@ -208,9 +228,9 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
 
     def embed_text(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -238,14 +258,15 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
     def embed_image(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -273,6 +294,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -1,9 +1,10 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common import OnnxProvider, ImageInput
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.late_interaction_multimodal.colpali import ColPali
+from fastembed.late_interaction_multimodal.colmodernvbert import ColModernVBERT
 
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
@@ -12,7 +13,10 @@ from fastembed.common.model_description import DenseModelDescription
 
 
 class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [ColPali]
+    EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [
+        ColPali,
+        ColModernVBERT,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -54,11 +58,11 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -104,7 +108,7 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -119,9 +123,9 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
 
     def embed_text(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -142,9 +146,9 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
 
     def embed_image(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -162,3 +166,24 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
             List of embeddings, one per image
         """
         yield from self.model.embed_image(images, batch_size, parallel, **kwargs)
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+            include_extension (bool): Whether to include tokens added by preprocessing
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(
+            texts, batch_size=batch_size, include_extension=include_extension, **kwargs
+        )

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 
 from fastembed.common import ImageInput
@@ -11,21 +11,21 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed_text(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -47,9 +47,9 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
 
     def embed_image(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -76,3 +76,11 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
     def embedding_size(self) -> int:
         """Returns embedding size for the current model"""
         raise NotImplementedError("Subclasses must implement this method")
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -2,7 +2,7 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from PIL import Image
@@ -11,19 +11,19 @@ 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.types import NumpyArray, Device
 from fastembed.common.utils import iter_batch
 from fastembed.image.transform.operators import Compose
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxMultimodalModel(OnnxModel[T]):
-    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+    ONNX_OUTPUT_NAMES: list[str] | None = None
 
     def __init__(self) -> None:
         super().__init__()
-        self.tokenizer: Optional[Tokenizer] = None
-        self.processor: Optional[Compose] = None
+        self.tokenizer: Tokenizer | None = None
+        self.processor: Compose | None = None
         self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_text_input(
@@ -60,10 +60,11 @@ class OnnxMultimodalModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -72,6 +73,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
@@ -114,14 +116,15 @@ class OnnxMultimodalModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -153,6 +156,9 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_text_worker_class(),
@@ -164,9 +170,11 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 yield from self._post_process_onnx_text_output(batch)  # type: ignore
 
     def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack():
+        with contextlib.ExitStack() as stack:
             image_files = [
-                Image.open(image) if not isinstance(image, Image.Image) else image
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
                 for image in images
             ]
             assert self.processor is not None, "Processor is not initialized"
@@ -181,14 +189,15 @@ class OnnxMultimodalModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[Iterable[ImageInput], ImageInput],
+        images: Iterable[ImageInput] | ImageInput,
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -220,6 +229,9 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_image_worker_class(),

fastembed/parallel_processor.py

@@ -8,8 +8,9 @@ from multiprocessing.context import BaseContext
 from multiprocessing.process import BaseProcess
 from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
-from typing import Any, Iterable, Optional, Type
+from typing import Any, Iterable, Type
 
+from fastembed.common.types import Device
 
 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -38,7 +39,7 @@ def _worker(
     output_queue: Queue,
     num_active_workers: BaseValue,
     worker_id: int,
-    kwargs: Optional[dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> None:
     """
     A worker that pulls data pints off the input queue, and places the execution result on the output queue.
@@ -93,21 +94,21 @@ class ParallelWorkerPool:
         self,
         num_workers: int,
         worker: Type[Worker],
-        start_method: Optional[str] = None,
-        device_ids: Optional[list[int]] = None,
-        cuda: bool = False,
+        start_method: str | None = None,
+        device_ids: list[int] | None = None,
+        cuda: bool | Device = Device.AUTO,
     ):
         self.worker_class = worker
         self.num_workers = num_workers
-        self.input_queue: Optional[Queue] = None
-        self.output_queue: Optional[Queue] = None
+        self.input_queue: Queue | None = None
+        self.output_queue: Queue | None = None
         self.ctx: BaseContext = get_context(start_method)
         self.processes: list[BaseProcess] = []
         self.queue_size = self.num_workers * max_internal_batch_size
         self.emergency_shutdown = False
         self.device_ids = device_ids
         self.cuda = cuda
-        self.num_active_workers: Optional[BaseValue] = None
+        self.num_active_workers: BaseValue | None = None
 
     def start(self, **kwargs: Any) -> None:
         self.input_queue = self.ctx.Queue(self.queue_size)
@@ -220,7 +221,7 @@ class ParallelWorkerPool:
                     f"Worker PID: {process.pid} terminated unexpectedly with code {process.exitcode}"
                 )
 
-    def join_or_terminate(self, timeout: Optional[int] = 1) -> None:
+    def join_or_terminate(self, timeout: int = 1) -> None:
         """
         Emergency shutdown
         @param timeout:

fastembed/postprocess/muvera.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import numpy as np
 
 from fastembed.common.types import NumpyArray
@@ -11,7 +9,7 @@ from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding
 )
 
 
-MultiVectorModel = Union[LateInteractionTextEmbeddingBase, LateInteractionMultimodalEmbeddingBase]
+MultiVectorModel = LateInteractionTextEmbeddingBase | LateInteractionMultimodalEmbeddingBase
 MAX_HAMMING_DISTANCE = 65  # 64 bits + 1
 POPCOUNT_LUT = np.array([bin(x).count("1") for x in range(256)], dtype=np.uint8)
 

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -1,7 +1,8 @@
-from typing import Optional, Sequence, Any
+from typing import Sequence, Any
 
 from fastembed.common import OnnxProvider
 from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.types import Device
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
 
 
@@ -11,14 +12,14 @@ class CustomTextCrossEncoder(OnnxTextCrossEncoder):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -1,9 +1,10 @@
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Sequence, Type
 
 from loguru import logger
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.rerank.cross_encoder.onnx_text_model import (
     OnnxCrossEncoderModel,
@@ -77,14 +78,14 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -96,10 +97,11 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -111,6 +113,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -123,7 +126,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             )
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -150,6 +153,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     def rerank(
@@ -178,7 +182,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         yield from self._rerank_pairs(
@@ -192,6 +196,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
@@ -204,6 +209,20 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     ) -> Iterable[float]:
         return (float(elem) for elem in output.model_output)
 
+    def token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the pairs.
+
+        Args:
+            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
+            batch_size: Batch size for tokenizing
+
+        Returns:
+            token count: overall number of tokens in the pairs
+        """
+        return self._token_count(pairs, batch_size=batch_size, **kwargs)
+
 
 class TextCrossEncoderWorker(TextRerankerWorker):
     def init_embedding(

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,7 +1,7 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from tokenizers import Encoding
@@ -12,14 +12,14 @@ from fastembed.common.onnx_model import (
     OnnxOutputContext,
     OnnxProvider,
 )
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxCrossEncoderModel(OnnxModel[float]):
-    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+    ONNX_OUTPUT_NAMES: list[str] | None = None
 
     @classmethod
     def _get_worker_class(cls) -> Type["TextRerankerWorker"]:
@@ -29,10 +29,11 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -41,6 +42,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
@@ -90,12 +92,13 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         cache_dir: str,
         pairs: Iterable[tuple[str, str]],
         batch_size: int,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -127,6 +130,9 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),
@@ -159,6 +165,20 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         """
         return onnx_input
 
+    def _token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **_: Any
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+
+        token_num = 0
+        assert self.tokenizer is not None
+        for batch in iter_batch(pairs, batch_size):
+            for tokens in self.tokenizer.encode_batch(batch):
+                token_num += sum(tokens.attention_mask)
+
+        return token_num
+
 
 class TextRerankerWorker(EmbeddingWorker[float]):
     def __init__(

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -1,7 +1,8 @@
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
 from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
 
@@ -53,11 +54,11 @@ class TextCrossEncoder(TextCrossEncoderBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -102,7 +103,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         """
@@ -140,7 +141,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         description: str = "",
         license: str = "",
         size_in_gb: float = 0.0,
-        additional_files: Optional[list[str]] = None,
+        additional_files: list[str] | None = None,
     ) -> None:
         registered_models = cls._list_supported_models()
         for registered_model in registered_models:
@@ -161,3 +162,17 @@ class TextCrossEncoder(TextCrossEncoderBase):
                 additional_files=additional_files or [],
             )
         )
+
+    def token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the pairs.
+
+        Args:
+            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
+            batch_size: Batch size for tokenizing
+
+        Returns:
+            token count: overall number of tokens in the pairs
+        """
+        return self.model.token_count(pairs, batch_size=batch_size, **kwargs)

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -1,4 +1,4 @@
-from typing import Any, Iterable, Optional
+from typing import Any, Iterable
 
 from fastembed.common.model_description import BaseModelDescription
 from fastembed.common.model_management import ModelManagement
@@ -8,8 +8,8 @@ class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
@@ -41,7 +41,7 @@ class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         """Rerank query-document pairs.
@@ -57,3 +57,7 @@ class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
             Iterable[float]: Scores for each individual pair
         """
         raise NotImplementedError("This method should be overridden by subclasses")
+
+    def token_count(self, pairs: Iterable[tuple[str, str]], **kwargs: Any) -> int:
+        """Returns the number of tokens in the pairs."""
+        raise NotImplementedError("This method should be overridden by subclasses")

fastembed/sparse/bm25.py

@@ -2,7 +2,7 @@ import os
 from collections import defaultdict
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Type, Union
+from typing import Any, Iterable, Type
 
 import mmh3
 import numpy as np
@@ -91,14 +91,14 @@ class Bm25(SparseTextEmbeddingBase):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
+        cache_dir: str | None = None,
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 256.0,
         language: str = "english",
         token_max_length: int = 40,
         disable_stemmer: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
@@ -158,11 +158,11 @@ class Bm25(SparseTextEmbeddingBase):
         self,
         model_name: str,
         cache_dir: str,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
     ) -> Iterable[SparseEmbedding]:
         is_small = False
 
@@ -205,9 +205,9 @@ class Bm25(SparseTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -268,6 +268,15 @@ class Bm25(SparseTextEmbeddingBase):
             embeddings.append(SparseEmbedding.from_dict(token_id2value))
         return embeddings
 
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        for text in texts:
+            document = remove_non_alphanumeric(text)
+            tokens = self.tokenizer.tokenize(document)
+            token_num += len(tokens)
+        return token_num
+
     def _term_frequency(self, tokens: list[str]) -> dict[int, float]:
         """Calculate the term frequency part of the BM25 formula.
 
@@ -302,9 +311,7 @@ class Bm25(SparseTextEmbeddingBase):
     def compute_token_id(cls, token: str) -> int:
         return abs(mmh3.hash(token))
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """To emulate BM25 behaviour, we don't need to use weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
         """

fastembed/sparse/bm42.py

@@ -1,7 +1,7 @@
 import math
 import string
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import mmh3
 import numpy as np
@@ -9,6 +9,7 @@ from py_rust_stemmers import SnowballStemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -31,9 +32,17 @@ supported_bm42_models: list[SparseModelDescription] = [
     ),
 ]
 
-MODEL_TO_LANGUAGE = {
+
+_MODEL_TO_LANGUAGE = {
     "Qdrant/bm42-all-minilm-l6-v2-attentions": "english",
 }
+MODEL_TO_LANGUAGE = {
+    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()
+}
+
+
+def get_language_by_model_name(model_name: str) -> str:
+    return MODEL_TO_LANGUAGE[model_name.lower()]
 
 
 class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
@@ -57,15 +66,15 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
         alpha: float = 0.5,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -79,10 +88,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             alpha (float, optional): Parameter, that defines the importance of the token weight in the document
                 versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
                 It is recommended to only change this parameter based on training data for a specific dataset.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -95,13 +105,14 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -124,7 +135,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.special_tokens_ids: set[int] = set()
         self.punctuation = set(string.punctuation)
         self.stopwords = set(self._load_stopwords(self._model_dir))
-        self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
+        self.stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))
         self.alpha = alpha
 
         if not self.lazy_load:
@@ -138,6 +149,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
         for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
@@ -272,9 +284,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -304,6 +316,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             alpha=self.alpha,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
         )
 
     @classmethod
@@ -314,9 +327,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             result[token_id] = 1.0
         return result
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         To emulate BM25 behaviour, we don't need to use smart weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
@@ -341,6 +352,13 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
         return Bm42TextEmbeddingWorker
 
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
 
 class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:

fastembed/sparse/minicoil.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from typing import Any, Optional, Sequence, Iterable, Union, Type
+from typing import Any, Sequence, Iterable, Type
 
 import numpy as np
 from numpy.typing import NDArray
@@ -10,6 +10,7 @@ from tokenizers import Tokenizer
 from fastembed.common.model_description import SparseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -46,9 +47,16 @@ supported_minicoil_models: list[SparseModelDescription] = [
     ),
 ]
 
-MODEL_TO_LANGUAGE = {
+_MODEL_TO_LANGUAGE = {
     "Qdrant/minicoil-v1": "english",
 }
+MODEL_TO_LANGUAGE = {
+    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()
+}
+
+
+def get_language_by_model_name(model_name: str) -> str:
+    return MODEL_TO_LANGUAGE[model_name.lower()]
 
 
 class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
@@ -65,17 +73,17 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 150.0,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -91,10 +99,11 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
                 Defaults to 0.75.
             avg_len (float, optional): The average length of the documents in the corpus. Defaults to 150.0.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -110,20 +119,22 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.device_ids = device_ids
         self.cuda = cuda
         self.device_id = device_id
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
         self.k = k
         self.b = b
         self.avg_len = avg_len
 
         # Initialize class attributes
-        self.tokenizer: Optional[Tokenizer] = None
+        self.tokenizer: Tokenizer | None = None
         self.invert_vocab: dict[int, str] = {}
         self.special_tokens: set[str] = set()
         self.special_tokens_ids: set[int] = set()
         self.stopwords: set[str] = set()
-        self.vocab_resolver: Optional[VocabResolver] = None
-        self.encoder: Optional[Encoder] = None
-        self.output_dim: Optional[int] = None
-        self.sparse_vector_converter: Optional[SparseVectorConverter] = None
+        self.vocab_resolver: VocabResolver | None = None
+        self.encoder: Encoder | None = None
+        self.output_dim: int | None = None
+        self.sparse_vector_converter: SparseVectorConverter | None = None
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
@@ -146,6 +157,7 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
         assert self.tokenizer is not None
@@ -156,7 +168,7 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.special_tokens_ids = set(self.special_token_to_id.values())
         self.stopwords = set(self._load_stopwords(self._model_dir))
 
-        stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[self.model_name])
+        stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))
 
         self.vocab_resolver = VocabResolver(
             tokenizer=VocabTokenizer(self.tokenizer),
@@ -177,11 +189,16 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             avg_len=self.avg_len,
         )
 
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -214,12 +231,11 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             is_query=False,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Encode a list of queries into list of embeddings.
         """

fastembed/sparse/sparse_embedding_base.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 import numpy as np
 from numpy.typing import NDArray
@@ -12,7 +12,7 @@ from fastembed.common.model_management import ModelManagement
 @dataclass
 class SparseEmbedding:
     values: NumpyArray
-    indices: Union[NDArray[np.int64], NDArray[np.int32]]
+    indices: NDArray[np.int64] | NDArray[np.int32]
 
     def as_object(self) -> dict[str, NumpyArray]:
         return {
@@ -35,8 +35,8 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
@@ -46,9 +46,9 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         raise NotImplementedError()
@@ -68,9 +68,7 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -86,3 +84,7 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
             yield from self.embed([query], **kwargs)
         else:
             yield from self.embed(query, **kwargs)
+
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/sparse/sparse_text_embedding.py

@@ -1,7 +1,8 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.bm42 import Bm42
 from fastembed.sparse.minicoil import MiniCOIL
@@ -53,11 +54,11 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -93,9 +94,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -115,9 +116,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -128,3 +127,17 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
             Iterable[SparseEmbedding]: The sparse embeddings.
         """
         yield from self.model.query_embed(query, **kwargs)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(texts, batch_size=batch_size, **kwargs)

fastembed/sparse/splade_pp.py

@@ -1,8 +1,9 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -53,6 +54,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             scores = row_scores[indices]
             yield SparseEmbedding(values=scores, indices=indices)
 
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
     @classmethod
     def _list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
@@ -65,14 +71,14 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -84,10 +90,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -99,13 +106,14 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -133,13 +141,14 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -168,6 +177,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/sparse/utils/sparse_vectors_converter.py

@@ -1,12 +1,11 @@
-from typing import Dict, List, Set
-from py_rust_stemmers import SnowballStemmer
-from fastembed.common.utils import get_all_punctuation, remove_non_alphanumeric
-import mmh3
 import copy
 from dataclasses import dataclass
 
+import mmh3
 import numpy as np
+from py_rust_stemmers import SnowballStemmer
 
+from fastembed.common.utils import get_all_punctuation, remove_non_alphanumeric
 from fastembed.sparse.sparse_embedding_base import SparseEmbedding
 
 GAP = 32000
@@ -16,16 +15,16 @@ INT32_MAX = 2**31 - 1
 @dataclass
 class WordEmbedding:
     word: str
-    forms: List[str]
+    forms: list[str]
     count: int
     word_id: int
-    embedding: List[float]
+    embedding: list[float]
 
 
 class SparseVectorConverter:
     def __init__(
         self,
-        stopwords: Set[str],
+        stopwords: set[str],
         stemmer: SnowballStemmer,
         k: float = 1.2,
         b: float = 0.75,
@@ -58,15 +57,15 @@ class SparseVectorConverter:
         return res
 
     @classmethod
-    def normalize_vector(cls, vector: List[float]) -> List[float]:
+    def normalize_vector(cls, vector: list[float]) -> list[float]:
         norm = sum([x**2 for x in vector]) ** 0.5
         if norm < 1e-8:
             return vector
         return [x / norm for x in vector]
 
     def clean_words(
-        self, sentence_embedding: Dict[str, WordEmbedding], token_max_length: int = 40
-    ) -> Dict[str, WordEmbedding]:
+        self, sentence_embedding: dict[str, WordEmbedding], token_max_length: int = 40
+    ) -> dict[str, WordEmbedding]:
         """
         Clean miniCOIL-produced sentence_embedding, as unknown to the miniCOIL's stemmer tokens should fully resemble
         our BM25 token representation.
@@ -85,7 +84,7 @@ class SparseVectorConverter:
                 }
         """
 
-        new_sentence_embedding: Dict[str, WordEmbedding] = {}
+        new_sentence_embedding: dict[str, WordEmbedding] = {}
 
         for word, embedding in sentence_embedding.items():
             # embedding = {
@@ -127,7 +126,7 @@ class SparseVectorConverter:
 
     def embedding_to_vector(
         self,
-        sentence_embedding: Dict[str, WordEmbedding],
+        sentence_embedding: dict[str, WordEmbedding],
         embedding_size: int,
         vocab_size: int,
     ) -> SparseEmbedding:
@@ -156,14 +155,14 @@ class SparseVectorConverter:
 
         """
 
-        indices: List[int] = []
-        values: List[float] = []
-        
+        indices: list[int] = []
+        values: list[float] = []
+
         # Example:
         # vocab_size = 10000
         # embedding_size = 4
         # GAP = 32000
-        # 
+        #
         # We want to start random words section from the bucket, that is guaranteed to not
         # include any vocab words.
         # We need (vocab_size * embedding_size) slots for vocab words.
@@ -171,9 +170,7 @@ class SparseVectorConverter:
         # Therefore, we can start random words from bucket (vocab_size * embedding_size) // GAP + 1 + 1
 
         # ID at which the scope of OOV words starts
-        unknown_words_shift = (
-            (vocab_size * embedding_size) // GAP + 2
-        ) * GAP
+        unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP
         sentence_embedding_cleaned = self.clean_words(sentence_embedding)
 
         # Calculate sentence length after cleaning
@@ -208,7 +205,7 @@ class SparseVectorConverter:
 
     def embedding_to_vector_query(
         self,
-        sentence_embedding: Dict[str, WordEmbedding],
+        sentence_embedding: dict[str, WordEmbedding],
         embedding_size: int,
         vocab_size: int,
     ) -> SparseEmbedding:
@@ -216,8 +213,8 @@ class SparseVectorConverter:
         Same as `embedding_to_vector`, but no TF
         """
 
-        indices: List[int] = []
-        values: List[float] = []
+        indices: list[int] = []
+        values: list[float] = []
 
         # ID at which the scope of OOV words starts
         unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP

fastembed/text/custom_text_embedding.py

@@ -1,5 +1,4 @@
-from typing import Optional, Sequence, Any, Iterable
-
+from typing import Sequence, Any, Iterable
 from dataclasses import dataclass
 
 import numpy as np
@@ -11,7 +10,7 @@ from fastembed.common.model_description import (
     DenseModelDescription,
 )
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common.utils import normalize, mean_pooling
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 
@@ -29,14 +28,14 @@ class CustomTextEmbedding(OnnxTextEmbedding):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(
@@ -64,7 +63,7 @@ class CustomTextEmbedding(OnnxTextEmbedding):
         return self._normalize(self._pool(output.model_output, output.attention_mask))
 
     def _pool(
-        self, embeddings: NumpyArray, attention_mask: Optional[NDArray[np.int64]] = None
+        self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None
     ) -> NumpyArray:
         if self._pooling == PoolingType.CLS:
             return embeddings[:, 0]

fastembed/text/multitask_embedding.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Type, Iterable, Union, Optional
+from typing import Any, Type, Iterable
 
 import numpy as np
 
@@ -45,11 +45,9 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
 
-    def __init__(self, *args: Any, task_id: Optional[int] = None, **kwargs: Any):
+    def __init__(self, *args: Any, task_id: int | None = None, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.default_task_id: Union[Task, int] = (
-            task_id if task_id is not None else self.PASSAGE_TASK
-        )
+        self.default_task_id: Task | int = task_id if task_id is not None else self.PASSAGE_TASK
 
     @classmethod
     def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
@@ -62,7 +60,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     def _preprocess_onnx_input(
         self,
         onnx_input: dict[str, NumpyArray],
-        task_id: Optional[Union[int, Task]] = None,
+        task_id: int | Task | None = None,
         **kwargs: Any,
     ) -> dict[str, NumpyArray]:
         if task_id is None:
@@ -72,10 +70,10 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        task_id: Optional[int] = None,
+        parallel: int | None = None,
+        task_id: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         task_id = (
@@ -83,7 +81,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         )  # required for multiprocessing
         yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
 
     def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:

fastembed/text/onnx_embedding.py

@@ -1,6 +1,6 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider, Device
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
@@ -199,14 +199,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -218,10 +218,11 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -233,13 +234,13 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -260,9 +261,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -291,6 +292,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
@@ -327,8 +329,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
 
 class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
     def init_embedding(

fastembed/text/onnx_text_model.py

@@ -1,13 +1,13 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider, Device
 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
@@ -15,7 +15,7 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxTextModel(OnnxModel[T]):
-    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+    ONNX_OUTPUT_NAMES: list[str] | None = None
 
     @classmethod
     def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
@@ -35,12 +35,12 @@ class OnnxTextModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.tokenizer: Optional[Tokenizer] = None
+        self.tokenizer: Tokenizer | None = None
         self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, Union[NumpyArray, NDArray[np.int64]]]:
+    ) -> dict[str, NumpyArray | NDArray[np.int64]]:
         """
         Preprocess the onnx input.
         """
@@ -50,10 +50,11 @@ class OnnxTextModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -62,6 +63,7 @@ class OnnxTextModel(OnnxModel[T]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
 
@@ -102,14 +104,15 @@ class OnnxTextModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -143,6 +146,9 @@ class OnnxTextModel(OnnxModel[T]):
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),
@@ -153,6 +159,19 @@ class OnnxTextModel(OnnxModel[T]):
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
 
+    def _token_count(self, texts: str | Iterable[str], batch_size: int = 1024, **_: Any) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+
+        token_num = 0
+        assert self.tokenizer is not None
+        texts = [texts] if isinstance(texts, str) else texts
+        for batch in iter_batch(texts, batch_size):
+            for tokens in self.tokenizer.encode_batch(batch):
+                token_num += sum(tokens.attention_mask)
+
+        return token_num
+
 
 class TextEmbeddingWorker(EmbeddingWorker[T]):
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:

fastembed/text/text_embedding.py

@@ -1,8 +1,8 @@
 import warnings
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider, Device
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 from fastembed.text.custom_text_embedding import CustomTextEmbedding
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
@@ -51,7 +51,7 @@ class TextEmbedding(TextEmbeddingBase):
         description: str = "",
         license: str = "",
         size_in_gb: float = 0.0,
-        additional_files: Optional[list[str]] = None,
+        additional_files: list[str] | None = None,
     ) -> None:
         registered_models = cls._list_supported_models()
         for registered_model in registered_models:
@@ -79,11 +79,11 @@ class TextEmbedding(TextEmbeddingBase):
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -149,7 +149,7 @@ class TextEmbedding(TextEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -164,9 +164,9 @@ class TextEmbedding(TextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -186,7 +186,7 @@ class TextEmbedding(TextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -212,3 +212,17 @@ class TextEmbedding(TextEmbeddingBase):
         """
         # This is model-specific, so that different models can have specialized implementations
         yield from self.model.passage_embed(texts, **kwargs)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(texts, batch_size=batch_size, **kwargs)

fastembed/text/text_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -9,21 +9,21 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
@@ -43,7 +43,7 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -69,3 +69,7 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def embedding_size(self) -> int:
         """Returns embedding size for the current model"""
         raise NotImplementedError("Subclasses must implement this method")
+
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

mkdocs.yml

@@ -13,6 +13,7 @@ copyright: |
 theme:
   name: material
   logo: assets/favicon.png
+  favicon: assets/favicon.png
   custom_dir: docs/overrides
   icon:
     repo: fontawesome/brands/github

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed"
-version = "0.7.3"
+name = "fastembed-gpu"
+version = "0.8.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"
@@ -11,24 +11,30 @@ repository = "https://github.com/qdrant/fastembed"
 keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
 
 [tool.poetry.dependencies]
-python = ">=3.9.0"
+python = ">=3.10.0"
 numpy = [
-    { version = ">=1.21", python = ">=3.10,<3.12" },
-    { version = ">=1.26", python = ">=3.12,<3.13" },
-    { version = ">=2.1.0", python = ">=3.13" },
-    { version = ">=1.21,<2.1.0", python = "<3.10" },
+    { version = ">=1.21,<2.3.0", python = "3.10" },
+    { version = ">=1.21", python = "3.11" },
+    { version = ">=1.26", python = "3.12" },
+    { version = ">=2.1.0", python = "3.13" },
+    { version = ">=2.3.0", python = ">=3.14" },
 ]
-onnxruntime = [
-    { version = ">=1.17.0,<1.20.0", python = "<3.10" },
-    { version = ">1.20.0", python = ">=3.13" },
-    { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
+onnxruntime-gpu = [
+    { version = ">=1.17.0,!=1.20.0,<1.24", python = "3.10" },
+    { version = ">=1.17.0,!=1.20.0,!=1.24.0,!=1.24.1", python = ">=3.11,<3.13" },
+    { version = ">1.21.0,!=1.24.0,!=1.24.1", python = "3.13" },
+    { version = ">=1.24.2", python = ">=3.14" },
 ]
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
-huggingface-hub = ">=0.20,<1.0"
+huggingface-hub = ">=0.20,<2.0"
 loguru = "^0.7.2"
-pillow = ">=10.3.0,<12.0.0"
+pillow = [
+    { version = ">=10.3.0,<13.0", python = ">=3.10,<3.13" },
+    { version = ">=11.0.0,<13.0", python = "3.13" },
+    { version = ">=12.0.0,<13.0", python = ">=3.14" },
+]
 mmh3 = ">=4.1.0,<6.0.0"
 py-rust-stemmers = "^0.1.0"
 
@@ -39,12 +45,16 @@ ruff = ">=0.3.1,<1.0"
 [tool.poetry.group.dev.dependencies]
 notebook = ">=7.0.2"
 pre-commit = "^3.6.2"
-onnx = ">=1.15.0"
+onnx = [
+    { version = ">=1.15.0", python = ">=3.10,<3.13" },
+    { version = ">=1.18.0", python = "3.13" },
+    { version = ">=1.20.0", python = ">=3.14" },
+]
 
 [tool.poetry.group.docs.dependencies]
 mkdocs-material = "^9.5.10"
 mkdocstrings = "^0.24.0"
-pillow = ">=10.3.0,<12.0.0"
+pillow = ">=10.3.0,<13.0.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 

tests/test_attention_embeddings.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 
 import numpy as np
 import pytest
@@ -7,98 +8,119 @@ from fastembed import SparseTextEmbedding
 from tests.utils import delete_model_cache
 
 
-@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
-def test_attention_embeddings(model_name: str) -> None:
+_MODELS_TO_CACHE = ("Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25")
+MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+
+
+@pytest.fixture(scope="module")
+def model_cache():
     is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name=model_name)
+    cache = {}
 
-    output = list(
-        model.query_embed(
-            [
-                "I must not fear. Fear is the mind-killer.",
-            ]
-        )
-    )
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = SparseTextEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            print("deleting model")
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
 
-    assert len(output) == 1
-
-    for result in output:
-        assert len(result.indices) == len(result.values)
-        assert np.allclose(result.values, np.ones(len(result.values)))
-
-    quotes = [
-        "I must not fear. Fear is the mind-killer.",
-        "All animals are equal, but some animals are more equal than others.",
-        "It was a pleasure to burn.",
-        "The sky above the port was the color of television, tuned to a dead channel.",
-        "In the beginning, the universe was created."
-        " This has made a lot of people very angry and been widely regarded as a bad move.",
-        "It's a truth universally acknowledged that a zombie in possession of brains must be in want of more brains.",
-        "War is peace. Freedom is slavery. Ignorance is strength.",
-        "We're not in Infinity; we're in the suburbs.",
-        "I was a thousand times more evil than thou!",
-        "History is merely a list of surprises... It can only prepare us to be surprised yet again.",
-        ".",  # Empty string
-    ]
-
-    output = list(model.embed(quotes))
-
-    assert len(output) == len(quotes)
-
-    for result in output[:-1]:
-        assert len(result.indices) == len(result.values)
-        assert len(result.indices) > 0
-
-    assert len(output[-1].indices) == 0
-
-    # Test support for unknown languages
-    output = list(
-        model.query_embed(
-            [
-                "привет мир!",
-            ]
-        )
-    )
-
-    assert len(output) == 1
-
-    for result in output:
-        assert len(result.indices) == len(result.values)
-        assert len(result.indices) == 2
+    yield get_model
 
     if is_ci:
-        delete_model_cache(model.model._model_dir)
+        for name, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
-def test_parallel_processing(model_name: str) -> None:
-    is_ci = os.getenv("CI")
+def test_attention_embeddings(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        output = list(
+            model.query_embed(
+                [
+                    "I must not fear. Fear is the mind-killer.",
+                ]
+            )
+        )
 
-    model = SparseTextEmbedding(model_name=model_name)
+        assert len(output) == 1
 
-    docs = ["hello world", "attention embedding", "Mangez-vous vraiment des grenouilles?"] * 100
-    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+        for result in output:
+            assert len(result.indices) == len(result.values)
+            assert np.allclose(result.values, np.ones(len(result.values)))
 
-    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+        quotes = [
+            "I must not fear. Fear is the mind-killer.",
+            "All animals are equal, but some animals are more equal than others.",
+            "It was a pleasure to burn.",
+            "The sky above the port was the color of television, tuned to a dead channel.",
+            "In the beginning, the universe was created."
+            " This has made a lot of people very angry and been widely regarded as a bad move.",
+            "It's a truth universally acknowledged that a zombie in possession of brains must be in want of more brains.",
+            "War is peace. Freedom is slavery. Ignorance is strength.",
+            "We're not in Infinity; we're in the suburbs.",
+            "I was a thousand times more evil than thou!",
+            "History is merely a list of surprises... It can only prepare us to be surprised yet again.",
+            ".",  # Empty string
+        ]
 
-    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+        output = list(model.embed(quotes))
 
-    assert len(embeddings) == len(docs)
+        assert len(output) == len(quotes)
 
-    for emb_1, emb_2, emb_3 in zip(embeddings, embeddings_2, embeddings_3):
-        assert np.allclose(emb_1.indices, emb_2.indices)
-        assert np.allclose(emb_1.indices, emb_3.indices)
-        assert np.allclose(emb_1.values, emb_2.values)
-        assert np.allclose(emb_1.values, emb_3.values)
+        for result in output[:-1]:
+            assert len(result.indices) == len(result.values)
+            assert len(result.indices) > 0
 
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        assert len(output[-1].indices) == 0
+
+        # Test support for unknown languages
+        output = list(
+            model.query_embed(
+                [
+                    "привет мир!",
+                ]
+            )
+        )
+
+        assert len(output) == 1
+
+        for result in output:
+            assert len(result.indices) == len(result.values)
+            assert len(result.indices) == 2
+
+
+@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
+def test_parallel_processing(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        docs = [
+            "hello world",
+            "attention embedding",
+            "Mangez-vous vraiment des grenouilles?",
+        ] * 100
+        embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+
+        embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+
+        embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+
+        assert len(embeddings) == len(docs)
+
+        for emb_1, emb_2, emb_3 in zip(embeddings, embeddings_2, embeddings_3):
+            assert np.allclose(emb_1.indices, emb_2.indices)
+            assert np.allclose(emb_1.indices, emb_3.indices)
+            assert np.allclose(emb_1.values, emb_2.values)
+            assert np.allclose(emb_1.values, emb_3.values)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
-def test_multilanguage(model_name: str) -> None:
-    is_ci = os.getenv("CI")
-
+def test_multilanguage(model_cache, model_name: str) -> None:
     docs = ["Mangez-vous vraiment des grenouilles?", "Je suis au lit"]
 
     model = SparseTextEmbedding(model_name=model_name, language="french")
@@ -109,39 +131,30 @@ def test_multilanguage(model_name: str) -> None:
     assert embeddings[1].values.shape == (1,)
     assert embeddings[1].indices.shape == (1,)
 
-    model = SparseTextEmbedding(model_name=model_name, language="english")
-    embeddings = list(model.embed(docs))[:2]
-    assert embeddings[0].values.shape == (5,)
-    assert embeddings[0].indices.shape == (5,)
+    with model_cache(model_name) as model:  # language = "english"
+        embeddings = list(model.embed(docs))[:2]
+        assert embeddings[0].values.shape == (5,)
+        assert embeddings[0].indices.shape == (5,)
 
-    assert embeddings[1].values.shape == (4,)
-    assert embeddings[1].indices.shape == (4,)
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        assert embeddings[1].values.shape == (4,)
+        assert embeddings[1].indices.shape == (4,)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
-def test_special_characters(model_name: str) -> None:
-    is_ci = os.getenv("CI")
-
-    docs = [
-        "Über den größten Flüssen Österreichs äußern sich Experten häufig: Öko-Systeme müssen geschützt werden!",
-        "L'élève français s'écrie : « Où est mon crayon ? J'ai besoin de finir cet exercice avant la récréation!",
-        "Într-o zi însorită, Ștefan și Ioana au mâncat mămăligă cu brânză și au băut țuică la cabană.",
-        "Üzgün öğretmen öğrencilere seslendi: Lütfen gürültü yapmayın, sınavınızı bitirmeye çalışıyorum!",
-        "Ο Ξενοφών είπε: «Ψάχνω για ένα ωραίο δώρο για τη γιαγιά μου. Ίσως ένα φυτό ή ένα βιβλίο;»",
-        "Hola! ¿Cómo estás? Estoy muy emocionado por el cumpleaños de mi hermano, ¡va a ser increíble! También quiero comprar un pastel de chocolate con fresas y un regalo especial: un libro titulado «Cien años de soledad",
-    ]
-
-    model = SparseTextEmbedding(model_name=model_name, language="english")
-    embeddings = list(model.embed(docs))
-    for idx, shape in enumerate([14, 18, 15, 10, 15]):
-        assert embeddings[idx].values.shape == (shape,)
-        assert embeddings[idx].indices.shape == (shape,)
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+def test_special_characters(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        docs = [
+            "Über den größten Flüssen Österreichs äußern sich Experten häufig: Öko-Systeme müssen geschützt werden!",
+            "L'élève français s'écrie : « Où est mon crayon ? J'ai besoin de finir cet exercice avant la récréation!",
+            "Într-o zi însorită, Ștefan și Ioana au mâncat mămăligă cu brânză și au băut țuică la cabană.",
+            "Üzgün öğretmen öğrencilere seslendi: Lütfen gürültü yapmayın, sınavınızı bitirmeye çalışıyorum!",
+            "Ο Ξενοφών είπε: «Ψάχνω για ένα ωραίο δώρο για τη γιαγιά μου. Ίσως ένα φυτό ή ένα βιβλίο;»",
+            "Hola! ¿Cómo estás? Estoy muy emocionado por el cumpleaños de mi hermano, ¡va a ser increíble! También quiero comprar un pastel de chocolate con fresas y un regalo especial: un libro titulado «Cien años de soledad",
+        ]
+        embeddings = list(model.embed(docs))
+        for idx, shape in enumerate([14, 18, 15, 10, 15]):
+            assert embeddings[idx].values.shape == (shape,)
+            assert embeddings[idx].indices.shape == (shape,)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])

tests/test_custom_models.py

@@ -70,9 +70,13 @@ def test_text_custom_model():
     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)
 
+    CustomTextEmbedding.SUPPORTED_MODELS.clear()
+    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
+
 
 def test_cross_encoder_custom_model():
     is_ci = os.getenv("CI")
@@ -110,6 +114,8 @@ def test_cross_encoder_custom_model():
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
+    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()
+
 
 def test_mock_add_custom_models():
     dim = 5
@@ -169,6 +175,9 @@ def test_mock_add_custom_models():
         )
         assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3)
 
+    CustomTextEmbedding.SUPPORTED_MODELS.clear()
+    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
+
 
 def test_do_not_add_existing_model():
     existing_base_model = "sentence-transformers/all-MiniLM-L6-v2"
@@ -203,6 +212,9 @@ def test_do_not_add_existing_model():
             size_in_gb=0.47,
         )
 
+    CustomTextEmbedding.SUPPORTED_MODELS.clear()
+    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
+
 
 def test_do_not_add_existing_cross_encoder():
     existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
@@ -227,3 +239,5 @@ def test_do_not_add_existing_cross_encoder():
             sources=ModelSource(hf=custom_model_name),
             size_in_gb=0.08,
         )
+
+    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()

tests/test_image_onnx_embeddings.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 from io import BytesIO
 
 import numpy as np
@@ -26,9 +27,37 @@ CANONICAL_VECTOR_VALUES = {
     ),
 }
 
+_MODELS_TO_CACHE = ("Qdrant/clip-ViT-B-32-vision",)
+MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+
+
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = ImageEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for name, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
 
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
-def test_embedding(model_name: str) -> None:
+def test_embedding(model_cache, model_name: str) -> None:
     is_ci = os.getenv("CI")
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
@@ -38,80 +67,69 @@ def test_embedding(model_name: str) -> None:
 
         dim = model_desc.dim
 
-        model = ImageEmbedding(model_name=model_desc.model)
+        with model_cache(model_desc.model) as model:
+            images = [
+                TEST_MISC_DIR / "image.jpeg",
+                str(TEST_MISC_DIR / "small_image.jpeg"),
+                Image.open((TEST_MISC_DIR / "small_image.jpeg")),
+                Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
+            ]
+            embeddings = list(model.embed(images))
+            embeddings = np.stack(embeddings, axis=0)
+            assert embeddings.shape == (len(images), dim)
 
-        images = [
+            canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+
+            assert np.allclose(
+                embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
+            ), model_desc.model
+
+            assert np.allclose(embeddings[1], embeddings[2]), model_desc.model
+
+
+@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
+def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        n_images = 32
+        test_images = [
             TEST_MISC_DIR / "image.jpeg",
             str(TEST_MISC_DIR / "small_image.jpeg"),
-            Image.open((TEST_MISC_DIR / "small_image.jpeg")),
-            Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
+            Image.open(TEST_MISC_DIR / "small_image.jpeg"),
         ]
-        embeddings = list(model.embed(images))
+        images = test_images * n_images
+
+        embeddings = list(model.embed(images, batch_size=10))
         embeddings = np.stack(embeddings, axis=0)
-        assert embeddings.shape == (len(images), dim)
+        assert np.allclose(embeddings[1], embeddings[2])
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_name]
 
-        assert np.allclose(
-            embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc.model
-
-        assert np.allclose(embeddings[1], embeddings[2]), model_desc.model
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+        assert embeddings.shape == (len(test_images) * n_images, n_dims)
+        assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
-def test_batch_embedding(n_dims: int, model_name: str) -> None:
-    is_ci = os.getenv("CI")
-    model = ImageEmbedding(model_name=model_name)
-    n_images = 32
-    test_images = [
-        TEST_MISC_DIR / "image.jpeg",
-        str(TEST_MISC_DIR / "small_image.jpeg"),
-        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
-    ]
-    images = test_images * n_images
+def test_parallel_processing(model_cache, n_dims: int, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        n_images = 32
+        test_images = [
+            TEST_MISC_DIR / "image.jpeg",
+            str(TEST_MISC_DIR / "small_image.jpeg"),
+            Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+        ]
+        images = test_images * n_images
+        embeddings = list(model.embed(images, batch_size=10, parallel=2))
+        embeddings = np.stack(embeddings, axis=0)
 
-    embeddings = list(model.embed(images, batch_size=10))
-    embeddings = np.stack(embeddings, axis=0)
-    assert np.allclose(embeddings[1], embeddings[2])
+        embeddings_2 = list(model.embed(images, batch_size=10, parallel=None))
+        embeddings_2 = np.stack(embeddings_2, axis=0)
 
-    canonical_vector = CANONICAL_VECTOR_VALUES[model_name]
+        embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
+        embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (len(test_images) * n_images, n_dims)
-    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
-@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
-def test_parallel_processing(n_dims: int, model_name: str) -> None:
-    is_ci = os.getenv("CI")
-    model = ImageEmbedding(model_name=model_name)
-
-    n_images = 32
-    test_images = [
-        TEST_MISC_DIR / "image.jpeg",
-        str(TEST_MISC_DIR / "small_image.jpeg"),
-        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
-    ]
-    images = test_images * n_images
-    embeddings = list(model.embed(images, batch_size=10, parallel=2))
-    embeddings = np.stack(embeddings, axis=0)
-
-    embeddings_2 = list(model.embed(images, batch_size=10, parallel=None))
-    embeddings_2 = np.stack(embeddings_2, axis=0)
-
-    embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
-    embeddings_3 = np.stack(embeddings_3, axis=0)
-
-    assert embeddings.shape == (n_images * len(test_images), n_dims)
-    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        assert embeddings.shape == (n_images * len(test_images), n_dims)
+        assert np.allclose(embeddings, embeddings_2, atol=1e-3)
+        assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
@@ -145,3 +163,13 @@ def test_embedding_size() -> None:
     assert model.embedding_size == 512
     if is_ci:
         delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
+def test_session_options(model_cache, model_name) -> None:
+    with model_cache(model_name) as default_model:
+        default_session_options = default_model.model.model.get_session_options()
+        assert default_session_options.enable_cpu_mem_arena is True
+        model = ImageEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
+        session_options = model.model.model.get_session_options()
+        assert session_options.enable_cpu_mem_arena is False

tests/test_late_interaction_embeddings.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 
 import pytest
 import numpy as np
@@ -150,42 +151,65 @@ CANONICAL_QUERY_VALUES = {
     ),
 }
 
+_MODELS_TO_CACHE = ("answerdotai/answerai-colbert-small-v1",)
+MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+
+
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = LateInteractionTextEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for name, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
+
 docs = ["Hello World"]
 
 
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_batch_embedding(model_name: str):
-    is_ci = os.getenv("CI")
+def test_batch_embedding(model_cache, model_name: str):
     docs_to_embed = docs * 10
 
-    model = LateInteractionTextEmbedding(model_name=model_name)
-    result = list(model.embed(docs_to_embed, batch_size=6))
-    expected_result = CANONICAL_COLUMN_VALUES[model_name]
+    with model_cache(model_name) as model:
+        result = list(model.embed(docs_to_embed, batch_size=6))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
 
-    for value in result:
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        for value in result:
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
 
 
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_batch_inference_size_same_as_single_inference(model_name: str):
-    is_ci = os.getenv("CI")
-
-    model = LateInteractionTextEmbedding(model_name=model_name)
-    docs_to_embed = ["short document", "A bit longer document, which should not affect the size"]
-    result = list(model.embed(docs_to_embed, batch_size=1))
-    result_2 = list(model.embed(docs_to_embed, batch_size=2))
-    assert len(result[0]) == len(result_2[0])
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+def test_batch_inference_size_same_as_single_inference(model_cache, model_name: str):
+    with model_cache(model_name) as model:
+        docs_to_embed = [
+            "short document",
+            "A bit longer document, which should not affect the size",
+        ]
+        result = list(model.embed(docs_to_embed, batch_size=1))
+        result_2 = list(model.embed(docs_to_embed, batch_size=2))
+        assert len(result[0]) == len(result_2[0])
 
 
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_single_embedding(model_name: str):
+def test_single_embedding(model_cache, model_name: str):
     is_ci = os.getenv("CI")
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     docs_to_embed = docs
@@ -195,20 +219,17 @@ def test_single_embedding(model_name: str):
             continue
 
         print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        whole_result = list(model.embed(docs_to_embed, batch_size=6))
-        assert len(whole_result) == 1
-        result = whole_result[0]
-        expected_result = CANONICAL_COLUMN_VALUES[model_name]
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+        with model_cache(model_desc.model) as model:
+            whole_result = list(model.embed(docs_to_embed, batch_size=6))
+            assert len(whole_result) == 1
+            result = whole_result[0]
+            expected_result = CANONICAL_COLUMN_VALUES[model_desc.model]
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
 
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_single_embedding_query(model_name: str):
+def test_single_embedding_query(model_cache, model_name: str):
     is_ci = os.getenv("CI")
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     queries_to_embed = docs
@@ -217,39 +238,34 @@ def test_single_embedding_query(model_name: str):
         if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue
 
-        print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        whole_result = list(model.query_embed(queries_to_embed))
-        assert len(whole_result) == 1
-        result = whole_result[0]
-        expected_result = CANONICAL_QUERY_VALUES[model_name]
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+        print("evaluating", model_desc.model)
+        with model_cache(model_desc.model) as model:
+            whole_result = list(model.query_embed(queries_to_embed))
+            assert len(whole_result) == 1
+            result = whole_result[0]
+            expected_result = CANONICAL_QUERY_VALUES[model_desc.model]
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
 
 @pytest.mark.parametrize("token_dim,model_name", [(96, "answerdotai/answerai-colbert-small-v1")])
-def test_parallel_processing(token_dim: int, model_name: str):
-    is_ci = os.getenv("CI")
-    model = LateInteractionTextEmbedding(model_name=model_name)
+def test_parallel_processing(model_cache, token_dim: int, model_name: str):
+    with model_cache(model_name) as model:
+        docs = ["hello world", "flag embedding"] * 100
+        embeddings = list(model.embed(docs, batch_size=10, parallel=2))
 
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+        embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
 
-    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+        # embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))  # inherits OnnxTextModel which
+        #         # is tested in TextEmbedding, disabling it here to reduce number of requests to hf
+        #         # multiprocessing is enough to test with `parallel=2`, and `parallel=None` is okay to tests since it reuses
+        #         # model from cache
 
-    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+        assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
 
-    assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
-
-    for i in range(len(embeddings)):
-        assert np.allclose(embeddings[i], embeddings_2[i], atol=1e-3)
-        assert np.allclose(embeddings[i], embeddings_3[i], atol=1e-3)
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        for i in range(len(embeddings)):
+            assert np.allclose(embeddings[i], embeddings_2[i], atol=1e-3)
+            # assert np.allclose(embeddings[i], embeddings_3[i], atol=1e-3)
 
 
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
@@ -292,3 +308,37 @@ def test_embedding_size():
     assert model.embedding_size == 96
     if is_ci:
         delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-ColBERT-small-v1"])
+def test_session_options(model_cache, model_name) -> None:
+    with model_cache(model_name) as default_model:
+        default_session_options = default_model.model.model.get_session_options()
+        assert default_session_options.enable_cpu_mem_arena is True
+        model = LateInteractionTextEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
+        session_options = model.model.model.get_session_options()
+        assert session_options.enable_cpu_mem_arena is False
+
+
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_token_count(model_cache, model_name) -> None:
+    with model_cache(model_name) as model:
+        documents = ["short doc", "it is a long document to check attention mask for paddings"]
+        short_doc_token_count = model.token_count(documents[0])
+        long_doc_token_count = model.token_count(documents[1])
+        documents_token_count = model.token_count(documents)
+        assert short_doc_token_count + long_doc_token_count == documents_token_count
+        # 2 is 2*DOC_MARKER_TOKEN_ID for each document
+        assert short_doc_token_count + long_doc_token_count + 2 == model.token_count(
+            documents, include_extension=True
+        )
+        assert short_doc_token_count + long_doc_token_count == model.token_count(
+            documents, batch_size=1
+        )
+        assert short_doc_token_count + long_doc_token_count == model.token_count(
+            documents, is_doc=False
+        )
+        # query min length is 32
+        assert model.token_count(documents, is_doc=False, include_extension=True) == 64
+        very_long_query = "It's a very long query which definitely contains more than 32 tokens and we're using it to check whether the method can handle large query properly without cutting it to 32 tokens"
+        assert model.token_count(very_long_query, is_doc=False, include_extension=True) > 32

tests/test_late_interaction_multimodal.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 
 import pytest
 from PIL import Image
@@ -6,7 +7,7 @@ import numpy as np
 
 from fastembed import LateInteractionMultimodalEmbedding
 from tests.config import TEST_MISC_DIR
-
+from tests.utils import delete_model_cache
 
 # vectors are abridged and rounded for brevity
 CANONICAL_IMAGE_VALUES = {
@@ -21,6 +22,17 @@ CANONICAL_IMAGE_VALUES = {
             [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
         ]
     ),
+    "Qdrant/colmodernvbert": np.array(
+        [
+            [0.11614, -0.15793, -0.11194, 0.0688, 0.08001, 0.10575, -0.07871],
+            [0.10094, -0.13301, -0.12069, 0.10932, 0.04645, 0.09884, 0.04048],
+            [0.13106, -0.18613, -0.13469, 0.10566, 0.03659, 0.07712, -0.03916],
+            [0.09754, -0.09596, -0.04839, 0.14991, 0.05692, 0.10569, -0.08349],
+            [0.02576, -0.15651, -0.09977, 0.09707, 0.13412, 0.09994, -0.09931],
+            [-0.06741, -0.1787, -0.19677, -0.07618, 0.13102, -0.02131, -0.02437],
+            [-0.02776, -0.10187, -0.13793, 0.03835, 0.04766, 0.04701, -0.15635],
+        ]
+    ),
 }
 
 CANONICAL_QUERY_VALUES = {
@@ -35,6 +47,17 @@ CANONICAL_QUERY_VALUES = {
             [-0.0165, -0.0106, 0.1672, -0.0768, 0.0389, -0.0038, 0.1137],
         ]
     ),
+    "Qdrant/colmodernvbert": np.array(
+        [
+            [0.05, 0.06557, 0.04026, 0.14981, 0.1842, 0.0263, -0.18706],
+            [-0.05664, -0.14028, 0.00649, -0.02849, 0.09034, -0.01494, 0.10693],
+            [-0.10147, -0.00716, 0.09084, -0.08236, -0.01849, -0.00972, -0.00461],
+            [-0.1233, -0.10814, -0.02337, -0.00329, 0.05984, 0.09934, 0.09846],
+            [-0.07053, -0.13119, -0.06487, 0.01508, 0.07459, 0.07655, 0.14821],
+            [0.00526, -0.13842, -0.05837, -0.02721, 0.13009, 0.05076, 0.17962],
+            [0.00924, -0.14383, -0.03057, -0.03691, 0.11718, 0.037, 0.13344],
+        ]
+    ),
 }
 
 queries = ["hello world", "flag embedding"]
@@ -44,43 +67,69 @@ images = [
     Image.open((TEST_MISC_DIR / "image.jpeg")),
 ]
 
+_MODELS_TO_CACHE = ("Qdrant/colmodernvbert",)
+MODELS_TO_CACHE = tuple(model_name.lower() for model_name in _MODELS_TO_CACHE)
 
-def test_batch_embedding():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
 
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = LateInteractionMultimodalEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for _, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
+
+def test_batch_embedding(model_cache):
     for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = list(model.embed_image(images, batch_size=2))
+        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
+            continue  # colpali is too large for ci
 
-        for value in result:
+        print("evaluating", model_name)
+        with model_cache(model_name) as model:
+            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)
+
+
+def test_single_embedding(model_cache):
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
+            continue  # colpali is too large for ci
+        print("evaluating", model_name)
+        with model_cache(model_name) as model:
+            result = next(iter(model.embed_image(images, batch_size=6)))
             token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
-def test_single_embedding():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
-    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        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
-        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_single_embedding_query():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
+def test_single_embedding_query(model_cache):
     for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
+            continue  # colpali is too large for ci
         print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = next(iter(model.embed_text(queries)))
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+        with model_cache(model_name) as model:
+            result = next(iter(model.embed_text(queries)))
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_get_embedding_size():
@@ -90,14 +139,27 @@ def test_get_embedding_size():
     model_name = "Qdrant/ColPali-v1.3-fp16"
     assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
 
+    model_name = "Qdrant/colmodernvbert"
+    assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
+
 
 def test_embedding_size():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-    model_name = "Qdrant/colpali-v1.3-fp16"
+    model_name = "Qdrant/colmodernvbert"
     model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
     assert model.embedding_size == 128
 
-    model_name = "Qdrant/ColPali-v1.3-fp16"
-    model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 128
+
+def test_token_count(model_cache) -> None:
+    model_name = "Qdrant/colmodernvbert"
+    with model_cache(model_name) as model:
+        documents = ["short doc", "it is a long document to check attention mask for paddings"]
+        short_doc_token_count = model.token_count(documents[0])
+        long_doc_token_count = model.token_count(documents[1])
+        documents_token_count = model.token_count(documents)
+        assert short_doc_token_count + long_doc_token_count == documents_token_count
+        assert short_doc_token_count + long_doc_token_count == model.token_count(
+            documents, batch_size=1
+        )
+        assert short_doc_token_count + long_doc_token_count < model.token_count(
+            documents, include_extension=True
+        )

tests/test_multi_gpu.py

@@ -1,5 +1,5 @@
 import pytest
-from typing import Optional
+
 from fastembed import (
     TextEmbedding,
     SparseTextEmbedding,
@@ -14,7 +14,7 @@ CACHE_DIR = "../model_cache"
 
 @pytest.mark.skip(reason="Requires a multi-gpu server")
 @pytest.mark.parametrize("device_id", [None, 0, 1])
-def test_gpu_via_providers(device_id: Optional[int]) -> None:
+def test_gpu_via_providers(device_id: int | None) -> None:
     docs = ["hello world", "flag embedding"]
 
     device_id = device_id if device_id is not None else 0
@@ -86,7 +86,7 @@ def test_gpu_via_providers(device_id: Optional[int]) -> None:
 
 @pytest.mark.skip(reason="Requires a multi-gpu server")
 @pytest.mark.parametrize("device_ids", [None, [0], [1], [0, 1]])
-def test_gpu_cuda_device_ids(device_ids: Optional[list[int]]) -> None:
+def test_gpu_cuda_device_ids(device_ids: list[int] | None) -> None:
     docs = ["hello world", "flag embedding"]
     device_id = device_ids[0] if device_ids else 0
     embedding_model = TextEmbedding(
@@ -171,7 +171,7 @@ def test_gpu_cuda_device_ids(device_ids: Optional[list[int]]) -> None:
 @pytest.mark.parametrize(
     "device_ids,parallel", [(None, None), (None, 2), ([1], None), ([1], 1), ([1], 2), ([0, 1], 2)]
 )
-def test_multi_gpu_parallel_inference(device_ids: Optional[list[int]], parallel: int) -> None:
+def test_multi_gpu_parallel_inference(device_ids: list[int] | None, parallel: int) -> None:
     docs = ["hello world", "flag embedding"] * 100
     batch_size = 5
 

tests/test_sparse_embeddings.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 
 import pytest
 import numpy as np
@@ -76,6 +77,40 @@ CANONICAL_QUERY_VALUES = {
 }
 
 
+_MODELS_TO_CACHE = (
+    "prithivida/Splade_PP_en_v1",
+    "Qdrant/minicoil-v1",
+    "Qdrant/bm25",
+    "Qdrant/bm42-all-minilm-l6-v2-attentions",
+)
+MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+
+
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = SparseTextEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for name, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
+
 docs = ["Hello World"]
 
 
@@ -83,23 +118,19 @@ docs = ["Hello World"]
     "model_name",
     ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
 )
-def test_batch_embedding(model_name: str) -> None:
-    is_ci = os.getenv("CI")
+def test_batch_embedding(model_cache, model_name: str) -> None:
     docs_to_embed = docs * 10
 
-    model = SparseTextEmbedding(model_name=model_name)
-    result = next(iter(model.embed(docs_to_embed, batch_size=6)))
-    expected_result = CANONICAL_COLUMN_VALUES[model_name]
-    assert result.indices.tolist() == expected_result["indices"]
+    with model_cache(model_name) as model:
+        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
+        assert result.indices.tolist() == expected_result["indices"]
 
-    for i, value in enumerate(result.values):
-        assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        for i, value in enumerate(result.values):
+            assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
 
 
-@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"])
-def test_single_embedding(model_name: str) -> None:
+def test_single_embedding(model_cache) -> None:
     is_ci = os.getenv("CI")
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
@@ -109,100 +140,106 @@ def test_single_embedding(model_name: str) -> None:
         ):  # attention models and bm25 are also parts of
             # SparseTextEmbedding, however, they have their own tests
             continue
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+        if not should_test_model(model_desc, model_desc.model, is_ci, is_manual):
             continue
 
-        model = SparseTextEmbedding(model_name=model_name)
+        with model_cache(model_desc.model) as model:
+            passage_result = next(iter(model.embed(docs, batch_size=6)))
+            query_result = next(iter(model.query_embed(docs)))
+            expected_result = CANONICAL_COLUMN_VALUES[model_desc.model]
+            expected_query_result = CANONICAL_QUERY_VALUES.get(model_desc.model, expected_result)
+            assert passage_result.indices.tolist() == expected_result["indices"]
+            for i, value in enumerate(passage_result.values):
+                assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
 
-        passage_result = next(iter(model.embed(docs, batch_size=6)))
-        query_result = next(iter(model.query_embed(docs)))
-        expected_result = CANONICAL_COLUMN_VALUES[model_name]
-        expected_query_result = CANONICAL_QUERY_VALUES.get(model_name, expected_result)
-        assert passage_result.indices.tolist() == expected_result["indices"]
-        for i, value in enumerate(passage_result.values):
-            assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
-
-        assert query_result.indices.tolist() == expected_query_result["indices"]
-        for i, value in enumerate(query_result.values):
-            assert pytest.approx(value, abs=0.001) == expected_query_result["values"][i]
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+            assert query_result.indices.tolist() == expected_query_result["indices"]
+            for i, value in enumerate(query_result.values):
+                assert pytest.approx(value, abs=0.001) == expected_query_result["values"][i]
 
 
 @pytest.mark.parametrize(
     "model_name",
     ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
 )
-def test_parallel_processing(model_name: str) -> None:
-    is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name=model_name)
-    docs = ["hello world", "flag embedding"] * 30
-    sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
-    sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))
-    sparse_embeddings = list(model.embed(docs, batch_size=10, parallel=None))
+def test_parallel_processing(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        docs = ["hello world", "flag embedding"] * 30
+        sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
+        # sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))  # inherits OnnxTextModel which
+        # is tested in TextEmbedding, disabling it here to reduce number of requests to hf
+        # multiprocessing is enough to test with `parallel=2`, and `parallel=None` is okay to tests since it reuses
+        # model from cache
+        sparse_embeddings = list(model.embed(docs, batch_size=10, parallel=None))
 
-    assert (
-        len(sparse_embeddings)
-        == len(sparse_embeddings_duo)
-        == len(sparse_embeddings_all)
-        == len(docs)
-    )
-
-    for sparse_embedding, sparse_embedding_duo, sparse_embedding_all in zip(
-        sparse_embeddings, sparse_embeddings_duo, sparse_embeddings_all
-    ):
         assert (
-            sparse_embedding.indices.tolist()
-            == sparse_embedding_duo.indices.tolist()
-            == sparse_embedding_all.indices.tolist()
+            len(sparse_embeddings)
+            == len(sparse_embeddings_duo)
+            # == len(sparse_embeddings_all)
+            == len(docs)
         )
-        assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3)
-        assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
 
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        for (
+            sparse_embedding,
+            sparse_embedding_duo,
+            # sparse_embedding_all
+        ) in zip(
+            sparse_embeddings,
+            sparse_embeddings_duo,
+            # sparse_embeddings_all
+        ):
+            assert (
+                sparse_embedding.indices.tolist() == sparse_embedding_duo.indices.tolist()
+                # == sparse_embedding_all.indices.tolist()
+            )
+            assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3)
+            # assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
 
 
-@pytest.fixture
-def bm25_instance() -> None:
-    ci = os.getenv("CI", True)
-    model = Bm25("Qdrant/bm25", language="english")
-    yield model
-    if ci:
-        delete_model_cache(model._model_dir)
+def test_stem_with_stopwords_and_punctuation(model_cache) -> None:
+    with model_cache("Qdrant/bm25") as model:
+        bm25_instance = model.model
+        # Setup
+        original_stopwords = bm25_instance.stopwords.copy()
+        original_punctuation = bm25_instance.punctuation.copy()
+
+        bm25_instance.stopwords = {"the", "is", "a"}
+        bm25_instance.punctuation = {".", ",", "!"}
+
+        # Test data
+        tokens = ["The", "quick", "brown", "fox", "is", "a", "test", "sentence", ".", "!"]
+
+        # Execute
+        result = bm25_instance._stem(tokens)
+
+        # Assert
+        expected = ["quick", "brown", "fox", "test", "sentenc"]
+        assert result == expected, f"Expected {expected}, but got {result}"
+
+        bm25_instance.stopwords = original_stopwords
+        bm25_instance.punctuation = original_punctuation
 
 
-def test_stem_with_stopwords_and_punctuation(bm25_instance: Bm25) -> None:
-    # Setup
-    bm25_instance.stopwords = {"the", "is", "a"}
-    bm25_instance.punctuation = {".", ",", "!"}
+def test_stem_case_insensitive_stopwords(model_cache) -> None:
+    with model_cache("Qdrant/bm25") as model:
+        bm25_instance = model.model
+        original_stopwords = bm25_instance.stopwords.copy()
+        original_punctuation = bm25_instance.punctuation.copy()
 
-    # Test data
-    tokens = ["The", "quick", "brown", "fox", "is", "a", "test", "sentence", ".", "!"]
+        # Setup
+        bm25_instance.stopwords = {"the", "is", "a"}
+        bm25_instance.punctuation = {".", ",", "!"}
 
-    # Execute
-    result = bm25_instance._stem(tokens)
+        # Test data
+        tokens = ["THE", "Quick", "Brown", "Fox", "IS", "A", "Test", "Sentence", ".", "!"]
 
-    # Assert
-    expected = ["quick", "brown", "fox", "test", "sentenc"]
-    assert result == expected, f"Expected {expected}, but got {result}"
+        # Execute
+        result = bm25_instance._stem(tokens)
 
-
-def test_stem_case_insensitive_stopwords(bm25_instance: Bm25) -> None:
-    # Setup
-    bm25_instance.stopwords = {"the", "is", "a"}
-    bm25_instance.punctuation = {".", ",", "!"}
-
-    # Test data
-    tokens = ["THE", "Quick", "Brown", "Fox", "IS", "A", "Test", "Sentence", ".", "!"]
-
-    # Execute
-    result = bm25_instance._stem(tokens)
-
-    # Assert
-    expected = ["quick", "brown", "fox", "test", "sentenc"]
-    assert result == expected, f"Expected {expected}, but got {result}"
+        # Assert
+        expected = ["quick", "brown", "fox", "test", "sentenc"]
+        assert result == expected, f"Expected {expected}, but got {result}"
+        bm25_instance.stopwords = original_stopwords
+        bm25_instance.punctuation = original_punctuation
 
 
 @pytest.mark.parametrize("disable_stemmer", [True, False])
@@ -244,3 +281,42 @@ def test_lazy_load(model_name: str) -> None:
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    [
+        "prithivida/Splade_PP_en_v1",
+        "Qdrant/minicoil-v1",
+        "Qdrant/bm42-all-minilm-l6-v2-attentions",
+    ],
+)
+def test_session_options(model_cache, model_name) -> None:
+    with model_cache(model_name) as default_model:
+        default_session_options = default_model.model.model.get_session_options()
+        assert default_session_options.enable_cpu_mem_arena is True
+        model = SparseTextEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
+        session_options = model.model.model.get_session_options()
+        assert session_options.enable_cpu_mem_arena is False
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    [
+        "prithivida/Splade_PP_en_v1",
+        "Qdrant/minicoil-v1",
+        "Qdrant/bm42-all-minilm-l6-v2-attentions",
+        "Qdrant/bm25",
+    ],
+)
+def test_token_count(model_cache, model_name) -> None:
+    with model_cache(model_name) as model:
+        documents = [
+            "Name me a couple of cities were the capitals of Germany?",
+            "Berlin is the current capital of Germany, Bonn is a former capital of Germany.",
+        ]
+        first_doc_token_count = model.token_count(documents[0])
+        second_doc_token_count = model.token_count(documents[1])
+        doc_token_count = model.token_count(documents)
+        assert first_doc_token_count + second_doc_token_count == doc_token_count
+        assert doc_token_count == model.token_count(documents, batch_size=1)

tests/test_text_cross_encoder.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 
 import numpy as np
 import pytest
@@ -16,8 +17,37 @@ CANONICAL_SCORE_VALUES = {
 }
 
 
+_MODELS_TO_CACHE = ("Xenova/ms-marco-MiniLM-L-6-v2",)
+MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+
+
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = TextCrossEncoder(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for name, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
+
 @pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_rerank(model_name: str) -> None:
+def test_rerank(model_cache, model_name: str) -> None:
     is_ci = os.getenv("CI")
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
@@ -25,11 +55,29 @@ def test_rerank(model_name: str) -> None:
         if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue
 
-        model = TextCrossEncoder(model_name=model_name)
+        with model_cache(model_desc.model) as model:
+            query = "What is the capital of France?"
+            documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
+            scores = np.array(list(model.rerank(query, documents)))
 
+            pairs = [(query, doc) for doc in documents]
+            scores2 = np.array(list(model.rerank_pairs(pairs)))
+            assert np.allclose(
+                scores, scores2, atol=1e-5
+            ), f"Model: {model_desc.model}, Scores: {scores}, Scores2: {scores2}"
+
+            canonical_scores = CANONICAL_SCORE_VALUES[model_desc.model]
+            assert np.allclose(
+                scores, canonical_scores, atol=1e-3
+            ), f"Model: {model_desc.model}, Scores: {scores}, Expected: {canonical_scores}"
+
+
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+def test_batch_rerank(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
         query = "What is the capital of France?"
-        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
-        scores = np.array(list(model.rerank(query, documents)))
+        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50
+        scores = np.array(list(model.rerank(query, documents, batch_size=10)))
 
         pairs = [(query, doc) for doc in documents]
         scores2 = np.array(list(model.rerank_pairs(pairs)))
@@ -37,38 +85,12 @@ def test_rerank(model_name: str) -> None:
             scores, scores2, atol=1e-5
         ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
 
-        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+        canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50)
+
+        assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}"
         assert np.allclose(
             scores, canonical_scores, atol=1e-3
         ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_batch_rerank(model_name: str) -> None:
-    is_ci = os.getenv("CI")
-
-    model = TextCrossEncoder(model_name=model_name)
-
-    query = "What is the capital of France?"
-    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50
-    scores = np.array(list(model.rerank(query, documents, batch_size=10)))
-
-    pairs = [(query, doc) for doc in documents]
-    scores2 = np.array(list(model.rerank_pairs(pairs)))
-    assert np.allclose(
-        scores, scores2, atol=1e-5
-    ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
-
-    canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50)
-
-    assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}"
-    assert np.allclose(
-        scores, canonical_scores, atol=1e-3
-    ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
@@ -86,21 +108,44 @@ def test_lazy_load(model_name: str) -> None:
 
 
 @pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_rerank_pairs_parallel(model_name: str) -> None:
-    is_ci = os.getenv("CI")
+def test_rerank_pairs_parallel(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        query = "What is the capital of France?"
+        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 10
+        pairs = [(query, doc) for doc in documents]
+        scores_parallel = np.array(list(model.rerank_pairs(pairs, parallel=2, batch_size=10)))
+        scores_sequential = np.array(list(model.rerank_pairs(pairs, batch_size=10)))
+        assert np.allclose(
+            scores_parallel, scores_sequential, atol=1e-5
+        ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Scores (Sequential): {scores_sequential}"
+        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+        assert np.allclose(
+            scores_parallel[: len(canonical_scores)], canonical_scores, atol=1e-3
+        ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Expected: {canonical_scores}"
 
-    model = TextCrossEncoder(model_name=model_name)
-    query = "What is the capital of France?"
-    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 10
-    pairs = [(query, doc) for doc in documents]
-    scores_parallel = np.array(list(model.rerank_pairs(pairs, parallel=2, batch_size=10)))
-    scores_sequential = np.array(list(model.rerank_pairs(pairs, batch_size=10)))
-    assert np.allclose(
-        scores_parallel, scores_sequential, atol=1e-5
-    ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Scores (Sequential): {scores_sequential}"
-    canonical_scores = CANONICAL_SCORE_VALUES[model_name]
-    assert np.allclose(
-        scores_parallel[: len(canonical_scores)], canonical_scores, atol=1e-3
-    ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Expected: {canonical_scores}"
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+def test_token_count(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        pairs = [
+            ("What is the capital of France?", "Paris is the capital of France."),
+            (
+                "Name me a couple of cities were the capitals of Germany?",
+                "Berlin is the current capital of Germany, Bonn is a former capital of Germany.",
+            ),
+        ]
+        first_pair_token_count = model.token_count([pairs[0]])
+        second_pair_token_count = model.token_count([pairs[1]])
+        pairs_token_count = model.token_count(pairs)
+        assert first_pair_token_count + second_pair_token_count == pairs_token_count
+        assert pairs_token_count == model.token_count(pairs, batch_size=1)
+
+
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+def test_session_options(model_cache, model_name) -> None:
+    with model_cache(model_name) as default_model:
+        default_session_options = default_model.model.model.get_session_options()
+        assert default_session_options.enable_cpu_mem_arena is True
+        model = TextCrossEncoder(model_name=model_name, enable_cpu_mem_arena=False)
+        session_options = model.model.model.get_session_options()
+        assert session_options.enable_cpu_mem_arena is False

tests/test_text_onnx_embeddings.py

@@ -1,5 +1,6 @@
 import os
 import platform
+from contextlib import contextmanager
 
 import numpy as np
 import pytest
@@ -71,9 +72,37 @@ CANONICAL_VECTOR_VALUES = {
 
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
+_MODELS_TO_CACHE = ("BAAI/bge-small-en-v1.5",)
+MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+
+
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = TextEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for name, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
 
 @pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
-def test_embedding(model_name: str) -> None:
+def test_embedding(model_cache, model_name: str) -> None:
     is_ci = os.getenv("CI")
     is_mac = platform.system() == "Darwin"
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
@@ -88,55 +117,44 @@ def test_embedding(model_name: str) -> None:
 
         dim = model_desc.dim
 
-        model = TextEmbedding(model_name=model_desc.model)
-        docs = ["hello world", "flag embedding"]
-        embeddings = list(model.embed(docs))
+        with model_cache(model_desc.model) as model:
+            docs = ["hello world", "flag embedding"]
+            embeddings = list(model.embed(docs))
+            embeddings = np.stack(embeddings, axis=0)
+            assert embeddings.shape == (2, dim)
+
+            canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+            assert np.allclose(
+                embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
+            ), model_desc.model
+
+
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
+def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        docs = ["hello world", "flag embedding"] * 100
+        embeddings = list(model.embed(docs, batch_size=10))
         embeddings = np.stack(embeddings, axis=0)
-        assert embeddings.shape == (2, dim)
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
-        assert np.allclose(
-            embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc.model
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+        assert embeddings.shape == (len(docs), n_dims)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
-def test_batch_embedding(n_dims: int, model_name: str) -> None:
-    is_ci = os.getenv("CI")
-    model = TextEmbedding(model_name=model_name)
+def test_parallel_processing(model_cache, n_dims: int, model_name: str) -> None:
+    with model_cache(model_name) as model:
+        docs = ["hello world", "flag embedding"] * 100
+        embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+        embeddings = np.stack(embeddings, axis=0)
 
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10))
-    embeddings = np.stack(embeddings, axis=0)
+        embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+        embeddings_2 = np.stack(embeddings_2, axis=0)
 
-    assert embeddings.shape == (len(docs), n_dims)
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+        embeddings_3 = np.stack(embeddings_3, axis=0)
 
-
-@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
-def test_parallel_processing(n_dims: int, model_name: str) -> None:
-    is_ci = os.getenv("CI")
-    model = TextEmbedding(model_name=model_name)
-
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
-    embeddings = np.stack(embeddings, axis=0)
-
-    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
-    embeddings_2 = np.stack(embeddings_2, axis=0)
-
-    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
-    embeddings_3 = np.stack(embeddings_3, axis=0)
-
-    assert embeddings.shape == (len(docs), n_dims)
-    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        assert embeddings.shape == (len(docs), n_dims)
+        assert np.allclose(embeddings, embeddings_2, atol=1e-3)
+        assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
 
 @pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
@@ -175,3 +193,27 @@ def test_embedding_size() -> None:
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["sentence-transformers/all-MiniLM-L6-v2"])
+def test_session_options(model_cache, model_name) -> None:
+    with model_cache(model_name) as default_model:
+        default_session_options = default_model.model.model.get_session_options()
+        assert default_session_options.enable_cpu_mem_arena is True
+        model = TextEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
+        session_options = model.model.model.get_session_options()
+        assert session_options.enable_cpu_mem_arena is False
+
+
+@pytest.mark.parametrize("model_name", ["sentence-transformers/all-MiniLM-L6-v2"])
+def test_token_count(model_cache, model_name) -> None:
+    with model_cache(model_name) as model:
+        documents = [
+            "Name me a couple of cities were the capitals of Germany?",
+            "Berlin is the current capital of Germany, Bonn is a former capital of Germany.",
+        ]
+        first_doc_token_count = model.token_count(documents[0])
+        second_doc_token_count = model.token_count(documents[1])
+        doc_token_count = model.token_count(documents)
+        assert first_doc_token_count + second_doc_token_count == doc_token_count
+        assert doc_token_count == model.token_count(documents, batch_size=1)

tests/utils.py

@@ -3,12 +3,12 @@ import traceback
 
 from pathlib import Path
 from types import TracebackType
-from typing import Union, Callable, Any, Type, Optional
+from typing import Callable, Any, Type
 
 from fastembed.common.model_description import BaseModelDescription
 
 
-def delete_model_cache(model_dir: Union[str, Path]) -> None:
+def delete_model_cache(model_dir: str | Path) -> None:
     """Delete the model cache directory.
 
     If a model was downloaded from the HuggingFace model hub, then _model_dir is the dir to snapshots, removing
@@ -42,14 +42,14 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
 def should_test_model(
     model_desc: BaseModelDescription,
     autotest_model_name: str,
-    is_ci: Optional[str],
+    is_ci: str | None,
     is_manual: bool,
 ):
     """Determine if a model should be tested based on environment
 
     Tests can be run either in ci or locally.
     Testing all models each time in ci is too long.
-    The testing scheme in ci and on a local machine are different, therefore, there are 3 possible scenarious.
+    The testing scheme in ci and on a local machine are different, therefore, there are 3 possible scenarios.
     1) Run lightweight tests in ci:
         - test only one model that has been manually chosen as a representative for a certain class family
     2) Run heavyweight (manual) tests in ci: