new: gpu package (#224)

* Implement MuveraEmbedding

* Add random generator parameter for reproducibility in MuveraEmbedding

* Document random_seed parameter

* Remove unnecessary module docstring from muvera_embedding.py

* refactor: clean up constructor parameters and improve formatting in MuveraEmbedding

* refactor: rename muvera_embedding.py to muvera.py and update related references

* feat: enhance MuveraEmbedding with multi-vector model support and improve parameter defaults

* feat: add embedding_size property to MuveraEmbedding

* feat: update MuveraPostprocessor to use model description for embedding size and add Jupyter notebook for MUVERA usage

* fix: fix types, doctest, rename variables, refactor (#545)

* fix: fix types, doctest, rename variables, refactor

* fix: fix python3.9 compatibility

* fix: make get_output_dimension protected

* Optimize muvera (#551)

* vectorize operations

* fix: fill empty clusters with dataset vectors

* rollback get_output_dimension

* fix: fix type hints

* fix: review comments

* tests: add tests

---------

Co-authored-by: George <george.panchuk@qdrant.tech>

.github/workflows/python-tests.yml

@@ -2,9 +2,10 @@ name: Tests
 run-name: Tests (gpu)
 
 on:
-  push:
-    branches: [ master, main, gpu ]
   pull_request:
+    branches: [ master, main, gpu ]
+  workflow_dispatch:
+  
 
 env:
   CARGO_TERM_COLOR: always
@@ -40,5 +41,7 @@ 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,8 +26,6 @@ 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 \

README.md

@@ -190,6 +190,23 @@ scores = list(encoder.rerank(query, documents))
 # [-11.48061752319336, 5.472434997558594]
 ```
 
+Text cross encoders can also be extended with models which are not in the list of supported models.
+
+```python
+from fastembed.rerank.cross_encoder import TextCrossEncoder 
+from fastembed.common.model_description import ModelSource
+
+TextCrossEncoder.add_custom_model(
+    model="Xenova/ms-marco-MiniLM-L-4-v2",
+    model_file="onnx/model.onnx",
+    sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-4-v2"),
+)
+model = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-4-v2")
+scores = list(model.rerank_pairs(
+    [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ..."),]
+))
+```
+
 ## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.
@@ -229,36 +246,36 @@ pip install qdrant-client[fastembed-gpu]
 You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
-from qdrant_client import QdrantClient
+from qdrant_client import QdrantClient, models
 
 # Initialize the client
 client = QdrantClient("localhost", port=6333) # For production
-# client = QdrantClient(":memory:") # For small experiments
+# client = QdrantClient(":memory:") # For experimentation
 
-# Prepare your documents, metadata, and IDs
-docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
-metadata = [
-    {"source": "Langchain-docs"},
-    {"source": "Llama-index-docs"},
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+payload = [
+    {"document": "Qdrant has Langchain integrations", "source": "Langchain-docs", },
+    {"document": "Qdrant also has Llama Index integrations", "source": "LlamaIndex-docs"},
 ]
+docs = [models.Document(text=data["document"], model=model_name) for data in payload]
 ids = [42, 2]
 
-# If you want to change the model:
-# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
-# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
-
-# Use the new add() instead of upsert()
-# This internally calls embed() of the configured embedding model
-client.add(
-    collection_name="demo_collection",
-    documents=docs,
-    metadata=metadata,
-    ids=ids
+client.create_collection(
+    "demo_collection",
+    vectors_config=models.VectorParams(
+        size=client.get_embedding_size(model_name), distance=models.Distance.COSINE)
 )
 
-search_result = client.query(
+client.upload_collection(
     collection_name="demo_collection",
-    query_text="This is a query document"
+    vectors=docs,
+    ids=ids,
+    payload=payload,
 )
+
+search_result = client.query_points(
+    collection_name="demo_collection",
+    query=models.Document(text="This is a query document", model=model_name)
+).points
 print(search_result)
 ```

fastembed/common/model_description.py

@@ -7,6 +7,11 @@ from typing import Optional, Any
 class ModelSource:
     hf: Optional[str] = None
     url: Optional[str] = None
+    _deprecated_tar_struct: bool = False
+
+    @property
+    def deprecated_tar_struct(self) -> bool:
+        return self._deprecated_tar_struct
 
     def __post_init__(self) -> None:
         if self.hf is None and self.url is None:

fastembed/common/model_management.py

@@ -330,9 +330,10 @@ class ModelManagement(Generic[T]):
         model_name: str,
         source_url: str,
         cache_dir: str,
+        deprecated_tar_struct: bool = False,
         local_files_only: bool = False,
     ) -> Path:
-        fast_model_name = f"fast-{model_name.split('/')[-1]}"
+        fast_model_name = f"{'fast-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
         cache_tmp_dir = Path(cache_dir) / "tmp"
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
@@ -438,6 +439,7 @@ class ModelManagement(Generic[T]):
                         model.model,
                         str(url_source),
                         str(cache_dir),
+                        deprecated_tar_struct=model.sources.deprecated_tar_struct,
                         local_files_only=local_files_only,
                     )
                 except Exception:

fastembed/common/onnx_model.py

@@ -28,7 +28,16 @@ class OnnxModel(Generic[T]):
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
@@ -68,15 +77,7 @@ class OnnxModel(Generic[T]):
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
-                # kSameAsRequested: Allocates only the requested memory, avoiding over-allocation.
-                # more precise than 'kNextPowerOfTwo', which grows memory aggressively.
-                # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-                onnx_providers = [
-                    (
-                        "CUDAExecutionProvider",
-                        {"device_id": device_id, "arena_extend_strategy": "kSameAsRequested"},
-                    )
-                ]
+                onnx_providers = [("CUDAExecutionProvider", {"device_id": device_id})]
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
@@ -140,7 +141,5 @@ class EmbeddingWorker(Worker, Generic[T]):
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError("Subclasses must implement this method")

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

fastembed/common/utils.py

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

fastembed/image/image_embedding.py

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

fastembed/image/image_embedding_base.py

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

fastembed/image/onnx_embedding.py

@@ -1,6 +1,5 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-import numpy as np
 
 from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
@@ -113,11 +112,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -178,6 +178,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 
@@ -194,8 +196,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
-        return normalize(output.model_output).astype(np.float32)
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return normalize(output.model_output)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):

fastembed/image/onnx_image_model.py

@@ -6,14 +6,13 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from PIL import Image
-import onnxruntime as ort
 
 from fastembed.image.transform.operators import Compose
 from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 # Holds type of the embedding result
@@ -24,7 +23,16 @@ class OnnxImageModel(OnnxModel[T]):
     def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
@@ -75,21 +83,7 @@ class OnnxImageModel(OnnxModel[T]):
             encoded = np.array(self.processor(image_files))
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(None, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -103,6 +97,8 @@ class OnnxImageModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -119,9 +115,7 @@ class OnnxImageModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed(batch), **kwargs)
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -131,6 +125,8 @@ class OnnxImageModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -142,13 +138,11 @@ class OnnxImageModel(OnnxModel[T]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
-                yield from self._post_process_onnx_output(batch)  # type: ignore
+                yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
 
 
 class ImageEmbeddingWorker(EmbeddingWorker[T]):
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed(batch, **kwargs)
+            embeddings = self.model.onnx_embed(batch)
             yield idx, embeddings

fastembed/image/transform/functional.py

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

fastembed/late_interaction/colbert.py

@@ -43,10 +43,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
-        self, output: OnnxOutputContext, is_doc: bool = True
+        self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any
     ) -> Iterable[NumpyArray]:
         if not is_doc:
