Merge branch 'main' into mean_pooling

.github/ISSUE_TEMPLATE/bug-report.yml

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

.github/workflows/ci.yml

@@ -21,5 +21,5 @@ jobs:
           path: .cache
           restore-keys: |
             mkdocs-material-
-      - run: pip install mkdocs-material mkdocstrings==0.27.0 pillow cairosvg mknotebooks
+      - run: pip install mkdocs-material mkdocstrings pillow cairosvg mknotebooks
       - run: mkdocs gh-deploy --force

.github/workflows/python-publish.yml

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

.github/workflows/python-tests.yml

@@ -1,11 +1,9 @@
 name: Tests
-run-name: Tests (gpu)
 
 on:
-  pull_request:
+  push:
     branches: [ master, main, gpu ]
-  workflow_dispatch:
-
+  pull_request:
 
 env:
   CARGO_TERM_COLOR: always
@@ -16,12 +14,15 @@ jobs:
     strategy:
       matrix:
         python-version:
+          - '3.9.x'
           - '3.10.x'
           - '3.11.x'
           - '3.12.x'
           - '3.13.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -40,7 +41,5 @@ jobs:
           poetry install --no-interaction --no-ansi --without dev,docs
 
       - name: Run pytest
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          poetry run pytest
+          poetry run pytest 

.github/workflows/type-checkers.yml

@@ -1,38 +0,0 @@
-name: type-checkers
-
-on: [push]
-
-jobs:
-  build:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      fail-fast: true
-      matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
-        os: [ubuntu-latest]
-
-    name: Python ${{ matrix.python-version }} test
-
-    steps:
-      - uses: actions/checkout@v1
-
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip poetry
-          poetry install --no-interaction --no-ansi --without dev,docs,test
-
-      - name: mypy
-        run: |
-            poetry run mypy fastembed \
-              --disallow-incomplete-defs \
-              --disallow-untyped-defs \
-              --disable-error-code=import-untyped
-
-      - name: pyright
-        run: |
-          poetry run pyright tests/type_stub.py

NOTICE

@@ -12,11 +12,5 @@ This distribution includes the following Jina AI models, each with its respectiv
 
 These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
 
-This distribution includes the following Google models, each with its respective license:
-- vidore/colpali-v1.3
-    - License: gemma
-
-Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms
-
 Additional Notes:
 This project also includes third-party libraries with their respective licenses. Please refer to the documentation of each library for details regarding its usage and licensing terms.

README.md

@@ -63,23 +63,6 @@ embeddings = list(model.embed(documents))
 
 ```
 
-Dense text embedding can also be extended with models which are not in the list of supported models.
-
-```python
-from fastembed import TextEmbedding
-from fastembed.common.model_description import PoolingType, ModelSource
-
-TextEmbedding.add_custom_model(
-    model="intfloat/multilingual-e5-small",
-    pooling=PoolingType.MEAN,
-    normalization=True,
-    sources=ModelSource(hf="intfloat/multilingual-e5-small"),  # can be used with an `url` to load files from a private storage
-    dim=384,
-    model_file="onnx/model.onnx",  # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
-)
-model = TextEmbedding(model_name="intfloat/multilingual-e5-small")
-embeddings = list(model.embed(documents))
-```
 
 
 ### 🔱 Sparse text embeddings
@@ -154,27 +137,6 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
-### Late interaction multimodal models (ColPali)
-
-```python
-from fastembed import LateInteractionMultimodalEmbedding
-
-doc_images = [
-    "./path/to/qdrant_pdf_doc_1_screenshot.jpg",
-    "./path/to/colpali_pdf_doc_2_screenshot.jpg",
-]
-
-query = "What is Qdrant?"
-
-model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
-doc_images_embeddings = list(model.embed_image(doc_images))
-# shape (2, 1030, 128)
-# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
-query_embedding = model.embed_text(query)
-# shape (1, 20, 128)
-# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
-```
-
 ### 🔄 Rerankers
 ```python
 from fastembed.rerank.cross_encoder import TextCrossEncoder
@@ -190,23 +152,6 @@ 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.
@@ -246,36 +191,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, models
+from qdrant_client import QdrantClient
 
 # Initialize the client
 client = QdrantClient("localhost", port=6333) # For production
-# client = QdrantClient(":memory:") # For experimentation
+# client = QdrantClient(":memory:") # For small experiments
 
-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"},
+# 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"},
 ]
-docs = [models.Document(text=data["document"], model=model_name) for data in payload]
 ids = [42, 2]
 
-client.create_collection(
-    "demo_collection",
-    vectors_config=models.VectorParams(
-        size=client.get_embedding_size(model_name), distance=models.Distance.COSINE)
+# 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.upload_collection(
+search_result = client.query(
     collection_name="demo_collection",
-    vectors=docs,
-    ids=ids,
-    payload=payload,
+    query_text="This is a query document"
 )
-
-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/__init__.py

@@ -2,7 +2,6 @@ import importlib.metadata
 
 from fastembed.image import ImageEmbedding
 from fastembed.late_interaction import LateInteractionTextEmbedding
-from fastembed.late_interaction_multimodal import LateInteractionMultimodalEmbedding
 from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
 from fastembed.text import TextEmbedding
 
@@ -18,5 +17,4 @@ __all__ = [
     "SparseEmbedding",
     "ImageEmbedding",
     "LateInteractionTextEmbedding",
-    "LateInteractionMultimodalEmbedding",
 ]

fastembed/common/__init__.py

@@ -1,3 +1,3 @@
-from fastembed.common.types import ImageInput, OnnxProvider, PathInput
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput, PilInput
 
-__all__ = ["OnnxProvider", "ImageInput", "PathInput"]
+__all__ = ["OnnxProvider", "ImageInput", "PathInput", "PilInput"]

fastembed/common/model_description.py

@@ -1,52 +0,0 @@
-from dataclasses import dataclass, field
-from enum import Enum
-from typing import Any
-
-
-@dataclass(frozen=True)
-class ModelSource:
-    hf: str | None = None
-    url: str | None = 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:
-            raise ValueError(
-                f"At least one source should be set, current sources: hf={self.hf}, url={self.url}"
-            )
-
-
-@dataclass(frozen=True)
-class BaseModelDescription:
-    model: str
-    sources: ModelSource
-    model_file: str
-    description: str
-    license: str
-    size_in_GB: float
-    additional_files: list[str] = field(default_factory=list)
-
-
-@dataclass(frozen=True)
-class DenseModelDescription(BaseModelDescription):
-    dim: int | None = None
-    tasks: dict[str, Any] | None = field(default_factory=dict)
-
-    def __post_init__(self) -> None:
-        assert self.dim is not None, "dim is required for dense model description"
-
-
-@dataclass(frozen=True)
-class SparseModelDescription(BaseModelDescription):
-    requires_idf: bool | None = None
-    vocab_size: int | None = None
-
-
-class PoolingType(str, Enum):
-    CLS = "CLS"
-    MEAN = "MEAN"
-    DISABLED = "DISABLED"

fastembed/common/model_management.py

@@ -3,9 +3,8 @@ import time
 import json
 import shutil
 import tarfile
-from copy import deepcopy
 from pathlib import Path
-from typing import Any, TypeVar, Generic
+from typing import Any, Optional
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -17,12 +16,9 @@ from huggingface_hub.utils import (
 )
 from loguru import logger
 from tqdm import tqdm
-from fastembed.common.model_description import BaseModelDescription
-
-T = TypeVar("T", bound=BaseModelDescription)
 
 
-class ModelManagement(Generic[T]):
+class ModelManagement:
     METADATA_FILE = "files_metadata.json"
 
     @classmethod
@@ -30,41 +26,12 @@ class ModelManagement(Generic[T]):
         """Lists the supported models.
 
         Returns:
-            list[T]: A list of dictionaries containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         raise NotImplementedError()
 
     @classmethod
-    def add_custom_model(
-        cls,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
-        """Add a custom model to the existing embedding classes based on the passed model descriptions
-
-        Model description dict should contain the fields same as in one of the model descriptions presented
-         in fastembed.common.model_description
-
-         E.g. for BaseModelDescription:
-              model: str
-              sources: ModelSource
-              model_file: str
-              description: str
-              license: str
-              size_in_GB: float
-              additional_files: list[str]
-
-        Returns:
-            None
-        """
-        raise NotImplementedError()
-
-    @classmethod
-    def _list_supported_models(cls) -> list[T]:
-        raise NotImplementedError()
-
-    @classmethod
-    def _get_model_description(cls, model_name: str) -> T:
+    def _get_model_description(cls, model_name: str) -> dict[str, Any]:
         """
         Gets the model description from the model_name.
 
@@ -75,10 +42,10 @@ class ModelManagement(Generic[T]):
             ValueError: If the model_name is not supported.
 
         Returns:
-            T: The model description.
+            dict[str, Any]: The model description.
         """
-        for model in cls._list_supported_models():
-            if model_name.lower() == model.model.lower():
+        for model in cls.list_supported_models():
+            if model_name.lower() == model["model"].lower():
                 return model
 
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
@@ -147,6 +114,7 @@ class ModelManagement(Generic[T]):
             extra_patterns (list[str]): extra patterns to allow in the snapshot download, typically
                 includes the required model files.
             local_files_only (bool, optional): Whether to only use local files. Defaults to False.
+            specific_model_path (Optional[str], optional): The path to the model dir already pooled from external source
         Returns:
             Path: The path to the model directory.
         """
@@ -180,8 +148,8 @@ class ModelManagement(Generic[T]):
 
         def _collect_file_metadata(
             model_dir: Path, repo_files: list[RepoFile]
-        ) -> dict[str, dict[str, int | str]]:
-            meta: dict[str, dict[str, int | str]] = {}
+        ) -> dict[str, dict[str, int]]:
+            meta = {}
             file_info_map = {f.path: f for f in repo_files}
             for file_path in model_dir.rglob("*"):
                 if file_path.is_file() and file_path.name != cls.METADATA_FILE:
@@ -193,7 +161,7 @@ class ModelManagement(Generic[T]):
                         }
             return meta
 
-        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) -> None:
+        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int]]) -> None:
             try:
                 if not model_dir.exists():
                     model_dir.mkdir(parents=True, exist_ok=True)
@@ -223,6 +191,11 @@ class ModelManagement(Generic[T]):
                     logger.warning(
                         "Local file sizes do not match the metadata."
                     )  # do not raise, still make an attempt to load the model
+            else:
+                logger.warning(
+                    "Metadata file not found. Proceeding without checking local files."
+                )  # if users have downloaded models from hf manually, or they're updating from previous versions of
+                # fastembed
             result = snapshot_download(
                 repo_id=hf_source_repo,
                 allow_patterns=allow_patterns,
@@ -320,14 +293,9 @@ class ModelManagement(Generic[T]):
 
     @classmethod
     def retrieve_model_gcs(
-        cls,
-        model_name: str,
-        source_url: str,
-        cache_dir: str,
-        deprecated_tar_struct: bool = False,
-        local_files_only: bool = False,
+        cls, model_name: str, source_url: str, cache_dir: str, local_files_only: bool = False
     ) -> Path:
-        fast_model_name = f"{'fast-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
+        fast_model_name = f"fast-{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
@@ -369,12 +337,14 @@ class ModelManagement(Generic[T]):
         return model_dir
 
     @classmethod
-    def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: Any) -> Path:
+    def download_model(
+        cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs: Any
+    ) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
         Args:
-            model (T): The model description.
+            model (dict[str, Any]): The model description.
                 Example:
                 ```
                 {
@@ -395,48 +365,26 @@ class ModelManagement(Generic[T]):
             Path: The path to the downloaded model directory.
         """
         local_files_only = kwargs.get("local_files_only", False)
-        hf_offline = os.environ.get("HF_HUB_OFFLINE", "").strip().upper()
-        if not local_files_only and hf_offline in {"1", "TRUE", "YES", "ON"}:
-            local_files_only = True
-            kwargs["local_files_only"] = True
-        specific_model_path: str | None = kwargs.pop("specific_model_path", None)
+        specific_model_path: Optional[str] = kwargs.pop("specific_model_path", None)
         if specific_model_path:
             return Path(specific_model_path)
         retries = 1 if local_files_only else retries
-        hf_source = model.sources.hf
-        url_source = model.sources.url
-
-        extra_patterns = [model.model_file]
-        extra_patterns.extend(model.additional_files)
-
-        if hf_source:
-            try:
-                cache_kwargs = deepcopy(kwargs)
-                cache_kwargs["local_files_only"] = True
-                return Path(
-                    cls.download_files_from_huggingface(
-                        hf_source,
-                        cache_dir=cache_dir,
-                        extra_patterns=extra_patterns,
-                        **cache_kwargs,
-                    )
-                )
-            except Exception:
-                pass
-            finally:
-                enable_progress_bars()
+        hf_source = model.get("sources", {}).get("hf")
+        url_source = model.get("sources", {}).get("url")
 
         sleep = 3.0
         while retries > 0:
             retries -= 1
 
-            if hf_source and not local_files_only:
-                # we have already tried loading with `local_files_only=True` via hf and we failed
+            if hf_source:
+                extra_patterns = [model["model_file"]]
+                extra_patterns.extend(model.get("additional_files", []))
+
                 try:
                     return Path(
                         cls.download_files_from_huggingface(
                             hf_source,
-                            cache_dir=cache_dir,
+                            cache_dir=str(cache_dir),
                             extra_patterns=extra_patterns,
                             **kwargs,
                         )
@@ -452,10 +400,9 @@ class ModelManagement(Generic[T]):
             if url_source or local_files_only:
                 try:
                     return cls.retrieve_model_gcs(
-                        model.model,
-                        str(url_source),
+                        model["model"],
+                        url_source,
                         str(cache_dir),
-                        deprecated_tar_struct=model.sources.deprecated_tar_struct,
                         local_files_only=local_files_only,
                     )
                 except Exception:
@@ -464,12 +411,11 @@ class ModelManagement(Generic[T]):
 
             if local_files_only:
                 logger.error("Could not find model in cache_dir")
-                break
             else:
                 logger.error(
                     f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
                 )
-                time.sleep(sleep)
-                sleep *= 3
+            time.sleep(sleep)
+            sleep *= 3
 
-        raise ValueError(f"Could not load model {model.model} from any source.")
+        raise ValueError(f"Could not load model {model['model']} from any source.")

fastembed/common/onnx_model.py

@@ -1,15 +1,12 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Generic, Iterable, Sequence, Type, TypeVar
+from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort
 
-from numpy.typing import NDArray
-from tokenizers import Tokenizer
-
-from fastembed.common.types import OnnxProvider, NumpyArray, Device
+from fastembed.common.types import OnnxProvider
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -18,38 +15,26 @@ T = TypeVar("T")
 
 @dataclass
 class OnnxOutputContext:
-    model_output: NumpyArray
-    attention_mask: NDArray[np.int64] | None = None
-    input_ids: NDArray[np.int64] | None = None
-    metadata: dict[str, Any] | None = None
+    model_output: np.ndarray
+    attention_mask: Optional[np.ndarray] = None
+    input_ids: Optional[np.ndarray] = None
 
 
 class OnnxModel(Generic[T]):
-    EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)
-
     @classmethod
-    def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
-        """Post-process the ONNX model output to convert it into a usable format.
-
-        Args:
-            output (OnnxOutputContext): The raw output from the ONNX model.
-            **kwargs: Additional keyword arguments that may be needed by specific implementations.
-
-        Returns:
-            Iterable[T]: Post-processed output as an iterable of type T.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
-        self.model: ort.InferenceSession | None = None
-        self.tokenizer: Tokenizer | None = None
+        self.model = None
+        self.tokenizer = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -59,30 +44,24 @@ class OnnxModel(Generic[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        available_providers = ort.get_available_providers()
-        cuda_available = "CUDAExecutionProvider" in available_providers
-        explicit_cuda = cuda is True or cuda == Device.CUDA
 
-        if explicit_cuda and providers is not None:
+        if cuda and providers is not None:
             warnings.warn(
-                f"`cuda` and `providers` are mutually exclusive parameters, "
-                f"cuda: {cuda}, providers: {providers}. If you'd like to use providers, cuda should be one of "
-                f"[False, Device.CPU, Device.AUTO].",
+                f"`cuda` and `providers` are mutually exclusive parameters, cuda: {cuda}, providers: {providers}",
                 category=UserWarning,
                 stacklevel=6,
             )
 
         if providers is not None:
             onnx_providers = list(providers)
-        elif explicit_cuda or (cuda == Device.AUTO and cuda_available):
+        elif cuda:
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
@@ -90,7 +69,8 @@ class OnnxModel(Generic[T]):
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
-        requested_provider_names: list[str] = []
+        available_providers = ort.get_available_providers()
+        requested_provider_names = []
         for provider in onnx_providers:
             # check providers available
             provider_name = provider if isinstance(provider, str) else provider[0]
@@ -107,14 +87,10 @@ class OnnxModel(Generic[T]):
             so.intra_op_num_threads = threads
             so.inter_op_num_threads = threads
 
-        if extra_session_options is not None:
-            self.add_extra_session_options(so, extra_session_options)
-
         self.model = ort.InferenceSession(
             str(model_path), providers=onnx_providers, sess_options=so
         )
         if "CUDAExecutionProvider" in requested_provider_names:
-            assert self.model is not None
             current_providers = self.model.get_providers()
             if "CUDAExecutionProvider" not in current_providers:
                 warnings.warn(
@@ -124,38 +100,6 @@ class OnnxModel(Generic[T]):
                     RuntimeWarning,
                 )
 
-    @classmethod
-    def _select_exposed_session_options(cls, model_kwargs: dict[str, Any]) -> dict[str, Any]:
-        """A convenience method to select the exposed session options in models
-
-        Args:
-            model_kwargs (dict[str, Any]): The model kwargs.
-
-        Returns:
-            dict[str, Any]: a dict with filtered exposed session options.
-        """
-        return {k: v for k, v in model_kwargs.items() if k in cls.EXPOSED_SESSION_OPTIONS}
-
-    @classmethod
-    def add_extra_session_options(
-        cls, session_options: ort.SessionOptions, extra_options: dict[str, Any]
-    ) -> None:
-        """Add extra session options to the existing options object in-place
-
-        Args:
-            session_options (ort.SessionOptions): The existing session options object.
-            extra_options (dict[str, Any]): The extra session options available in cls.EXPOSED_SESSION_OPTIONS.
-
-        Returns:
-            None
-        """
-        for option in extra_options:
-            assert (
-                option in cls.EXPOSED_SESSION_OPTIONS
-            ), f"{option} is unknown or not exposed (exposed options: {cls.EXPOSED_SESSION_OPTIONS})"
-        if "enable_cpu_mem_arena" in extra_options:
-            session_options.enable_cpu_mem_arena = extra_options["enable_cpu_mem_arena"]
-
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
@@ -163,13 +107,13 @@ class OnnxModel(Generic[T]):
         raise NotImplementedError("Subclasses must implement this method")
 
 
-class EmbeddingWorker(Worker, Generic[T]):
+class EmbeddingWorker(Worker):
     def init_embedding(
         self,
         model_name: str,
         cache_dir: str,
         **kwargs: Any,
-    ) -> OnnxModel[T]:
+    ) -> OnnxModel:
         raise NotImplementedError()
 
     def __init__(
@@ -181,7 +125,7 @@ class EmbeddingWorker(Worker, Generic[T]):
         self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
     @classmethod
-    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:

fastembed/common/pooling.py

@@ -0,0 +1,38 @@
+import numpy as np
+from numpy import ufunc
+
+
+class LateInteractionPooler(object):
+    def __init__(self, agg_type="row", agg="mean"):
+        self.agg = agg
+        self.type = agg_type
+
+    def _pick_operation(self) -> ufunc:
+        if self.agg == "mean":
+            return np.mean
+        elif self.agg == "max":
+            return np.max
+        elif self.agg == "min":
+            return np.min
+        else:
+            raise NotImplementedError(
+                f"LateInteractionPooler only supports agg=mean,min,max, provided {self.agg}"
+            )
+
+    def pool(self, embeddings_batch) -> np.array:
+        if isinstance(embeddings_batch, np.ndarray) and len(embeddings_batch.shape) == 2:
+            embeddings_batch = [embeddings_batch]
+
+        if self.type == "row":
+            pooled_embedding = self.pool_row(embeddings_batch)
+        elif self.type == "col":
+            pooled_embedding = self.pool_col(embeddings_batch)
+        else:
+            raise ValueError("type must be 'row' or 'col'")
+        return pooled_embedding
+
+    def pool_row(self, embeddings_batch) -> np.array:
+        return self._pick_operation()(embeddings_batch, axis=-1)
+
+    def pool_col(self, embeddings_batch) -> np.array:
+        return self._pick_operation()(embeddings_batch, axis=-2)

fastembed/common/preprocessor_utils.py

@@ -1,5 +1,4 @@
 import json
-from typing import Any
 from pathlib import Path
 
 from tokenizers import AddedToken, Tokenizer
@@ -7,7 +6,7 @@ from tokenizers import AddedToken, Tokenizer
 from fastembed.image.transform.operators import Compose
 
 
-def load_special_tokens(model_dir: Path) -> dict[str, Any]:
+def load_special_tokens(model_dir: Path) -> dict:
     tokens_map_path = model_dir / "special_tokens_map.json"
     if not tokens_map_path.exists():
         raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
@@ -18,7 +17,7 @@ def load_special_tokens(model_dir: Path) -> dict[str, Any]:
     return tokens_map
 
 
-def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
+def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict]:
     config_path = model_dir / "config.json"
     if not config_path.exists():
         raise ValueError(f"Could not find config.json in {model_dir}")
@@ -36,9 +35,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:
@@ -50,10 +49,9 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=max_context)
-    if not tokenizer.padding:
-        tokenizer.enable_padding(
-            pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
-        )
+    tokenizer.enable_padding(
+        pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
+    )
 
     for token in tokens_map.values():
         if isinstance(token, str):
@@ -61,7 +59,7 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
         elif isinstance(token, dict):
             tokenizer.add_special_tokens([AddedToken(**token)])
 
-    special_token_to_id: dict[str, int] = {}
+    special_token_to_id = {}
 
     for token in tokens_map.values():
         if isinstance(token, str):

fastembed/common/types.py

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

fastembed/common/utils.py

@@ -5,17 +5,14 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, TypeVar
+from typing import Iterable, Optional, TypeVar
 
 import numpy as np
-from numpy.typing import NDArray
-
-from fastembed.common.types import NumpyArray
 
 T = TypeVar("T")
 
 
-def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
+def normalize(input_array: np.ndarray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> np.ndarray:
     # Calculate the Lp norm along the specified dimension
     norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
     norm = np.maximum(norm, eps)  # Avoid division by zero
@@ -23,15 +20,6 @@ def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e
     return normalized_array
 
 
-def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
-    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
-    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
-    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
-    sum_mask = np.sum(input_mask_expanded, axis=1)
-    pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
-    return pooled_embeddings
-
-
 def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))
@@ -45,7 +33,7 @@ def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
         yield b
 
 
-def define_cache_dir(cache_dir: str | None = None) -> Path:
+def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
     """
     Define the cache directory for fastembed
     """

fastembed/embedding.py

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

fastembed/image/image_embedding.py

@@ -1,11 +1,10 @@
-from typing import Any, Iterable, Sequence, Type
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
 
-from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
-from fastembed.common.model_description import DenseModelDescription
 
 
 class ImageEmbedding(ImageEmbeddingBase):
@@ -36,30 +35,26 @@ class ImageEmbedding(ImageEmbeddingBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -77,49 +72,16 @@ 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: int | None = 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: ImageInput | Iterable[ImageInput],
+        images: ImageInput,
         batch_size: int = 16,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
-        Encode a list of images into list of embeddings.
+        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:
             images: Iterator of image paths or single image path to embed

fastembed/image/image_embedding_base.py

@@ -1,32 +1,31 @@
-from typing import Iterable, Any
+from typing import Iterable, Optional, Any
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 from fastembed.common.types import ImageInput
 
 
-class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
+class ImageEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: int | None = None
 
     def embed(
         self,
-        images: ImageInput | Iterable[ImageInput],
+        images: ImageInput,
         batch_size: int = 16,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
         Embeds a list of images into a list of embeddings.
 
@@ -40,16 +39,6 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.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,77 +1,83 @@
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type
 
-
-from fastembed.common.types import NumpyArray, Device
+import numpy as np
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
 
-from fastembed.common.model_description import DenseModelDescription, ModelSource
-
-supported_onnx_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="Qdrant/clip-ViT-B-32-vision",
-        dim=512,
-        description="Image embeddings, Multimodal (text&image), 2021 year",
-        license="mit",
-        size_in_GB=0.34,
-        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-vision"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="Qdrant/resnet50-onnx",
-        dim=2048,
-        description="Image embeddings, Unimodal (image), 2016 year",
-        license="apache-2.0",
-        size_in_GB=0.1,
-        sources=ModelSource(hf="Qdrant/resnet50-onnx"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="Qdrant/Unicom-ViT-B-16",
-        dim=768,
-        description="Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
-        license="apache-2.0",
-        size_in_GB=0.82,
-        sources=ModelSource(hf="Qdrant/Unicom-ViT-B-16"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="Qdrant/Unicom-ViT-B-32",
-        dim=512,
-        description="Image embeddings, Multimodal (text&image), 2023 year",
-        license="apache-2.0",
-        size_in_GB=0.48,
-        sources=ModelSource(hf="Qdrant/Unicom-ViT-B-32"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-clip-v1",
-        dim=768,
-        description="Image embeddings, Multimodal (text&image), 2024 year",
-        license="apache-2.0",
-        size_in_GB=0.34,
-        sources=ModelSource(hf="jinaai/jina-clip-v1"),
-        model_file="onnx/vision_model.onnx",
-    ),
+supported_onnx_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-vision",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2021 year",
+        "license": "mit",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-vision",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/resnet50-onnx",
+        "dim": 2048,
+        "description": "Image embeddings, Unimodal (image), 2016 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.1,
+        "sources": {
+            "hf": "Qdrant/resnet50-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-16",
+        "dim": 768,
+        "description": "Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.82,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-16",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-32",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.48,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-32",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "jinaai/jina-clip-v1",
+        "dim": 768,
+        "description": "Image embeddings, Multimodal (text&image), 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "jinaai/jina-clip-v1",
+        },
+        "model_file": "onnx/vision_model.onnx",
+    },
 ]
 
 
-class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
+class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
     def __init__(
         self,
         model_name: str,
-
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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: int | None = None,
-        specific_model_path: str | None = None,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -83,11 +89,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
+            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` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
+                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.
@@ -100,27 +105,26 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = 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.cache_dir = define_cache_dir(cache_dir)
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -132,31 +136,30 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         """
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
         )
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
     def embed(
         self,
-        images: ImageInput | Iterable[ImageInput],
+        images: ImageInput,
         batch_size: int = 16,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of images into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -182,32 +185,27 @@ 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,
-            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
         return OnnxImageEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
 
         return onnx_input
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
-        return normalize(output.model_output)
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        return normalize(output.model_output).astype(np.float32)
 
 
-class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
         return OnnxImageEmbedding(
             model_name=model_name,

fastembed/image/onnx_image_model.py

@@ -2,13 +2,11 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
 from PIL import Image
 
-from fastembed.image.transform.operators import Compose
-from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
@@ -20,28 +18,19 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 class OnnxImageModel(OnnxModel[T]):
     @classmethod
-    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
-        """Post-process the ONNX model output to convert it into a usable format.
-
-        Args:
-            output (OnnxOutputContext): The raw output from the ONNX model.
-            **kwargs: Additional keyword arguments that may be needed by specific implementations.
-
-        Returns:
-            Iterable[T]: Post-processed output as an iterable of type T.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def __init__(self) -> None:
+    def __init__(self):
         super().__init__()
