fastembed-gpu 0.7.2__tar.gz → 0.7.4__tar.gz

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/PKG-INFO

@@ -1,8 +1,10 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: fastembed-gpu
-Version: 0.7.2
+Version: 0.7.4
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache License
+License-File: LICENSE
+License-File: NOTICE
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
 Author: Qdrant Team
 Author-email: info@qdrant.tech
@@ -14,17 +16,22 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Requires-Dist: huggingface-hub (>=0.20,<1.0)
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: huggingface-hub (>=0.20,<2.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: mmh3 (>=4.1.0,<6.0.0)
-Requires-Dist: numpy (>=1.21) ; python_version >= "3.10" and python_version < "3.12"
+Requires-Dist: numpy (>=1.21) ; python_version == "3.11"
 Requires-Dist: numpy (>=1.21,<2.1.0) ; python_version < "3.10"
+Requires-Dist: numpy (>=1.21,<2.3.0) ; python_version == "3.10"
 Requires-Dist: numpy (>=1.26) ; python_version == "3.12"
-Requires-Dist: numpy (>=2.1.0) ; python_version >= "3.13"
+Requires-Dist: numpy (>=2.1.0) ; python_version == "3.13"
+Requires-Dist: numpy (>=2.3.0) ; python_version >= "3.14"
 Requires-Dist: onnxruntime-gpu (>1.20.0) ; python_version >= "3.13"
 Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0) ; python_version >= "3.10" and python_version < "3.13"
 Requires-Dist: onnxruntime-gpu (>=1.17.0,<1.20.0) ; python_version < "3.10"
-Requires-Dist: pillow (>=10.3.0,<12.0.0)
+Requires-Dist: pillow (>=10.3.0,<11.0) ; python_version < "3.10"
+Requires-Dist: pillow (>=10.3.0,<12.0) ; python_version >= "3.10" and python_version < "3.13"
+Requires-Dist: pillow (>=11.0.0,<12.0) ; python_version >= "3.13"
 Requires-Dist: py-rust-stemmers (>=0.1.0,<0.2.0)
 Requires-Dist: requests (>=2.31,<3.0)
 Requires-Dist: tokenizers (>=0.15,<1.0)

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/common/model_management.py

@@ -3,6 +3,7 @@ import time
 import json
 import shutil
 import tarfile
+from copy import deepcopy
 from pathlib import Path
 from typing import Any, Optional, Union, TypeVar, Generic
 
@@ -224,11 +225,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,
@@ -408,14 +404,32 @@ class ModelManagement(Generic[T]):
         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 +462,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/common/onnx_model.py

@@ -24,6 +24,8 @@ class OnnxOutputContext:
 
 
 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")
@@ -60,6 +62,7 @@ class OnnxModel(Generic[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
@@ -99,6 +102,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 +119,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/image/onnx_embedding.py

@@ -98,6 +98,7 @@ 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
@@ -134,6 +135,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
@@ -180,6 +182,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/image/onnx_image_model.py

@@ -55,6 +55,7 @@ class OnnxImageModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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)
 
@@ -99,6 +101,7 @@ class OnnxImageModel(OnnxModel[T]):
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -130,6 +133,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction/colbert.py

@@ -2,12 +2,13 @@ import string
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from tokenizers import Encoding
+from tokenizers import Encoding, Tokenizer
 
+from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.types import NumpyArray
 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,
 )
@@ -48,24 +49,24 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         if not is_doc:
             for embedding in output.model_output:
                 yield embedding
+        else:
+            if output.input_ids is None or output.attention_mask is None:
+                raise ValueError(
+                    "input_ids and attention_mask must be provided for document post-processing"
+                )
 
-        if output.input_ids is None or output.attention_mask is None:
-            raise ValueError(
-                "input_ids and attention_mask must be provided for document post-processing"
-            )
-
-        for i, token_sequence in enumerate(output.input_ids):
-            for j, token_id in enumerate(token_sequence):  # type: ignore
-                if token_id in self.skip_list or token_id == self.pad_token_id:
-                    output.attention_mask[i, j] = 0
+            for i, token_sequence in enumerate(output.input_ids):
+                for j, token_id in enumerate(token_sequence):  # type: ignore
+                    if token_id in self.skip_list or token_id == self.pad_token_id:
+                        output.attention_mask[i, j] = 0
 
-        output.model_output *= np.expand_dims(output.attention_mask, 2)
-        norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
-        norm_clamped = np.maximum(norm, 1e-12)
-        output.model_output /= norm_clamped
+            output.model_output *= np.expand_dims(output.attention_mask, 2)
+            norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
+            norm_clamped = np.maximum(norm, 1e-12)
+            output.model_output /= norm_clamped
 
-        for embedding, attention_mask in zip(output.model_output, output.attention_mask):
-            yield embedding[attention_mask == 1]
+            for embedding, attention_mask in zip(output.model_output, output.attention_mask):
+                yield embedding[attention_mask == 1]
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
@@ -87,29 +88,46 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     def _tokenize_query(self, query: str) -> list[Encoding]:
-        assert self.tokenizer is not None
-        encoded = self.tokenizer.encode_batch([query])
-        # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
-        if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
-            prev_padding = None
-            if self.tokenizer.padding:
-                prev_padding = self.tokenizer.padding
-            self.tokenizer.enable_padding(
-                pad_token=self.MASK_TOKEN,
-                pad_id=self.mask_token_id,
-                length=self.MIN_QUERY_LENGTH,
-            )
-            encoded = self.tokenizer.encode_batch([query])
-            if prev_padding is None:
-                self.tokenizer.no_padding()
-            else:
-                self.tokenizer.enable_padding(**prev_padding)
+        assert self.query_tokenizer is not None
+        encoded = self.query_tokenizer.encode_batch([query])
         return encoded
 
     def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
         encoded = self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
         return encoded
 