-            return output.model_output.astype(np.float32)
+            for embedding in output.model_output:
+                yield embedding
 
         if output.input_ids is None or output.attention_mask is None:
             raise ValueError(
@@ -58,11 +59,13 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
                 if token_id in self.skip_list or token_id == self.pad_token_id:
                     output.attention_mask[i, j] = 0
 
-        output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32)
+        output.model_output *= np.expand_dims(output.attention_mask, 2)
         norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
         norm_clamped = np.maximum(norm, 1e-12)
         output.model_output /= norm_clamped
-        return output.model_output.astype(np.float32)
+
+        for embedding, attention_mask in zip(output.model_output, output.attention_mask):
+            yield embedding[attention_mask == 1]
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
@@ -169,11 +172,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
         self.mask_token_id: Optional[int] = None
         self.pad_token_id: Optional[int] = None
@@ -233,6 +237,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

fastembed/late_interaction/late_interaction_embedding_base.py

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

fastembed/late_interaction/late_interaction_text_embedding.py

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

fastembed/late_interaction/token_embeddings.py

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

fastembed/late_interaction_multimodal/colpali.py

@@ -95,11 +95,12 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
         self.mask_token_id = None
         self.pad_token_id = None
@@ -142,7 +143,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         assert self.model_description.dim is not None, "Model dim is not defined"
         return output.model_output.reshape(
             output.model_output.shape[0], -1, self.model_description.dim
-        ).astype(np.float32)
+        )
 
     def _post_process_onnx_text_output(
         self,
@@ -157,7 +158,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         Returns:
             Iterable[NumpyArray]: Post-processed output as NumPy arrays.
         """
-        return output.model_output.astype(np.float32)
+        return output.model_output
 
     def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         texts_query: list[str] = []
@@ -174,7 +175,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
     ) -> dict[str, NumpyArray]:
         onnx_input["input_ids"] = np.array(
             [
-                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()
+                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()  # type: ignore[index]
                 for input_ids in onnx_input["input_ids"]
             ]
         )
@@ -235,6 +236,8 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 
@@ -268,6 +271,8 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

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

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

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

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -6,14 +6,13 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from PIL import Image
-import onnxruntime as ort
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common import OnnxProvider, ImageInput
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer, load_preprocessor
 from fastembed.common.types import NumpyArray
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.image.transform.operators import Compose
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -104,21 +103,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             )
 
         onnx_input = self._preprocess_onnx_text_input(onnx_input, **kwargs)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -135,6 +120,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -151,9 +138,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_text_output(
-                    self.onnx_embed_text(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_text_output(self.onnx_embed_text(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -163,6 +148,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -186,21 +173,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             encoded = np.array(self.processor(image_files))
         onnx_input = {"pixel_values": encoded}
         onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(None, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -214,6 +187,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -230,9 +205,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_image_output(
-                    self.onnx_embed_image(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_image_output(self.onnx_embed_image(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -242,6 +215,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -274,11 +249,9 @@ class TextEmbeddingWorker(EmbeddingWorker[T]):
     ) -> OnnxMultimodalModel:
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed_text(batch, **kwargs)
+            onnx_output = self.model.onnx_embed_text(batch)
             yield idx, onnx_output
 
 
@@ -300,9 +273,7 @@ class ImageEmbeddingWorker(EmbeddingWorker[T]):
     ) -> OnnxMultimodalModel:
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed_image(batch, **kwargs)
+            embeddings = self.model.onnx_embed_image(batch)
             yield idx, embeddings

fastembed/parallel_processor.py

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

fastembed/postprocess/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.postprocess.muvera import Muvera
+
+__all__ = ["Muvera"]

fastembed/postprocess/muvera.py

@@ -0,0 +1,364 @@
+from typing import Union
+
+import numpy as np
+
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+
+
+MultiVectorModel = Union[LateInteractionTextEmbeddingBase, LateInteractionMultimodalEmbeddingBase]
+MAX_HAMMING_DISTANCE = 65  # 64 bits + 1
+POPCOUNT_LUT = np.array([bin(x).count("1") for x in range(256)], dtype=np.uint8)
+
+
+def hamming_distance_matrix(ids: np.ndarray) -> np.ndarray:
+    """Compute full Hamming distance matrix
+
+    Args:
+    ids: shape (n,) - array of ids, only size of the array matters
+
+    Return:
+        np.ndarray (n, n) - hamming distance matrix
+    """
+    n = len(ids)
+    xor_vals = np.bitwise_xor(ids[:, None], ids[None, :])  # (n, n) uint64
+    bytes_view = xor_vals.view(np.uint8).reshape(n, n, 8)  # (n, n, 8)
+    return POPCOUNT_LUT[bytes_view].sum(axis=2)
+
+
+class SimHashProjection:
+    """
+    SimHash projection component for MUVERA clustering.
+
+    This class implements locality-sensitive hashing using random hyperplanes
+    to partition the vector space into 2^k_sim clusters. Each vector is assigned
+    to a cluster based on which side of k_sim random hyperplanes it falls on.
+
+    Attributes:
+        k_sim (int): Number of SimHash functions (hyperplanes)
+        dim (int): Dimensionality of input vectors
+        simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (dim, k_sim)
+    """
+
+    def __init__(self, k_sim: int, dim: int, random_generator: np.random.Generator):
+        """
+        Initialize SimHash projection with random hyperplanes.
+
+        Args:
+            k_sim (int): Number of SimHash functions, determines 2^k_sim clusters
+            dim (int): Dimensionality of input vectors
+            random_generator (np.random.Generator): Random number generator for reproducibility
+        """
+        self.k_sim = k_sim
+        self.dim = dim
+        # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution
+        self.simhash_vectors = random_generator.normal(size=(dim, k_sim))
+
+    def get_cluster_ids(self, vectors: np.ndarray) -> np.ndarray:
+        """
+        Compute the cluster IDs for a given vector using SimHash.
+
+        The cluster ID is determined by computing the dot product of the vector
+        with each hyperplane normal vector, taking the sign, and interpreting
+        the resulting binary string as an integer.
+
+        Args:
+            vectors (np.ndarray): Input vectors of shape (n, dim,)
+
+        Returns:
+            np.ndarray: Cluster IDs in range [0, 2^k_sim - 1]
+
+        Raises:
+            AssertionError: If a vector shape doesn't match expected dimensionality
+        """
+        dot_product = (
+            vectors @ self.simhash_vectors
+        )  # (token_num, dim) x (dim, k_sim) -> (token_num, k_sim)
+        cluster_ids = (dot_product > 0) @ (1 << np.arange(self.k_sim))
+        return cluster_ids
+
+
+class Muvera:
+    """
+    MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation.
+
+    This class creates Fixed Dimensional Encodings (FDEs) from variable-length
+    sequences of vectors by using SimHash clustering and random projections.
+    The process involves:
+    1. Clustering vectors using multiple SimHash projections
+    2. Computing cluster centers (with different strategies for docs vs queries)
+    3. Applying random projections for dimensionality reduction
+    4. Concatenating results from all projections
+
+    Attributes:
+        k_sim (int): Number of SimHash functions per projection
+        dim (int): Input vector dimensionality
+        dim_proj (int): Output dimensionality after random projection
+        r_reps (int): Number of random projection repetitions
+        random_seed (int): Random seed for consistent random matrix generation
+        simhash_projections (List[SimHashProjection]): SimHash instances for clustering
+        dim_reduction_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj)
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        k_sim: int = 5,
+        dim_proj: int = 16,
+        r_reps: int = 20,
+        random_seed: int = 42,
+    ):
+        """
+        Initialize MUVERA algorithm with specified parameters.
+
+        Args:
+            dim (int): Dimensionality of individual input vectors
+            k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters).
+                                   Defaults to 5.
+            dim_proj (int, optional): Dimensionality after random projection (must be <= dim).
+                                    Defaults to 16.
+            r_reps (int, optional): Number of random projection repetitions for robustness.
+                                    Defaults to 20.
+            random_seed (int, optional): Seed for random number generator to ensure
+                                         reproducible results. Defaults to 42.
+
+        Raises:
+            ValueError: If dim_proj > dim (cannot project to higher dimensionality)
+        """
+        if dim_proj > dim:
+            raise ValueError(
+                f"Cannot project to a higher dimensionality (dim_proj={dim_proj} > dim={dim})"
+            )
+
+        self.k_sim = k_sim
+        self.dim = dim
+        self.dim_proj = dim_proj
+        self.r_reps = r_reps
+        # Create r_reps independent SimHash projections for robustness
+        generator = np.random.default_rng(random_seed)
+        self.simhash_projections = [
+            SimHashProjection(k_sim=self.k_sim, dim=self.dim, random_generator=generator)
+            for _ in range(r_reps)
+        ]
+        # Random projection matrices with entries from {-1, +1} for each repetition
+        self.dim_reduction_projections = generator.choice([-1, 1], size=(r_reps, dim, dim_proj))
+
+    @classmethod
+    def from_multivector_model(
+        cls,
+        model: MultiVectorModel,
+        k_sim: int = 5,
+        dim_proj: int = 16,
+        r_reps: int = 20,  # noqa[naming]
+        random_seed: int = 42,
+    ) -> "Muvera":
+        """
+        Create a Muvera instance from a multi-vector embedding model.
+
+        This class method provides a convenient way to initialize a MUVERA
+        that is compatible with a given multi-vector model by automatically extracting
+        the embedding dimensionality from the model.
+
+        Args:
+            model (MultiVectorModel): A late interaction text or multimodal embedding model
+                                    that provides multi-vector embeddings. Must have an
+                                    `embedding_size` attribute specifying the dimensionality
+                                    of individual vectors.
+            k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters).
+                                   Defaults to 5.
+            dim_proj (int, optional): Dimensionality after random projection (must be <= model's
+                                    embedding_size). Defaults to 16.
+            r_reps (int, optional): Number of random projection repetitions for robustness.
+                                    Defaults to 20.
+            random_seed (int, optional): Seed for random number generator to ensure
+                                         reproducible results. Defaults to 42.
+
+        Returns:
+            Muvera: A configured MUVERA instance ready to process embeddings from the given model.
+
+        Raises:
+            ValueError: If dim_proj > model.embedding_size (cannot project to higher dimensionality)
+
+        Example:
+            >>> from fastembed import LateInteractionTextEmbedding
+            >>> model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+            >>> muvera = Muvera.from_multivector_model(
+            ...     model=model,
+            ...     k_sim=6,
+            ...     dim_proj=32
+            ... )
+            >>> # Now use postprocessor with embeddings from the model
+            >>> embeddings = np.array(list(model.embed(["sample text"])))
+            >>> fde = muvera.process_document(embeddings[0])
+        """
+        return cls(
+            dim=model.embedding_size,
+            k_sim=k_sim,
+            dim_proj=dim_proj,
+            r_reps=r_reps,
+            random_seed=random_seed,
+        )
+
+    def _get_output_dimension(self) -> int:
+        """
+        Get the output dimension of the MUVERA algorithm.
+
+        Returns:
+            int: Output dimension (r_reps * num_partitions * dim_proj) where b = 2^k_sim
+        """
+        num_partitions = 2**self.k_sim
+        return self.r_reps * num_partitions * self.dim_proj
+
+    @property
+    def embedding_size(self) -> int:
+        return self._get_output_dimension()
+
+    def process_document(self, vectors: NumpyArray) -> NumpyArray:
+        """
+        Encode a document's vectors into a Fixed Dimensional Encoding (FDE).
+
+        Uses document-specific settings: normalizes cluster centers by vector count
+        and fills empty clusters using Hamming distance-based selection.
+
+        Args:
+            vectors (NumpyArray): Document vectors of shape (n_tokens, dim)
+
+        Returns:
+            NumpyArray: Fixed dimensional encodings of shape (r_reps * b * dim_proj,)
+        """
+        return self.process(vectors, fill_empty_clusters=True, normalize_by_count=True)
+
+    def process_query(self, vectors: NumpyArray) -> NumpyArray:
+        """
+        Encode a query's vectors into a Fixed Dimensional Encoding (FDE).
+
+        Uses query-specific settings: no normalization by count and no empty
+        cluster filling to preserve query vector magnitudes.
+
+        Args:
+            vectors (NumpyArray]): Query vectors of shape (n_tokens, dim)
+
+        Returns:
+            NumpyArray: Fixed dimensional encoding of shape (r_reps * b * dim_proj,)
+        """
+        return self.process(vectors, fill_empty_clusters=False, normalize_by_count=False)
+
+    def process(
+        self,
+        vectors: NumpyArray,
+        fill_empty_clusters: bool = True,
+        normalize_by_count: bool = True,
+    ) -> NumpyArray:
+        """
+        Core encoding method that transforms variable-length vector sequences into FDEs.
+
+        The encoding process:
+        1. For each of r_reps random projections:
+           a. Assign vectors to clusters using SimHash
+           b. Compute cluster centers (sum of vectors in each cluster)
+           c. Optionally normalize by cluster size
+           d. Fill empty clusters using Hamming distance if requested
+           e. Apply random projection for dimensionality reduction
+           f. Flatten cluster centers into a vector
+        2. Concatenate all projection results
+
+        Args:
+            vectors (np.ndarray): Input vectors of shape (n_vectors, dim)
+            fill_empty_clusters (bool): Whether to fill empty clusters using nearest
+                                      vectors based on Hamming distance of cluster IDs
+            normalize_by_count (bool): Whether to normalize cluster centers by the
+                                     number of vectors assigned to each cluster
+
+        Returns:
+            np.ndarray: Fixed dimensional encoding of shape (r_reps * b * dim_proj)
+                        where B = 2^k_sim is the number of clusters
+
+        Raises:
+            AssertionError: If input vectors don't have expected dimensionality
+        """
+        assert (
+            vectors.shape[1] == self.dim
+        ), f"Expected vectors of shape (n, {self.dim}), got {vectors.shape}"
+
+        # Store results from each random projection
+        output_vectors = []
+
+        # num of space partitions in SimHash
+        num_partitions = 2**self.k_sim
+        cluster_center_ids = np.arange(num_partitions)
+        precomputed_hamming_matrix = (
+            hamming_distance_matrix(cluster_center_ids) if fill_empty_clusters else None
+        )
+
+        for projection_index, simhash in enumerate(self.simhash_projections):
+            # Initialize cluster centers and count vectors assigned to each cluster
+            cluster_centers = np.zeros((num_partitions, self.dim))
+            cluster_center_id_to_vectors: dict[int, list[int]] = {
+                cluster_center_id: [] for cluster_center_id in cluster_center_ids
+            }
+            cluster_vector_counts = None
+            empty_mask = None
+
+            # Assign each vector to its cluster and accumulate cluster centers
+            vector_cluster_ids = simhash.get_cluster_ids(vectors)
+            for cluster_id, (vec_idx, vec) in zip(vector_cluster_ids, enumerate(vectors)):
+                cluster_centers[cluster_id] += vec
+                cluster_center_id_to_vectors[cluster_id].append(vec_idx)
+
+            if normalize_by_count or fill_empty_clusters:
+                cluster_vector_counts = np.bincount(vector_cluster_ids, minlength=num_partitions)
+                empty_mask = cluster_vector_counts == 0
+
+            if normalize_by_count:
+                assert empty_mask is not None
+                assert cluster_vector_counts is not None
+                non_empty_mask = ~empty_mask
+                cluster_centers[non_empty_mask] /= cluster_vector_counts[non_empty_mask][:, None]
+
+            # Fill empty clusters using vectors with minimum Hamming distance
+            if fill_empty_clusters:
+                assert empty_mask is not None
+                assert precomputed_hamming_matrix is not None
+                masked_hamming = np.where(
+                    empty_mask[None, :], MAX_HAMMING_DISTANCE, precomputed_hamming_matrix
+                )
+                nearest_non_empty = np.argmin(masked_hamming, axis=1)
+                fill_vectors = np.array(
+                    [
+                        vectors[cluster_center_id_to_vectors[cluster_id][0]]
+                        for cluster_id in nearest_non_empty[empty_mask]
+                    ]
+                ).reshape(-1, self.dim)
+                cluster_centers[empty_mask] = fill_vectors
+
+            # Apply random projection for dimensionality reduction if needed
+            if self.dim_proj < self.dim:
+                dim_reduction_projection = self.dim_reduction_projections[
+                    projection_index
+                ]  # Get projection matrix for this repetition
+                projected_centers = (1 / np.sqrt(self.dim_proj)) * (
+                    cluster_centers @ dim_reduction_projection
+                )
+
+                # Flatten cluster centers into a single vector and add to output
+                output_vectors.append(projected_centers.flatten())
+                continue
+
+            # If no projection needed (dim_proj == dim), use original cluster centers
+            output_vectors.append(cluster_centers.flatten())
+
+        # Concatenate results from all R_reps projections into final FDE
+        return np.concatenate(output_vectors)
+
+
+if __name__ == "__main__":
+    v_arrs = np.random.randn(10, 100, 128)
+    muvera = Muvera(128, 4, 8, 20, 42)
+
+    for v_arr in v_arrs:
+        muvera.process(v_arr)  # type: ignore

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -0,0 +1,46 @@
+from typing import Optional, Sequence, Any
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import BaseModelDescription
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+
+
+class CustomTextCrossEncoder(OnnxTextCrossEncoder):
+    SUPPORTED_MODELS: list[BaseModelDescription] = []
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_ids=device_ids,
+            lazy_load=lazy_load,
+            device_id=device_id,
+            specific_model_path=specific_model_path,
+            **kwargs,
+        )
+
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        return cls.SUPPORTED_MODELS
+
+    @classmethod
+    def add_model(
+        cls,
+        model_description: BaseModelDescription,
+    ) -> None:
+        cls.SUPPORTED_MODELS.append(model_description)

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -131,11 +131,12 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -189,6 +190,8 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 
@@ -196,7 +199,9 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     def _get_worker_class(cls) -> Type[TextRerankerWorker]:
         return TextCrossEncoderWorker
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
         return (float(elem) for elem in output.model_output)
 
 

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -4,7 +4,6 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
-import onnxruntime as ort
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -15,7 +14,7 @@ from fastembed.common.onnx_model import (
 )
 from fastembed.common.types import NumpyArray
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
@@ -72,21 +71,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         tokenized_input = self.tokenize(pairs, **kwargs)
         inputs = self._build_onnx_input(tokenized_input)
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # Enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # Helps prevent excessive memory retention, especially for dynamic workloads.
-            # Source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         relevant_output = outputs[0]
         scores: NumpyArray = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