-        self.processor: Compose | None = None
+        self.processor = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -51,11 +40,10 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -64,30 +52,25 @@ class OnnxImageModel(OnnxModel[T]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
-            extra_session_options=extra_session_options,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _build_onnx_input(self, encoded: NumpyArray) -> dict[str, NumpyArray]:
-        input_name = self.model.get_inputs()[0].name  # type: ignore[union-attr]
-        return {input_name: encoded}
+    def _build_onnx_input(self, encoded: np.ndarray) -> dict[str, np.ndarray]:
+        return {node.name: encoded for node in self.model.get_inputs()}
 
     def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack() as stack:
+        with contextlib.ExitStack():
             image_files = [
-                stack.enter_context(Image.open(image))
-                if not isinstance(image, Image.Image)
-                else image
+                Image.open(image) if not isinstance(image, Image.Image) else image
                 for image in images
             ]
-            assert self.processor is not None, "Processor is not initialized"
-            encoded = np.array(self.processor(image_files))
+            encoded = self.processor(image_files)
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
-        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+        model_output = self.model.run(None, onnx_input)
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -95,15 +78,12 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: ImageInput | Iterable[ImageInput],
+        images: ImageInput,
         batch_size: int = 256,
-        parallel: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        local_files_only: bool = False,
-        specific_model_path: str | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -120,7 +100,7 @@ class OnnxImageModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_output(self.onnx_embed(batch), **kwargs)
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -130,14 +110,9 @@ 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,
             }
 
-            if extra_session_options is not None:
-                params.update(extra_session_options)
-
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),
@@ -146,10 +121,10 @@ class OnnxImageModel(OnnxModel[T]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
-                yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
+                yield from self._post_process_onnx_output(batch)
 
 
-class ImageEmbeddingWorker(EmbeddingWorker[T]):
+class ImageEmbeddingWorker(EmbeddingWorker):
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             embeddings = self.model.onnx_embed(batch)

fastembed/image/transform/functional.py

@@ -1,8 +1,8 @@
+from typing import Sized, Union
+
 import numpy as np
 from PIL import Image
 
-from fastembed.common.types import NumpyArray
-
 
 def convert_to_rgb(image: Image.Image) -> Image.Image:
     if image.mode == "RGB":
@@ -13,9 +13,9 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Image.Image | NumpyArray,
+    image: Union[Image.Image, np.ndarray],
     size: tuple[int, int],
-) -> NumpyArray:
+) -> np.ndarray:
     if isinstance(image, np.ndarray):
         _, orig_height, orig_width = image.shape
     else:
@@ -40,7 +40,7 @@ def center_crop(
     new_height = max(crop_height, orig_height)
     new_width = max(crop_width, orig_width)
     new_shape = image.shape[:-2] + (new_height, new_width)
-    new_image = np.zeros_like(image, shape=new_shape, dtype=np.float32)
+    new_image = np.zeros_like(image, shape=new_shape)
 
     top_pad = (new_height - orig_height) // 2
     bottom_pad = top_pad + orig_height
@@ -61,41 +61,44 @@ def center_crop(
 
 
 def normalize(
-    image: NumpyArray,
-    mean: float | list[float],
-    std: float | list[float],
-) -> NumpyArray:
+    image: np.ndarray,
+    mean: Union[float, np.ndarray],
+    std: Union[float, np.ndarray],
+) -> np.ndarray:
+    if not isinstance(image, np.ndarray):
+        raise ValueError("image must be a numpy array")
+
     num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
 
     if not np.issubdtype(image.dtype, np.floating):
         image = image.astype(np.float32)
 
-    mean_list = mean if isinstance(mean, list) else [mean] * num_channels
+    if isinstance(mean, Sized):
+        if len(mean) != num_channels:
+            raise ValueError(
+                f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}"
+            )
+    else:
+        mean = [mean] * num_channels
+    mean = np.array(mean, dtype=image.dtype)
 
-    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_list)}"
-        )
+    if isinstance(std, Sized):
+        if len(std) != num_channels:
+            raise ValueError(
+                f"std must have {num_channels} elements if it is an iterable, got {len(std)}"
+            )
+    else:
+        std = [std] * num_channels
+    std = np.array(std, dtype=image.dtype)
 
-    mean_arr = np.array(mean_list, dtype=np.float32)
-
-    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_list)}"
-        )
-
-    std_arr = np.array(std_list, dtype=np.float32)
-
-    image_upd = ((image.T - mean_arr) / std_arr).T
-    return image_upd
+    image = ((image.T - mean) / std).T
+    return image
 
 
 def resize(
     image: Image.Image,
-    size: int | tuple[int, int],
-    resample: int | Image.Resampling = Image.Resampling.BILINEAR,
+    size: Union[int, tuple[int, int]],
+    resample: Union[int, Image.Resampling] = Image.Resampling.BILINEAR,
 ) -> Image.Image:
     if isinstance(size, tuple):
         return image.resize(size, resample)
@@ -111,11 +114,11 @@ def resize(
     return image.resize(new_size, resample)
 
 
-def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyArray:
+def rescale(image: np.ndarray, scale: float, dtype: type = np.float32) -> np.ndarray:
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Image.Image | NumpyArray) -> NumpyArray:
+def pil2ndarray(image: Union[Image.Image, np.ndarray]) -> np.ndarray:
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image
@@ -124,7 +127,7 @@ def pil2ndarray(image: Image.Image | NumpyArray) -> NumpyArray:
 def pad2square(
     image: Image.Image,
     size: int,
-    fill_color: str | int | tuple[int, ...] = 0,
+    fill_color: Union[str, int, tuple[int, ...]] = 0,
 ) -> Image.Image:
     height, width = image.height, image.width
 
@@ -145,77 +148,3 @@ def pad2square(
     new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
     new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
     return new_image
-
-
-def resize_longest_edge(
-    image: Image.Image,
-    max_size: int,
-    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
-) -> Image.Image:
-    height, width = image.height, image.width
-    aspect_ratio = width / height
-
-    if width >= height:
-        # Width is longer
-        new_width = max_size
-        new_height = int(new_width / aspect_ratio)
-    else:
-        # Height is longer
-        new_height = max_size
-        new_width = int(new_height * aspect_ratio)
-
-    # Ensure even dimensions
-    if new_height % 2 != 0:
-        new_height += 1
-    if new_width % 2 != 0:
-        new_width += 1
-
-    return image.resize((new_width, new_height), resample)
-
-
-def crop_ndarray(
-    image: NumpyArray,
-    x1: int,
-    y1: int,
-    x2: int,
-    y2: int,
-    channel_first: bool = True,
-) -> NumpyArray:
-    if channel_first:
-        # (C, H, W) format
-        return image[:, y1:y2, x1:x2]
-    else:
-        # (H, W, C) format
-        return image[y1:y2, x1:x2, :]
-
-
-def resize_ndarray(
-    image: NumpyArray,
-    size: tuple[int, int],
-    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
-    channel_first: bool = True,
-) -> NumpyArray:
-    # Convert to PIL-friendly format (H, W, C)
-    if channel_first:
-        img_hwc = image.transpose((1, 2, 0))
-    else:
-        img_hwc = image
-
-    # Handle different dtypes
-    if img_hwc.dtype == np.float32 or img_hwc.dtype == np.float64:
-        # Assume normalized, scale to 0-255 for PIL
-        img_hwc_scaled = (img_hwc * 255).astype(np.uint8)
-        pil_img = Image.fromarray(img_hwc_scaled, mode="RGB")
-        resized = pil_img.resize(size, resample)
-        result = np.array(resized).astype(np.float32) / 255.0
-    else:
-        # uint8 or similar
-        pil_img = Image.fromarray(img_hwc.astype(np.uint8), mode="RGB")
-        resized = pil_img.resize(size, resample)
-        result = np.array(resized)
-
-    # Convert back to original format
-    if channel_first:
-        result = result.transpose((2, 0, 1))
-
-    return result

fastembed/image/transform/operators.py

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

fastembed/late_interaction/colbert.py

@@ -1,83 +1,78 @@
 import string
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from tokenizers import Encoding, Tokenizer
+from tokenizers import Encoding
 
-from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.types import NumpyArray, Device
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.common.utils import define_cache_dir
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_colbert_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="colbert-ir/colbertv2.0",
-        dim=128,
-        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2023 year",
-        license="mit",
-        size_in_GB=0.44,
-        sources=ModelSource(hf="colbert-ir/colbertv2.0"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="answerdotai/answerai-colbert-small-v1",
-        dim=96,
-        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2024 year",
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="answerdotai/answerai-colbert-small-v1"),
-        model_file="vespa_colbert.onnx",
-    ),
+
+supported_colbert_models = [
+    {
+        "model": "colbert-ir/colbertv2.0",
+        "dim": 128,
+        "description": "Late interaction model",
+        "license": "mit",
+        "size_in_GB": 0.44,
+        "sources": {
+            "hf": "colbert-ir/colbertv2.0",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "answerdotai/answerai-colbert-small-v1",
+        "dim": 96,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "answerdotai/answerai-colbert-small-v1",
+        },
+        "model_file": "vespa_colbert.onnx",
+    },
 ]
 
 
-class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
+class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
     QUERY_MARKER_TOKEN_ID = 1
     DOCUMENT_MARKER_TOKEN_ID = 2
     MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
-        self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+        self, output: OnnxOutputContext, is_doc: bool = True
+    ) -> Iterable[np.ndarray]:
         if not is_doc:
-            for embedding in output.model_output:
-                yield embedding
-        else:
-            if output.input_ids is None or output.attention_mask is None:
-                raise ValueError(
-                    "input_ids and attention_mask must be provided for document post-processing"
-                )
+            return output.model_output.astype(np.float32)
 
-            for i, token_sequence in enumerate(output.input_ids):
-                for j, token_id in enumerate(token_sequence):  # type: ignore
-                    if token_id in self.skip_list or token_id == self.pad_token_id:
-                        output.attention_mask[i, j] = 0
+        if output.input_ids is None or output.attention_mask is None:
+            raise ValueError(
+                "input_ids and attention_mask must be provided for document post-processing"
+            )
 
-            output.model_output *= np.expand_dims(output.attention_mask, 2)
-            norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
-            norm_clamped = np.maximum(norm, 1e-12)
-            output.model_output /= norm_clamped
+        for i, token_sequence in enumerate(output.input_ids):
+            for j, token_id in enumerate(token_sequence):
+                if token_id in self.skip_list or token_id == self.pad_token_id:
+                    output.attention_mask[i, j] = 0
 
-            for embedding, attention_mask in zip(output.model_output, output.attention_mask):
-                yield embedding[attention_mask == 1]
+        output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32)
+        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)
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
-        onnx_input["input_ids"] = np.insert(
-            onnx_input["input_ids"].astype(np.int64), 1, marker_token, axis=1
-        )
-        onnx_input["attention_mask"] = np.insert(
-            onnx_input["attention_mask"].astype(np.int64), 1, 1, axis=1
-        )
+        onnx_input["input_ids"] = np.insert(onnx_input["input_ids"], 1, marker_token, axis=1)
+        onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
         return onnx_input
 
     def tokenize(self, documents: list[str], is_doc: bool = True, **kwargs: Any) -> list[Encoding]:
@@ -88,66 +83,48 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     def _tokenize_query(self, query: str) -> list[Encoding]:
-        assert self.query_tokenizer is not None
-        encoded = self.query_tokenizer.encode_batch([query])
+        encoded = self.tokenizer.encode_batch([query])
+        # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
+        if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
+            prev_padding = None
+            if self.tokenizer.padding:
+                prev_padding = self.tokenizer.padding
+            self.tokenizer.enable_padding(
+                pad_token=self.MASK_TOKEN,
+                pad_id=self.mask_token_id,
+                length=self.MIN_QUERY_LENGTH,
+            )
+            encoded = self.tokenizer.encode_batch([query])
+            if prev_padding is None:
+                self.tokenizer.no_padding()
+            else:
+                self.tokenizer.enable_padding(**prev_padding)
         return encoded
 
     def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
-        encoded = self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+        encoded = self.tokenizer.encode_batch(documents)
         return encoded
 
