improve: Added pool of shared memory and reuse mechanisim

.github/workflows/python-tests.yml

@@ -42,7 +42,5 @@ jobs:
           poetry install --no-interaction --no-ansi --without dev,docs
 
       - name: Run pytest
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          poetry run pytest
+          poetry run pytest 

.github/workflows/type-checkers.yml

@@ -26,6 +26,8 @@ jobs:
           python -m pip install --upgrade pip poetry
           poetry install --no-interaction --no-ansi --without dev,docs,test
 
+          poetry run pip install "numpy<2.0.0"  # https://github.com/python/mypy/issues/17396
+
       - name: mypy
         run: |
             poetry run mypy fastembed \

fastembed/common/onnx_model.py

@@ -28,16 +28,7 @@ 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, **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.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:

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/common/types.py

@@ -16,7 +16,6 @@ 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/image/image_embedding.py

@@ -77,40 +77,6 @@ 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/image/image_embedding_base.py

@@ -18,7 +18,6 @@ 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,
@@ -43,13 +42,3 @@ 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/image/onnx_embedding.py

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

fastembed/image/onnx_image_model.py

@@ -23,16 +23,7 @@ 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, **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.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
@@ -97,8 +88,6 @@ 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
@@ -115,7 +104,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), **kwargs)
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -125,8 +114,6 @@ 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,
             }
 
@@ -138,7 +125,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, **kwargs)  # type: ignore
+                yield from self._post_process_onnx_output(batch)  # type: ignore
 
 
 class ImageEmbeddingWorker(EmbeddingWorker[T]):

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_list = mean if isinstance(mean, list) else [mean] * num_channels
+    mean = mean if isinstance(mean, list) else [mean] * num_channels
 
-    if len(mean_list) != num_channels:
+    if len(mean) != 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_list)}"
+            f"{len(mean)}"
         )
 
-    mean_arr = np.array(mean_list, dtype=np.float32)
+    mean_arr = np.array(mean, dtype=np.float32)
 
-    std_list = std if isinstance(std, list) else [std] * num_channels
-    if len(std_list) != num_channels:
+    std = std if isinstance(std, list) else [std] * num_channels
+    if len(std) != num_channels:
         raise ValueError(
-            f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std_list)}"
+            f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std)}"
         )
 
-    std_arr = np.array(std_list, dtype=np.float32)
+    std_arr = np.array(std, dtype=np.float32)
 
-    image_upd = ((image.T - mean_arr) / std_arr).T
-    return image_upd
+    image = ((image.T - mean_arr) / std_arr).T
+    return image
 
 
 def resize(

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, **kwargs: Any
+        self, output: OnnxOutputContext, is_doc: bool = True
     ) -> Iterable[NumpyArray]:
         if not is_doc:
             return output.model_output
@@ -169,12 +169,11 @@ 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=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
         self.mask_token_id: Optional[int] = None
         self.pad_token_id: Optional[int] = None
@@ -234,8 +233,6 @@ 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/late_interaction/late_interaction_embedding_base.py

@@ -17,7 +17,6 @@ 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,
@@ -59,13 +58,3 @@ 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/late_interaction/late_interaction_text_embedding.py

@@ -80,40 +80,6 @@ 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/late_interaction/token_embeddings.py

@@ -1,83 +0,0 @@
-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/late_interaction_multimodal/colpali.py

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

@@ -83,40 +83,6 @@ 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/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -19,7 +19,6 @@ 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,
@@ -66,13 +65,3 @@ 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/late_interaction_multimodal/onnx_multimodal_model.py

@@ -120,8 +120,6 @@ 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
@@ -148,8 +146,6 @@ 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,
             }
 
@@ -187,8 +183,6 @@ 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
@@ -215,8 +209,6 @@ 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/parallel_processor.py

@@ -3,6 +3,8 @@ import os
 from collections import defaultdict
 from copy import deepcopy
 from enum import Enum
+from dataclasses import dataclass
+from multiprocessing import shared_memory, Manager, Lock
 from multiprocessing import Queue, get_context
 from multiprocessing.context import BaseContext
 from multiprocessing.process import BaseProcess
@@ -10,6 +12,11 @@ from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
 from typing import Any, Iterable, Optional, Type
 
+import numpy as np
+from numpy.typing import NDArray
+
+from fastembed.common.types import NumpyArray
+
 
 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -17,6 +24,13 @@ processing_timeout = 10 * 60  # seconds
 max_internal_batch_size = 200
 
 
+@dataclass
+class OnnxOutputContext:
+    model_output: NumpyArray
+    attention_mask: Optional[NDArray[np.int64]] = None
+    input_ids: Optional[NDArray[np.int64]] = None
+
+
 class QueueSignals(str, Enum):
     stop = "stop"
     confirm = "confirm"
@@ -32,12 +46,54 @@ class Worker:
         raise NotImplementedError()
 
 