@@ -109,6 +94,8 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -125,9 +112,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(pairs, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed_pairs(batch, cuda=cuda, providers=providers, **kwargs)
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed_pairs(batch, **kwargs))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -137,6 +122,8 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -150,7 +137,18 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             for batch in pool.ordered_map(iter_batch(pairs, batch_size), **params):
                 yield from self._post_process_onnx_output(batch)  # type: ignore
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[float]: Post-processed output as an iterable of float values.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def _preprocess_onnx_input(
@@ -180,9 +178,7 @@ class TextRerankerWorker(EmbeddingWorker[float]):
     ) -> OnnxCrossEncoderModel:
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed_pairs(batch, **kwargs)
+            onnx_output = self.model.onnx_embed_pairs(batch)
             yield idx, onnx_output

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -3,13 +3,19 @@ from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.model_description import (
+    ModelSource,
+    BaseModelDescription,
+)
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
     CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
         OnnxTextCrossEncoder,
+        CustomTextCrossEncoder,
     ]
 
     @classmethod
@@ -124,3 +130,34 @@ class TextCrossEncoder(TextCrossEncoderBase):
         yield from self.model.rerank_pairs(
             pairs, batch_size=batch_size, parallel=parallel, **kwargs
         )
+
+    @classmethod
+    def add_custom_model(
+        cls,
+        model: str,
+        sources: ModelSource,
+        model_file: str = "onnx/model.onnx",
+        description: str = "",
+        license: str = "",
+        size_in_gb: float = 0.0,
+        additional_files: Optional[list[str]] = None,
+    ) -> None:
+        registered_models = cls._list_supported_models()
+        for registered_model in registered_models:
+            if model == registered_model.model:
+                raise ValueError(
+                    f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, "
+                    f"please use another model name"
+                )
+
+        CustomTextCrossEncoder.add_model(
+            BaseModelDescription(
+                model=model,
+                sources=sources,
+                model_file=model_file,
+                description=description,
+                license=license,
+                size_in_GB=size_in_gb,
+                additional_files=additional_files or [],
+            )
+        )