-    def token_count(
-        self,
-        texts: str | Iterable[str],
-        batch_size: int = 1024,
-        is_doc: bool = True,
-        include_extension: bool = False,
-        **kwargs: Any,
-    ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-        token_num = 0
-        texts = [texts] if isinstance(texts, str) else texts
-        tokenizer = self.tokenizer if is_doc else self.query_tokenizer
-        assert tokenizer is not None
-        for batch in iter_batch(texts, batch_size):
-            for tokens in tokenizer.encode_batch(batch):
-                if is_doc:
-                    token_num += sum(tokens.attention_mask)
-                else:
-                    attend_count = sum(tokens.attention_mask)
-                    if include_extension:
-                        token_num += max(attend_count, self.MIN_QUERY_LENGTH)
-
-                    else:
-                        token_num += attend_count
-            if include_extension:
-                token_num += len(
-                    batch
-                )  # add 1 for each cls.DOC_MARKER_TOKEN_ID or cls.QUERY_MARKER_TOKEN_ID
-
-        return token_num
-
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_colbert_models
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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: int | None = None,
-        specific_model_path: str | None = None,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -159,11 +136,10 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
+            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` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
+                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.
@@ -176,34 +152,31 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
-        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
-        self.mask_token_id: int | None = None
-        self.pad_token_id: int | None = None
-        self.skip_list: set[int] = set()
-
-        self.query_tokenizer: Tokenizer | None = None
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.skip_list = set()
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -211,16 +184,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
         )
-        self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir)
-
-        assert self.tokenizer is not None
         self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
         self.skip_list = {
@@ -230,20 +199,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         current_max_length = self.tokenizer.truncation["max_length"]
         # ensure not to overflow after adding document-marker
         self.tokenizer.enable_truncation(max_length=current_max_length - 1)
-        self.query_tokenizer.enable_truncation(max_length=current_max_length - 1)
-        self.query_tokenizer.enable_padding(
-            pad_token=self.MASK_TOKEN,
-            pad_id=self.mask_token_id,
-            length=self.MIN_QUERY_LENGTH,
-        )
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -268,13 +231,10 @@ 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,
-            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
         if isinstance(query, str):
             query = [query]
 
@@ -287,11 +247,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return ColbertEmbeddingWorker
 
 
-class ColbertEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+class ColbertEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
         return Colbert(
             model_name=model_name,

fastembed/late_interaction/jina_colbert.py

@@ -1,20 +1,24 @@
 from typing import Any, Type
 
-from fastembed.common.types import NumpyArray
-from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
+import numpy as np
 
-supported_jina_colbert_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="jinaai/jina-colbert-v2",
-        dim=128,
-        description="New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
-        license="cc-by-nc-4.0",
-        size_in_GB=2.24,
-        sources=ModelSource(hf="jinaai/jina-colbert-v2"),
-        model_file="onnx/model.onnx",
-        additional_files=["onnx/model.onnx_data"],
-    )
+from fastembed.late_interaction.colbert import Colbert
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+supported_jina_colbert_models = [
+    {
+        "model": "jinaai/jina-colbert-v2",
+        "dim": 128,
+        "description": "New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
+        "license": "cc-by-nc-4.0",
+        "size_in_GB": 2.24,
+        "sources": {
+            "hf": "jinaai/jina-colbert-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "additional_files": ["onnx/model.onnx_data"],
+    },
 ]
 
 
@@ -25,21 +29,21 @@ class JinaColbert(Colbert):
     MASK_TOKEN = "<mask>"
 
     @classmethod
-    def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return JinaColbertEmbeddingWorker
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_jina_colbert_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
 
         # the attention mask for jina-colbert-v2 is always 1 in queries
@@ -48,7 +52,7 @@ class JinaColbert(Colbert):
         return onnx_input
 
 
-class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):
+class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
         return JinaColbert(
             model_name=model_name,

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -1,34 +1,33 @@
-from typing import Iterable, Any
+from typing import Iterable, Optional, Union, Any
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
+class LateInteractionTextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: int | None = None
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -37,13 +36,13 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -51,30 +50,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        else:
+        if isinstance(query, Iterable):
             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")
-
-    def token_count(
-        self,
-        texts: str | Iterable[str],
-        batch_size: int = 1024,
-        **kwargs: Any,
-    ) -> int:
-        """Returns the number of tokens in the texts."""
-        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,8 +1,7 @@
-from typing import Any, Iterable, Sequence, Type
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type, Union
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray, Device
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -39,30 +38,26 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -80,47 +75,13 @@ 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: int | None = 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: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -138,7 +99,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -146,35 +107,8 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.model.query_embed(query, **kwargs)
-
-    def token_count(
-        self,
-        texts: str | Iterable[str],
-        batch_size: int = 1024,
-        is_doc: bool = True,
-        include_extension: bool = False,
-        **kwargs: Any,
-    ) -> int:
-        """Returns the number of tokens in the texts.
-
-        Args:
-            texts (str | Iterable[str]): The list of texts to embed.
-            batch_size (int): Batch size for encoding
-            is_doc (bool): Whether the texts are documents (disable embedding a query with include_mask=True).
-            include_extension (bool): Turn on to count DOC / QUERY marker tokens, and [MASK] token in query mode.
-
-        Returns:
-            int: Sum of number of tokens in the texts.
-        """
-        return self.model.token_count(
-            texts,
-            batch_size=batch_size,
-            is_doc=is_doc,
-            include_extension=include_extension,
-            **kwargs,
-        )

fastembed/late_interaction/token_embeddings.py

@@ -1,83 +0,0 @@
-from dataclasses import asdict
-from typing import Iterable, 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: str | Iterable[str],
-        batch_size: int = 256,
-        parallel: int | None = 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/__init__.py

@@ -1,5 +0,0 @@
-from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding import (
-    LateInteractionMultimodalEmbedding,
-)
-
-__all__ = ["LateInteractionMultimodalEmbedding"]

fastembed/late_interaction_multimodal/colmodernvbert.py

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

fastembed/late_interaction_multimodal/colpali.py

@@ -1,327 +0,0 @@
-from typing import Any, Iterable, Sequence, Type
-
-import numpy as np
-from tokenizers import Encoding
-
-from fastembed.common import OnnxProvider, ImageInput
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray, Device
-from fastembed.common.utils import define_cache_dir, iter_batch
-from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
-    LateInteractionMultimodalEmbeddingBase,
-)
-from fastembed.late_interaction_multimodal.onnx_multimodal_model import (
-    OnnxMultimodalModel,
-    TextEmbeddingWorker,
-    ImageEmbeddingWorker,
-)
-from fastembed.common.model_description import DenseModelDescription, ModelSource
-
-supported_colpali_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="Qdrant/colpali-v1.3-fp16",
-        dim=128,
-        description="Text embeddings, Multimodal (text&image), English, 50 tokens query length truncation, 2024.",
-        license="mit",
-        size_in_GB=6.5,
-        sources=ModelSource(hf="Qdrant/colpali-v1.3-fp16"),
-        additional_files=["model.onnx_data"],
-        model_file="model.onnx",
-    ),
-]
-
-
-class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyArray]):
-    QUERY_PREFIX = "Query: "
-    BOS_TOKEN = "<s>"
-    PAD_TOKEN = "<pad>"
-    QUERY_MARKER_TOKEN_ID = [2, 5098]
-    IMAGE_PLACEHOLDER_SIZE = (3, 448, 448)
-    EMPTY_TEXT_PLACEHOLDER = np.array(
-        [257152] * 1024 + [2, 50721, 573, 2416, 235265, 108]
-    )  # This is a tokenization of '<image>' * 1024 + '<bos>Describe the image.\n' line which is used as placeholder
-    # while processing an image
-    EVEN_ATTENTION_MASK = np.array([1] * 1030)
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        lazy_load: bool = False,
-        device_id: int | None = None,
-        specific_model_path: str | None = None,
-        **kwargs: Any,
-    ):
-        """
-        Args:
-            model_name (str): The name of the model to use.
-            cache_dir (str, optional): The path to the cache directory.
-                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
-                                       Defaults to `fastembed_cache` in the system's temp directory.
-            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
-            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
-                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
-            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
-            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
-                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
-            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
-
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-
-        super().__init__(model_name, cache_dir, threads, **kwargs)
-        self.providers = providers
-        self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
-
-        # List of device ids, that can be used for data parallel processing in workers
-        self.device_ids = device_ids
-        self.cuda = cuda
-
-        # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
-        if device_id is not None:
-            self.device_id = device_id
-        elif self.device_ids is not None:
-            self.device_id = self.device_ids[0]
-
-        self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
-
-        self._specific_model_path = specific_model_path
-        self._model_dir = self.download_model(
-            self.model_description,
-            self.cache_dir,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-        )
-        self.mask_token_id = None
-        self.pad_token_id = None
-
-        if not self.lazy_load:
-            self.load_onnx_model()
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        """Lists the supported models.
-
-        Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
-        """
-        return supported_colpali_models
-
-    def load_onnx_model(self) -> None:
-        self._load_onnx_model(
-            model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
-            threads=self.threads,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
-        )
-
-    def _post_process_onnx_image_output(
-        self,
-        output: OnnxOutputContext,
-    ) -> Iterable[NumpyArray]:
-        """
-        Post-process the ONNX model output to convert it into a usable format.
-
-        Args:
-            output (OnnxOutputContext): The raw output from the ONNX model.
-
-        Returns:
-            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
-        """
-        assert self.model_description.dim is not None, "Model dim is not defined"
-        return output.model_output.reshape(
-            output.model_output.shape[0], -1, self.model_description.dim
-        )
-
-    def _post_process_onnx_text_output(
-        self,
-        output: OnnxOutputContext,
-    ) -> Iterable[NumpyArray]:
-        """
-        Post-process the ONNX model output to convert it into a usable format.
-
-        Args:
-            output (OnnxOutputContext): The raw output from the ONNX model.
-
-        Returns:
-            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
-        """
-        return output.model_output
-
-    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
-        texts_query: list[str] = []
-        for query in documents:
-            query = self.BOS_TOKEN + self.QUERY_PREFIX + query + self.PAD_TOKEN * 10
-            query += "\n"
-
-            texts_query.append(query)
-        encoded = self.tokenizer.encode_batch(texts_query)  # type: ignore[union-attr]
-        return encoded
-
-    def token_count(
-        self,
-        texts: str | Iterable[str],
-        batch_size: int = 1024,
-        include_extension: bool = False,
-        **kwargs: Any,
-    ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-        token_num = 0
-        texts = [texts] if isinstance(texts, str) else texts
-        assert self.tokenizer is not None
-        tokenize_func = self.tokenize if include_extension else self.tokenizer.encode_batch
-        for batch in iter_batch(texts, batch_size):
-            token_num += sum([sum(encoding.attention_mask) for encoding in tokenize_func(batch)])
-        return token_num
-
-    def _preprocess_onnx_text_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
-        onnx_input["input_ids"] = np.array(
-            [
-                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()  # type: ignore[index]
-                for input_ids in onnx_input["input_ids"]
-            ]
-        )
-        empty_image_placeholder: NumpyArray = np.zeros(
-            self.IMAGE_PLACEHOLDER_SIZE, dtype=np.float32
-        )
-        onnx_input["pixel_values"] = np.array(
-            [empty_image_placeholder for _ in onnx_input["input_ids"]],
-        )
-        return onnx_input
-
-    def _preprocess_onnx_image_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
-        """
-        Add placeholders for text input when processing image data for ONNX.
-        Args:
-            onnx_input (Dict[str, NumpyArray]): Preprocessed image inputs.
-            **kwargs: Additional arguments.
-        Returns:
-            Dict[str, NumpyArray]: ONNX input with text placeholders.
-        """
-        onnx_input["input_ids"] = np.array(
-            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["pixel_values"]]
-        )
-        onnx_input["attention_mask"] = np.array(
-            [self.EVEN_ATTENTION_MASK for _ in onnx_input["pixel_values"]]
-        )
-        return onnx_input
-
-    def embed_text(
-        self,
-        documents: str | Iterable[str],
-        batch_size: int = 256,
-        parallel: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        """
-        Encode a list of documents into list of embeddings.
-
-        Args:
-            documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per document
-        """
-        yield from self._embed_documents(
-            model_name=self.model_name,
-            cache_dir=str(self.cache_dir),
-            documents=documents,
-            batch_size=batch_size,
-            parallel=parallel,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_ids=self.device_ids,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            extra_session_options=self._extra_session_options,
-            **kwargs,
-        )
-
-    def embed_image(
-        self,
-        images: ImageInput | Iterable[ImageInput],
-        batch_size: int = 16,
-        parallel: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        """
-        Encode a list of images into list of embeddings.
-
-        Args:
-            images: Iterator of image paths or single image path to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per document
-        """
-        yield from self._embed_images(
-            model_name=self.model_name,
-            cache_dir=str(self.cache_dir),
-            images=images,
-            batch_size=batch_size,
-            parallel=parallel,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_ids=self.device_ids,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            extra_session_options=self._extra_session_options,
-            **kwargs,
-        )
-
-    @classmethod
-    def _get_text_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
-        return ColPaliTextEmbeddingWorker
-
-    @classmethod
-    def _get_image_worker_class(cls) -> Type[ImageEmbeddingWorker[NumpyArray]]:
-        return ColPaliImageEmbeddingWorker
-
-
-class ColPaliTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColPali:
-        return ColPali(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=1,
-            **kwargs,
-        )
-
-
-class ColPaliImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColPali:
-        return ColPali(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=1,
-            **kwargs,
-        )

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -1,189 +0,0 @@
-from typing import Any, Iterable, Sequence, Type
-from dataclasses import asdict
-
-from fastembed.common import OnnxProvider, ImageInput
-from fastembed.common.types import NumpyArray, Device
-from fastembed.late_interaction_multimodal.colpali import ColPali
-from fastembed.late_interaction_multimodal.colmodernvbert import ColModernVBERT
-
-from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
-    LateInteractionMultimodalEmbeddingBase,
-)
-from fastembed.common.model_description import DenseModelDescription
-
-
-class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [
-        ColPali,
-        ColModernVBERT,
-    ]
-
-    @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.
-
-            Example:
-                ```
-                [
-                    {
-                         "model": "Qdrant/colpali-v1.3-fp16",
-                         "dim": 128,
-                         "description": "Text embeddings, Unimodal (text), Aligned to image latent space, ColBERT-compatible, 512 tokens max, 2024.",
-                         "license": "mit",
-                         "size_in_GB": 6.06,
-                         "sources": {
-                            "hf": "Qdrant/colpali-v1.3-fp16",
-                         },
-                         "additional_files": [
-                            "model.onnx_data",
-                         ],
-                        "model_file": "model.onnx",
-                    },
-                ]
-                ```
-        """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
-        for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
-        return result
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        lazy_load: bool = False,
-        **kwargs: Any,
-    ):
-        super().__init__(model_name, cache_dir, threads, **kwargs)
-        for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
-                self.model = EMBEDDING_MODEL_TYPE(
-                    model_name,
-                    cache_dir,
-                    threads=threads,
-                    providers=providers,
-                    cuda=cuda,
-                    device_ids=device_ids,
-                    lazy_load=lazy_load,
-                    **kwargs,
-                )
-                return
-
-        raise ValueError(
-            f"Model {model_name} is not supported in LateInteractionMultimodalEmbedding."
-            "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: int | None = 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: str | Iterable[str],
-        batch_size: int = 256,
-        parallel: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        """
-        Encode a list of documents into list of embeddings.
-
-        Args:
-            documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per document
-        """
-        yield from self.model.embed_text(documents, batch_size, parallel, **kwargs)
-
-    def embed_image(
-        self,
-        images: ImageInput | Iterable[ImageInput],
-        batch_size: int = 16,
-        parallel: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        """
-        Encode a list of images into list of embeddings.
-
-        Args:
-            images: Iterator of image paths or single image path to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per image
-        """
-        yield from self.model.embed_image(images, batch_size, parallel, **kwargs)
-
-    def token_count(
-        self,
-        texts: str | Iterable[str],
-        batch_size: int = 1024,
-        include_extension: bool = False,
-        **kwargs: Any,
-    ) -> int:
-        """Returns the number of tokens in the texts.
-
-        Args:
-            texts (str | Iterable[str]): The list of texts to embed.
-            batch_size (int): Batch size for encoding
-            include_extension (bool): Whether to include tokens added by preprocessing
-
-        Returns:
-            int: Sum of number of tokens in the texts.
-        """
-        return self.model.token_count(
-            texts, batch_size=batch_size, include_extension=include_extension, **kwargs
-        )

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -1,86 +0,0 @@
-from typing import Iterable, Any
-
-
-from fastembed.common import ImageInput
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.model_management import ModelManagement
-from fastembed.common.types import NumpyArray
-
-
-class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescription]):
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        **kwargs: Any,
-    ):
-        self.model_name = model_name
-        self.cache_dir = cache_dir
-        self.threads = threads
-        self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: int | None = None
-
-    def embed_text(
-        self,
-        documents: str | Iterable[str],
-        batch_size: int = 256,
-        parallel: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        """
-        Embeds a list of documents into a list of embeddings.
-
-        Args:
-            documents (Iterable[str]): The list of texts 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.
-            **kwargs: Additional keyword argument to pass to the embed method.
-
-        Yields:
-            Iterable[NumpyArray]: The embeddings.
-        """
-        raise NotImplementedError()
-
-    def embed_image(
-        self,
-        images: ImageInput | Iterable[ImageInput],
-        batch_size: int = 16,
-        parallel: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        """
-        Encode a list of images into list of embeddings.
-        Args:
-            images: Iterator of image paths or single image path to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per 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")
-
-    def token_count(
-        self,
-        texts: str | Iterable[str],
-        **kwargs: Any,
-    ) -> int:
-        """Returns the number of tokens in the texts."""
-        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -1,291 +0,0 @@
-import contextlib
-import os
-from multiprocessing import get_all_start_methods
-from pathlib import Path
-from typing import Any, Iterable, Sequence, Type
-
-import numpy as np
-from PIL import Image
-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, Device
-from fastembed.common.utils import iter_batch
-from fastembed.image.transform.operators import Compose
-from fastembed.parallel_processor import ParallelWorkerPool
-
-
-class OnnxMultimodalModel(OnnxModel[T]):
-    ONNX_OUTPUT_NAMES: list[str] | None = None
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.tokenizer: Tokenizer | None = None
-        self.processor: Compose | None = None
-        self.special_token_to_id: dict[str, int] = {}
-
-    def _preprocess_onnx_text_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
-        """
-        Preprocess the onnx input.
-        """
-        return onnx_input
-
-    def _preprocess_onnx_image_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
-        """
-        Preprocess the onnx input.
-        """
-        return onnx_input
-
-    @classmethod
-    def _get_text_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
-        raise NotImplementedError("Subclasses must implement this method")
-
-    @classmethod
-    def _get_image_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
-        raise NotImplementedError("Subclasses must implement this method")
-
-    def _post_process_onnx_image_output(self, output: OnnxOutputContext) -> Iterable[T]:
-        raise NotImplementedError("Subclasses must implement this method")
-
-    def _post_process_onnx_text_output(self, output: OnnxOutputContext) -> Iterable[T]:
-        raise NotImplementedError("Subclasses must implement this method")
-
-    def _load_onnx_model(
-        self,
-        model_dir: Path,
-        model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
-    ) -> None:
-        super()._load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_file,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_id=device_id,
-            extra_session_options=extra_session_options,
-        )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
-        assert self.tokenizer is not None
-        self.processor = load_preprocessor(model_dir=model_dir)
-
-    def load_onnx_model(self) -> None:
-        raise NotImplementedError("Subclasses must implement this method")
-
-    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
-        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
-
-    def onnx_embed_text(
-        self,
-        documents: list[str],
-        **kwargs: Any,
-    ) -> OnnxOutputContext:
-        encoded = self.tokenize(documents, **kwargs)
-        input_ids = np.array([e.ids for e in encoded])
-        attention_mask = np.array([e.attention_mask for e in encoded])  # type: ignore[union-attr]
-        input_names = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
-        onnx_input: dict[str, NumpyArray] = {
-            "input_ids": np.array(input_ids, dtype=np.int64),
-        }
-        if "attention_mask" in input_names:
-            onnx_input["attention_mask"] = np.array(attention_mask, dtype=np.int64)
-        if "token_type_ids" in input_names:
-            onnx_input["token_type_ids"] = np.array(
-                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
-            )
-
-        onnx_input = self._preprocess_onnx_text_input(onnx_input, **kwargs)
-        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),
-            input_ids=onnx_input.get("input_ids", input_ids),
-        )
-
-    def _embed_documents(
-        self,
-        model_name: str,
-        cache_dir: str,
-        documents: str | Iterable[str],
-        batch_size: int = 256,
-        parallel: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        local_files_only: bool = False,
-        specific_model_path: str | None = None,
-        extra_session_options: dict[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> Iterable[T]:
-        is_small = False
-
-        if isinstance(documents, str):
-            documents = [documents]
-            is_small = True
-
-        if isinstance(documents, list):
-            if len(documents) < batch_size:
-                is_small = True
-
-        if parallel is None or is_small:
-            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))
-        else:
-            if parallel == 0:
-                parallel = os.cpu_count()
-
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "model_name": model_name,
-                "cache_dir": cache_dir,
-                "providers": providers,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
-                **kwargs,
-            }
-
-            if extra_session_options is not None:
-                params.update(extra_session_options)
-
-            pool = ParallelWorkerPool(
-                num_workers=parallel or 1,
-                worker=self._get_text_worker_class(),
-                cuda=cuda,
-                device_ids=device_ids,
-                start_method=start_method,
-            )
-            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                yield from self._post_process_onnx_text_output(batch)  # type: ignore
-
-    def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack() as stack:
-            image_files = [
-                stack.enter_context(Image.open(image))
-                if not isinstance(image, Image.Image)
-                else image
-                for image in images
-            ]
-            assert self.processor is not None, "Processor is not initialized"
-            encoded = np.array(self.processor(image_files))
-        onnx_input = {"pixel_values": encoded}
-        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
-        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)
-
-    def _embed_images(
-        self,
-        model_name: str,
-        cache_dir: str,
-        images: Iterable[ImageInput] | ImageInput,
-        batch_size: int = 256,
-        parallel: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        local_files_only: bool = False,
-        specific_model_path: str | None = None,
-        extra_session_options: dict[str, Any] | None = None,
-        **kwargs: Any,
-    ) -> Iterable[T]:
-        is_small = False
-
-        if isinstance(images, (str, Path, Image.Image)):
-            images = [images]
-            is_small = True
-
-        if isinstance(images, list) and len(images) < batch_size:
-            is_small = True
-
-        if parallel is None or is_small:
-            if not hasattr(self, "model") or self.model is None:
-                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))
-        else:
-            if parallel == 0:
-                parallel = os.cpu_count()
-
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "model_name": model_name,
-                "cache_dir": cache_dir,
-                "providers": providers,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
-                **kwargs,
-            }
-
-            if extra_session_options is not None:
-                params.update(extra_session_options)
-
-            pool = ParallelWorkerPool(
-                num_workers=parallel or 1,
-                worker=self._get_image_worker_class(),
-                cuda=cuda,
-                device_ids=device_ids,
-                start_method=start_method,
-            )
-            for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
-                yield from self._post_process_onnx_image_output(batch)  # type: ignore
-
-
-class TextEmbeddingWorker(EmbeddingWorker[T]):
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ):
-        self.model: OnnxMultimodalModel
-        super().__init__(model_name, cache_dir, **kwargs)
-
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ) -> OnnxMultimodalModel:
-        raise NotImplementedError()
-
-    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)
-            yield idx, onnx_output
-
-
-class ImageEmbeddingWorker(EmbeddingWorker[T]):
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ):
-        self.model: OnnxMultimodalModel
-        super().__init__(model_name, cache_dir, **kwargs)
-
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ) -> OnnxMultimodalModel:
-        raise NotImplementedError()
-
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
-        for idx, batch in items:
-            embeddings = self.model.onnx_embed_image(batch)
-            yield idx, embeddings

fastembed/parallel_processor.py

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

fastembed/postprocess/__init__.py

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

fastembed/postprocess/muvera.py

@@ -1,362 +0,0 @@
-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 = 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/py.typed

@@ -1 +0,0 @@
-partial

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -1,47 +0,0 @@
-from typing import Sequence, Any
-
-from fastembed.common import OnnxProvider
-from fastembed.common.model_description import BaseModelDescription
-from fastembed.common.types import Device
-from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
-
-
-class CustomTextCrossEncoder(OnnxTextCrossEncoder):
-    SUPPORTED_MODELS: list[BaseModelDescription] = []
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        lazy_load: bool = False,
-        device_id: int | None = None,
-        specific_model_path: str | None = 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

@@ -1,91 +1,101 @@
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type
 
 from loguru import logger
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.rerank.cross_encoder.onnx_text_model import (
     OnnxCrossEncoderModel,
     TextRerankerWorker,
 )
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import BaseModelDescription, ModelSource
 
-supported_onnx_models: list[BaseModelDescription] = [
-    BaseModelDescription(
-        model="Xenova/ms-marco-MiniLM-L-6-v2",
-        description="MiniLM-L-6-v2 model optimized for re-ranking tasks.",
-        license="apache-2.0",
-        size_in_GB=0.08,
-        sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-6-v2"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="Xenova/ms-marco-MiniLM-L-12-v2",
-        description="MiniLM-L-12-v2 model optimized for re-ranking tasks.",
-        license="apache-2.0",
-        size_in_GB=0.12,
-        sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-12-v2"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="BAAI/bge-reranker-base",
-        description="BGE reranker base model for cross-encoder re-ranking.",
-        license="mit",
-        size_in_GB=1.04,
-        sources=ModelSource(hf="BAAI/bge-reranker-base"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="jinaai/jina-reranker-v1-tiny-en",
-        description="Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="jinaai/jina-reranker-v1-tiny-en"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="jinaai/jina-reranker-v1-turbo-en",
-        description="Designed for blazing-fast re-ranking with 8K context length.",
-        license="apache-2.0",
-        size_in_GB=0.15,
-        sources=ModelSource(hf="jinaai/jina-reranker-v1-turbo-en"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="jinaai/jina-reranker-v2-base-multilingual",
-        description="A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
-        license="cc-by-nc-4.0",
-        size_in_GB=1.11,
-        sources=ModelSource(hf="jinaai/jina-reranker-v2-base-multilingual"),
-        model_file="onnx/model.onnx",
-    ),
+supported_onnx_models = [
+    {
+        "model": "Xenova/ms-marco-MiniLM-L-6-v2",
+        "size_in_GB": 0.08,
+        "sources": {
+            "hf": "Xenova/ms-marco-MiniLM-L-6-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "Xenova/ms-marco-MiniLM-L-12-v2",
+        "size_in_GB": 0.12,
+        "sources": {
+            "hf": "Xenova/ms-marco-MiniLM-L-12-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "MiniLM-L-12-v2 model optimized for re-ranking tasks.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "BAAI/bge-reranker-base",
+        "size_in_GB": 1.04,
+        "sources": {
+            "hf": "BAAI/bge-reranker-base",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "BGE reranker base model for cross-encoder re-ranking.",
+        "license": "mit",
+    },
+    {
+        "model": "jinaai/jina-reranker-v1-tiny-en",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-tiny-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v1-turbo-en",
+        "size_in_GB": 0.15,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-turbo-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v2-base-multilingual",
+        "size_in_GB": 1.11,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v2-base-multilingual",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
+        "license": "cc-by-nc-4.0",
+    },
 ]
 
 
 class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     @classmethod
-    def _list_supported_models(cls) -> list[BaseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[BaseModelDescription]: A list of BaseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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: int | None = None,
-        specific_model_path: str | None = None,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -97,11 +107,10 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
+            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` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
+                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.
@@ -113,7 +122,6 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -126,20 +134,20 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             )
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = 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.cache_dir = define_cache_dir(cache_dir)
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -148,12 +156,11 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
         )
 
     def rerank(
@@ -182,7 +189,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         yield from self._rerank_pairs(
@@ -194,9 +201,6 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
@@ -204,25 +208,9 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     def _get_worker_class(cls) -> Type[TextRerankerWorker]:
         return TextCrossEncoderWorker
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[float]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
         return (float(elem) for elem in output.model_output)
 
-    def token_count(
-        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        """Returns the number of tokens in the pairs.
-
-        Args:
-            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
-            batch_size: Batch size for tokenizing
-
-        Returns:
-            token count: overall number of tokens in the pairs
-        """
-        return self._token_count(pairs, batch_size=batch_size, **kwargs)
-
 
 class TextCrossEncoderWorker(TextRerankerWorker):
     def init_embedding(

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,9 +1,10 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
+from numpy.typing import NDArray
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -12,14 +13,13 @@ from fastembed.common.onnx_model import (
     OnnxOutputContext,
     OnnxProvider,
 )
-from fastembed.common.types import NumpyArray, Device
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxCrossEncoderModel(OnnxModel[float]):
-    ONNX_OUTPUT_NAMES: list[str] | None = None
+    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
 
     @classmethod
     def _get_worker_class(cls) -> Type["TextRerankerWorker"]:
@@ -29,11 +29,10 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -42,17 +41,17 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
-            extra_session_options=extra_session_options,
         )
         self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
-        assert self.tokenizer is not None
 
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
-        return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
+        return self.tokenizer.encode_batch(pairs)
 
-    def _build_onnx_input(self, tokenized_input: list[Encoding]) -> dict[str, NumpyArray]:
-        input_names: set[str] = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
-        inputs: dict[str, NumpyArray] = {
+    def _build_onnx_input(
+        self, tokenized_input
+    ) -> dict[str, NDArray[Union[np.float32, np.int64]]]:
+        input_names = {node.name for node in self.model.get_inputs()}
+        inputs = {
             "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
         }
         if "token_type_ids" in input_names:
@@ -73,9 +72,9 @@ 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)
-        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
         relevant_output = outputs[0]
-        scores: NumpyArray = relevant_output[:, 0]
+        scores = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
 
     def _rerank_documents(
@@ -92,13 +91,10 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         cache_dir: str,
         pairs: Iterable[tuple[str, str]],
         batch_size: int,
-        parallel: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        local_files_only: bool = False,
-        specific_model_path: str | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -125,14 +121,9 @@ 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,
             }
 
-            if extra_session_options is not None:
-                params.update(extra_session_options)
-
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),
@@ -141,63 +132,21 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(pairs, batch_size), **params):
-                yield from self._post_process_onnx_output(batch)  # type: ignore
+                yield from self._post_process_onnx_output(batch)
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[float]:
-        """Post-process the ONNX model output to convert it into a usable format.
-
-        Args:
-            output (OnnxOutputContext): The raw output from the ONNX model.
-            **kwargs: Additional keyword arguments that may be needed by specific implementations.
-
-        Returns:
-            Iterable[float]: Post-processed output as an iterable of float values.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
-    def _token_count(
-        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **_: Any
-    ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-
-        token_num = 0
-        assert self.tokenizer is not None
-        for batch in iter_batch(pairs, batch_size):
-            for tokens in self.tokenizer.encode_batch(batch):
-                token_num += sum(tokens.attention_mask)
-
-        return token_num
-
-
-class TextRerankerWorker(EmbeddingWorker[float]):
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ):
-        self.model: OnnxCrossEncoderModel
-        super().__init__(model_name, cache_dir, **kwargs)
-
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ) -> OnnxCrossEncoderModel:
-        raise NotImplementedError()
 
+class TextRerankerWorker(EmbeddingWorker):
     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)

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -1,22 +1,13 @@
-from typing import Any, Iterable, Sequence, Type
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type
 
 from fastembed.common import OnnxProvider
-from fastembed.common.types import Device
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
-from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
-
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import (
-    ModelSource,
-    BaseModelDescription,
-)
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
     CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
         OnnxTextCrossEncoder,
-        CustomTextCrossEncoder,
     ]
 
     @classmethod
@@ -24,7 +15,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         """Lists the supported models.
 
         Returns:
-            list[BaseModelDescription]: A list of dictionaries containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
 
             Example:
                 ```
@@ -42,31 +33,27 @@ class TextCrossEncoder(TextCrossEncoderBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[BaseModelDescription]:
-        result: list[BaseModelDescription] = []
+        result = []
         for encoder in cls.CROSS_ENCODER_REGISTRY:
-            result.extend(encoder._list_supported_models())
+            result.extend(encoder.list_supported_models())
         return result
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         for CROSS_ENCODER_TYPE in self.CROSS_ENCODER_REGISTRY:
-            supported_models = CROSS_ENCODER_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = CROSS_ENCODER_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = CROSS_ENCODER_TYPE(
                     model_name=model_name,
                     cache_dir=cache_dir,
@@ -103,7 +90,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         """
@@ -131,48 +118,3 @@ 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: list[str] | None = 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 [],
-            )
-        )
-
-    def token_count(
-        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        """Returns the number of tokens in the pairs.
-
-        Args:
-            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
-            batch_size: Batch size for tokenizing
-
-        Returns:
-            token count: overall number of tokens in the pairs
-        """
-        return self.model.token_count(pairs, batch_size=batch_size, **kwargs)

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

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

fastembed/sparse/bm25.py

@@ -2,7 +2,7 @@ import os
 from collections import defaultdict
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Type
+from typing import Any, Iterable, Optional, Type, Union
 
 import mmh3
 import numpy as np
@@ -19,11 +19,14 @@ from fastembed.sparse.sparse_embedding_base import (
     SparseTextEmbeddingBase,
 )
 from fastembed.sparse.utils.tokenizer import SimpleTokenizer
-from fastembed.common.model_description import SparseModelDescription, ModelSource
-
 
 supported_languages = [
     "arabic",
+    "azerbaijani",
+    "basque",
+    "bengali",
+    "catalan",
+    "chinese",
     "danish",
     "dutch",
     "english",
@@ -31,30 +34,37 @@ supported_languages = [
     "french",
     "german",
     "greek",
+    "hebrew",
+    "hinglish",
     "hungarian",
+    "indonesian",
     "italian",
+    "kazakh",
+    "nepali",
     "norwegian",
     "portuguese",
     "romanian",
     "russian",
+    "slovene",
     "spanish",
     "swedish",
-    "tamil",
+    "tajik",
     "turkish",
 ]
 
-supported_bm25_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="Qdrant/bm25",
-        vocab_size=0,
-        description="BM25 as sparse embeddings meant to be used with Qdrant",
-        license="apache-2.0",
-        size_in_GB=0.01,
-        sources=ModelSource(hf="Qdrant/bm25"),
-        additional_files=[f"{lang}.txt" for lang in supported_languages],
-        requires_idf=True,
-        model_file="mock.file",
-    ),
+supported_bm25_models = [
+    {
+        "model": "Qdrant/bm25",
+        "description": "BM25 as sparse embeddings meant to be used with Qdrant",
+        "license": "apache-2.0",
+        "size_in_GB": 0.01,
+        "sources": {
+            "hf": "Qdrant/bm25",
+        },
+        "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
+        "additional_files": [f"{lang}.txt" for lang in supported_languages],
+        "requires_idf": True,
+    },
 ]
 
 
@@ -91,14 +101,14 @@ class Bm25(SparseTextEmbeddingBase):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
+        cache_dir: Optional[str] = None,
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 256.0,
         language: str = "english",
         token_max_length: int = 40,
         disable_stemmer: bool = False,
-        specific_model_path: str | None = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
@@ -113,14 +123,13 @@ class Bm25(SparseTextEmbeddingBase):
         self.avg_len = avg_len
 
         model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
-        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         self.token_max_length = token_max_length
@@ -128,7 +137,7 @@ class Bm25(SparseTextEmbeddingBase):
         self.disable_stemmer = disable_stemmer
 
         if disable_stemmer:
-            self.stopwords: set[str] = set()
+            self.stopwords = set()
             self.stemmer = None
         else:
             self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
@@ -137,11 +146,11 @@ class Bm25(SparseTextEmbeddingBase):
         self.tokenizer = SimpleTokenizer
 
     @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_bm25_models
 
@@ -158,11 +167,9 @@ class Bm25(SparseTextEmbeddingBase):
         self,
         model_name: str,
         cache_dir: str,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
-        local_files_only: bool = False,
-        specific_model_path: str | None = None,
+        parallel: Optional[int] = None,
     ) -> Iterable[SparseEmbedding]:
         is_small = False
 
@@ -191,8 +198,6 @@ class Bm25(SparseTextEmbeddingBase):
                 "language": self.language,
                 "token_max_length": self.token_max_length,
                 "disable_stemmer": self.disable_stemmer,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
             }
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
@@ -201,13 +206,13 @@ class Bm25(SparseTextEmbeddingBase):
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 for record in batch:
-                    yield record  # type: ignore
+                    yield record
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -231,12 +236,10 @@ 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]:
-        stemmed_tokens: list[str] = []
+        stemmed_tokens = []
         for token in tokens:
             lower_token = token.lower()
 
@@ -259,7 +262,7 @@ class Bm25(SparseTextEmbeddingBase):
         self,
         documents: list[str],
     ) -> list[SparseEmbedding]:
-        embeddings: list[SparseEmbedding] = []
+        embeddings = []
         for document in documents:
             document = remove_non_alphanumeric(document)
             tokens = self.tokenizer.tokenize(document)
@@ -268,15 +271,6 @@ class Bm25(SparseTextEmbeddingBase):
             embeddings.append(SparseEmbedding.from_dict(token_id2value))
         return embeddings
 
-    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
-        token_num = 0
-        texts = [texts] if isinstance(texts, str) else texts
-        for text in texts:
-            document = remove_non_alphanumeric(text)
-            tokens = self.tokenizer.tokenize(document)
-            token_num += len(tokens)
-        return token_num
-
     def _term_frequency(self, tokens: list[str]) -> dict[int, float]:
         """Calculate the term frequency part of the BM25 formula.
 
@@ -292,8 +286,8 @@ class Bm25(SparseTextEmbeddingBase):
         Returns:
             dict[int, float]: The token_id to term frequency mapping.
         """
-        tf_map: dict[int, float] = {}
-        counter: defaultdict[str, int] = defaultdict(int)
+        tf_map = {}
+        counter = defaultdict(int)
         for stemmed_token in tokens:
             counter[stemmed_token] += 1
 
@@ -311,7 +305,9 @@ class Bm25(SparseTextEmbeddingBase):
     def compute_token_id(cls, token: str) -> int:
         return abs(mmh3.hash(token))
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """To emulate BM25 behaviour, we don't need to use weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
         """
@@ -347,9 +343,7 @@ class Bm25Worker(Worker):
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
-    def process(
-        self, items: Iterable[tuple[int, Any]]
-    ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             onnx_output = self.model.raw_embed(batch)
             yield idx, onnx_output

fastembed/sparse/bm42.py

@@ -1,7 +1,7 @@
 import math
 import string
 from pathlib import Path
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import mmh3
 import numpy as np
@@ -9,40 +9,32 @@ from py_rust_stemmers import SnowballStemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import SparseModelDescription, ModelSource
 
-supported_bm42_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="Qdrant/bm42-all-minilm-l6-v2-attentions",
-        vocab_size=30522,
-        description="Light sparse embedding model, which assigns an importance score to each token in the text",
-        license="apache-2.0",
-        size_in_GB=0.09,
-        sources=ModelSource(hf="Qdrant/all_miniLM_L6_v2_with_attentions"),
-        model_file="model.onnx",
-        additional_files=["stopwords.txt"],
-        requires_idf=True,
-    ),
+supported_bm42_models = [
+    {
+        "model": "Qdrant/bm42-all-minilm-l6-v2-attentions",
+        "vocab_size": 30522,
+        "description": "Light sparse embedding model, which assigns an importance score to each token in the text",
+        "license": "apache-2.0",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "Qdrant/all_miniLM_L6_v2_with_attentions",
+        },
+        "model_file": "model.onnx",
+        "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
+    },
 ]
 
-
-_MODEL_TO_LANGUAGE = {
+MODEL_TO_LANGUAGE = {
     "Qdrant/bm42-all-minilm-l6-v2-attentions": "english",
 }
-MODEL_TO_LANGUAGE = {
-    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()
-}
-
-
-def get_language_by_model_name(model_name: str) -> str:
-    return MODEL_TO_LANGUAGE[model_name.lower()]
 
 
 class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
@@ -66,15 +58,15 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         alpha: float = 0.5,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        device_id: int | None = None,
-        specific_model_path: str | None = None,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -88,11 +80,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             alpha (float, optional): Parameter, that defines the importance of the token weight in the document
                 versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
                 It is recommended to only change this parameter based on training data for a specific dataset.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
+            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` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
+                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.
@@ -105,37 +96,36 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
-        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
-        self.invert_vocab: dict[int, str] = {}
+        self.invert_vocab = {}
 
-        self.special_tokens: set[str] = set()
-        self.special_tokens_ids: set[int] = set()
+        self.special_tokens = set()
+        self.special_tokens_ids = set()
         self.punctuation = set(string.punctuation)
         self.stopwords = set(self._load_stopwords(self._model_dir))
-        self.stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))
+        self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
         self.alpha = alpha
 
         if not self.lazy_load:
@@ -144,22 +134,20 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
         )
-
-        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
+        for token, idx in self.tokenizer.get_vocab().items():
             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))
 
     def _filter_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
-        result: list[tuple[str, Any]] = []
+        result = []
         for token, value in tokens:
             if token in self.stopwords or token in self.punctuation:
                 continue
@@ -167,7 +155,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return result
 
     def _stem_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
-        result: list[tuple[str, Any]] = []
+        result = []
         for token, value in tokens:
             processed_token = self.stemmer.stem_word(token)
             result.append((processed_token, value))
@@ -177,7 +165,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _aggregate_weights(
         cls, tokens: list[tuple[str, list[int]]], weights: list[float]
     ) -> list[tuple[str, float]]:
-        result: list[tuple[str, float]] = []
+        result = []
         for token, idxs in tokens:
             sum_weight = sum(weights[idx] for idx in idxs)
             result.append((token, sum_weight))
@@ -186,11 +174,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _reconstruct_bpe(
         self, bpe_tokens: Iterable[tuple[int, str]]
     ) -> list[tuple[str, list[int]]]:
-        result: list[tuple[str, list[int]]] = []
-        acc: str = ""
-        acc_idx: list[int] = []
+        result = []
+        acc = ""
+        acc_idx = []
 
-        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix  # type: ignore[union-attr]
+        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix
         continuing_subword_prefix_len = len(continuing_subword_prefix)
 
         for idx, token in bpe_tokens:
@@ -218,7 +206,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         So that the scoring doesn't depend on absolute values assigned by the model, but on the relative importance.
         """
 
-        new_vector: dict[int, float] = {}
+        new_vector = {}
 
         for token, value in vector.items():
             token_id = abs(mmh3.hash(token))
@@ -230,13 +218,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
         return new_vector
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
         if output.input_ids is None:
             raise ValueError("input_ids must be provided for document post-processing")
 
-        token_ids_batch = output.input_ids.astype(int)
+        token_ids_batch = output.input_ids
 
         # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
         pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask
@@ -255,7 +241,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
             weighted = self._aggregate_weights(stemmed, attention_value)
 
-            max_token_weight: dict[str, float] = {}
+            max_token_weight = {}
 
             for token, weight in weighted:
                 max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
@@ -265,11 +251,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding.from_dict(rescored)
 
     @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_bm42_models
 
@@ -284,9 +270,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -314,20 +300,19 @@ 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,
-            extra_session_options=self._extra_session_options,
         )
 
     @classmethod
     def _query_rehash(cls, tokens: Iterable[str]) -> dict[int, float]:
-        result: dict[int, float] = {}
+        result = {}
         for token in tokens:
             token_id = abs(mmh3.hash(token))
             result[token_id] = 1.0
         return result
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """
         To emulate BM25 behaviour, we don't need to use smart weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
@@ -340,7 +325,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             self.load_onnx_model()
 
         for text in query:
-            encoded = self.tokenizer.encode(text)  # type: ignore[union-attr]
+            encoded = self.tokenizer.encode(text)
             document_tokens_with_ids = enumerate(encoded.tokens)
             reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
             filtered = self._filter_pair_tokens(reconstructed)
@@ -349,18 +334,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding.from_dict(self._query_rehash(token for token, _ in stemmed))
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return Bm42TextEmbeddingWorker
 
-    def token_count(
-        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-        return self._token_count(texts, batch_size=batch_size, **kwargs)
 
-
-class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
         return Bm42(
             model_name=model_name,

fastembed/sparse/minicoil.py

@@ -1,372 +0,0 @@
-from pathlib import Path
-
-from typing import Any, Sequence, Iterable, 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.types import Device
-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",
-}
-MODEL_TO_LANGUAGE = {
-    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()
-}
-
-
-def get_language_by_model_name(model_name: str) -> str:
-    return MODEL_TO_LANGUAGE[model_name.lower()]
-
-
-class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
-    """
-        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: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        k: float = 1.2,
-        b: float = 0.75,
-        avg_len: float = 150.0,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        lazy_load: bool = False,
-        device_id: int | None = None,
-        specific_model_path: str | None = 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 (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
-            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
-            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
-                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
-            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
-            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._extra_session_options = self._select_exposed_session_options(kwargs)
-
-        self.k = k
-        self.b = b
-        self.avg_len = avg_len
-
-        # Initialize class attributes
-        self.tokenizer: Tokenizer | None = None
-        self.invert_vocab: dict[int, str] = {}
-        self.special_tokens: set[str] = set()
-        self.special_tokens_ids: set[int] = set()
-        self.stopwords: set[str] = set()
-        self.vocab_resolver: VocabResolver | None = None
-        self.encoder: Encoder | None = None
-        self.output_dim: int | None = None
-        self.sparse_vector_converter: SparseVectorConverter | None = None
-
-        self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
-        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,
-            extra_session_options=self._extra_session_options,
-        )
-
-        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(get_language_by_model_name(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 token_count(
-        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        return self._token_count(texts, batch_size=batch_size, **kwargs)
-
-    def embed(
-        self,
-        documents: str | Iterable[str],
-        batch_size: int = 256,
-        parallel: int | None = 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,
-            extra_session_options=self._extra_session_options,
-            **kwargs,
-        )
-
-    def query_embed(self, query: 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_embedding_base.py

@@ -1,27 +1,24 @@
 from dataclasses import dataclass
-from typing import Iterable, Any
+from typing import Iterable, Optional, Union, Any
 
 import numpy as np
-from numpy.typing import NDArray
 
-from fastembed.common.model_description import SparseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
 @dataclass
 class SparseEmbedding:
-    values: NumpyArray
-    indices: NDArray[np.int64] | NDArray[np.int32]
+    values: np.ndarray
+    indices: np.ndarray
 
-    def as_object(self) -> dict[str, NumpyArray]:
+    def as_object(self) -> dict[str, np.ndarray]:
         return {
             "values": self.values,
             "indices": self.indices,
         }
 
     def as_dict(self) -> dict[int, float]:
-        return {int(i): float(v) for i, v in zip(self.indices, self.values)}  # type: ignore
+        return {i: v for i, v in zip(self.indices, self.values)}
 
     @classmethod
     def from_dict(cls, data: dict[int, float]) -> "SparseEmbedding":
@@ -31,12 +28,12 @@ class SparseEmbedding:
         return cls(values=np.array(values), indices=np.array(indices))
 
 
-class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
+class SparseTextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
@@ -46,9 +43,9 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         raise NotImplementedError()
@@ -68,7 +65,9 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -82,9 +81,5 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        else:
+        if isinstance(query, Iterable):
             yield from self.embed(query, **kwargs)
-
-    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
-        """Returns the number of tokens in the texts."""
-        raise NotImplementedError("Subclasses must implement this method")

fastembed/sparse/sparse_text_embedding.py

@@ -1,22 +1,18 @@
-from typing import Any, Iterable, Sequence, Type
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 from fastembed.common import OnnxProvider
-from fastembed.common.types import Device
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.bm42 import Bm42
-from fastembed.sparse.minicoil import MiniCOIL
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
 from fastembed.sparse.splade_pp import SpladePP
 import warnings
-from fastembed.common.model_description import SparseModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25, MiniCOIL]
+    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -42,28 +38,24 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
-        result: list[SparseModelDescription] = []
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "prithvida/Splade_PP_en_v1".lower():
+        if model_name == "prithvida/Splade_PP_en_v1":
             warnings.warn(
                 "The right spelling is prithivida/Splade_PP_en_v1. "
                 "Support of this name will be removed soon, please fix the model_name",
@@ -73,8 +65,8 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
             model_name = "prithivida/Splade_PP_en_v1"
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -94,9 +86,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -116,7 +108,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -127,17 +121,3 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
             Iterable[SparseEmbedding]: The sparse embeddings.
         """
         yield from self.model.query_embed(query, **kwargs)
-
-    def token_count(
-        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        """Returns the number of tokens in the texts.
-
-        Args:
-            texts (str | Iterable[str]): The list of texts to embed.
-            batch_size (int): Batch size for encoding
-
-        Returns:
-            int: Sum of number of tokens in the texts.
-        """
-        return self.model.token_count(texts, batch_size=batch_size, **kwargs)

fastembed/sparse/splade_pp.py

@@ -1,43 +1,43 @@
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import SparseModelDescription, ModelSource
 
-supported_splade_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="prithivida/Splade_PP_en_v1",
-        vocab_size=30522,
-        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"),
-        model_file="model.onnx",
-    ),
-    SparseModelDescription(
-        model="prithvida/Splade_PP_en_v1",
-        vocab_size=30522,
-        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"),
-        model_file="model.onnx",
-    ),
+supported_splade_models = [
+    {
+        "model": "prithivida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "prithvida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
 
 class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 
@@ -54,31 +54,26 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             scores = row_scores[indices]
             yield SparseEmbedding(values=scores, indices=indices)
 
-    def token_count(
-        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        return self._token_count(texts, batch_size=batch_size, **kwargs)
-
     @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_splade_models
 
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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: int | None = None,
-        specific_model_path: str | None = None,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -90,11 +85,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.
+            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` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
+                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.
@@ -106,28 +100,27 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
-        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -136,19 +129,18 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
         )
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -175,18 +167,15 @@ 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,
-            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return SpladePPEmbeddingWorker
 
 
-class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+class SpladePPEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
         return SpladePP(
             model_name=model_name,

fastembed/sparse/utils/minicoil_encoder.py

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

fastembed/sparse/utils/sparse_vectors_converter.py

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

fastembed/sparse/utils/vocab_resolver.py

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

fastembed/text/clip_embedding.py

@@ -1,43 +1,41 @@
 from typing import Any, Iterable, Type
 
-from fastembed.common.types import NumpyArray
+import numpy as np
+
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
 
-supported_clip_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="Qdrant/clip-ViT-B-32-text",
-        dim=512,
-        description=(
-            "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2021 year"
-        ),
-        license="mit",
-        size_in_GB=0.25,
-        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-text"),
-        model_file="model.onnx",
-    ),
+supported_clip_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-text",
+        "dim": 512,
+        "description": "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
+        "license": "mit",
+        "size_in_GB": 0.25,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-text",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
 
 class CLIPOnnxEmbedding(OnnxTextEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return CLIPEmbeddingWorker
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_clip_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return output.model_output
 
 

fastembed/text/custom_text_embedding.py

@@ -1,97 +0,0 @@
-from typing import Sequence, Any, Iterable
-from dataclasses import dataclass
-
-import numpy as np
-from numpy.typing import NDArray
-
-from fastembed.common import OnnxProvider
-from fastembed.common.model_description import (
-    PoolingType,
-    DenseModelDescription,
-)
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray, Device
-from fastembed.common.utils import normalize, mean_pooling
-from fastembed.text.onnx_embedding import OnnxTextEmbedding
-
-
-@dataclass(frozen=True)
-class PostprocessingConfig:
-    pooling: PoolingType
-    normalization: bool
-
-
-class CustomTextEmbedding(OnnxTextEmbedding):
-    SUPPORTED_MODELS: list[DenseModelDescription] = []
-    POSTPROCESSING_MAPPING: dict[str, PostprocessingConfig] = {}
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        lazy_load: bool = False,
-        device_id: int | None = None,
-        specific_model_path: str | None = 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,
-        )
-        self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
-        self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        return cls.SUPPORTED_MODELS
-
-    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(
-        self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None
-    ) -> NumpyArray:
-        if self._pooling == PoolingType.CLS:
-            return embeddings[:, 0]
-
-        if self._pooling == PoolingType.MEAN:
-            if attention_mask is None:
-                raise ValueError("attention_mask must be provided for mean pooling")
-            return mean_pooling(embeddings, attention_mask)
-
-        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
-
-    @classmethod
-    def add_model(
-        cls,
-        model_description: DenseModelDescription,
-        pooling: PoolingType,
-        normalization: bool,
-    ) -> None:
-        cls.SUPPORTED_MODELS.append(model_description)
-        cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
-            pooling=pooling, normalization=normalization
-        )

fastembed/text/multitask_embedding.py

@@ -1,35 +1,32 @@
 from enum import Enum
-from typing import Any, Type, Iterable
+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
-from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
 
-supported_multitask_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v3",
-        dim=1024,
-        tasks={
+supported_multitask_models = [
+    {
+        "model": "jinaai/jina-embeddings-v3",
+        "dim": 1024,
+        "tasks": {
             "retrieval.query": 0,
             "retrieval.passage": 1,
             "separation": 2,
             "classification": 3,
             "text-matching": 4,
         },
-        description=(
-            "Multi-task unimodal (text) embedding model, multi-lingual (~100), "
-            "1024 tokens truncation, and 8192 sequence length. Prefixes for queries/documents: not necessary, 2024 year."
-        ),
-        license="cc-by-nc-4.0",
-        size_in_GB=2.29,
-        sources=ModelSource(hf="jinaai/jina-embeddings-v3"),
-        model_file="onnx/model.onnx",
-        additional_files=["onnx/model.onnx_data"],
-    ),
+        "description": "Multi-task unimodal (text) embedding model, multi-lingual (~100), 1024 tokens truncation, and 8192 sequence length. Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "cc-by-nc-4.0",
+        "size_in_GB": 2.29,
+        "sources": {
+            "hf": "jinaai/jina-embeddings-v3",
+        },
+        "model_file": "onnx/model.onnx",
+        "additional_files": ["onnx/model.onnx_data"],
+    },
 ]
 
 
@@ -45,47 +42,43 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
 
-    def __init__(self, *args: Any, task_id: int | None = None, **kwargs: Any):
+    def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.default_task_id: Task | int = task_id if task_id is not None else self.PASSAGE_TASK
+        self._current_task_id = self.PASSAGE_TASK
 
     @classmethod
-    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
         return JinaEmbeddingV3Worker
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         return supported_multitask_models
 
     def _preprocess_onnx_input(
-        self,
-        onnx_input: dict[str, NumpyArray],
-        task_id: int | Task | None = None,
-        **kwargs: Any,
-    ) -> dict[str, NumpyArray]:
-        if task_id is None:
-            raise ValueError(f"task_id must be provided for JinaEmbeddingV3, got <{task_id}>")
-        onnx_input["task_id"] = np.array(task_id, dtype=np.int64)
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
+        onnx_input["task_id"] = np.array(self._current_task_id, dtype=np.int64)
         return onnx_input
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
-        task_id: int | None = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
-        task_id = (
-            task_id if task_id is not None else self.default_task_id
-        )  # required for multiprocessing
-        yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
+        parallel: Optional[int] = None,
+        task_id: int = PASSAGE_TASK,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        self._current_task_id = task_id
+        kwargs["task_id"] = task_id
+        yield from super().embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
-        yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
+        self._current_task_id = self.QUERY_TASK
+        yield from super().embed(query, **kwargs)
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
-        yield from super().embed(texts, task_id=self.PASSAGE_TASK, **kwargs)
+    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
+        self._current_task_id = self.PASSAGE_TASK
+        yield from super().embed(texts, **kwargs)
 
 
 class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
@@ -93,17 +86,13 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
+        **kwargs,
     ) -> JinaEmbeddingV3:
-        return JinaEmbeddingV3(
+        model = JinaEmbeddingV3(
             model_name=model_name,
             cache_dir=cache_dir,
             threads=1,
             **kwargs,
         )
-
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
-        self.model: JinaEmbeddingV3  # mypy complaints `self.model` does not have `default_task_id`
-        for idx, batch in items:
-            onnx_output = self.model.onnx_embed(batch, task_id=self.model.default_task_id)
-            yield idx, onnx_output
+        model._current_task_id = kwargs["task_id"]
+        return model

fastembed/text/onnx_embedding.py

@@ -1,212 +1,199 @@
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-from fastembed.common.types import NumpyArray, OnnxProvider, Device
+import numpy as np
+
+from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 from fastembed.text.text_embedding_base import TextEmbeddingBase
-from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_onnx_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="BAAI/bge-base-en",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.42,
-        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",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-base-en-v1.5",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.21,
-        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",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-large-en-v1.5",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=1.20,
-        sources=ModelSource(hf="qdrant/bge-large-en-v1.5-onnx"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-small-en",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.13,
-        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",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-small-en-v1.5",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.067,
-        sources=ModelSource(hf="qdrant/bge-small-en-v1.5-onnx-q"),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-small-zh-v1.5",
-        dim=512,
-        description=(
-            "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.09,
-        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",
-    ),
-    DenseModelDescription(
-        model="mixedbread-ai/mxbai-embed-large-v1",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.64,
-        sources=ModelSource(hf="mixedbread-ai/mxbai-embed-large-v1"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-xs",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.09,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-xs"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-s",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-s"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-m",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.43,
-        sources=ModelSource(hf="Snowflake/snowflake-arctic-embed-m"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-m-long",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.54,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-m-long"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-l",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=1.02,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-l"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-clip-v1",
-        dim=768,
-        description=(
-            "Text embeddings, Multimodal (text&image), English, Prefixes for queries/documents: "
-            "not necessary, 2024 year"
-        ),
-        license="apache-2.0",
-        size_in_GB=0.55,
-        sources=ModelSource(hf="jinaai/jina-clip-v1"),
-        model_file="onnx/text_model.onnx",
-    ),
+supported_onnx_models = [
+    {
+        "model": "BAAI/bge-base-en",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.42,
+        "sources": {
+            "hf": "Qdrant/fast-bge-base-en",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-base-en-v1.5",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.21,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            "hf": "qdrant/bge-base-en-v1.5-onnx-q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 1.20,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-en",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "Qdrant/bge-small-en",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-en-v1.5",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.067,
+        "sources": {
+            "hf": "qdrant/bge-small-en-v1.5-onnx-q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-zh-v1.5",
+        "dim": 512,
+        "description": "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "Qdrant/bge-small-zh-v1.5",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "thenlper/gte-large",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 1.20,
+        "sources": {
+            "hf": "qdrant/gte-large-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "mixedbread-ai/mxbai-embed-large-v1",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.64,
+        "sources": {
+            "hf": "mixedbread-ai/mxbai-embed-large-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-xs",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-xs",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-s",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-s",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.43,
+        "sources": {
+            "hf": "Snowflake/snowflake-arctic-embed-m",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m-long",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-m-long",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-l",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 1.02,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-l",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-clip-v1",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text&image), English, Prefixes for queries/documents: not necessary, 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.55,
+        "sources": {
+            "hf": "jinaai/jina-clip-v1",
+        },
+        "model_file": "onnx/text_model.onnx",
+    },
 ]
 
 
-class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
+class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
     """Implementation of the Flag Embedding model."""
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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: int | None = None,
-        specific_model_path: str | None = None,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -218,11 +205,10 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to Device.AUTO.
+            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` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
-                with `providers`. Defaults to None.
+                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.
@@ -234,26 +220,26 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
-        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = 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.cache_dir = define_cache_dir(cache_dir)
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -261,11 +247,11 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -290,55 +276,43 @@ 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,
-            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
         return OnnxTextEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         embeddings = output.model_output
-
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
             processed_embeddings = embeddings[:, 0]
         elif embeddings.ndim == 2:  # (batch_size, embedding_dim)
             processed_embeddings = embeddings
         else:
             raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
-        return normalize(processed_embeddings)
+        return normalize(processed_embeddings).astype(np.float32)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
-            extra_session_options=self._extra_session_options,
         )
 
-    def token_count(
-        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        return self._token_count(texts, batch_size=batch_size, **kwargs)
 
-
-class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(
         self,
         model_name: str,

fastembed/text/onnx_text_model.py

@@ -1,13 +1,12 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from numpy.typing import NDArray
-from tokenizers import Encoding, Tokenizer
+from tokenizers import Encoding
 
-from fastembed.common.types import NumpyArray, OnnxProvider, Device
+from fastembed.common import 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
@@ -15,32 +14,23 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxTextModel(OnnxModel[T]):
-    ONNX_OUTPUT_NAMES: list[str] | None = None
+    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
-        """Post-process the ONNX model output to convert it into a usable format.
-
-        Args:
-            output (OnnxOutputContext): The raw output from the ONNX model.
-            **kwargs: Additional keyword arguments that may be needed by specific implementations.
-
-        Returns:
-            Iterable[T]: Post-processed output as an iterable of type T.
-        """
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def __init__(self) -> None:
+    def __init__(self):
         super().__init__()
-        self.tokenizer: Tokenizer | None = None
-        self.special_token_to_id: dict[str, int] = {}
+        self.tokenizer = None
+        self.special_token_to_id = {}
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray | NDArray[np.int64]]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -50,11 +40,11 @@ class OnnxTextModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -63,7 +53,6 @@ class OnnxTextModel(OnnxModel[T]):
             providers=providers,
             cuda=cuda,
             device_id=device_id,
-            extra_session_options=extra_session_options,
         )
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
 
@@ -71,7 +60,7 @@ class OnnxTextModel(OnnxModel[T]):
         raise NotImplementedError("Subclasses must implement this method")
 
     def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
-        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+        return self.tokenizer.encode_batch(documents)
 
     def onnx_embed(
         self,
@@ -81,8 +70,8 @@ class OnnxTextModel(OnnxModel[T]):
         encoded = self.tokenize(documents, **kwargs)
         input_ids = np.array([e.ids for e in encoded])
         attention_mask = np.array([e.attention_mask for e in encoded])
-        input_names = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
-        onnx_input: dict[str, NumpyArray] = {
+        input_names = {node.name for node in self.model.get_inputs()}
+        onnx_input = {
             "input_ids": np.array(input_ids, dtype=np.int64),
         }
         if "attention_mask" in input_names:
@@ -93,7 +82,7 @@ class OnnxTextModel(OnnxModel[T]):
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -104,15 +93,12 @@ class OnnxTextModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        local_files_only: bool = False,
-        specific_model_path: str | None = None,
-        extra_session_options: dict[str, Any] | None = None,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -129,9 +115,7 @@ class OnnxTextModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, **kwargs), **kwargs
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -141,14 +125,9 @@ 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,
             }
 
-            if extra_session_options is not None:
-                params.update(extra_session_options)
-
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),
@@ -157,24 +136,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, **kwargs)  # type: ignore
-
-    def _token_count(self, texts: str | Iterable[str], batch_size: int = 1024, **_: Any) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-
-        token_num = 0
-        assert self.tokenizer is not None
-        texts = [texts] if isinstance(texts, str) else texts
-        for batch in iter_batch(texts, batch_size):
-            for tokens in self.tokenizer.encode_batch(batch):
-                token_num += sum(tokens.attention_mask)
-
-        return token_num
+                yield from self._post_process_onnx_output(batch)
 
 
-class TextEmbeddingWorker(EmbeddingWorker[T]):
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
+class TextEmbeddingWorker(EmbeddingWorker):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             onnx_output = self.model.onnx_embed(batch)
             yield idx, onnx_output

fastembed/text/pooled_embedding.py

@@ -1,124 +1,115 @@
 from typing import Any, Iterable, Type
 
 import numpy as np
-from numpy.typing import NDArray
 
-from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import mean_pooling
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
 
-supported_pooled_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="nomic-ai/nomic-embed-text-v1.5",
-        dim=768,
-        description=(
-            "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.52,
-        sources=ModelSource(hf="nomic-ai/nomic-embed-text-v1.5"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="nomic-ai/nomic-embed-text-v1.5-Q",
-        dim=768,
-        description=(
-            "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="nomic-ai/nomic-embed-text-v1.5"),
-        model_file="onnx/model_quantized.onnx",
-    ),
-    DenseModelDescription(
-        model="nomic-ai/nomic-embed-text-v1",
-        dim=768,
-        description=(
-            "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.52,
-        sources=ModelSource(hf="nomic-ai/nomic-embed-text-v1"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2019 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.22,
-        sources=ModelSource(hf="qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q"),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2021 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=1.00,
-        sources=ModelSource(hf="xenova/paraphrase-multilingual-mpnet-base-v2"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="intfloat/multilingual-e5-large",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="mit",
-        size_in_GB=2.24,
-        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"],
-    ),
+supported_pooled_models = [
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2019 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.22,
+        "sources": {
+            "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 1.00,
+        "sources": {
+            "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "intfloat/multilingual-e5-large",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "mit",
+        "size_in_GB": 2.24,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+            "hf": "qdrant/multilingual-e5-large-onnx",
+        },
+        "model_file": "model.onnx",
+        "additional_files": ["model.onnx_data"],
+    },
 ]
 
 
 class PooledEmbedding(OnnxTextEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return PooledEmbeddingWorker
 
     @classmethod
-    def mean_pooling(
-        cls, model_output: NumpyArray, attention_mask: NDArray[np.int64]
-    ) -> NumpyArray:
-        return mean_pooling(model_output, attention_mask)
+    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
+        input_mask_expanded = input_mask_expanded.astype(float)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_pooled_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         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)
+        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
 
 
 class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/pooled_normalized_embedding.py

@@ -1,152 +1,113 @@
 from typing import Any, Iterable, Type
 
+import numpy as np
 
-from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import normalize
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
 from fastembed.text.pooled_embedding import PooledEmbedding
-from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_pooled_normalized_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="sentence-transformers/all-MiniLM-L6-v2",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 256 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2021 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.09,
-        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",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v2-base-en",
-        dim=768,
-        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.52,
-        sources=ModelSource(hf="xenova/jina-embeddings-v2-base-en"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v2-small-en",
-        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",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v2-base-de",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.32,
-        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-de"),
-        model_file="onnx/model_fp16.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v2-base-code",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), "
-            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.64,
-        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-code"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v2-base-zh",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), supports mixed Chinese-English input text, "
-            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.64,
-        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-zh"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-embeddings-v2-base-es",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), supports mixed Spanish-English input text, "
-            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.64,
-        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-es"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="thenlper/gte-base",
-        dim=768,
-        description=(
-            "General text embeddings, Unimodal (text), supports English only input text, "
-            "512 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
-        ),
-        license="mit",
-        size_in_GB=0.44,
-        sources=ModelSource(hf="thenlper/gte-base"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="thenlper/gte-large",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=1.20,
-        sources=ModelSource(hf="qdrant/gte-large-onnx"),
-        model_file="model.onnx",
-    ),
+supported_pooled_normalized_models = [
+    {
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 256 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-en",
+        "dim": 768,
+        "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.52,
+        "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-small-en",
+        "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": {"hf": "xenova/jina-embeddings-v2-small-en"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-de",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.32,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
+        "model_file": "onnx/model_fp16.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-code",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.64,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-code"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-zh",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), supports mixed Chinese-English input text, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.64,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-zh"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-es",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), supports mixed Spanish-English input text, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.64,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-es"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "thenlper/gte-base",
+        "dim": 768,
+        "description": "General text embeddings, Unimodal (text), supports English only input text, 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
+        "license": "mit",
+        "size_in_GB": 0.44,
+        "sources": {"hf": "thenlper/gte-base"},
+        "model_file": "onnx/model.onnx",
+    },
 ]
 
 
 class PooledNormalizedEmbedding(PooledEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return PooledNormalizedEmbeddingWorker
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_pooled_normalized_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         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))
+        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
 
 
 class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/text_embedding.py

@@ -1,16 +1,14 @@
 import warnings
-from typing import Any, Iterable, Sequence, Type
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-from fastembed.common.types import NumpyArray, OnnxProvider, Device
+import numpy as np
+from fastembed.common import OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
-from fastembed.text.custom_text_embedding import CustomTextEmbedding
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.multitask_embedding import JinaEmbeddingV3
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
-from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
 
 
 class TextEmbedding(TextEmbeddingBase):
@@ -20,97 +18,78 @@ class TextEmbedding(TextEmbeddingBase):
         PooledNormalizedEmbedding,
         PooledEmbedding,
         JinaEmbeddingV3,
-        CustomTextEmbedding,
     ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
-        """Lists the supported models.
+        """
+        Lists the supported models.
 
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "intfloat/multilingual-e5-large",
+                        "dim": 1024,
+                        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+                        "license": "mit",
+                        "size_in_GB": 2.24,
+                        "sources": {
+                            "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+                            "hf": "qdrant/multilingual-e5-large-onnx",
+                        }
+                    }
+                ]
+                ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
-    @classmethod
-    def add_custom_model(
-        cls,
-        model: str,
-        pooling: PoolingType,
-        normalization: bool,
-        sources: ModelSource,
-        dim: int,
-        model_file: str = "onnx/model.onnx",
-        description: str = "",
-        license: str = "",
-        size_in_gb: float = 0.0,
-        additional_files: list[str] | None = None,
-    ) -> None:
-        registered_models = cls._list_supported_models()
-        for registered_model in registered_models:
-            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"
-                )
-
-        CustomTextEmbedding.add_model(
-            DenseModelDescription(
-                model=model,
-                sources=sources,
-                dim=dim,
-                model_file=model_file,
-                description=description,
-                license=license,
-                size_in_GB=size_in_gb,
-                additional_files=additional_files or [],
-            ),
-            pooling=pooling,
-            normalization=normalization,
-        )
-
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
+        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,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
+        if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
             warnings.warn(
-                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
-                "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
+                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. "
+                "Please review the latest documentation and release notes to ensure compatibility with your workflow. ",
                 UserWarning,
                 stacklevel=2,
             )