+class SharedMemoryPool:
+    def __init__(self, lock: Lock):
+        self._lock = lock
+        self._pool: dict[str, tuple[shared_memory.SharedMemory, int, np.dtype]] = {}
+        self._free_buffers: list[str] = []
+
+    def allocate(self, size: int, dtype: np.dtype) -> tuple[shared_memory.SharedMemory, str]:
+        best_match = None
+        best_size = float("inf")
+        for buf_name in self._free_buffers:
+            shm, buf_size, buf_dtype = self._pool[buf_name]
+            # get best match for needed size
+            if buf_size >= size and buf_dtype == dtype and buf_size < best_size:
+                best_match = buf_name
+                best_size = buf_size
+            if best_match:
+                self._free_buffers.remove(best_match)
+                return self._pool[best_match][0], best_match
+            shm = shared_memory.SharedMemory(create=True, size=size)
+            self._pool[shm.name] = (shm, size, dtype)
+            return shm, shm.name
+        # if no match found, create new buffer
+        shm = shared_memory.SharedMemory(create=True, size=size)
+        self._pool[shm.name] = (shm, size, dtype)
+        return shm, shm.name
+
+    def release(self, name: str) -> None:
+        if name in self._pool and name not in self._free_buffers:
+            self._free_buffers.append(name)
+
+    def cleanup(self) -> None:
+        for shm, _, _ in self._pool.values():
+            shm.close()
+            shm.unlink()
+        self._pool.clear()
+        self._free_buffers.clear()
+
+    def __del__(self):
+        self.cleanup()
+
+
 def _worker(
     worker_class: Type[Worker],
     input_queue: Queue,
     output_queue: Queue,
     num_active_workers: BaseValue,
     worker_id: int,
+    shared_pool: SharedMemoryPool,
     kwargs: Optional[dict[str, Any]] = None,
 ) -> None:
     """
@@ -55,7 +111,6 @@ def _worker(
     try:
         worker = worker_class.start(**kwargs)
 
-        # Keep going until you get an item that's None.
         def input_queue_iterable() -> Iterable[Any]:
             while True:
                 item = input_queue.get()
@@ -64,7 +119,24 @@ def _worker(
                 yield item
 
         for processed_item in worker.process(input_queue_iterable()):
-            output_queue.put(processed_item)
+            idx, output_context = processed_item
+            output_metadata = {}
+            for field in ["model_output", "attention_mask", "input_ids"]:
+                array = getattr(output_context, field, None)
+                if array is not None:
+                    shm, shm_name = shared_pool.allocate(array.nbytes, array.dtype)
+                    shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf)
+                    np.copyto(shm_array, array)
+                    output_metadata[field] = {
+                        "name": shm_name,
+                        "shape": array.shape,
+                        "dtype": array.dtype.str,
+                    }
+                    shm.close()
+            output_queue.put((idx, output_metadata))
+            for field in output_metadata:  # mark release to reuse
+                shared_pool.release(output_metadata[field]["name"])
+
     except Exception as e:  # pylint: disable=broad-except
         logging.exception(e)
         output_queue.put(QueueSignals.error)
@@ -108,6 +180,8 @@ class ParallelWorkerPool:
         self.device_ids = device_ids
         self.cuda = cuda
         self.num_active_workers: Optional[BaseValue] = None
+        self.manager = Manager()
+        self.shared_pool = SharedMemoryPool(self.manager.Lock())
 
     def start(self, **kwargs: Any) -> None:
         self.input_queue = self.ctx.Queue(self.queue_size)
@@ -133,6 +207,7 @@ class ParallelWorkerPool:
                     self.output_queue,
                     self.num_active_workers,
                     worker_id,
+                    self.shared_pool,
                     worker_kwargs,
                 ),
             )
@@ -178,7 +253,17 @@ class ParallelWorkerPool:
                     if out_item == QueueSignals.error:
                         self.join_or_terminate()
                         raise RuntimeError("Thread unexpectedly terminated")
-                    yield out_item
+
+                    idx, output_metadata = out_item
+                    output_arrays = {}
+                    for field, meta in output_metadata.items():
+                        shm = shared_memory.SharedMemory(name=meta["name"])
+                        array = np.ndarray(
+                            meta["shape"], dtype=meta["dtype"], buffer=shm.buf
+                        ).copy()
+                        output_arrays[field] = array
+                        shm.close()
+                    yield (idx, OnnxOutputContext(**output_arrays))
                     read += 1
 
                 self.input_queue.put((idx, item))
@@ -193,9 +278,18 @@ class ParallelWorkerPool:
                 if out_item == QueueSignals.error:
                     self.join_or_terminate()
                     raise RuntimeError("Thread unexpectedly terminated")
-                yield out_item
+
+                idx, output_metadata = out_item
+                output_arrays = {}
+                for field, meta in output_metadata.items():
+                    shm = shared_memory.SharedMemory(name=meta["name"])
+                    array = np.ndarray(meta["shape"], dtype=meta["dtype"], buffer=shm.buf).copy()
+                    output_arrays[field] = array
+                    shm.close()
+                yield (idx, OnnxOutputContext(**output_arrays))
                 read += 1
         finally:
+            self.shared_pool.cleanup()
             assert self.input_queue is not None, "Input queue is None"
             assert self.output_queue is not None, "Output queue is None"
             self.join()
@@ -231,11 +325,13 @@ class ParallelWorkerPool:
             if process.is_alive():
                 process.terminate()
         self.processes.clear()
+        self.shared_pool.cleanup()
 
     def join(self) -> None:
         for process in self.processes:
             process.join()
         self.processes.clear()
+        self.shared_pool.cleanup()
 
     def __del__(self) -> None:
         """
@@ -250,3 +346,4 @@ class ParallelWorkerPool:
         for process in self.processes:
             if process.is_alive():
                 process.terminate()
+        self.shared_pool.cleanup()

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

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

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -94,8 +94,6 @@ 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
@@ -122,8 +120,6 @@ 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,
             }
 
@@ -137,18 +133,7 @@ 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, **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.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _preprocess_onnx_input(

fastembed/sparse/bm25.py

@@ -21,9 +21,13 @@ from fastembed.sparse.sparse_embedding_base import (
 from fastembed.sparse.utils.tokenizer import SimpleTokenizer
 from fastembed.common.model_description import SparseModelDescription, ModelSource
 
-
 supported_languages = [
     "arabic",
+    "azerbaijani",
+    "basque",
+    "bengali",
+    "catalan",
+    "chinese",
     "danish",
     "dutch",
     "english",
@@ -31,15 +35,21 @@ supported_languages = [
     "french",
     "german",
     "greek",
+    "hebrew",
+    "hinglish",
     "hungarian",
+    "indonesian",
     "italian",
+    "kazakh",
+    "nepali",
     "norwegian",
     "portuguese",
     "romanian",
     "russian",
+    "slovene",
     "spanish",
     "swedish",
-    "tamil",
+    "tajik",
     "turkish",
 ]
 
@@ -115,12 +125,11 @@ class Bm25(SparseTextEmbeddingBase):
         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(
             model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         self.token_max_length = token_max_length
@@ -161,8 +170,6 @@ class Bm25(SparseTextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
     ) -> Iterable[SparseEmbedding]:
         is_small = False
 
@@ -191,8 +198,6 @@ class Bm25(SparseTextEmbeddingBase):
                 "language": self.language,
                 "token_max_length": self.token_max_length,
                 "disable_stemmer": self.disable_stemmer,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
             }
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
@@ -231,8 +236,6 @@ class Bm25(SparseTextEmbeddingBase):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
         )
 
     def _stem(self, tokens: list[str]) -> list[str]:

fastembed/sparse/bm42.py

@@ -110,12 +110,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         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,
+            specific_model_path=specific_model_path,
         )
 
         self.invert_vocab: dict[int, str] = {}
@@ -218,9 +217,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
         return new_vector
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
         if output.input_ids is None:
             raise ValueError("input_ids must be provided for document post-processing")
 
@@ -302,8 +299,6 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             cuda=self.cuda,
             device_ids=self.device_ids,
             alpha=self.alpha,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
         )
 
     @classmethod

fastembed/sparse/minicoil.py