fastembed/sparse/bm25.py

@@ -21,13 +21,9 @@ 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",
@@ -35,21 +31,15 @@ supported_languages = [
     "french",
     "german",
     "greek",
-    "hebrew",
-    "hinglish",
     "hungarian",
-    "indonesian",
     "italian",
-    "kazakh",
-    "nepali",
     "norwegian",
     "portuguese",
     "romanian",
     "russian",
-    "slovene",
     "spanish",
     "swedish",
-    "tajik",
+    "tamil",
     "turkish",
 ]
 
@@ -125,11 +115,12 @@ 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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         self.token_max_length = token_max_length
@@ -170,6 +161,8 @@ 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
 
@@ -198,6 +191,8 @@ 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,
@@ -236,6 +231,8 @@ 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]:
@@ -344,7 +341,7 @@ class Bm25Worker(Worker):
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+        self, items: Iterable[tuple[int, Any]]
     ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
         for idx, batch in items:
             onnx_output = self.model.raw_embed(batch)

fastembed/sparse/bm42.py

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

@@ -0,0 +1,356 @@
+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,6 +4,7 @@ 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,
@@ -14,7 +15,7 @@ from fastembed.common.model_description import SparseModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
+    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25, MiniCOIL]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -61,7 +62,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name == "prithvida/Splade_PP_en_v1":
+        if model_name.lower() == "prithvida/Splade_PP_en_v1".lower():
             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,14 +27,16 @@ 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) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 