-        if model_name.lower() in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2".lower(),
-            "thenlper/gte-large".lower(),
-            "intfloat/multilingual-e5-large".lower(),
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2".lower(),
+        if model_name == "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":
+            warnings.warn(
+                "The model 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2' has been updated to "
+                "include a mean pooling layer. Please ensure your usage aligns with the new functionality. "
+                "Support for the previous version without mean pooling will be removed as of version 0.5.2.",
+                UserWarning,
+                stacklevel=2,
+            )
+        if model_name in {
+            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+            "intfloat/multilingual-e5-large",
         }:
             warnings.warn(
-                f"The model {model_name} now uses mean pooling instead of CLS embedding. "
-                f"In order to preserve the previous behaviour, consider either pinning fastembed version to 0.5.1 or "
-                "using `add_custom_model` functionality.",
+                f"{model_name} has been updated as of fastembed 0.5.2, outputs are now average pooled.",
                 UserWarning,
                 stacklevel=2,
             )
+
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name=model_name,
                     cache_dir=cache_dir,
@@ -128,47 +107,13 @@ 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: int | None = 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: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -186,7 +131,7 @@ class TextEmbedding(TextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -194,12 +139,12 @@ class TextEmbedding(TextEmbeddingBase):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[NumpyArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
         # This is model-specific, so that different models can have specialized implementations
         yield from self.model.query_embed(query, **kwargs)
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -212,17 +157,3 @@ class TextEmbedding(TextEmbeddingBase):
         """
         # This is model-specific, so that different models can have specialized implementations
         yield from self.model.passage_embed(texts, **kwargs)
-
-    def token_count(
-        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
-    ) -> int:
-        """Returns the number of tokens in the texts.
-
-        Args:
-            texts (str | Iterable[str]): The list of texts to embed.
-            batch_size (int): Batch size for encoding
-
-        Returns:
-            int: Sum of number of tokens in the texts.
-        """
-        return self.model.token_count(texts, batch_size=batch_size, **kwargs)

fastembed/text/text_embedding_base.py

@@ -1,34 +1,33 @@
-from typing import Iterable, Any
+from typing import Iterable, Optional, Union, Any
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
+class TextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: int | None = None
 
     def embed(
         self,
-        documents: str | Iterable[str],
+        documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int | None = None,
+        parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -37,13 +36,14 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[NumpyArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+
         """
         Embeds queries
 
@@ -51,25 +51,11 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[NumpyArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        else:
+        if isinstance(query, Iterable):
             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")
-
-    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
-        """Returns the number of tokens in the texts."""
-        raise NotImplementedError("Subclasses must implement this method")

mkdocs.yml

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

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed-gpu"
-version = "0.8.0"
+name = "fastembed"
+version = "0.5.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -11,31 +11,25 @@ repository = "https://github.com/qdrant/fastembed"
 keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
 
 [tool.poetry.dependencies]
-python = ">=3.10.0"
+python = ">=3.9.0"
 numpy = [
-    { version = ">=1.21,<2.3.0", python = "3.10" },
-    { version = ">=1.21", python = "3.11" },
-    { version = ">=1.26", python = "3.12" },
-    { version = ">=2.1.0", python = "3.13" },
-    { version = ">=2.3.0", python = ">=3.14" },
+    { version = ">=1.21,<2.1.0", python = "<3.10" },
+    { version = ">=1.21", python = ">=3.10,<3.12" },
+    { version = ">=1.26", python = ">=3.12,<3.13" },
+    { version = ">=2.1.0", python = ">=3.13" }
 ]
-onnxruntime-gpu = [
-    { version = ">=1.17.0,!=1.20.0,<1.24", python = "3.10" },
-    { version = ">=1.17.0,!=1.20.0,!=1.24.0,!=1.24.1", python = ">=3.11,<3.13" },
-    { version = ">1.21.0,!=1.24.0,!=1.24.1", python = "3.13" },
-    { version = ">=1.24.2", python = ">=3.14" },
+onnxruntime = [
+    { version = ">=1.17.0,<1.20.0", python = "<3.10" },
+    { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
+    { version = ">1.20.0", python = ">=3.13" }
 ]
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
-huggingface-hub = ">=0.20,<2.0"
+huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
-pillow = [
-    { version = ">=10.3.0,<13.0", python = ">=3.10,<3.13" },
-    { version = ">=11.0.0,<13.0", python = "3.13" },
-    { version = ">=12.0.0,<13.0", python = ">=3.14" },
-]
-mmh3 = ">=4.1.0,<6.0.0"
+pillow = "^10.3.0"
+mmh3 = "^4.1.0"
 py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.test.dependencies]
@@ -45,16 +39,12 @@ ruff = ">=0.3.1,<1.0"
 [tool.poetry.group.dev.dependencies]
 notebook = ">=7.0.2"
 pre-commit = "^3.6.2"
-onnx = [
-    { version = ">=1.15.0", python = ">=3.10,<3.13" },
-    { version = ">=1.18.0", python = "3.13" },
-    { version = ">=1.20.0", python = ">=3.14" },
-]
+onnx = ">=1.15.0"
 
 [tool.poetry.group.docs.dependencies]
 mkdocs-material = "^9.5.10"
 mkdocstrings = "^0.24.0"
-pillow = ">=10.3.0,<13.0.0"
+pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 

tests/test_attention_embeddings.py

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

tests/test_common.py

@@ -1,30 +0,0 @@
-from fastembed import (
-    TextEmbedding,
-    SparseTextEmbedding,
-    ImageEmbedding,
-    LateInteractionMultimodalEmbedding,
-    LateInteractionTextEmbedding,
-)
-
-
-def test_text_list_supported_models():
-    for model_type in [
-        TextEmbedding,
-        SparseTextEmbedding,
-        ImageEmbedding,
-        LateInteractionMultimodalEmbedding,
-        LateInteractionTextEmbedding,
-    ]:
-        supported_models = model_type.list_supported_models()
-        assert isinstance(supported_models, list)
-        description = supported_models[0]
-        assert isinstance(description, dict)
-
-        assert "model" in description and description["model"]
-        if model_type != SparseTextEmbedding:
-            assert "dim" in description and description["dim"]
-        assert "license" in description and description["license"]
-        assert "size_in_GB" in description and description["size_in_GB"]
-        assert "model_file" in description and description["model_file"]
-        assert "sources" in description and description["sources"]
-        assert "hf" in description["sources"] or "url" in description["sources"]

tests/test_custom_models.py

@@ -1,243 +0,0 @@
-import itertools
-import os
-import numpy as np
-import pytest
-
-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
-
-
-@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():
-    is_ci = os.getenv("CI")
-    custom_model_name = "intfloat/multilingual-e5-small"
-    canonical_vector = np.array(
-        [3.1317e-02, 3.0939e-02, -3.5117e-02, -6.7274e-02, 8.5084e-02], dtype=np.float32
-    )
-    pooling = PoolingType.MEAN
-    normalization = True
-    dim = 384
-    size_in_gb = 0.47
-    source = ModelSource(hf=custom_model_name)
-
-    TextEmbedding.add_custom_model(
-        custom_model_name,
-        pooling=pooling,
-        normalization=normalization,
-        sources=source,
-        dim=dim,
-        size_in_gb=size_in_gb,
-    )
-
-    assert CustomTextEmbedding.SUPPORTED_MODELS[0] == DenseModelDescription(
-        model=custom_model_name,
-        sources=source,
-        model_file="onnx/model.onnx",
-        description="",
-        license="",
-        size_in_GB=size_in_gb,
-        additional_files=[],
-        dim=dim,
-        tasks={},
-    )
-    assert CustomTextEmbedding.POSTPROCESSING_MAPPING[custom_model_name] == PostprocessingConfig(
-        pooling=pooling, normalization=normalization
-    )
-
-    model = TextEmbedding(custom_model_name)
-    docs = ["hello world", "flag embedding"]
-    embeddings = list(model.embed(docs))
-    embeddings = np.stack(embeddings, axis=0)
-    assert embeddings.shape == (2, dim)
-
-    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-    CustomTextEmbedding.SUPPORTED_MODELS.clear()
-    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
-
-
-def test_cross_encoder_custom_model():
-    is_ci = os.getenv("CI")
-    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)
-
-    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()
-
-
-def test_mock_add_custom_models():
-    dim = 5
-    size_in_gb = 0.1
-    source = ModelSource(hf="artificial")
-
-    num_tokens = 10
-    dummy_pooled_embedding = np.random.random((1, dim)).astype(np.float32)
-    dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(np.float32)
-    dummy_attention_mask = np.ones((1, num_tokens)).astype(np.int64)
-
-    dummy_token_output = OnnxOutputContext(
-        model_output=dummy_token_embedding, attention_mask=dummy_attention_mask
-    )
-    dummy_pooled_output = OnnxOutputContext(model_output=dummy_pooled_embedding)
-    input_data = {
-        f"{PoolingType.MEAN.lower()}-normalized": dummy_token_output,
-        f"{PoolingType.MEAN.lower()}": dummy_token_output,
-        f"{PoolingType.CLS.lower()}-normalized": dummy_token_output,
-        f"{PoolingType.CLS.lower()}": dummy_token_output,
-        f"{PoolingType.DISABLED.lower()}-normalized": dummy_pooled_output,
-        f"{PoolingType.DISABLED.lower()}": dummy_pooled_output,
-    }
-
-    expected_output = {
-        f"{PoolingType.MEAN.lower()}-normalized": normalize(
-            mean_pooling(dummy_token_embedding, dummy_attention_mask)
-        ),
-        f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
-        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),
-        f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
-    }
-
-    for pooling, normalization in itertools.product(
-        (PoolingType.MEAN, PoolingType.CLS, PoolingType.DISABLED), (True, False)
-    ):
-        model_name = f"{pooling.name.lower()}{'-normalized' if normalization else ''}"
-        TextEmbedding.add_custom_model(
-            model_name,
-            pooling=pooling,
-            normalization=normalization,
-            sources=source,
-            dim=dim,
-            size_in_gb=size_in_gb,
-        )
-
-        custom_text_embedding = CustomTextEmbedding(
-            model_name,
-            lazy_load=True,
-            specific_model_path="./",  # disable model downloading and loading
-        )
-
-        post_processed_output = next(
-            iter(custom_text_embedding._post_process_onnx_output(input_data[model_name]))
-        )
-        assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3)
-
-    CustomTextEmbedding.SUPPORTED_MODELS.clear()
-    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
-
-
-def test_do_not_add_existing_model():
-    existing_base_model = "sentence-transformers/all-MiniLM-L6-v2"
-    custom_model_name = "intfloat/multilingual-e5-small"
-
-    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
-        TextEmbedding.add_custom_model(
-            existing_base_model,
-            pooling=PoolingType.MEAN,
-            normalization=True,
-            sources=ModelSource(hf=existing_base_model),
-            dim=384,
-            size_in_gb=0.47,
-        )
-
-    TextEmbedding.add_custom_model(
-        custom_model_name,
-        pooling=PoolingType.MEAN,
-        normalization=False,
-        sources=ModelSource(hf=existing_base_model),
-        dim=384,
-        size_in_gb=0.47,
-    )
-
-    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
-        TextEmbedding.add_custom_model(
-            custom_model_name,
-            pooling=PoolingType.MEAN,
-            normalization=True,
-            sources=ModelSource(hf=custom_model_name),
-            dim=384,
-            size_in_gb=0.47,
-        )
-
-    CustomTextEmbedding.SUPPORTED_MODELS.clear()
-    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
-
-
-def test_do_not_add_existing_cross_encoder():
-    existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
-    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,
-        )
-
-    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()