@@ -1,356 +0,0 @@
-from pathlib import Path
-
-from typing import Any, Optional, Sequence, Iterable, Union, Type
-
-import numpy as np
-from numpy.typing import NDArray
-from py_rust_stemmers import SnowballStemmer
-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.utils import define_cache_dir
-from fastembed.sparse.sparse_embedding_base import (
-    SparseEmbedding,
-    SparseTextEmbeddingBase,
-)
-from fastembed.sparse.utils.minicoil_encoder import Encoder
-from fastembed.sparse.utils.sparse_vectors_converter import SparseVectorConverter, WordEmbedding
-from fastembed.sparse.utils.vocab_resolver import VocabResolver, VocabTokenizer
-from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-
-
-MINICOIL_MODEL_FILE = "minicoil.triplet.model.npy"
-MINICOIL_VOCAB_FILE = "minicoil.triplet.model.vocab"
-STOPWORDS_FILE = "stopwords.txt"
-
-
-supported_minicoil_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="Qdrant/minicoil-v1",
-        vocab_size=19125,
-        description="Sparse embedding model, that resolves semantic meaning of the words, "
-        "while keeping exact keyword match behavior. "
-        "Based on jinaai/jina-embeddings-v2-small-en-tokens",
-        license="apache-2.0",
-        size_in_GB=0.09,
-        sources=ModelSource(hf="Qdrant/minicoil-v1"),
-        model_file="onnx/model.onnx",
-        additional_files=[
-            STOPWORDS_FILE,
-            MINICOIL_MODEL_FILE,
-            MINICOIL_VOCAB_FILE,
-        ],
-        requires_idf=True,
-    ),
-]
-
-MODEL_TO_LANGUAGE = {
-    "Qdrant/minicoil-v1": "english",
-}
-
-
-class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
-    """
-        MiniCOIL is a sparse embedding model, that resolves semantic meaning of the words,
-        while keeping exact keyword match behavior.
-
-        Each vocabulary token is converted into 4d component of a sparse vector, which is then weighted by the token frequency in the corpus.
-        If the token is not found in the corpus, it is treated exactly like in BM25.
-    `
-        The model is based on `jinaai/jina-embeddings-v2-small-en-tokens`
-    """
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        k: float = 1.2,
-        b: float = 0.75,
-        avg_len: float = 150.0,
-        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 providers to use for onnxruntime.
-            k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.
-                I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
-            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.
-            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.
-            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
-
-        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.device_ids = device_ids
-        self.cuda = cuda
-        self.device_id = device_id
-        self.k = k
-        self.b = b
-        self.avg_len = avg_len
-
-        # Initialize class attributes
-        self.tokenizer: Optional[Tokenizer] = 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.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,
-        )
-
-        if not self.lazy_load:
-            self.load_onnx_model()
-
-    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,
-        )
-
-        assert self.tokenizer is not None
-
-        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
-            self.invert_vocab[idx] = token
-        self.special_tokens = set(self.special_token_to_id.keys())
-        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])
-
-        self.vocab_resolver = VocabResolver(
-            tokenizer=VocabTokenizer(self.tokenizer),
-            stopwords=self.stopwords,
-            stemmer=stemmer,
-        )
-        self.vocab_resolver.load_json_vocab(str(self._model_dir / MINICOIL_VOCAB_FILE))
-
-        weights = np.load(str(self._model_dir / MINICOIL_MODEL_FILE), mmap_mode="r")
-        self.encoder = Encoder(weights)
-        self.output_dim = self.encoder.output_dim
-
-        self.sparse_vector_converter = SparseVectorConverter(
-            stopwords=self.stopwords,
-            stemmer=stemmer,
-            k=self.k,
-            b=self.b,
-            avg_len=self.avg_len,
-        )
-
-    def embed(
-        self,
-        documents: Union[str, Iterable[str]],
-        batch_size: int = 256,
-        parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[SparseEmbedding]:
-        """
-        Encode a list of documents into list of embeddings.
-        We use mean pooling with attention so that the model can handle variable-length inputs.
-
-        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,
-            k=self.k,
-            b=self.b,
-            avg_len=self.avg_len,
-            is_query=False,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            **kwargs,
-        )
-
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
-        """
-        Encode a list of queries into list of embeddings.
-        """
-        yield from self._embed_documents(
-            model_name=self.model_name,
-            cache_dir=str(self.cache_dir),
-            documents=query,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_ids=self.device_ids,
-            k=self.k,
-            b=self.b,
-            avg_len=self.avg_len,
-            is_query=True,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            **kwargs,
-        )
-
-    @classmethod
-    def _load_stopwords(cls, model_dir: Path) -> list[str]:
-        stopwords_path = model_dir / STOPWORDS_FILE
-        if not stopwords_path.exists():
-            return []
-
-        with open(stopwords_path, "r") as f:
-            return f.read().splitlines()
-
-    @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
-        """Lists the supported models.
-
-        Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
-        """
-        return supported_minicoil_models
-
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, is_query: bool = False, **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
-        if output.input_ids is None:
-            raise ValueError("input_ids must be provided for document post-processing")
-
-        assert self.vocab_resolver is not None
-        assert self.encoder is not None
-        assert self.sparse_vector_converter is not None
-
-        # 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
-
-        vocab_size = self.vocab_resolver.vocab_size()
-        embedding_size = self.encoder.output_dim
-
-        # For each document we only select those embeddings that are not masked out
-
-        for i in range(embeddings.shape[0]):
-            # Size: (sequence_length, hidden_size)
-            token_embeddings = embeddings[i, masks[i] == 1]
-
-            # Size: (sequence_length)
-            token_ids: NDArray[np.int64] = output.input_ids[i, masks[i] == 1]
-
-            word_ids_array, counts, oov, forms = self.vocab_resolver.resolve_tokens(token_ids)
-
-            # Size: (1, words)
-            word_ids_array_expanded: NDArray[np.int64] = np.expand_dims(word_ids_array, axis=0)
-
-            # Size: (1, words, embedding_size)
-            token_embeddings_array: NDArray[np.float32] = np.expand_dims(token_embeddings, axis=0)
-
-            assert word_ids_array_expanded.shape[1] == token_embeddings_array.shape[1]
-
-            # Size of word_ids_mapping: (unique_words, 2) - [vocab_id, batch_id]
-            # Size of embeddings: (unique_words, embedding_size)
-            ids_mapping, minicoil_embeddings = self.encoder.forward(
-                word_ids_array_expanded, token_embeddings_array
-            )
-
-            # Size of counts: (unique_words)
-            words_ids: list[int] = ids_mapping[:, 0].tolist()  # type: ignore[assignment]
-
-            sentence_result: dict[str, WordEmbedding] = {}
-
-            words = [self.vocab_resolver.lookup_word(word_id) for word_id in words_ids]
-
-            for word, word_id, emb in zip(words, words_ids, minicoil_embeddings.tolist()):  # type: ignore[arg-type]
-                if word_id == 0:
-                    continue
-
-                sentence_result[word] = WordEmbedding(
-                    word=word,
-                    forms=forms[word],
-                    count=int(counts[word_id]),
-                    word_id=int(word_id),
-                    embedding=emb,  # type: ignore[arg-type]
-                )
-
-            for oov_word, count in oov.items():
-                # {
-                #     "word": oov_word,
-                #     "forms": [oov_word],
-                #     "count": int(count),
-                #     "word_id": -1,
-                #     "embedding": [1]
-                # }
-                sentence_result[oov_word] = WordEmbedding(
-                    word=oov_word, forms=[oov_word], count=int(count), word_id=-1, embedding=[1]
-                )
-
-            if not is_query:
-                yield self.sparse_vector_converter.embedding_to_vector(
-                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
-                )
-            else:
-                yield self.sparse_vector_converter.embedding_to_vector_query(
-                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
-                )
-
-    @classmethod
-    def _get_worker_class(cls) -> Type["MiniCoilTextEmbeddingWorker"]:
-        return MiniCoilTextEmbeddingWorker
-
-
-class MiniCoilTextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> MiniCOIL:
-        return MiniCOIL(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=1,
-            **kwargs,
-        )

fastembed/sparse/sparse_text_embedding.py

@@ -4,7 +4,6 @@ from dataclasses import asdict
 from fastembed.common import OnnxProvider
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.bm42 import Bm42
-from fastembed.sparse.minicoil import MiniCOIL
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
@@ -15,7 +14,7 @@ from fastembed.common.model_description import SparseModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25, MiniCOIL]
+    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -62,7 +61,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "prithvida/Splade_PP_en_v1".lower():
+        if model_name == "prithvida/Splade_PP_en_v1":
             warnings.warn(
                 "The right spelling is prithivida/Splade_PP_en_v1. "
                 "Support of this name will be removed soon, please fix the model_name",

fastembed/sparse/splade_pp.py

@@ -18,7 +18,7 @@ supported_splade_models: list[SparseModelDescription] = [
         description="Independent Implementation of SPLADE++ Model for English.",
         license="apache-2.0",
         size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
+        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
         model_file="model.onnx",
     ),
     SparseModelDescription(
@@ -27,16 +27,14 @@ supported_splade_models: list[SparseModelDescription] = [
         description="Independent Implementation of SPLADE++ Model for English.",
         license="apache-2.0",
         size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
+        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
         model_file="model.onnx",
     ),
 ]
 
 
 class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 
@@ -114,12 +112,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         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,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -166,8 +163,6 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             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/sparse/utils/minicoil_encoder.py

@@ -1,146 +0,0 @@
-"""
-Pure numpy implementation of encoder model for a single word.
-
-This model is not trainable, and should only be used for inference.
-"""
-
-import numpy as np
-from fastembed.common.types import NumpyArray
-
-
-class Encoder:
-    """
-    Encoder(768, 4, 10000)
-
-    Will look like this:
-
-
-                                         Per-word
-                                         Encoder Matrix
-     ┌─────────────────────┐
-     │ Token Embedding(768)├──────┐      (10k, 768, 4)
-     └─────────────────────┘      │         ┌─────────┐
-                                  │         │         │
-     ┌─────────────────────┐      │       ┌─┴───────┐ │
-     │                     │      │       │         │ │
-     └─────────────────────┘      │     ┌─┴───────┐ │ │      ┌─────────┐
-                                  └────►│         │ │ ├─────►│Tanh     │
-     ┌─────────────────────┐            │         │ │ │      └─────────┘
-     │                     │            │         │ ├─┘
-     └─────────────────────┘            │         ├─┘
-                                        │         │
-     ┌─────────────────────┐            └─────────┘
-     │                     │
-     └─────────────────────┘
-
-     Final linear transformation is accompanied by a non-linear activation function: Tanh.
-
-     Tanh is used to ensure that the output is in the range [-1, 1].
-     It would be easier to visually interpret the output of the model, assuming that each dimension
-     would need to encode a type of semantic cluster.
-    """
-
-    def __init__(
-        self,
-        weights: NumpyArray,
-    ):
-        self.weights = weights
-        self.vocab_size, self.input_dim, self.output_dim = weights.shape
-
-        self.encoder_weights: NumpyArray = weights
-
-        # Activation function
-        self.activation = np.tanh
-
-    @staticmethod
-    def convert_vocab_ids(vocab_ids: NumpyArray) -> NumpyArray:
-        """
-        Convert vocab_ids of shape (batch_size, seq_len) into (batch_size, seq_len, 2)
-        by appending batch_id alongside each vocab_id.
-        """
-        batch_size, seq_len = vocab_ids.shape
-        batch_ids = np.arange(batch_size, dtype=vocab_ids.dtype).reshape(batch_size, 1)
-        batch_ids = np.repeat(batch_ids, seq_len, axis=1)
-        # Stack vocab_ids and batch_ids along the last dimension
-        combined: NumpyArray = np.stack((vocab_ids, batch_ids), axis=2).astype(np.int32)
-        return combined
-
-    @classmethod
-    def avg_by_vocab_ids(
-        cls, vocab_ids: NumpyArray, embeddings: NumpyArray
-    ) -> tuple[NumpyArray, NumpyArray]:
-        """
-        Takes:
-            vocab_ids: (batch_size, seq_len) int array
-            embeddings: (batch_size, seq_len, input_dim) float array
-
-        Returns:
-            unique_flattened_vocab_ids: (total_unique, 2) array of [vocab_id, batch_id]
-            unique_flattened_embeddings: (total_unique, input_dim) averaged embeddings
-        """
-        input_dim = embeddings.shape[2]
-
-        # Flatten vocab_ids and embeddings
-        # flattened_vocab_ids: (batch_size*seq_len, 2)
-        flattened_vocab_ids = cls.convert_vocab_ids(vocab_ids).reshape(-1, 2)
-
-        # flattened_embeddings: (batch_size*seq_len, input_dim)
-        flattened_embeddings = embeddings.reshape(-1, input_dim)
-
-        # Find unique (vocab_id, batch_id) pairs
-        unique_flattened_vocab_ids, inverse_indices = np.unique(
-            flattened_vocab_ids, axis=0, return_inverse=True
-        )
-
-        # Prepare arrays to accumulate sums
-        unique_count = unique_flattened_vocab_ids.shape[0]
-        unique_flattened_embeddings = np.zeros((unique_count, input_dim), dtype=np.float32)
-        unique_flattened_count = np.zeros(unique_count, dtype=np.int32)
-
-        # Use np.add.at to accumulate sums based on inverse indices
-        np.add.at(unique_flattened_embeddings, inverse_indices, flattened_embeddings)
-        np.add.at(unique_flattened_count, inverse_indices, 1)
-
-        # Compute averages
-        unique_flattened_embeddings /= unique_flattened_count[:, None]
-
-        return unique_flattened_vocab_ids.astype(np.int32), unique_flattened_embeddings.astype(
-            np.float32
-        )
-
-    def forward(
-        self, vocab_ids: NumpyArray, embeddings: NumpyArray
-    ) -> tuple[NumpyArray, NumpyArray]:
-        """
-        Args:
-            vocab_ids: (batch_size, seq_len) int array
-            embeddings: (batch_size, seq_len, input_dim) float array
-
-        Returns:
-            unique_flattened_vocab_ids_and_batch_ids: (total_unique, 2)
-            unique_flattened_encoded: (total_unique, output_dim)
-        """
-        # Average embeddings for duplicate vocab_ids
-        unique_flattened_vocab_ids_and_batch_ids, unique_flattened_embeddings = (
-            self.avg_by_vocab_ids(vocab_ids, embeddings)
-        )
-
-        # Select the encoder weights for each unique vocab_id
-        unique_flattened_vocab_ids = unique_flattened_vocab_ids_and_batch_ids[:, 0].astype(
-            np.int32
-        )
-
-        # unique_encoder_weights: (total_unique, input_dim, output_dim)
-        unique_encoder_weights = self.encoder_weights[unique_flattened_vocab_ids]
-
-        # Compute linear transform: (total_unique, output_dim)
-        # Using Einstein summation for matrix multiplication:
-        # 'bi,bio->bo' means: for each "b" (batch element), multiply embeddings (b,i) by weights (b,i,o) -> (b,o)
-        unique_flattened_encoded = np.einsum(
-            "bi,bio->bo", unique_flattened_embeddings, unique_encoder_weights
-        )
-
-        # Apply Tanh activation and ensure float32 type
-        unique_flattened_encoded = self.activation(unique_flattened_encoded).astype(np.float32)
-
-        return unique_flattened_vocab_ids_and_batch_ids.astype(np.int32), unique_flattened_encoded

fastembed/sparse/utils/sparse_vectors_converter.py

@@ -1,247 +0,0 @@
-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 numpy as np
-
-from fastembed.sparse.sparse_embedding_base import SparseEmbedding
-
-GAP = 32000
-INT32_MAX = 2**31 - 1
-
-
-@dataclass
-class WordEmbedding:
-    word: str
-    forms: List[str]
-    count: int
-    word_id: int
-    embedding: List[float]
-
-
-class SparseVectorConverter:
-    def __init__(
-        self,
-        stopwords: Set[str],
-        stemmer: SnowballStemmer,
-        k: float = 1.2,
-        b: float = 0.75,
-        avg_len: float = 150.0,
-    ):
-        punctuation = set(get_all_punctuation())
-        special_tokens = {"[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"}
-
-        self.stemmer = stemmer
-        self.unwanted_tokens = punctuation | special_tokens | stopwords
-
-        self.k = k
-        self.b = b
-        self.avg_len = avg_len
-
-    @classmethod
-    def unkn_word_token_id(
-        cls, word: str, shift: int
-    ) -> int:  # 2-3 words can collide in 1 index with this mapping, not considering mm3 collisions
-        token_hash = abs(mmh3.hash(word))
-
-        range_size = INT32_MAX - shift
-        remapped_hash = shift + (token_hash % range_size)
-
-        return remapped_hash
-
-    def bm25_tf(self, num_occurrences: int, sentence_len: int) -> float:
-        res = num_occurrences * (self.k + 1)
-        res /= num_occurrences + self.k * (1 - self.b + self.b * sentence_len / self.avg_len)
-        return res
-
-    @classmethod
-    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]:
-        """
-        Clean miniCOIL-produced sentence_embedding, as unknown to the miniCOIL's stemmer tokens should fully resemble
-        our BM25 token representation.
-
-        sentence_embedding = {"9°": {"word": "9°", "word_id": -1, "count": 2, "embedding": [1], "forms": ["9°"]},
-                "9": {"word": "9", "word_id": -1, "count": 2, "embedding": [1], "forms": ["9"]},
-                "bat": {"word": "bat", "word_id": 2, "count": 3, "embedding": [0.2, 0.1, -0.2, -0.2], "forms": ["bats", "bat"]},
-                "9°9": {"word": "9°9", "word_id": -1, "count": 1, "embedding": [1], "forms": ["9°9"]},
-                "screech": {"word": "screech", "word_id": -1, "count": 1, "embedding": [1], "forms": ["screech"]},
-                "screeched": {"word": "screeched", "word_id": -1, "count": 1, "embedding": [1], "forms": ["screeched"]}
-                }
-        cleaned_embedding_ground_truth = {
-                "9": {"word": "9", "word_id": -1, "count": 6, "embedding": [1], "forms": ["9°", "9", "9°9", "9°9"]},
-                "bat": {"word": "bat", "word_id": 2, "count": 3, "embedding": [0.2, 0.1, -0.2, -0.2], "forms": ["bats", "bat"]},
-                "screech": {"word": "screech", "word_id": -1, "count": 2, "embedding": [1], "forms": ["screech", "screeched"]}
-                }
-        """
-
-        new_sentence_embedding: Dict[str, WordEmbedding] = {}
-
-        for word, embedding in sentence_embedding.items():
-            # embedding = {
-            #     "word": "vector",
-            #     "forms": ["vector", "vectors"],
-            #     "count": 2,
-            #     "word_id": 1231,
-            #     "embedding": [0.1, 0.2, 0.3, 0.4]
-            # }
-            if embedding.word_id > 0:
-                # Known word, no need to clean
-                new_sentence_embedding[word] = embedding
-            else:
-                # Unknown word
-                if word in self.unwanted_tokens:
-                    continue
-
-                # Example complex word split:
-                # word = `word^vec`
-                word_cleaned = remove_non_alphanumeric(word).strip()
-                # word_cleaned = `word vec`
-
-                if len(word_cleaned) > 0:
-                    # Subwords: ['word', 'vec']
-                    for subword in word_cleaned.split():
-                        stemmed_subword: str = self.stemmer.stem_word(subword)
-                        if (
-                            len(stemmed_subword) <= token_max_length
-                            and stemmed_subword not in self.unwanted_tokens
-                        ):
-                            if stemmed_subword not in new_sentence_embedding:
-                                new_sentence_embedding[stemmed_subword] = copy.deepcopy(embedding)
-                                new_sentence_embedding[stemmed_subword].word = stemmed_subword
-                            else:
-                                new_sentence_embedding[stemmed_subword].count += embedding.count
-                                new_sentence_embedding[stemmed_subword].forms += embedding.forms
-
-        return new_sentence_embedding
-
-    def embedding_to_vector(
-        self,
-        sentence_embedding: Dict[str, WordEmbedding],
-        embedding_size: int,
-        vocab_size: int,
-    ) -> SparseEmbedding:
-        """
-        Convert miniCOIL sentence embedding to Qdrant sparse vector
-
-        Example input:
-
-        ```
-        {
-            "vector": WordEmbedding({ // Vocabulary word, encoded with miniCOIL normally
-                "word": "vector",
-                "forms": ["vector", "vectors"],
-                "count": 2,
-                "word_id": 1231,
-                "embedding": [0.1, 0.2, 0.3, 0.4]
-            }),
-            "axiotic": WordEmbedding({ // Out-of-vocabulary word, fallback to BM25
-                "word": "axiotic",
-                "forms": ["axiotics"],
-                "count": 1,
-                "word_id": -1,
-            })
-        }
-        ```
-
-        """
-
-        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.
-        # Therefore we need (vocab_size * embedding_size) // GAP + 1 buckets for vocab words.
-        # 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
-        sentence_embedding_cleaned = self.clean_words(sentence_embedding)
-
-        # Calculate sentence length after cleaning
-        sentence_len = 0
-        for embedding in sentence_embedding_cleaned.values():
-            sentence_len += embedding.count
-
-        for embedding in sentence_embedding_cleaned.values():
-            word_id = embedding.word_id
-            num_occurrences = embedding.count
-            tf = self.bm25_tf(num_occurrences, sentence_len)
-            if (
-                word_id > 0
-            ):  # miniCOIL starts with ID 1, we generally won't have word_id == 0 (UNK), as we don't add
-                # these words to sentence_embedding
-                embedding_values = embedding.embedding
-                normalized_embedding = self.normalize_vector(embedding_values)
-
-                for val_id, value in enumerate(normalized_embedding):
-                    indices.append(
-                        word_id * embedding_size + val_id
-                    )  # since miniCOIL IDs start with 1
-                    values.append(value * tf)
-            else:
-                indices.append(self.unkn_word_token_id(embedding.word, unknown_words_shift))
-                values.append(tf)
-
-        return SparseEmbedding(
-            indices=np.array(indices, dtype=np.int32),
-            values=np.array(values, dtype=np.float32),
-        )
-
-    def embedding_to_vector_query(
-        self,
-        sentence_embedding: Dict[str, WordEmbedding],
-        embedding_size: int,
-        vocab_size: int,
-    ) -> SparseEmbedding:
-        """
-        Same as `embedding_to_vector`, but no TF
-        """
-
-        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
-
-        sentence_embedding_cleaned = self.clean_words(sentence_embedding)
-
-        for embedding in sentence_embedding_cleaned.values():
-            word_id = embedding.word_id
-            tf = 1.0
-
-            if word_id >= 0:  # miniCOIL starts with ID 1
-                embedding_values = embedding.embedding
-                normalized_embedding = self.normalize_vector(embedding_values)
-
-                for val_id, value in enumerate(normalized_embedding):
-                    indices.append(
-                        word_id * embedding_size + val_id
-                    )  # since miniCOIL IDs start with 1
-                    values.append(value * tf)
-            else:
-                indices.append(self.unkn_word_token_id(embedding.word, unknown_words_shift))
-                values.append(tf)
-
-        return SparseEmbedding(
-            indices=np.array(indices, dtype=np.int32),
-            values=np.array(values, dtype=np.float32),
-        )

fastembed/sparse/utils/vocab_resolver.py

@@ -1,202 +0,0 @@
-from collections import defaultdict
-from typing import Iterable
-
-from py_rust_stemmers import SnowballStemmer
-import numpy as np
-from tokenizers import Tokenizer
-from numpy.typing import NDArray
-
-from fastembed.common.types import NumpyArray
-
-
-class VocabTokenizerBase:
-    def tokenize(self, sentence: str) -> NumpyArray:
-        raise NotImplementedError()
-
-    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
-        raise NotImplementedError()
-
-
-class VocabTokenizer(VocabTokenizerBase):
-    def __init__(self, tokenizer: Tokenizer):
-        self.tokenizer = tokenizer
-
-    def tokenize(self, sentence: str) -> NumpyArray:
-        return np.array(self.tokenizer.encode(sentence).ids)
-
-    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
-        return [self.tokenizer.id_to_token(token_id) for token_id in token_ids]
-
-
-class VocabResolver:
-    def __init__(self, tokenizer: VocabTokenizerBase, stopwords: set[str], stemmer: SnowballStemmer):
-        # Word to id mapping
-        self.vocab: dict[str, int] = {}
-        # Id to word mapping
-        self.words: list[str] = []
-        # Lemma to word mapping
-        self.stem_mapping: dict[str, str] = {}
-        self.tokenizer: VocabTokenizerBase = tokenizer
-        self.stemmer = stemmer
-        self.stopwords: set[str] = stopwords
-
-    def tokenize(self, sentence: str) -> NumpyArray:
-        return self.tokenizer.tokenize(sentence)
-
-    def lookup_word(self, word_id: int) -> str:
-        if word_id == 0:
-            return "UNK"
-        return self.words[word_id - 1]
-
-    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
-        return self.tokenizer.convert_ids_to_tokens(token_ids)
-
-    def vocab_size(self) -> int:
-        # We need +1 for UNK token
-        return len(self.vocab) + 1
-
-    def save_vocab(self, path: str) -> None:
-        with open(path, "w") as f:
-            for word in self.words:
-                f.write(word + "\n")
-
-    def save_json_vocab(self, path: str) -> None:
-        import json
-
-        with open(path, "w") as f:
-            json.dump({"vocab": self.words, "stem_mapping": self.stem_mapping}, f, indent=2)
-
-    def load_json_vocab(self, path: str) -> None:
-        import json
-
-        with open(path, "r") as f:
-            data = json.load(f)
-            self.words = data["vocab"]
-            self.vocab = {word: idx + 1 for idx, word in enumerate(self.words)}
-            self.stem_mapping = data["stem_mapping"]
-
-    def add_word(self, word: str) -> None:
-        if word not in self.vocab:
-            self.vocab[word] = len(self.vocab) + 1
-            self.words.append(word)
-            stem = self.stemmer.stem_word(word)
-            if stem not in self.stem_mapping:
-                self.stem_mapping[stem] = word
-            else:
-                existing_word = self.stem_mapping[stem]
-                if len(existing_word) > len(word):
-                    # Prefer shorter words for the same stem
-                    # Example: "swim" is preferred over "swimming"
-                    self.stem_mapping[stem] = word
-
-    def load_vocab(self, path: str) -> None:
-        with open(path, "r") as f:
-            for line in f:
-                self.add_word(line.strip())
-
-    @classmethod
-    def _reconstruct_bpe(
-        cls, bpe_tokens: Iterable[tuple[int, str]]
-    ) -> list[tuple[str, list[int]]]:
-        result: list[tuple[str, list[int]]] = []
-        acc: str = ""
-        acc_idx: list[int] = []
-
-        continuing_subword_prefix = "##"
-        continuing_subword_prefix_len = len(continuing_subword_prefix)
-
-        for idx, token in bpe_tokens:
-            if token.startswith(continuing_subword_prefix):
-                acc += token[continuing_subword_prefix_len:]
-                acc_idx.append(idx)
-            else:
-                if acc:
-                    result.append((acc, acc_idx))
-                    acc_idx = []
-                acc = token
-                acc_idx.append(idx)
-
-        if acc:
-            result.append((acc, acc_idx))
-        return result
-
-    def resolve_tokens(
-        self, token_ids: NDArray[np.int64]
-    ) -> tuple[NDArray[np.int64], dict[int, int], dict[str, int], dict[str, list[str]]]:
-        """
-        Mark known tokens (including composed tokens) with vocab ids.
-
-        Args:
-            token_ids: (seq_len) - list of ids of tokens
-                Example:
-                    [
-                        101,  3897, 19332, 12718, 23348,
-                        1010,  1996,  7151,  2296, 4845,
-                        2359,  2005,  4234,  1010,  4332,
-                        2871,  3191,  2062, 102
-                    ]
-
-            returns:
-                - token_ids with vocab ids
-                    [
-                        0,  151, 151, 0, 0,
-                        912,  0,  0,  0, 332,
-                        332,  332,  0,  7121,  191,
-                        0,  0,  332, 0
-                    ]
-                - counts of each token
-                    {
-                        151: 1,
-                        332: 3,
-                        7121: 1,
-                        191: 1,
-                        912: 1
-                    }
-                - oov counts of each token
-                    {
-                        "the": 1,
-                        "a": 1,
-                        "[CLS]": 1,
-                        "[SEP]": 1,
-                        ...
-                    }
-                - forms of each token
-                    {
-                        "hello": ["hello"],
-                        "world": ["worlds", "world", "worlding"],
-                    }
-
-        """
-        tokens = self.convert_ids_to_tokens(token_ids)
-        tokens_mapping = self._reconstruct_bpe(enumerate(tokens))
-
-        counts: dict[int, int] = defaultdict(int)
-        oov_count: dict[str, int] = defaultdict(int)
-
-        forms: dict[str, list[str]] = defaultdict(list)
-
-        for token, mapped_token_ids in tokens_mapping:
-            vocab_id = 0
-            if token in self.stopwords:
-                vocab_id = 0
-            elif token in self.vocab:
-                vocab_id = self.vocab[token]
-                forms[token].append(token)
-            elif token in self.stem_mapping:
-                vocab_id = self.vocab[self.stem_mapping[token]]
-                forms[self.stem_mapping[token]].append(token)
-            else:
-                stem = self.stemmer.stem_word(token)
-                if stem in self.stem_mapping:
-                    vocab_id = self.vocab[self.stem_mapping[stem]]
-                    forms[self.stem_mapping[stem]].append(token)
-
-            for token_id in mapped_token_ids:
-                token_ids[token_id] = vocab_id
-
-            if vocab_id == 0:
-                oov_count[token] += 1
-            else:
-                counts[vocab_id] += 1
-        return token_ids, counts, oov_count, forms
-

fastembed/text/clip_embedding.py

@@ -35,9 +35,7 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
         """
         return supported_clip_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         return output.model_output
 
 

fastembed/text/custom_text_embedding.py

@@ -58,9 +58,7 @@ class CustomTextEmbedding(OnnxTextEmbedding):
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         return cls.SUPPORTED_MODELS
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         return self._normalize(self._pool(output.model_output, output.attention_mask))
 
     def _pool(
@@ -77,11 +75,6 @@ class CustomTextEmbedding(OnnxTextEmbedding):
         if self._pooling == PoolingType.DISABLED:
             return embeddings
 
-        raise ValueError(
-            f"Unsupported pooling type {self._pooling}. "
-            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, {PoolingType.DISABLED}."
-        )
-
     def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
         return normalize(embeddings) if self._normalization else embeddings
 

fastembed/text/multitask_embedding.py

@@ -3,7 +3,6 @@ from typing import Any, Type, Iterable, Union, Optional
 
 import numpy as np
 
-from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.types import NumpyArray
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
@@ -45,11 +44,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, **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.current_task_id: Union[Task, int] = self.PASSAGE_TASK
 
     @classmethod
     def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
@@ -60,14 +57,9 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         return supported_multitask_models
 
     def _preprocess_onnx_input(
-        self,
-        onnx_input: dict[str, NumpyArray],
-        task_id: Optional[Union[int, Task]] = None,
-        **kwargs: Any,
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
     ) -> dict[str, NumpyArray]:
-        if task_id is None:
-            raise ValueError(f"task_id must be provided for JinaEmbeddingV3, got <{task_id}>")
-        onnx_input["task_id"] = np.array(task_id, dtype=np.int64)
+        onnx_input["task_id"] = np.array(self.current_task_id, dtype=np.int64)
         return onnx_input
 
     def embed(
@@ -75,19 +67,20 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        task_id: Optional[int] = None,
+        task_id: int = PASSAGE_TASK,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
-        task_id = (
-            task_id if task_id is not None else self.default_task_id
-        )  # required for multiprocessing
-        yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
+        self.current_task_id = task_id
+        kwargs["task_id"] = task_id
+        yield from super().embed(documents, batch_size, parallel, **kwargs)
 
     def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
-        yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
+        self.current_task_id = self.QUERY_TASK
+        yield from super().embed(query, **kwargs)
 
     def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
-        yield from super().embed(texts, task_id=self.PASSAGE_TASK, **kwargs)
+        self.current_task_id = self.PASSAGE_TASK
+        yield from super().embed(texts, **kwargs)
 
 
 class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
@@ -97,15 +90,11 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
         cache_dir: str,
         **kwargs: Any,
     ) -> JinaEmbeddingV3:
-        return JinaEmbeddingV3(
+        model = JinaEmbeddingV3(
             model_name=model_name,
             cache_dir=cache_dir,
             threads=1,
             **kwargs,
         )
-
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
-        self.model: JinaEmbeddingV3  # mypy complaints `self.model` does not have `default_task_id`
-        for idx, batch in items:
-            onnx_output = self.model.onnx_embed(batch, task_id=self.model.default_task_id)
-            yield idx, onnx_output
+        model.current_task_id = kwargs["task_id"]
+        return model

fastembed/text/onnx_embedding.py

@@ -247,12 +247,11 @@ class OnnxTextEmbedding(TextEmbeddingBase, 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=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -289,8 +288,6 @@ class OnnxTextEmbedding(TextEmbeddingBase, 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,
         )
 
@@ -306,9 +303,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         """
         return onnx_input
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         embeddings = output.model_output
 
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)

fastembed/text/onnx_text_model.py

@@ -21,16 +21,7 @@ class OnnxTextModel(OnnxModel[T]):
     def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    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.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
@@ -108,8 +99,6 @@ class OnnxTextModel(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
@@ -126,9 +115,7 @@ class OnnxTextModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, **kwargs), **kwargs
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -138,8 +125,6 @@ class OnnxTextModel(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,
             }
 
@@ -151,7 +136,7 @@ class OnnxTextModel(OnnxModel[T]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
+                yield from self._post_process_onnx_output(batch)  # type: ignore
 
 
 class TextEmbeddingWorker(EmbeddingWorker[T]):

fastembed/text/pooled_embedding.py

@@ -110,9 +110,7 @@ class PooledEmbedding(OnnxTextEmbedding):
         """
         return supported_pooled_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 

fastembed/text/pooled_normalized_embedding.py

@@ -138,9 +138,7 @@ class PooledNormalizedEmbedding(PooledEmbedding):
         """
         return supported_pooled_normalized_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 

fastembed/text/text_embedding.py

@@ -55,7 +55,7 @@ class TextEmbedding(TextEmbeddingBase):
     ) -> None:
         registered_models = cls._list_supported_models()
         for registered_model in registered_models:
-            if model.lower() == registered_model.model.lower():
+            if model == registered_model.model:
                 raise ValueError(
                     f"Model {model} is already registered in TextEmbedding, if you still want to add this model, "
                     f"please use another model name"
@@ -88,18 +88,18 @@ class TextEmbedding(TextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
+        if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
             warnings.warn(
                 "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
                 "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
                 UserWarning,
                 stacklevel=2,
             )
-        if model_name.lower() in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2".lower(),
-            "thenlper/gte-large".lower(),
-            "intfloat/multilingual-e5-large".lower(),
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2".lower(),
+        if model_name in {
+            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+            "thenlper/gte-large",
+            "intfloat/multilingual-e5-large",
+            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
         }:
             warnings.warn(
                 f"The model {model_name} now uses mean pooling instead of CLS embedding. "
@@ -128,40 +128,6 @@ class TextEmbedding(TextEmbeddingBase):
             "Please check the supported models using `TextEmbedding.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/text/text_embedding_base.py

@@ -17,7 +17,6 @@ class TextEmbeddingBase(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,
@@ -59,13 +58,3 @@ class TextEmbeddingBase(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 passed 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")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.7.1"
+version = "0.6.0"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -19,8 +19,8 @@ numpy = [
     { version = ">=1.21,<2.1.0", python = "<3.10" },
 ]
 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" },
     { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
 ]
 tqdm = "^4.66"

tests/test_image_onnx_embeddings.py

@@ -127,21 +127,3 @@ def test_lazy_load(model_name: str) -> None:
     assert hasattr(model.model, "model")
     if is_ci:
         delete_model_cache(model.model._model_dir)
-
-
-def test_get_embedding_size() -> None:
-    assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-ViT-B-32-vision") == 512
-    assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-vit-b-32-vision") == 512
-
-
-def test_embedding_size() -> None:
-    is_ci = os.getenv("CI")
-    model_name = "Qdrant/clip-ViT-B-32-vision"
-    model = ImageEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 512
-
-    model_name = "Qdrant/clip-vit-b-32-vision"
-    model = ImageEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 512
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/test_late_interaction_embeddings.py

@@ -254,24 +254,3 @@ def test_lazy_load(model_name: str):
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
-
-
-def test_get_embedding_size():
-    model_name = "answerdotai/answerai-colbert-small-v1"
-    assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 96
-
-    model_name = "answerdotai/answerai-ColBERT-small-v1"
-    assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 96
-
-
-def test_embedding_size():
-    is_ci = os.getenv("CI")
-    model_name = "answerdotai/answerai-colbert-small-v1"
-    model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 96
-
-    model_name = "answerdotai/answerai-ColBERT-small-v1"
-    model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 96
-    if is_ci:
-        delete_model_cache(model.model._model_dir)

tests/test_late_interaction_multimodal.py

@@ -81,23 +81,3 @@ def test_single_embedding_query():
         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():
-    model_name = "Qdrant/colpali-v1.3-fp16"
-    assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
-
-    model_name = "Qdrant/ColPali-v1.3-fp16"
-    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 = 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

tests/test_sparse_embeddings.py

@@ -43,46 +43,13 @@ CANONICAL_COLUMN_VALUES = {
             2.1904349327087402,
             1.0531445741653442,
         ],
-    },
-    "Qdrant/minicoil-v1": {
-        "indices": [80, 81, 82, 83, 6664, 6665, 6666, 6667],
-        "values": [
-            0.52634597,
-            0.8711344,
-            1.2264385,
-            0.52123857,
-            0.974713,
-            -0.97803956,
-            -0.94312465,
-            -0.12508166,
-        ],
-    },
+    }
 }
 
-CANONICAL_QUERY_VALUES = {
-    "Qdrant/minicoil-v1": {
-        "indices": [80, 81, 82, 83, 6664, 6665, 6666, 6667],
-        "values": [
-            0.31389374,
-            0.5195128,
-            0.7314033,
-            0.3108479,
-            0.5812834,
-            -0.5832673,
-            -0.5624452,
-            -0.0745942,
-        ],
-    },
-}
-
-
 docs = ["Hello World"]
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
-)
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
 def test_batch_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
@@ -98,7 +65,7 @@ def test_batch_embedding(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"])
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
 def test_single_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
@@ -117,23 +84,16 @@ def test_single_embedding(model_name: str) -> None:
         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]
+        for result in [passage_result, query_result]:
+            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)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
-)
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
 def test_parallel_processing(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name)

tests/test_text_multitask_embeddings.py

@@ -109,25 +109,9 @@ def test_single_embedding():
 
             canonical_vector = task["vectors"]
             assert np.allclose(
-                embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+                embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
             ), model_desc.model
 
-        classification_embeddings = list(model.embed(documents=docs, task_id=Task.CLASSIFICATION))
-        classification_embeddings = np.stack(classification_embeddings, axis=0)
-
-        assert classification_embeddings.shape == (len(docs), dim)
-
-        model = TextEmbedding(model_name=model_name, task_id=Task.CLASSIFICATION)
-        default_embeddings = list(model.embed(documents=docs))
-        default_embeddings = np.stack(default_embeddings, axis=0)
-
-        assert default_embeddings.shape == (len(docs), dim)
-
-        assert np.allclose(
-            classification_embeddings,
-            default_embeddings,
-            atol=1e-4,
-        ), model_desc.model
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
@@ -156,7 +140,7 @@ def test_single_embedding_query():
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
-            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+            embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
@@ -188,7 +172,7 @@ def test_single_embedding_passage():
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
-            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+            embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
@@ -223,6 +207,27 @@ def test_parallel_processing(dim: int, model_name: str):
         delete_model_cache(model.model._model_dir)
 
 
+def test_task_assignment():
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    if is_ci and not is_manual:
+        pytest.skip("Skipping in CI non-manual mode")
+
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
+        model_name = model_desc.model
+
+        model = TextEmbedding(model_name=model_name)
+
+        for i, task_id in enumerate(Task):
+            _ = list(model.embed(documents=docs, batch_size=1, task_id=i))
+            assert model.model.current_task_id == task_id
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
 @pytest.mark.parametrize("model_name", ["jinaai/jina-embeddings-v3"])
 def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")

tests/test_text_onnx_embeddings.py

@@ -156,22 +156,3 @@ def test_lazy_load(model_name: str) -> None:
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
-
-
-def test_get_embedding_size() -> None:
-    assert TextEmbedding.get_embedding_size("sentence-transformers/all-MiniLM-L6-v2") == 384
-    assert TextEmbedding.get_embedding_size("sentence-transformers/all-minilm-l6-v2") == 384
-
-
-def test_embedding_size() -> None:
-    is_ci = os.getenv("CI")
-    model_name = "sentence-transformers/all-MiniLM-L6-v2"
-    model = TextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 384
-
-    model_name = "sentence-transformers/all-minilm-l6-v2"
-    model = TextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 384
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)