@@ -112,11 +114,12 @@ 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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -163,6 +166,8 @@ 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

@@ -0,0 +1,146 @@
+"""
+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

@@ -0,0 +1,247 @@
+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

@@ -0,0 +1,202 @@
+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,7 +35,9 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
         """
         return supported_clip_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
         return output.model_output
 
 

fastembed/text/custom_text_embedding.py

@@ -58,7 +58,9 @@ class CustomTextEmbedding(OnnxTextEmbedding):
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         return cls.SUPPORTED_MODELS
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
         return self._normalize(self._pool(output.model_output, output.attention_mask))
 
     def _pool(
@@ -75,6 +77,11 @@ 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,6 +3,7 @@ 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
@@ -44,9 +45,11 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
 
-    def __init__(self, *args: Any, **kwargs: Any):
+    def __init__(self, *args: Any, task_id: Optional[int] = None, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.current_task_id: Union[Task, int] = self.PASSAGE_TASK
+        self.default_task_id: Union[Task, int] = (
+            task_id if task_id is not None else self.PASSAGE_TASK
+        )
 
     @classmethod
     def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
@@ -57,9 +60,14 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         return supported_multitask_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+        self,
+        onnx_input: dict[str, NumpyArray],
+        task_id: Optional[Union[int, Task]] = None,
+        **kwargs: Any,
     ) -> dict[str, NumpyArray]:
-        onnx_input["task_id"] = np.array(self.current_task_id, dtype=np.int64)
+        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)
         return onnx_input
 
     def embed(
@@ -67,20 +75,19 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        task_id: int = PASSAGE_TASK,
+        task_id: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
-        self.current_task_id = task_id
-        kwargs["task_id"] = task_id
-        yield from super().embed(documents, batch_size, parallel, **kwargs)
+        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)
 
     def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
-        self.current_task_id = self.QUERY_TASK
-        yield from super().embed(query, **kwargs)
+        yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
 
     def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
-        self.current_task_id = self.PASSAGE_TASK
-        yield from super().embed(texts, **kwargs)
+        yield from super().embed(texts, task_id=self.PASSAGE_TASK, **kwargs)
 
 
 class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
@@ -90,11 +97,15 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
         cache_dir: str,
         **kwargs: Any,
     ) -> JinaEmbeddingV3:
-        model = JinaEmbeddingV3(
+        return JinaEmbeddingV3(
             model_name=model_name,
             cache_dir=cache_dir,
             threads=1,
             **kwargs,
         )
-        model.current_task_id = kwargs["task_id"]
-        return model
+
+    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

fastembed/text/onnx_embedding.py

@@ -1,6 +1,5 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-import numpy as np
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
@@ -21,6 +20,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="Qdrant/fast-bge-base-en",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -36,6 +36,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="qdrant/bge-base-en-v1.5-onnx-q",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -63,6 +64,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="Qdrant/bge-small-en",
             url="https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -90,6 +92,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="Qdrant/bge-small-zh-v1.5",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -244,11 +247,12 @@ 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=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -285,6 +289,8 @@ 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,
         )
 
@@ -300,7 +306,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         """
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
         embeddings = output.model_output
 
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
@@ -309,7 +317,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             processed_embeddings = embeddings
         else:
             raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
-        return normalize(processed_embeddings).astype(np.float32)
+        return normalize(processed_embeddings)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(

fastembed/text/onnx_text_model.py

@@ -4,14 +4,13 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-import onnxruntime as ort
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
@@ -22,7 +21,16 @@ 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) -> Iterable[T]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
@@ -83,21 +91,7 @@ class OnnxTextModel(OnnxModel[T]):
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -114,6 +108,8 @@ 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
@@ -131,7 +127,7 @@ class OnnxTextModel(OnnxModel[T]):
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
                 yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, cuda=cuda, providers=providers)
+                    self.onnx_embed(batch, **kwargs), **kwargs
                 )
         else:
             if parallel == 0:
@@ -142,6 +138,8 @@ 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,
             }
 
@@ -153,13 +151,11 @@ 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)  # type: ignore
+                yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
 
 
 class TextEmbeddingWorker(EmbeddingWorker[T]):
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, OnnxOutputContext]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed(batch, **kwargs)
+            onnx_output = self.model.onnx_embed(batch)
             yield idx, onnx_output

fastembed/text/pooled_embedding.py

@@ -82,6 +82,7 @@ supported_pooled_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="qdrant/multilingual-e5-large-onnx",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model.onnx",
         additional_files=["model.onnx_data"],
@@ -109,13 +110,15 @@ class PooledEmbedding(OnnxTextEmbedding):
         """
         return supported_pooled_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
+        return self.mean_pooling(embeddings, attn_mask)
 
 
 class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/pooled_normalized_embedding.py

@@ -1,6 +1,5 @@
 from typing import Any, Iterable, Type
 
-import numpy as np
 
 from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -22,6 +21,7 @@ supported_pooled_normalized_models: list[DenseModelDescription] = [
         sources=ModelSource(
             url="https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
             hf="qdrant/all-MiniLM-L6-v2-onnx",
+            _deprecated_tar_struct=True,
         ),
         model_file="model.onnx",
     ),