tests/test_image_onnx_embeddings.py

@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 from io import BytesIO
 
 import numpy as np
@@ -9,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, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_VECTOR_VALUES = {
     "Qdrant/clip-ViT-B-32-vision": np.array([-0.0098, 0.0128, -0.0274, 0.002, -0.0059]),
@@ -27,109 +26,86 @@ CANONICAL_VECTOR_VALUES = {
     ),
 }
 
-_MODELS_TO_CACHE = ("Qdrant/clip-ViT-B-32-vision",)
-MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
 
-
-@pytest.fixture(scope="module")
-def model_cache():
+def test_embedding() -> None:
     is_ci = os.getenv("CI")
-    cache = {}
 
-    @contextmanager
-    def get_model(model_name: str):
-        lowercase_model_name = model_name.lower()
-        if lowercase_model_name not in cache:
-            cache[lowercase_model_name] = ImageEmbedding(lowercase_model_name)
-        yield cache[lowercase_model_name]
-        if lowercase_model_name not in MODELS_TO_CACHE:
-            model_inst = cache.pop(lowercase_model_name)
-            if is_ci:
-                delete_model_cache(model_inst.model._model_dir)
-            del model_inst
-
-    yield get_model
-
-    if is_ci:
-        for name, model in cache.items():
-            delete_model_cache(model.model._model_dir)
-    cache.clear()
-
-
-@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
-def test_embedding(model_cache, 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 should_test_model(model_desc, model_name, is_ci, is_manual):
+    for model_desc in ImageEmbedding.list_supported_models():
+        if not is_ci and model_desc["size_in_GB"] > 1:
             continue
 
-        dim = model_desc.dim
+        dim = model_desc["dim"]
 
-        with model_cache(model_desc.model) as model:
-            images = [
-                TEST_MISC_DIR / "image.jpeg",
-                str(TEST_MISC_DIR / "small_image.jpeg"),
-                Image.open((TEST_MISC_DIR / "small_image.jpeg")),
-                Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
-            ]
-            embeddings = list(model.embed(images))
-            embeddings = np.stack(embeddings, axis=0)
-            assert embeddings.shape == (len(images), dim)
+        model = ImageEmbedding(model_name=model_desc["model"])
 
-            canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+        images = [
+            TEST_MISC_DIR / "image.jpeg",
+            str(TEST_MISC_DIR / "small_image.jpeg"),
+            Image.open((TEST_MISC_DIR / "small_image.jpeg")),
+            Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
+        ]
+        embeddings = list(model.embed(images))
+        embeddings = np.stack(embeddings, axis=0)
+        assert embeddings.shape == (len(images), dim)
 
-            assert np.allclose(
-                embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-            ), model_desc.model
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
 
-            assert np.allclose(embeddings[1], embeddings[2]), model_desc.model
+        assert np.allclose(
+            embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
+        ), model_desc["model"]
+
+        assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
-def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        n_images = 32
-        test_images = [
-            TEST_MISC_DIR / "image.jpeg",
-            str(TEST_MISC_DIR / "small_image.jpeg"),
-            Image.open(TEST_MISC_DIR / "small_image.jpeg"),
-        ]
-        images = test_images * n_images
+def test_batch_embedding(n_dims: int, model_name: str) -> None:
+    is_ci = os.getenv("CI")
+    model = ImageEmbedding(model_name=model_name)
+    n_images = 32
+    test_images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    images = test_images * n_images
 
-        embeddings = list(model.embed(images, batch_size=10))
-        embeddings = np.stack(embeddings, axis=0)
-        assert np.allclose(embeddings[1], embeddings[2])
+    embeddings = list(model.embed(images, batch_size=10))
+    embeddings = np.stack(embeddings, axis=0)
 
-        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)
+    assert embeddings.shape == (len(test_images) * n_images, n_dims)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
-def test_parallel_processing(model_cache, n_dims: int, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        n_images = 32
-        test_images = [
-            TEST_MISC_DIR / "image.jpeg",
-            str(TEST_MISC_DIR / "small_image.jpeg"),
-            Image.open(TEST_MISC_DIR / "small_image.jpeg"),
-        ]
-        images = test_images * n_images
-        embeddings = list(model.embed(images, batch_size=10, parallel=2))
-        embeddings = np.stack(embeddings, axis=0)
+def test_parallel_processing(n_dims: int, model_name: str) -> None:
+    is_ci = os.getenv("CI")
+    model = ImageEmbedding(model_name=model_name)
 
-        embeddings_2 = list(model.embed(images, batch_size=10, parallel=None))
-        embeddings_2 = np.stack(embeddings_2, axis=0)
+    n_images = 32
+    test_images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    images = test_images * n_images
+    embeddings = list(model.embed(images, batch_size=10, parallel=2))
+    embeddings = np.stack(embeddings, axis=0)
 
-        embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
-        embeddings_3 = np.stack(embeddings_3, axis=0)
+    embeddings_2 = list(model.embed(images, batch_size=10, parallel=None))
+    embeddings_2 = np.stack(embeddings_2, axis=0)
 
-        assert embeddings.shape == (n_images * len(test_images), n_dims)
-        assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-        assert np.allclose(embeddings, embeddings_3, atol=1e-3)
+    embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
+    embeddings_3 = np.stack(embeddings_3, axis=0)
+
+    assert embeddings.shape == (n_images * len(test_images), n_dims)
+    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
+    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
@@ -145,31 +121,3 @@ def test_lazy_load(model_name: str) -> None:
     assert hasattr(model.model, "model")
     if is_ci:
         delete_model_cache(model.model._model_dir)
-
-
-def test_get_embedding_size() -> None:
-    assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-ViT-B-32-vision") == 512
-    assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-vit-b-32-vision") == 512
-
-
-def test_embedding_size() -> None:
-    is_ci = os.getenv("CI")
-    model_name = "Qdrant/clip-ViT-B-32-vision"
-    model = ImageEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 512
-
-    model_name = "Qdrant/clip-vit-b-32-vision"
-    model = ImageEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 512
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
-@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
-def test_session_options(model_cache, model_name) -> None:
-    with model_cache(model_name) as default_model:
-        default_session_options = default_model.model.model.get_session_options()
-        assert default_session_options.enable_cpu_mem_arena is True
-        model = ImageEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
-        session_options = model.model.model.get_session_options()
-        assert session_options.enable_cpu_mem_arena is False

tests/test_late_interaction_embeddings.py

@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 
 import pytest
 import numpy as np
@@ -7,7 +6,7 @@ import numpy as np
 from fastembed.late_interaction.late_interaction_text_embedding import (
     LateInteractionTextEmbedding,
 )
-from tests.utils import delete_model_cache, should_test_model
+from tests.utils import delete_model_cache
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -151,124 +150,82 @@ CANONICAL_QUERY_VALUES = {
     ),
 }
 
-_MODELS_TO_CACHE = ("answerdotai/answerai-colbert-small-v1",)
-MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
-
-
-@pytest.fixture(scope="module")
-def model_cache():
-    is_ci = os.getenv("CI")
-    cache = {}
-
-    @contextmanager
-    def get_model(model_name: str):
-        lowercase_model_name = model_name.lower()
-        if lowercase_model_name not in cache:
-            cache[lowercase_model_name] = LateInteractionTextEmbedding(lowercase_model_name)
-        yield cache[lowercase_model_name]
-        if lowercase_model_name not in MODELS_TO_CACHE:
-            model_inst = cache.pop(lowercase_model_name)
-            if is_ci:
-                delete_model_cache(model_inst.model._model_dir)
-            del model_inst
-
-    yield get_model
-
-    if is_ci:
-        for name, model in cache.items():
-            delete_model_cache(model.model._model_dir)
-    cache.clear()
-
-
 docs = ["Hello World"]
 
 
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_batch_embedding(model_cache, model_name: str):
+def test_batch_embedding():
+    is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    with model_cache(model_name) as model:
+    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))
-        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)
 
-
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_batch_inference_size_same_as_single_inference(model_cache, model_name: str):
-    with model_cache(model_name) as model:
-        docs_to_embed = [
-            "short document",
-            "A bit longer document, which should not affect the size",
-        ]
-        result = list(model.embed(docs_to_embed, batch_size=1))
-        result_2 = list(model.embed(docs_to_embed, batch_size=2))
-        assert len(result[0]) == len(result_2[0])
+        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_cache, model_name: str):
+def test_single_embedding():
     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
-
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         print("evaluating", model_name)
-        with model_cache(model_desc.model) as model:
-            whole_result = list(model.embed(docs_to_embed, batch_size=6))
-            assert len(whole_result) == 1
-            result = whole_result[0]
-            expected_result = CANONICAL_COLUMN_VALUES[model_desc.model]
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+        model = LateInteractionTextEmbedding(model_name=model_name)
+        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        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_cache, model_name: str):
+def test_single_embedding_query():
     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
+    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)
 
-        print("evaluating", model_desc.model)
-        with model_cache(model_desc.model) as model:
-            whole_result = list(model.query_embed(queries_to_embed))
-            assert len(whole_result) == 1
-            result = whole_result[0]
-            expected_result = CANONICAL_QUERY_VALUES[model_desc.model]
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
+        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(model_cache, token_dim: int, model_name: str):
-    with model_cache(model_name) as model:
-        docs = ["hello world", "flag embedding"] * 100
-        embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+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 = 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))  # inherits OnnxTextModel which
-        #         # is tested in TextEmbedding, disabling it here to reduce number of requests to hf
-        #         # multiprocessing is enough to test with `parallel=2`, and `parallel=None` is okay to tests since it reuses
-        #         # model from cache
+    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+    embeddings_3 = np.stack(embeddings_3, axis=0)
 
-        assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
+    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)
 
-        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"])
+@pytest.mark.parametrize(
+    "model_name",
+    ["colbert-ir/colbertv2.0"],
+)
 def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")
 
@@ -287,58 +244,3 @@ def test_lazy_load(model_name: str):
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
-
-
-def test_get_embedding_size():
-    model_name = "answerdotai/answerai-colbert-small-v1"
-    assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 96
-
-    model_name = "answerdotai/answerai-ColBERT-small-v1"
-    assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 96
-
-
-def test_embedding_size():
-    is_ci = os.getenv("CI")
-    model_name = "answerdotai/answerai-colbert-small-v1"
-    model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 96
-
-    model_name = "answerdotai/answerai-ColBERT-small-v1"
-    model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 96
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-ColBERT-small-v1"])
-def test_session_options(model_cache, model_name) -> None:
-    with model_cache(model_name) as default_model:
-        default_session_options = default_model.model.model.get_session_options()
-        assert default_session_options.enable_cpu_mem_arena is True
-        model = LateInteractionTextEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
-        session_options = model.model.model.get_session_options()
-        assert session_options.enable_cpu_mem_arena is False
-
-
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_token_count(model_cache, model_name) -> None:
-    with model_cache(model_name) as model:
-        documents = ["short doc", "it is a long document to check attention mask for paddings"]
-        short_doc_token_count = model.token_count(documents[0])
-        long_doc_token_count = model.token_count(documents[1])
-        documents_token_count = model.token_count(documents)
-        assert short_doc_token_count + long_doc_token_count == documents_token_count
-        # 2 is 2*DOC_MARKER_TOKEN_ID for each document
-        assert short_doc_token_count + long_doc_token_count + 2 == model.token_count(
-            documents, include_extension=True
-        )
-        assert short_doc_token_count + long_doc_token_count == model.token_count(
-            documents, batch_size=1
-        )
-        assert short_doc_token_count + long_doc_token_count == model.token_count(
-            documents, is_doc=False
-        )
-        # query min length is 32
-        assert model.token_count(documents, is_doc=False, include_extension=True) == 64
-        very_long_query = "It's a very long query which definitely contains more than 32 tokens and we're using it to check whether the method can handle large query properly without cutting it to 32 tokens"
-        assert model.token_count(very_long_query, is_doc=False, include_extension=True) > 32

tests/test_late_interaction_multimodal.py

@@ -1,165 +0,0 @@
-import os
-from contextlib import contextmanager
-
-import pytest
-from PIL import Image
-import numpy as np
-
-from fastembed import LateInteractionMultimodalEmbedding
-from tests.config import TEST_MISC_DIR
-from tests.utils import delete_model_cache
-
-# vectors are abridged and rounded for brevity
-CANONICAL_IMAGE_VALUES = {
-    "Qdrant/colpali-v1.3-fp16": np.array(
-        [
-            [-0.0345, -0.022, 0.0567, -0.0518, -0.0782, 0.1714, -0.1738],
-            [-0.1181, -0.099, 0.0268, 0.0774, 0.0228, 0.0563, -0.1021],
-            [-0.117, -0.0683, 0.0371, 0.0921, 0.0107, 0.0659, -0.0666],
-            [-0.1393, -0.0948, 0.037, 0.0951, -0.0126, 0.0678, -0.087],
-            [-0.0957, -0.081, 0.0404, 0.052, 0.0409, 0.0335, -0.064],
-            [-0.0626, -0.0445, 0.056, 0.0592, -0.0229, 0.0409, -0.0301],
-            [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
-        ]
-    ),
-    "Qdrant/colmodernvbert": np.array(
-        [
-            [0.11614, -0.15793, -0.11194, 0.0688, 0.08001, 0.10575, -0.07871],
-            [0.10094, -0.13301, -0.12069, 0.10932, 0.04645, 0.09884, 0.04048],
-            [0.13106, -0.18613, -0.13469, 0.10566, 0.03659, 0.07712, -0.03916],
-            [0.09754, -0.09596, -0.04839, 0.14991, 0.05692, 0.10569, -0.08349],
-            [0.02576, -0.15651, -0.09977, 0.09707, 0.13412, 0.09994, -0.09931],
-            [-0.06741, -0.1787, -0.19677, -0.07618, 0.13102, -0.02131, -0.02437],
-            [-0.02776, -0.10187, -0.13793, 0.03835, 0.04766, 0.04701, -0.15635],
-        ]
-    ),
-}
-
-CANONICAL_QUERY_VALUES = {
-    "Qdrant/colpali-v1.3-fp16": np.array(
-        [
-            [-0.0023, 0.1477, 0.1594, 0.046, -0.0196, 0.0554, 0.1567],
-            [-0.0139, -0.0057, 0.0932, 0.0052, -0.0678, 0.0131, 0.0537],
-            [0.0054, 0.0364, 0.2078, -0.074, 0.0355, 0.061, 0.1593],
-            [-0.0076, -0.0154, 0.2266, 0.0103, 0.0089, -0.024, 0.098],
-            [-0.0274, 0.0098, 0.2106, -0.0634, 0.0616, -0.0021, 0.0708],
-            [0.0074, 0.0025, 0.1631, -0.0802, 0.0418, -0.0219, 0.1022],
-            [-0.0165, -0.0106, 0.1672, -0.0768, 0.0389, -0.0038, 0.1137],
-        ]
-    ),
-    "Qdrant/colmodernvbert": np.array(
-        [
-            [0.05, 0.06557, 0.04026, 0.14981, 0.1842, 0.0263, -0.18706],
-            [-0.05664, -0.14028, 0.00649, -0.02849, 0.09034, -0.01494, 0.10693],
-            [-0.10147, -0.00716, 0.09084, -0.08236, -0.01849, -0.00972, -0.00461],
-            [-0.1233, -0.10814, -0.02337, -0.00329, 0.05984, 0.09934, 0.09846],
-            [-0.07053, -0.13119, -0.06487, 0.01508, 0.07459, 0.07655, 0.14821],
-            [0.00526, -0.13842, -0.05837, -0.02721, 0.13009, 0.05076, 0.17962],
-            [0.00924, -0.14383, -0.03057, -0.03691, 0.11718, 0.037, 0.13344],
-        ]
-    ),
-}
-
-queries = ["hello world", "flag embedding"]
-images = [
-    TEST_MISC_DIR / "image.jpeg",
-    str(TEST_MISC_DIR / "image.jpeg"),
-    Image.open((TEST_MISC_DIR / "image.jpeg")),
-]
-
-_MODELS_TO_CACHE = ("Qdrant/colmodernvbert",)
-MODELS_TO_CACHE = tuple(model_name.lower() for model_name in _MODELS_TO_CACHE)
-
-
-@pytest.fixture(scope="module")
-def model_cache():
-    is_ci = os.getenv("CI")
-    cache = {}
-
-    @contextmanager
-    def get_model(model_name: str):
-        lowercase_model_name = model_name.lower()
-        if lowercase_model_name not in cache:
-            cache[lowercase_model_name] = LateInteractionMultimodalEmbedding(lowercase_model_name)
-        yield cache[lowercase_model_name]
-        if lowercase_model_name not in MODELS_TO_CACHE:
-            model_inst = cache.pop(lowercase_model_name)
-            if is_ci:
-                delete_model_cache(model_inst.model._model_dir)
-            del model_inst
-
-    yield get_model
-
-    if is_ci:
-        for _, model in cache.items():
-            delete_model_cache(model.model._model_dir)
-    cache.clear()
-
-
-def test_batch_embedding(model_cache):
-    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
-            continue  # colpali is too large for ci
-
-        print("evaluating", model_name)
-        with model_cache(model_name) as model:
-            result = list(model.embed_image(images, batch_size=2))
-
-            for value in result:
-                token_num, abridged_dim = expected_result.shape
-                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_single_embedding(model_cache):
-    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
-            continue  # colpali is too large for ci
-        print("evaluating", model_name)
-        with model_cache(model_name) as model:
-            result = next(iter(model.embed_image(images, batch_size=6)))
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_single_embedding_query(model_cache):
-    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
-        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
-            continue  # colpali is too large for ci
-        print("evaluating", model_name)
-        with model_cache(model_name) as model:
-            result = next(iter(model.embed_text(queries)))
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_get_embedding_size():
-    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
-
-    model_name = "Qdrant/colmodernvbert"
-    assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
-
-
-def test_embedding_size():
-    model_name = "Qdrant/colmodernvbert"
-    model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 128
-
-
-def test_token_count(model_cache) -> None:
-    model_name = "Qdrant/colmodernvbert"
-    with model_cache(model_name) as model:
-        documents = ["short doc", "it is a long document to check attention mask for paddings"]
-        short_doc_token_count = model.token_count(documents[0])
-        long_doc_token_count = model.token_count(documents[1])
-        documents_token_count = model.token_count(documents)
-        assert short_doc_token_count + long_doc_token_count == documents_token_count
-        assert short_doc_token_count + long_doc_token_count == model.token_count(
-            documents, batch_size=1
-        )
-        assert short_doc_token_count + long_doc_token_count < model.token_count(
-            documents, include_extension=True
-        )

tests/test_multi_gpu.py

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

tests/test_pooling.py

@@ -0,0 +1,34 @@
+import os
+
+import numpy as np
+
+from fastembed.late_interaction.late_interaction_text_embedding import (
+    LateInteractionTextEmbedding,
+)
+from fastembed.common.pooling import LateInteractionPooler
+
+from tests.utils import delete_model_cache
+
+CANONICAL_COLUMN_VALUES = {
+    "colbert-ir/colbertv2.0": np.array(
+        [4.0727495e-03, -2.4026826e-03, -6.8204990e-04, -7.1383954e-05, 4.4963313e-03]
+    ),
+}
+
+docs = ["Hello World"]
+
+
+def test_batch_embedding():
+    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)
+        pooler = LateInteractionPooler()
+        result = list(model.embed(docs_to_embed, batch_size=6))
+        pooled_result = pooler.pool(result)
+        assert np.allclose(pooled_result[0], expected_result, atol=2e-3)
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)