+    def token_count(
+        self,
+        texts: Union[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.
@@ -157,6 +175,7 @@ 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
@@ -183,6 +202,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self.pad_token_id: Optional[int] = None
         self.skip_list: set[int] = set()
 
+        self.query_tokenizer: Optional[Tokenizer] = None
+
         if not self.lazy_load:
             self.load_onnx_model()
 
@@ -194,7 +215,10 @@ 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)
+
         assert self.tokenizer is not None
         self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
@@ -205,6 +229,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         current_max_length = self.tokenizer.truncation["max_length"]
         # ensure not to overflow after adding document-marker
         self.tokenizer.enable_truncation(max_length=current_max_length - 1)
+        self.query_tokenizer.enable_truncation(max_length=current_max_length - 1)
+        self.query_tokenizer.enable_padding(
+            pad_token=self.MASK_TOKEN,
+            pad_id=self.mask_token_id,
+            length=self.MIN_QUERY_LENGTH,
+        )
 
     def embed(
         self,
@@ -239,6 +269,7 @@ 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,
         )
 

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction/late_interaction_embedding_base.py

@@ -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: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction/late_interaction_text_embedding.py

@@ -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: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction_multimodal/colpali.py

@@ -6,7 +6,7 @@ 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.utils import define_cache_dir, iter_batch
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
 )
@@ -80,6 +80,7 @@ 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
@@ -125,6 +126,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 +172,23 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         encoded = self.tokenizer.encode_batch(texts_query)  # type: ignore[union-attr]
         return encoded
 
+    def token_count(
+        self,
+        texts: Union[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]:
@@ -238,6 +257,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,
         )
 
@@ -273,6 +293,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -162,3 +162,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: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -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: Union[str, Iterable[str]],
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -64,6 +64,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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
@@ -122,6 +124,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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(),
@@ -189,6 +195,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -220,6 +227,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -111,6 +111,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
@@ -150,6 +151,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(
@@ -192,6 +194,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 +207,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -33,6 +33,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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
@@ -96,6 +98,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -161,3 +161,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/sparse/bm25.py

@@ -268,6 +268,15 @@ class Bm25(SparseTextEmbeddingBase):
             embeddings.append(SparseEmbedding.from_dict(token_id2value))
         return embeddings
 
+    def token_count(self, texts: Union[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.
 

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/sparse/bm42.py

@@ -31,9 +31,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]):
@@ -95,6 +103,7 @@ 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
@@ -124,7 +133,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 +147,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]
@@ -304,6 +314,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
@@ -341,6 +352,13 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
         return Bm42TextEmbeddingWorker
 
+    def token_count(
+        self, texts: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/sparse/minicoil.py

@@ -46,9 +46,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]):
@@ -110,6 +117,8 @@ 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
@@ -146,6 +155,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 +166,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,6 +187,11 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             avg_len=self.avg_len,
         )
 
+    def token_count(
+        self, texts: Union[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]],
@@ -214,6 +229,7 @@ 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,
         )
 

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/sparse/sparse_embedding_base.py

@@ -86,3 +86,7 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
             yield from self.embed([query], **kwargs)
         else:
             yield from self.embed(query, **kwargs)
+
+    def token_count(self, texts: Union[str, Iterable[str]], **kwargs: Any) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/sparse/sparse_text_embedding.py

@@ -128,3 +128,17 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
             Iterable[SparseEmbedding]: The sparse embeddings.
         """
         yield from self.model.query_embed(query, **kwargs)
+
+    def token_count(
+        self, texts: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/sparse/splade_pp.py

@@ -53,6 +53,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             scores = row_scores[indices]
             yield SparseEmbedding(values=scores, indices=indices)
 
+    def token_count(
+        self, texts: Union[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.
@@ -99,6 +104,7 @@ 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
@@ -133,6 +139,7 @@ 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(
@@ -168,6 +175,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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/text/onnx_embedding.py

@@ -233,7 +233,7 @@ 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
@@ -291,6 +291,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 +328,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: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/text/onnx_text_model.py

@@ -54,6 +54,7 @@ class OnnxTextModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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)
 
@@ -110,6 +112,7 @@ class OnnxTextModel(OnnxModel[T]):
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = 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,21 @@ 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: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/text/text_embedding.py

@@ -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: Union[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_gpu-0.7.2 → fastembed_gpu-0.7.4}/fastembed/text/text_embedding_base.py

@@ -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: Union[str, Iterable[str]], **kwargs: Any) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed-gpu"
-version = "0.7.2"
+version = "0.7.4"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -13,10 +13,12 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 [tool.poetry.dependencies]
 python = ">=3.9.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,<3.11" },
+    { version = ">=1.21", python = ">=3.11,<3.12" },
+    { version = ">=1.26", python = ">=3.12,<3.13" },
+    { version = ">=2.1.0", python = ">=3.13,<3.14" },
+    { version = ">=2.3.0", python = ">=3.14" },
 ]
 onnxruntime-gpu = [
     { version = ">=1.17.0,<1.20.0", python = "<3.10" },
@@ -26,9 +28,13 @@ onnxruntime-gpu = [
 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,<11.0", python = "<3.10" },
+    { version = ">=10.3.0,<12.0", python = ">=3.10,<3.13" },
+    { version = ">=11.0.0,<12.0", python = ">=3.13" },
+]
 mmh3 = ">=4.1.0,<6.0.0"
 py-rust-stemmers = "^0.1.0"
 
@@ -39,12 +45,15 @@ 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.13" },
+    { version = ">=1.18.0", python = ">=3.13" },
+]
 
 [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"