@@ -138,13 +138,15 @@ class PooledNormalizedEmbedding(PooledEmbedding):
         """
         return supported_pooled_normalized_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+        return normalize(self.mean_pooling(embeddings, attn_mask))
 
 
 class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):

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 == registered_model.model:
+            if model.lower() == registered_model.model.lower():
                 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 == "nomic-ai/nomic-embed-text-v1.5-Q":
+        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
             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 in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-            "thenlper/gte-large",
-            "intfloat/multilingual-e5-large",
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+        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(),
         }:
             warnings.warn(
                 f"The model {model_name} now uses mean pooling instead of CLS embedding. "
@@ -128,6 +128,40 @@ 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,6 +17,7 @@ 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,
@@ -58,3 +59,13 @@ 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-gpu"
-version = "0.6.0"
+version = "0.7.2"
 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-gpu = [
-    { version = ">1.20.0", python = ">=3.13" },
     { version = ">=1.17.0,<1.20.0", python = "<3.10" },
+    { version = ">1.20.0", python = ">=3.13" },
     { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
 ]
 tqdm = "^4.66"

tests/test_custom_models.py

@@ -3,10 +3,17 @@ import os
 import numpy as np
 import pytest
 
-from fastembed.common.model_description import PoolingType, ModelSource, DenseModelDescription
+from fastembed.common.model_description import (
+    PoolingType,
+    ModelSource,
+    DenseModelDescription,
+    BaseModelDescription,
+)
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import normalize, mean_pooling
 from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+from fastembed.rerank.cross_encoder import TextCrossEncoder
 from fastembed.text.text_embedding import TextEmbedding
 from tests.utils import delete_model_cache
 
@@ -14,8 +21,10 @@ from tests.utils import delete_model_cache
 @pytest.fixture(autouse=True)
 def restore_custom_models_fixture():
     CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextCrossEncoder.SUPPORTED_MODELS = []
     yield
     CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextCrossEncoder.SUPPORTED_MODELS = []
 
 
 def test_text_custom_model():
@@ -65,6 +74,43 @@ def test_text_custom_model():
         delete_model_cache(model.model._model_dir)
 
 
+def test_cross_encoder_custom_model():
+    is_ci = os.getenv("CI")
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+    size_in_gb = 0.08
+    source = ModelSource(hf=custom_model_name)
+    canonical_vector = np.array([-5.7170815, -11.112114], dtype=np.float32)
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        model_file="onnx/model.onnx",
+        sources=source,
+        size_in_gb=size_in_gb,
+    )
+
+    assert CustomTextCrossEncoder.SUPPORTED_MODELS[0] == BaseModelDescription(
+        model=custom_model_name,
+        sources=source,
+        model_file="onnx/model.onnx",
+        description="",
+        license="",
+        size_in_GB=size_in_gb,
+    )
+
+    model = TextCrossEncoder(custom_model_name)
+    pairs = [
+        ("What is AI?", "Artificial intelligence is ..."),
+        ("What is ML?", "Machine learning is ..."),
+    ]
+    scores = list(model.rerank_pairs(pairs))
+
+    embeddings = np.stack(scores, axis=0)
+    assert embeddings.shape == (2,)
+    assert np.allclose(embeddings, canonical_vector, atol=1e-3)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
 def test_mock_add_custom_models():
     dim = 5
     size_in_gb = 0.1
@@ -91,15 +137,11 @@ def test_mock_add_custom_models():
     expected_output = {
         f"{PoolingType.MEAN.lower()}-normalized": normalize(
             mean_pooling(dummy_token_embedding, dummy_attention_mask)
-        ).astype(np.float32),
+        ),
         f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
-        f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]).astype(
-            np.float32
-        ),
+        f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]),
         f"{PoolingType.CLS.lower()}": dummy_token_embedding[:, 0],
-        f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding).astype(
-            np.float32
-        ),
+        f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding),
         f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
     }
 
@@ -160,3 +202,28 @@ def test_do_not_add_existing_model():
             dim=384,
             size_in_gb=0.47,
         )
+
+
+def test_do_not_add_existing_cross_encoder():
+    existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+
+    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
+        TextCrossEncoder.add_custom_model(
+            existing_base_model,
+            sources=ModelSource(hf=existing_base_model),
+            size_in_gb=0.08,
+        )
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        sources=ModelSource(hf=existing_base_model),
+        size_in_gb=0.08,
+    )
+
+    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
+        TextCrossEncoder.add_custom_model(
+            custom_model_name,
+            sources=ModelSource(hf=custom_model_name),
+            size_in_gb=0.08,
+        )

tests/test_image_onnx_embeddings.py

@@ -8,7 +8,7 @@ from PIL import Image
 
 from fastembed import ImageEmbedding
 from tests.config import TEST_MISC_DIR
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_VECTOR_VALUES = {
     "Qdrant/clip-ViT-B-32-vision": np.array([-0.0098, 0.0128, -0.0274, 0.002, -0.0059]),
@@ -27,11 +27,13 @@ CANONICAL_VECTOR_VALUES = {
 }
 
 
-def test_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
+def test_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
     for model_desc in ImageEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue
 
         dim = model_desc.dim
@@ -74,8 +76,12 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
 
     embeddings = list(model.embed(images, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
+    assert np.allclose(embeddings[1], embeddings[2])
+
+    canonical_vector = CANONICAL_VECTOR_VALUES[model_name]
 
     assert embeddings.shape == (len(test_images) * n_images, n_dims)
+    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
@@ -121,3 +127,21 @@ 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

@@ -6,7 +6,7 @@ import numpy as np
 from fastembed.late_interaction.late_interaction_text_embedding import (
     LateInteractionTextEmbedding,
 )
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -153,79 +153,105 @@ CANONICAL_QUERY_VALUES = {
 docs = ["Hello World"]
 
 
-def test_batch_embedding():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_batch_embedding(model_name: str):
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        result = list(model.embed(docs_to_embed, batch_size=6))
+    model = LateInteractionTextEmbedding(model_name=model_name)
+    result = list(model.embed(docs_to_embed, batch_size=6))
+    expected_result = CANONICAL_COLUMN_VALUES[model_name]
 
-        for value in result:
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-
-def test_single_embedding():
-    is_ci = os.getenv("CI")
-    docs_to_embed = docs
-
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+    for value in result:
         token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-
-def test_single_embedding_query():
-    is_ci = os.getenv("CI")
-    queries_to_embed = docs
-
-    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        result = next(iter(model.query_embed(queries_to_embed)))
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-
-def test_parallel_processing():
-    is_ci = os.getenv("CI")
-    model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
-    token_dim = 128
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
-    embeddings = np.stack(embeddings, axis=0)
-
-    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
-    embeddings_2 = np.stack(embeddings_2, axis=0)
-
-    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
-    embeddings_3 = np.stack(embeddings_3, axis=0)
-
-    assert embeddings.shape[0] == len(docs) and embeddings.shape[-1] == token_dim
-    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
+        assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["colbert-ir/colbertv2.0"],
-)
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_batch_inference_size_same_as_single_inference(model_name: str):
+    is_ci = os.getenv("CI")
+
+    model = LateInteractionTextEmbedding(model_name=model_name)
+    docs_to_embed = [
+        "short document",
+        "A bit longer document, which should not affect the size"
+    ]
+    result = list(model.embed(docs_to_embed, batch_size=1))
+    result_2 = list(model.embed(docs_to_embed, batch_size=2))
+    assert len(result[0]) == len(result_2[0])
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_single_embedding(model_name: str):
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    docs_to_embed = docs
+
+    for model_desc in LateInteractionTextEmbedding._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
+        print("evaluating", model_name)
+        model = LateInteractionTextEmbedding(model_name=model_name)
+        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_single_embedding_query(model_name: str):
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    queries_to_embed = docs
+
+    for model_desc in LateInteractionTextEmbedding._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
+        print("evaluating", model_name)
+        model = LateInteractionTextEmbedding(model_name=model_name)
+        result = next(iter(model.query_embed(queries_to_embed)))
+        expected_result = CANONICAL_QUERY_VALUES[model_name]
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("token_dim,model_name", [(96, "answerdotai/answerai-colbert-small-v1")])
+def test_parallel_processing(token_dim: int, model_name: str):
+    is_ci = os.getenv("CI")
+    model = LateInteractionTextEmbedding(model_name=model_name)
+
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+
+    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+
+    assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
+
+    for i in range(len(embeddings)):
+        assert np.allclose(embeddings[i], embeddings_2[i], atol=1e-3)
+        assert np.allclose(embeddings[i], embeddings_3[i], atol=1e-3)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
 def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")
 
@@ -244,3 +270,24 @@ 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

@@ -1,5 +1,6 @@
 import os
 
+import pytest
 from PIL import Image
 import numpy as np
 
@@ -45,38 +46,58 @@ images = [
 
 
 def test_batch_embedding():
-    is_ci = os.getenv("CI")
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
 
-    if not is_ci:
-        for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-            print("evaluating", model_name)
-            model = LateInteractionMultimodalEmbedding(model_name=model_name)
-            result = list(model.embed_image(images, batch_size=2))
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = list(model.embed_image(images, batch_size=2))
 
-            for value in result:
-                token_num, abridged_dim = expected_result.shape
-                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+        for value in result:
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_single_embedding():
-    is_ci = os.getenv("CI")
-    if not is_ci:
-        for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-            print("evaluating", model_name)
-            model = LateInteractionMultimodalEmbedding(model_name=model_name)
-            result = next(iter(model.embed_image(images, batch_size=6)))
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
+
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = next(iter(model.embed_image(images, batch_size=6)))
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_single_embedding_query():
-    is_ci = os.getenv("CI")
-    if not is_ci:
-        queries_to_embed = queries
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
 
-        for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
-            print("evaluating", model_name)
-            model = LateInteractionMultimodalEmbedding(model_name=model_name)
-            result = next(iter(model.embed_text(queries_to_embed)))
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = next(iter(model.embed_text(queries)))
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+
+
+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_postprocess.py

@@ -0,0 +1,38 @@
+import numpy as np
+
+from fastembed import LateInteractionTextEmbedding
+from fastembed.postprocess import Muvera
+
+CANONICAL_VALUES = [-2.61810007e-04, 1.89005750e00, -2.32070747e00]
+CANONICAL_QUERY_VALUES = [
+    -0.85783903,
+    1.1077204,
+    -0.09522747,
+]  # part of the values are zeros, should be compared with the result of nonzero mask
+
+DIM = 128
+K_SIM = 5
+DIM_PROJ = 16
+R_REPS = 20
+
+
+def test_single_input():
+    model = LateInteractionTextEmbedding("colbert-ir/colbertv2.0", lazy_load=True)
+    random_generator = np.random.default_rng(42)
+    multivector = random_generator.random((10, 128))
+
+    for muvera in (
+        Muvera(dim=DIM, k_sim=K_SIM, dim_proj=DIM_PROJ, r_reps=R_REPS, random_seed=42),
+        Muvera.from_multivector_model(model, k_sim=K_SIM, dim_proj=DIM_PROJ, r_reps=R_REPS),
+    ):
+        fde = muvera.process(multivector)
+        assert fde.shape[0] == muvera.embedding_size
+        assert np.allclose(fde[:3], CANONICAL_VALUES)
+
+        fde_doc = muvera.process_document(multivector)
+        assert fde_doc.shape[0] == muvera.embedding_size
+        assert np.allclose(fde, fde_doc)
+
+        fde_query = muvera.process_query(multivector)
+        assert fde_query.shape[0] == muvera.embedding_size
+        assert np.allclose(fde_query[np.nonzero(fde_query)][:3], CANONICAL_QUERY_VALUES)

tests/test_sparse_embeddings.py

@@ -5,10 +5,10 @@ import numpy as np
 
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_COLUMN_VALUES = {
-    "prithvida/Splade_PP_en_v1": {
+    "prithivida/Splade_PP_en_v1": {
         "indices": [
             2040,
             2047,
@@ -43,46 +43,100 @@ 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"]
 
 
-def test_batch_embedding() -> None:
+@pytest.mark.parametrize(
+    "model_name",
+    ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
+)
+def test_batch_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        model = SparseTextEmbedding(model_name=model_name)
-        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
-        assert result.indices.tolist() == expected_result["indices"]
+    model = SparseTextEmbedding(model_name=model_name)
+    result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+    expected_result = CANONICAL_COLUMN_VALUES[model_name]
+    assert result.indices.tolist() == expected_result["indices"]
 
-        for i, value in enumerate(result.values):
-            assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    for i, value in enumerate(result.values):
+        assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"])
+def test_single_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    for model_desc in SparseTextEmbedding._list_supported_models():
+        if (
+            model_desc.model not in CANONICAL_COLUMN_VALUES
+        ):  # attention models and bm25 are also parts of
+            # SparseTextEmbedding, however, they have their own tests
+            continue
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         model = SparseTextEmbedding(model_name=model_name)
 
         passage_result = next(iter(model.embed(docs, batch_size=6)))
         query_result = next(iter(model.query_embed(docs)))
-        for result in [passage_result, query_result]:
-            assert result.indices.tolist() == expected_result["indices"]
+        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 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)
 
 
-def test_parallel_processing() -> None:
+@pytest.mark.parametrize(
+    "model_name",
+    ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
+)
+def test_parallel_processing(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+    model = SparseTextEmbedding(model_name=model_name)
     docs = ["hello world", "flag embedding"] * 30
     sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
     sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))
@@ -172,10 +226,7 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["prithivida/Splade_PP_en_v1"],
-)
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)

tests/test_text_cross_encoder.py

@@ -4,7 +4,7 @@ import numpy as np
 import pytest
 
 from fastembed.rerank.cross_encoder import TextCrossEncoder
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_SCORE_VALUES = {
     "Xenova/ms-marco-MiniLM-L-6-v2": np.array([8.500708, -2.541011]),
@@ -15,44 +15,37 @@ CANONICAL_SCORE_VALUES = {
     "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]),
 }
 
-SELECTED_MODELS = {
-    "Xenova": "Xenova/ms-marco-MiniLM-L-6-v2",
-    "BAAI": "BAAI/bge-reranker-base",
-    "jinaai": "jinaai/jina-reranker-v1-tiny-en",
-}
 
-
-@pytest.mark.parametrize(
-    "model_name",
-    [model_name for model_name in CANONICAL_SCORE_VALUES],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
-    model = TextCrossEncoder(model_name=model_name)
+    for model_desc in TextCrossEncoder._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
 
-    query = "What is the capital of France?"
-    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
-    scores = np.array(list(model.rerank(query, documents)))
+        model = TextCrossEncoder(model_name=model_name)
 
-    pairs = [(query, doc) for doc in documents]
-    scores2 = np.array(list(model.rerank_pairs(pairs)))
-    assert np.allclose(
-        scores, scores2, atol=1e-5
-    ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
+        query = "What is the capital of France?"
+        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
+        scores = np.array(list(model.rerank(query, documents)))
 
-    canonical_scores = CANONICAL_SCORE_VALUES[model_name]
-    assert np.allclose(
-        scores, canonical_scores, atol=1e-3
-    ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        pairs = [(query, doc) for doc in documents]
+        scores2 = np.array(list(model.rerank_pairs(pairs)))
+        assert np.allclose(
+            scores, scores2, atol=1e-5
+        ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
+
+        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+        assert np.allclose(
+            scores, canonical_scores, atol=1e-3
+        ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    [model_name for model_name in SELECTED_MODELS.values()],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_batch_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
 
@@ -78,10 +71,7 @@ def test_batch_rerank(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["Xenova/ms-marco-MiniLM-L-6-v2"],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextCrossEncoder(model_name=model_name, lazy_load=True)
@@ -95,10 +85,7 @@ def test_lazy_load(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    [model_name for model_name in SELECTED_MODELS.values()],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_rerank_pairs_parallel(model_name: str) -> None:
     is_ci = os.getenv("CI")
 

tests/test_text_multitask_embeddings.py

@@ -4,7 +4,7 @@ import numpy as np
 import pytest
 
 from fastembed import TextEmbedding
-from fastembed.text.multitask_embedding import Task
+from fastembed.text.multitask_embedding import JinaEmbeddingV3, Task
 from tests.utils import delete_model_cache
 
 
@@ -60,52 +60,43 @@ CANONICAL_VECTOR_VALUES = {
 docs = ["Hello World", "Follow the white rabbit."]
 
 
-def test_batch_embedding():
+@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
+def test_batch_embedding(dim: int, model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
     docs_to_embed = docs * 10
     default_task = Task.RETRIEVAL_PASSAGE
 
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
+    model = TextEmbedding(model_name=model_name)
 
-        model_name = model_desc.model
-        dim = model_desc.dim
+    embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
+    embeddings = np.stack(embeddings, axis=0)
 
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
+    assert embeddings.shape == (len(docs_to_embed), dim)
 
-        model = TextEmbedding(model_name=model_name)
+    canonical_vector = CANONICAL_VECTOR_VALUES[model_name][default_task]["vectors"]
+    assert np.allclose(
+        embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+    ), model_name
 
-        print(f"evaluating {model_name} default task")
-
-        embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
-        embeddings = np.stack(embeddings, axis=0)
-
-        assert embeddings.shape == (len(docs_to_embed), dim)
-
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_name][default_task]["vectors"]
-        assert np.allclose(
-            embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-        ), model_desc.model
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 def test_single_embedding():
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
 
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
+    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
         dim = model_desc.dim
 
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
         model = TextEmbedding(model_name=model_name)
 
         for task in CANONICAL_VECTOR_VALUES[model_name]:
@@ -118,27 +109,42 @@ def test_single_embedding():
 
             canonical_vector = task["vectors"]
             assert np.allclose(
-                embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+                embeddings[:, : 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)
 
 
 def test_single_embedding_query():
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
     task_id = Task.RETRIEVAL_QUERY
 
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
+    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
         dim = model_desc.dim
 
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
         model = TextEmbedding(model_name=model_name)
 
         print(f"evaluating {model_name} query_embed task_id: {task_id}")
@@ -150,7 +156,7 @@ def test_single_embedding_query():
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
-            embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
@@ -159,18 +165,18 @@ def test_single_embedding_query():
 
 def test_single_embedding_passage():
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
     task_id = Task.RETRIEVAL_PASSAGE
 
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
+    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
         dim = model_desc.dim
 
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
         model = TextEmbedding(model_name=model_name)
 
         print(f"evaluating {model_name} passage_embed task_id: {task_id}")
@@ -182,21 +188,22 @@ def test_single_embedding_passage():
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
-            embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
-def test_parallel_processing():
+@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
+def test_parallel_processing(dim: int, model_name: str):
     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")
 
     docs = ["Hello World", "Follow the white rabbit."] * 10
 
-    model_name = "jinaai/jina-embeddings-v3"
-    dim = 1024
-
     model = TextEmbedding(model_name=model_name)
 
     task_id = Task.SEPARATION
@@ -216,33 +223,14 @@ def test_parallel_processing():
         delete_model_cache(model.model._model_dir)
 
 
-def test_task_assignment():
-    is_ci = os.getenv("CI")
-
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
-        model_name = model_desc.model
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
-        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"],
-)
+@pytest.mark.parametrize("model_name", ["jinaai/jina-embeddings-v3"])
 def test_lazy_load(model_name: str):
     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")
+
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
 

tests/test_text_onnx_embeddings.py

@@ -5,7 +5,7 @@ import numpy as np
 import pytest
 
 from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
@@ -72,17 +72,19 @@ CANONICAL_VECTOR_VALUES = {
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
 
-def test_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+def test_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     is_mac = platform.system() == "Darwin"
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
     for model_desc in TextEmbedding._list_supported_models():
-        if (
-            (not is_ci and model_desc.size_in_GB > 1)
-            or model_desc.model in MULTI_TASK_MODELS
-            or (is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q")
+        if model_desc.model in MULTI_TASK_MODELS or (
+            is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q"
         ):
             continue
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
 
         dim = model_desc.dim
 
@@ -95,15 +97,12 @@ def test_embedding() -> None:
         canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc["model"]
+        ), model_desc.model
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "n_dims,model_name",
-    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
-)
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_batch_embedding(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
@@ -112,15 +111,12 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
     embeddings = list(model.embed(docs, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (200, n_dims)
+    assert embeddings.shape == (len(docs), n_dims)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "n_dims,model_name",
-    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
-)
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_parallel_processing(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
@@ -135,7 +131,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
     embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (200, n_dims)
+    assert embeddings.shape == (len(docs), n_dims)
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
@@ -143,10 +139,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["BAAI/bge-small-en-v1.5"],
-)
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name, lazy_load=True)
@@ -163,3 +156,22 @@ 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)

tests/utils.py

@@ -3,7 +3,9 @@ import traceback
 
 from pathlib import Path
 from types import TracebackType
-from typing import Union, Callable, Any, Type
+from typing import Union, Callable, Any, Type, Optional
+
+from fastembed.common.model_description import BaseModelDescription
 
 
 def delete_model_cache(model_dir: Union[str, Path]) -> None:
@@ -35,3 +37,31 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
     if model_dir.exists():
         # todo: PermissionDenied is raised on blobs removal in Windows, with blobs > 2GB
         shutil.rmtree(model_dir, onerror=on_error)
+
+
+def should_test_model(
+    model_desc: BaseModelDescription,
+    autotest_model_name: str,
+    is_ci: Optional[str],
+    is_manual: bool,
+):
+    """Determine if a model should be tested based on environment
+
+    Tests can be run either in ci or locally.
+    Testing all models each time in ci is too long.
+    The testing scheme in ci and on a local machine are different, therefore, there are 3 possible scenarious.
+    1) Run lightweight tests in ci:
+        - test only one model that has been manually chosen as a representative for a certain class family
+    2) Run heavyweight (manual) tests in ci:
+        - test all models
+        Running tests in ci each time is too expensive, however, it's fine to run it one time with a manual dispatch
+    3) Run tests locally:
+        - test all models, which are not too heavy, since network speed might be a bottleneck
+
+    """
+    if not is_ci:
+        if model_desc.size_in_GB > 1:
+            return False
+    elif not is_manual and model_desc.model != autotest_model_name:
+        return False
+    return True