fastembed-gpu 0.6.1__tar.gz → 0.7.1__tar.gz

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: fastembed-gpu
-Version: 0.6.1
+Version: 0.7.1
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache License
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/common/onnx_model.py

@@ -28,7 +28,16 @@ class OnnxModel(Generic[T]):
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/common/preprocessor_utils.py

@@ -36,9 +36,9 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
-        assert (
-            "model_max_length" in tokenizer_config or "max_length" in tokenizer_config
-        ), "Models without model_max_length or max_length are not supported."
+        assert "model_max_length" in tokenizer_config or "max_length" in tokenizer_config, (
+            "Models without model_max_length or max_length are not supported."
+        )
         if "model_max_length" not in tokenizer_config:
             max_context = tokenizer_config["max_length"]
         elif "max_length" not in tokenizer_config:

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/common/types.py

@@ -16,6 +16,7 @@ ImageInput: TypeAlias = Union[PathInput, Image.Image]
 
 OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]
 NumpyArray = Union[
+    NDArray[np.float64],
     NDArray[np.float32],
     NDArray[np.float16],
     NDArray[np.int8],

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/image/image_embedding.py

@@ -77,6 +77,40 @@ class ImageEmbedding(ImageEmbeddingBase):
             "Please check the supported models using `ImageEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed(
         self,
         images: Union[ImageInput, Iterable[ImageInput]],

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/image/image_embedding_base.py

@@ -18,6 +18,7 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -42,3 +43,13 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
             Iterable[NdArray]: The embeddings.
         """
         raise NotImplementedError()
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/image/onnx_embedding.py

@@ -112,11 +112,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
         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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -177,6 +178,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             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,
             **kwargs,
         )
 
@@ -193,7 +196,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
         return normalize(output.model_output)
 
 

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/image/onnx_image_model.py

@@ -23,7 +23,16 @@ class OnnxImageModel(OnnxModel[T]):
     def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
@@ -88,6 +97,8 @@ class OnnxImageModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -104,7 +115,7 @@ class OnnxImageModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+                yield from self._post_process_onnx_output(self.onnx_embed(batch), **kwargs)
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -114,6 +125,8 @@ class OnnxImageModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -125,7 +138,7 @@ class OnnxImageModel(OnnxModel[T]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
-                yield from self._post_process_onnx_output(batch)  # type: ignore
+                yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
 
 
 class ImageEmbeddingWorker(EmbeddingWorker[T]):

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/image/transform/functional.py

@@ -72,26 +72,26 @@ def normalize(
     if not np.issubdtype(image.dtype, np.floating):
         image = image.astype(np.float32)
 
-    mean = mean if isinstance(mean, list) else [mean] * num_channels
+    mean_list = mean if isinstance(mean, list) else [mean] * num_channels
 
-    if len(mean) != num_channels:
+    if len(mean_list) != num_channels:
         raise ValueError(
             f"mean must have the same number of channels as the image, image has {num_channels} channels, got "
-            f"{len(mean)}"
+            f"{len(mean_list)}"
         )
 
-    mean_arr = np.array(mean, dtype=np.float32)
+    mean_arr = np.array(mean_list, dtype=np.float32)
 
-    std = std if isinstance(std, list) else [std] * num_channels
-    if len(std) != num_channels:
+    std_list = std if isinstance(std, list) else [std] * num_channels
+    if len(std_list) != num_channels:
         raise ValueError(
-            f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std)}"
+            f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std_list)}"
         )
 
-    std_arr = np.array(std, dtype=np.float32)
+    std_arr = np.array(std_list, dtype=np.float32)
 
-    image = ((image.T - mean_arr) / std_arr).T
-    return image
+    image_upd = ((image.T - mean_arr) / std_arr).T
+    return image_upd
 
 
 def resize(

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction/colbert.py

@@ -43,7 +43,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
-        self, output: OnnxOutputContext, is_doc: bool = True
+        self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any
     ) -> Iterable[NumpyArray]:
         if not is_doc:
             return output.model_output
@@ -169,11 +169,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
         self.mask_token_id: Optional[int] = None
         self.pad_token_id: Optional[int] = None
@@ -233,6 +234,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             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,
             **kwargs,
         )
 

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction/late_interaction_embedding_base.py

@@ -17,6 +17,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -58,3 +59,13 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
             yield from self.embed([query], **kwargs)
         else:
             yield from self.embed(query, **kwargs)
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction/late_interaction_text_embedding.py

@@ -80,6 +80,40 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
             "Please check the supported models using `LateInteractionTextEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed(
         self,
         documents: Union[str, Iterable[str]],

fastembed_gpu-0.7.1/fastembed/late_interaction/token_embeddings.py

@@ -0,0 +1,83 @@
+from dataclasses import asdict
+from typing import Union, Iterable, Optional, Any, Type
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+supported_token_embeddings_models = [
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-small-en-tokens",
+        dim=512,
+        description="Text embeddings, Unimodal (text), English, 8192 input tokens truncation,"
+        " Prefixes for queries/documents: not necessary, 2023 year.",
+        license="apache-2.0",
+        size_in_GB=0.12,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-small-en"),
+        model_file="onnx/model.onnx",
+    ),
+]
+
+
+class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
+    @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_token_embeddings_models
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return TokensEmbeddingWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        # Size: (batch_size, sequence_length, hidden_size)
+        embeddings = output.model_output
+        # Size: (batch_size, sequence_length)
+        assert output.attention_mask is not None
+        masks = output.attention_mask
+
+        # For each document we only select those embeddings that are not masked out
+        for i in range(embeddings.shape[0]):
+            yield embeddings[i, masks[i] == 1]
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        yield from super().embed(documents, batch_size=batch_size, parallel=parallel, **kwargs)
+
+
+class TokensEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs: Any
+    ) -> TokenEmbeddingsModel:
+        return TokenEmbeddingsModel(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/colpali.py

@@ -95,11 +95,12 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
         self.mask_token_id = None
         self.pad_token_id = None
@@ -174,7 +175,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
     ) -> dict[str, NumpyArray]:
         onnx_input["input_ids"] = np.array(
             [
-                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()
+                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()  # type: ignore[index]
                 for input_ids in onnx_input["input_ids"]
             ]
         )
@@ -235,6 +236,8 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             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,
             **kwargs,
         )
 
@@ -268,6 +271,8 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             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,
             **kwargs,
         )
 

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -83,6 +83,40 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
             "Please check the supported models using `LateInteractionMultimodalEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed_text(
         self,
         documents: Union[str, Iterable[str]],

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -19,6 +19,7 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed_text(
         self,
@@ -65,3 +66,13 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
             List of embeddings, one per image
         """
         raise NotImplementedError()
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -120,6 +120,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -146,6 +148,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -183,6 +187,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -209,6 +215,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -131,11 +131,12 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
 
         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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -189,6 +190,8 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             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,
             **kwargs,
         )
 
@@ -196,7 +199,9 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     def _get_worker_class(cls) -> Type[TextRerankerWorker]:
         return TextCrossEncoderWorker
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
         return (float(elem) for elem in output.model_output)
 
 

{fastembed_gpu-0.6.1 → fastembed_gpu-0.7.1}/fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -94,6 +94,8 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -120,6 +122,8 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -133,7 +137,18 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             for batch in pool.ordered_map(iter_batch(pairs, batch_size), **params):
                 yield from self._post_process_onnx_output(batch)  # type: ignore
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[float]: Post-processed output as an iterable of float values.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def _preprocess_onnx_input(