tests/test_postprocess.py

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

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

tests/test_text_cross_encoder.py

@@ -1,11 +1,10 @@
 import os
-from contextlib import contextmanager
 
 import numpy as np
 import pytest
 
 from fastembed.rerank.cross_encoder import TextCrossEncoder
-from tests.utils import delete_model_cache, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_SCORE_VALUES = {
     "Xenova/ms-marco-MiniLM-L-6-v2": np.array([8.500708, -2.541011]),
@@ -16,84 +15,73 @@ CANONICAL_SCORE_VALUES = {
     "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]),
 }
 
-
-_MODELS_TO_CACHE = ("Xenova/ms-marco-MiniLM-L-6-v2",)
-MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
+SELECTED_MODELS = {
+    "Xenova": "Xenova/ms-marco-MiniLM-L-6-v2",
+    "BAAI": "BAAI/bge-reranker-base",
+    "jinaai": "jinaai/jina-reranker-v1-tiny-en",
+}
 
 
-@pytest.fixture(scope="module")
-def model_cache():
+@pytest.mark.parametrize(
+    "model_name",
+    [model_name for model_name in CANONICAL_SCORE_VALUES],
+)
+def test_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    cache = {}
 
-    @contextmanager
-    def get_model(model_name: str):
-        lowercase_model_name = model_name.lower()
-        if lowercase_model_name not in cache:
-            cache[lowercase_model_name] = TextCrossEncoder(lowercase_model_name)
-        yield cache[lowercase_model_name]
-        if lowercase_model_name not in MODELS_TO_CACHE:
-            model_inst = cache.pop(lowercase_model_name)
-            if is_ci:
-                delete_model_cache(model_inst.model._model_dir)
-            del model_inst
+    model = TextCrossEncoder(model_name=model_name)
 
-    yield get_model
+    query = "What is the capital of France?"
+    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
+    scores = np.array(list(model.rerank(query, documents)))
 
+    pairs = [(query, doc) for doc in documents]
+    scores2 = np.array(list(model.rerank_pairs(pairs)))
+    assert np.allclose(
+        scores, scores2, atol=1e-5
+    ), f"Model: {model_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:
-        for name, model in cache.items():
-            delete_model_cache(model.model._model_dir)
-    cache.clear()
+        delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_rerank(model_cache, model_name: str) -> None:
+@pytest.mark.parametrize(
+    "model_name",
+    [model_name for model_name in SELECTED_MODELS.values()],
+)
+def test_batch_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
-    for model_desc in TextCrossEncoder._list_supported_models():
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
+    model = TextCrossEncoder(model_name=model_name)
 
-        with model_cache(model_desc.model) as model:
-            query = "What is the capital of France?"
-            documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
-            scores = np.array(list(model.rerank(query, documents)))
+    query = "What is the capital of France?"
+    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50
+    scores = np.array(list(model.rerank(query, documents, batch_size=10)))
 
-            pairs = [(query, doc) for doc in documents]
-            scores2 = np.array(list(model.rerank_pairs(pairs)))
-            assert np.allclose(
-                scores, scores2, atol=1e-5
-            ), f"Model: {model_desc.model}, Scores: {scores}, Scores2: {scores2}"
+    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_desc.model]
-            assert np.allclose(
-                scores, canonical_scores, atol=1e-3
-            ), f"Model: {model_desc.model}, Scores: {scores}, Expected: {canonical_scores}"
+    canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50)
+
+    assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}"
+    assert np.allclose(
+        scores, canonical_scores, atol=1e-3
+    ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_batch_rerank(model_cache, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        query = "What is the capital of France?"
-        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50
-        scores = np.array(list(model.rerank(query, documents, batch_size=10)))
-
-        pairs = [(query, doc) for doc in documents]
-        scores2 = np.array(list(model.rerank_pairs(pairs)))
-        assert np.allclose(
-            scores, scores2, atol=1e-5
-        ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
-
-        canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50)
-
-        assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}"
-        assert np.allclose(
-            scores, canonical_scores, atol=1e-3
-        ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
-
-
-@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)
@@ -107,45 +95,25 @@ def test_lazy_load(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_rerank_pairs_parallel(model_cache, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        query = "What is the capital of France?"
-        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 10
-        pairs = [(query, doc) for doc in documents]
-        scores_parallel = np.array(list(model.rerank_pairs(pairs, parallel=2, batch_size=10)))
-        scores_sequential = np.array(list(model.rerank_pairs(pairs, batch_size=10)))
-        assert np.allclose(
-            scores_parallel, scores_sequential, atol=1e-5
-        ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Scores (Sequential): {scores_sequential}"
-        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
-        assert np.allclose(
-            scores_parallel[: len(canonical_scores)], canonical_scores, atol=1e-3
-        ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Expected: {canonical_scores}"
+@pytest.mark.parametrize(
+    "model_name",
+    [model_name for model_name in SELECTED_MODELS.values()],
+)
+def test_rerank_pairs_parallel(model_name: str) -> None:
+    is_ci = os.getenv("CI")
 
-
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_token_count(model_cache, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        pairs = [
-            ("What is the capital of France?", "Paris is the capital of France."),
-            (
-                "Name me a couple of cities were the capitals of Germany?",
-                "Berlin is the current capital of Germany, Bonn is a former capital of Germany.",
-            ),
-        ]
-        first_pair_token_count = model.token_count([pairs[0]])
-        second_pair_token_count = model.token_count([pairs[1]])
-        pairs_token_count = model.token_count(pairs)
-        assert first_pair_token_count + second_pair_token_count == pairs_token_count
-        assert pairs_token_count == model.token_count(pairs, batch_size=1)
-
-
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
-def test_session_options(model_cache, model_name) -> None:
-    with model_cache(model_name) as default_model:
-        default_session_options = default_model.model.model.get_session_options()
-        assert default_session_options.enable_cpu_mem_arena is True
-        model = TextCrossEncoder(model_name=model_name, enable_cpu_mem_arena=False)
-        session_options = model.model.model.get_session_options()
-        assert session_options.enable_cpu_mem_arena is False
+    model = TextCrossEncoder(model_name=model_name)
+    query = "What is the capital of France?"
+    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 10
+    pairs = [(query, doc) for doc in documents]
+    scores_parallel = np.array(list(model.rerank_pairs(pairs, parallel=2, batch_size=10)))
+    scores_sequential = np.array(list(model.rerank_pairs(pairs, batch_size=10)))
+    assert np.allclose(
+        scores_parallel, scores_sequential, atol=1e-5
+    ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Scores (Sequential): {scores_sequential}"
+    canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+    assert np.allclose(
+        scores_parallel[: len(canonical_scores)], canonical_scores, atol=1e-3
+    ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Expected: {canonical_scores}"
+    if is_ci:
+        delete_model_cache(model.model._model_dir)

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 JinaEmbeddingV3, Task
+from fastembed.text.multitask_embedding import Task
 from tests.utils import delete_model_cache
 
 
@@ -60,42 +60,51 @@ CANONICAL_VECTOR_VALUES = {
 docs = ["Hello World", "Follow the white rabbit."]
 
 
-@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
-def test_batch_embedding(dim: int, model_name: str):
+def test_batch_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")
-
     docs_to_embed = docs * 10
     default_task = Task.RETRIEVAL_PASSAGE
 
-    model = TextEmbedding(model_name=model_name)
+    for model_desc in TextEmbedding.list_supported_models():
+        if not is_ci and model_desc["size_in_GB"] > 1:
+            continue
 
-    embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
-    embeddings = np.stack(embeddings, axis=0)
+        model_name = model_desc["model"]
+        dim = model_desc["dim"]
 
-    assert embeddings.shape == (len(docs_to_embed), dim)
+        if model_name not in CANONICAL_VECTOR_VALUES.keys():
+            continue
 
-    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
+        model = TextEmbedding(model_name=model_name)
 
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        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)
 
 
 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 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
+    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"]
+        dim = model_desc["dim"]
+
+        if model_name not in CANONICAL_VECTOR_VALUES.keys():
+            continue
 
         model = TextEmbedding(model_name=model_name)
 
@@ -109,41 +118,26 @@ def test_single_embedding():
 
             canonical_vector = task["vectors"]
             assert np.allclose(
-                embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-            ), model_desc.model
+                embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+            ), model_desc["model"]
 
-        classification_embeddings = list(model.embed(documents=docs, task_id=Task.CLASSIFICATION))
-        classification_embeddings = np.stack(classification_embeddings, axis=0)
-
-        assert classification_embeddings.shape == (len(docs), dim)
-
-        model = TextEmbedding(model_name=model_name, task_id=Task.CLASSIFICATION)
-        default_embeddings = list(model.embed(documents=docs))
-        default_embeddings = np.stack(default_embeddings, axis=0)
-
-        assert default_embeddings.shape == (len(docs), dim)
-
-        assert np.allclose(
-            classification_embeddings,
-            default_embeddings,
-            atol=1e-4,
-        ), model_desc.model
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
 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 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
+    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"]
+        dim = model_desc["dim"]
+
+        if model_name not in CANONICAL_VECTOR_VALUES.keys():
+            continue
 
         model = TextEmbedding(model_name=model_name)
 
@@ -156,8 +150,8 @@ def test_single_embedding_query():
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
-            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-        ), model_desc.model
+            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)
@@ -165,17 +159,17 @@ 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 JinaEmbeddingV3._list_supported_models():
-        # todo: once we add more models, we should not test models >1GB size locally
+    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
-        dim = model_desc.dim
+        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)
 
@@ -188,22 +182,21 @@ def test_single_embedding_passage():
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
-            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-        ), model_desc.model
+            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)
 
 
-@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
-def test_parallel_processing(dim: int, model_name: str):
+def test_parallel_processing():
     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
@@ -223,14 +216,33 @@ def test_parallel_processing(dim: int, model_name: str):
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["jinaai/jina-embeddings-v3"])
+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"],
+)
 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

@@ -1,12 +1,11 @@
 import os
 import platform
-from contextlib import contextmanager
 
 import numpy as np
 import pytest
 
 from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
@@ -53,7 +52,7 @@ CANONICAL_VECTOR_VALUES = {
         [0.0802303, 0.3700881, -4.3053818, 0.4431803, -0.271572]
     ),
     "thenlper/gte-large": np.array(
-        [-0.00986551, -0.00018734, 0.00605892, -0.03289612, -0.0387564],
+        [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
     ),
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]
@@ -72,92 +71,82 @@ CANONICAL_VECTOR_VALUES = {
 
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
-_MODELS_TO_CACHE = ("BAAI/bge-small-en-v1.5",)
-MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
 
-
-@pytest.fixture(scope="module")
-def model_cache():
-    is_ci = os.getenv("CI")
-    cache = {}
-
-    @contextmanager
-    def get_model(model_name: str):
-        lowercase_model_name = model_name.lower()
-        if lowercase_model_name not in cache:
-            cache[lowercase_model_name] = TextEmbedding(lowercase_model_name)
-        yield cache[lowercase_model_name]
-        if lowercase_model_name not in MODELS_TO_CACHE:
-            model_inst = cache.pop(lowercase_model_name)
-            if is_ci:
-                delete_model_cache(model_inst.model._model_dir)
-            del model_inst
-
-    yield get_model
-
-    if is_ci:
-        for name, model in cache.items():
-            delete_model_cache(model.model._model_dir)
-    cache.clear()
-
-
-@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
-def test_embedding(model_cache, model_name: str) -> None:
+def test_embedding() -> 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 model_desc.model in MULTI_TASK_MODELS or (
-            is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q"
+    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")
         ):
             continue
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
 
-        dim = model_desc.dim
+        dim = model_desc["dim"]
 
-        with model_cache(model_desc.model) as model:
-            docs = ["hello world", "flag embedding"]
-            embeddings = list(model.embed(docs))
-            embeddings = np.stack(embeddings, axis=0)
-            assert embeddings.shape == (2, dim)
-
-            canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
-            assert np.allclose(
-                embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-            ), model_desc.model
-
-
-@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
-def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        docs = ["hello world", "flag embedding"] * 100
-        embeddings = list(model.embed(docs, batch_size=10))
+        model = TextEmbedding(model_name=model_desc["model"])
+        docs = ["hello world", "flag embedding"]
+        embeddings = list(model.embed(docs))
         embeddings = np.stack(embeddings, axis=0)
+        assert embeddings.shape == (2, dim)
 
-        assert embeddings.shape == (len(docs), n_dims)
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
+        assert np.allclose(
+            embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
+        ), model_desc["model"]
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
-def test_parallel_processing(model_cache, n_dims: int, model_name: str) -> None:
-    with model_cache(model_name) as model:
-        docs = ["hello world", "flag embedding"] * 100
-        embeddings = list(model.embed(docs, batch_size=10, parallel=2))
-        embeddings = np.stack(embeddings, axis=0)
+@pytest.mark.parametrize(
+    "n_dims,model_name",
+    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
+)
+def test_batch_embedding(n_dims: int, model_name: str) -> None:
+    is_ci = os.getenv("CI")
+    model = TextEmbedding(model_name=model_name)
 
-        embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
-        embeddings_2 = np.stack(embeddings_2, axis=0)
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10))
+    embeddings = np.stack(embeddings, axis=0)
 
-        embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
-        embeddings_3 = np.stack(embeddings_3, axis=0)
-
-        assert embeddings.shape == (len(docs), n_dims)
-        assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-        assert np.allclose(embeddings, embeddings_3, atol=1e-3)
+    assert embeddings.shape == (200, n_dims)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+@pytest.mark.parametrize(
+    "n_dims,model_name",
+    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
+)
+def test_parallel_processing(n_dims: int, model_name: str) -> None:
+    is_ci = os.getenv("CI")
+    model = TextEmbedding(model_name=model_name)
+
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+    embeddings = np.stack(embeddings, axis=0)
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+    embeddings_2 = np.stack(embeddings_2, axis=0)
+
+    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+    embeddings_3 = np.stack(embeddings_3, axis=0)
+
+    assert embeddings.shape == (200, n_dims)
+    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
+    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+@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)
@@ -174,46 +163,3 @@ def test_lazy_load(model_name: str) -> None:
 
     if is_ci:
         delete_model_cache(model.model._model_dir)
-
-
-def test_get_embedding_size() -> None:
-    assert TextEmbedding.get_embedding_size("sentence-transformers/all-MiniLM-L6-v2") == 384
-    assert TextEmbedding.get_embedding_size("sentence-transformers/all-minilm-l6-v2") == 384
-
-
-def test_embedding_size() -> None:
-    is_ci = os.getenv("CI")
-    model_name = "sentence-transformers/all-MiniLM-L6-v2"
-    model = TextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 384
-
-    model_name = "sentence-transformers/all-minilm-l6-v2"
-    model = TextEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 384
-
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
-@pytest.mark.parametrize("model_name", ["sentence-transformers/all-MiniLM-L6-v2"])
-def test_session_options(model_cache, model_name) -> None:
-    with model_cache(model_name) as default_model:
-        default_session_options = default_model.model.model.get_session_options()
-        assert default_session_options.enable_cpu_mem_arena is True
-        model = TextEmbedding(model_name=model_name, enable_cpu_mem_arena=False)
-        session_options = model.model.model.get_session_options()
-        assert session_options.enable_cpu_mem_arena is False
-
-
-@pytest.mark.parametrize("model_name", ["sentence-transformers/all-MiniLM-L6-v2"])
-def test_token_count(model_cache, model_name) -> None:
-    with model_cache(model_name) as model:
-        documents = [
-            "Name me a couple of cities were the capitals of Germany?",
-            "Berlin is the current capital of Germany, Bonn is a former capital of Germany.",
-        ]
-        first_doc_token_count = model.token_count(documents[0])
-        second_doc_token_count = model.token_count(documents[1])
-        doc_token_count = model.token_count(documents)
-        assert first_doc_token_count + second_doc_token_count == doc_token_count
-        assert doc_token_count == model.token_count(documents, batch_size=1)

tests/type_stub.py

@@ -1,56 +0,0 @@
-from fastembed import TextEmbedding, LateInteractionTextEmbedding, SparseTextEmbedding
-from fastembed.sparse.bm25 import Bm25
-from fastembed.rerank.cross_encoder import TextCrossEncoder
-
-
-text_embedder = TextEmbedding(cache_dir="models")
-late_interaction_embedder = LateInteractionTextEmbedding(model_name="", cache_dir="models")
-reranker = TextCrossEncoder(model_name="", cache_dir="models")
-sparse_embedder = SparseTextEmbedding(model_name="", cache_dir="models")
-bm25_embedder = Bm25(
-    model_name="",
-    k=1.0,
-    b=1.0,
-    avg_len=1.0,
-    language="",
-    token_max_length=1,
-    disable_stemmer=False,
-    specific_model_path="models",
-)
-
-text_embedder.list_supported_models()
-text_embedder.embed(documents=[""], batch_size=1, parallel=1)
-text_embedder.embed(documents="", parallel=None, task_id=1)
-text_embedder.query_embed(query=[""], batch_size=1, parallel=1)
-text_embedder.query_embed(query="", parallel=None)
-text_embedder.passage_embed(texts=[""], batch_size=1, parallel=1)
-text_embedder.passage_embed(texts=[""], parallel=None)
-
-late_interaction_embedder.list_supported_models()
-late_interaction_embedder.embed(documents=[""], batch_size=1, parallel=1)
-late_interaction_embedder.embed(documents="", parallel=None)
-late_interaction_embedder.query_embed(query=[""], batch_size=1, parallel=1)
-late_interaction_embedder.query_embed(query="", parallel=None)
-late_interaction_embedder.passage_embed(texts=[""], batch_size=1, parallel=1)
-late_interaction_embedder.passage_embed(texts=[""], parallel=None)
-
-reranker.list_supported_models()
-reranker.rerank(query="", documents=[""], batch_size=1, parallel=1)
-reranker.rerank(query="", documents=[""], parallel=None)
-reranker.rerank_pairs(pairs=[("", "")], batch_size=1, parallel=1)
-reranker.rerank_pairs(pairs=[("", "")], parallel=None)
-
-sparse_embedder.list_supported_models()
-sparse_embedder.embed(documents=[""], batch_size=1, parallel=1)
-sparse_embedder.embed(documents="", batch_size=1, parallel=None)
-sparse_embedder.query_embed(query=[""], batch_size=1, parallel=1)
-sparse_embedder.query_embed(query="", batch_size=1, parallel=None)
-sparse_embedder.passage_embed(texts=[""], batch_size=1, parallel=1)
-sparse_embedder.passage_embed(texts=[""], batch_size=1, parallel=None)
-
-bm25_embedder.list_supported_models()
-bm25_embedder.embed(documents=[""], batch_size=1, parallel=1)
-bm25_embedder.embed(documents="", batch_size=1, parallel=None)
-bm25_embedder.query_embed(query=[""], batch_size=1, parallel=1)
-bm25_embedder.query_embed(query="", batch_size=1, parallel=None)
-bm25_embedder.raw_embed(documents=[""])

tests/utils.py

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