Clear supported_models list to keep only supported models

.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-tests.yml

@@ -1,10 +1,9 @@
 name: Tests
 
 on:
-  pull_request:
+  push:
     branches: [ master, main, gpu ]
-  workflow_dispatch:
-  
+  pull_request:
 
 env:
   CARGO_TERM_COLOR: always
@@ -42,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.9", "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 Optional, Any
-
-
-@dataclass(frozen=True)
-class ModelSource:
-    hf: Optional[str] = None
-    url: Optional[str] = None
-    _deprecated_tar_struct: bool = False
-
-    @property
-    def deprecated_tar_struct(self) -> bool:
-        return self._deprecated_tar_struct
-
-    def __post_init__(self) -> None:
-        if self.hf is None and self.url is None:
-            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: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = 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: Optional[bool] = None
-    vocab_size: Optional[int] = None
-
-
-class PoolingType(str, Enum):
-    CLS = "CLS"
-    MEAN = "MEAN"
-    DISABLED = "DISABLED"

fastembed/common/model_management.py

@@ -4,7 +4,7 @@ import json
 import shutil
 import tarfile
 from pathlib import Path
-from typing import Any, Optional, Union, TypeVar, Generic
+from typing import Any, Optional
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -16,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
@@ -29,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.
 
@@ -74,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__}.")
@@ -146,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.
         """
@@ -179,8 +148,8 @@ class ModelManagement(Generic[T]):
 
         def _collect_file_metadata(
             model_dir: Path, repo_files: list[RepoFile]
-        ) -> dict[str, dict[str, Union[int, str]]]:
-            meta: dict[str, dict[str, Union[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:
@@ -192,9 +161,7 @@ class ModelManagement(Generic[T]):
                         }
             return meta
 
-        def _save_file_metadata(
-            model_dir: Path, meta: dict[str, dict[str, Union[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)
@@ -326,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
@@ -375,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:
                 ```
                 {
@@ -405,22 +369,22 @@ class ModelManagement(Generic[T]):
         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
+        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:
-                extra_patterns = [model.model_file]
-                extra_patterns.extend(model.additional_files)
+                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,
                         )
@@ -436,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:
@@ -455,4 +418,4 @@ class ModelManagement(Generic[T]):
             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

@@ -6,10 +6,7 @@ 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
+from fastembed.common.types import OnnxProvider
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -18,35 +15,26 @@ T = TypeVar("T")
 
 @dataclass
 class OnnxOutputContext:
-    model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
+    model_output: np.ndarray
+    attention_mask: Optional[np.ndarray] = None
+    input_ids: Optional[np.ndarray] = None
 
 
 class OnnxModel(Generic[T]):
     @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: Optional[ort.InferenceSession] = None
-        self.tokenizer: Optional[Tokenizer] = 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.
         """
@@ -82,7 +70,7 @@ class OnnxModel(Generic[T]):
             onnx_providers = ["CPUExecutionProvider"]
 
         available_providers = ort.get_available_providers()
-        requested_provider_names: list[str] = []
+        requested_provider_names = []
         for provider in onnx_providers:
             # check providers available
             provider_name = provider if isinstance(provider, str) else provider[0]
@@ -103,7 +91,6 @@ class OnnxModel(Generic[T]):
             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(
@@ -120,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__(
@@ -138,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/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:
@@ -60,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,9 +1,7 @@
 from pathlib import Path
 import sys
 from PIL import Image
-from typing import Any, Union
-import numpy as np
-from numpy.typing import NDArray
+from typing import Any, Iterable, Union
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -12,14 +10,7 @@ else:
 
 
 PathInput: TypeAlias = Union[str, Path]
-ImageInput: TypeAlias = Union[PathInput, Image.Image]
+PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
+ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
 
 OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]
-NumpyArray = Union[
-    NDArray[np.float64],
-    NDArray[np.float32],
-    NDArray[np.float16],
-    NDArray[np.int8],
-    NDArray[np.int64],
-    NDArray[np.int32],
-]

fastembed/common/utils.py

@@ -8,14 +8,11 @@ from itertools import islice
 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))

fastembed/image/image_embedding.py

@@ -1,11 +1,10 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
 
-from fastembed.common.types import NumpyArray
 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,13 +35,9 @@ 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__(
@@ -58,8 +53,8 @@ class ImageEmbedding(ImageEmbeddingBase):
     ):
         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: Optional[int] = None
-        for description in descriptions:
-            if description.model.lower() == model_name.lower():
-                embedding_size = description.dim
-                break
-        if embedding_size is None:
-            model_names = [description.model for description in descriptions]
-            raise ValueError(
-                f"Embedding size for model {model_name} was None. "
-                f"Available model names: {model_names}"
-            )
-        return embedding_size
-
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 16,
         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,12 +1,12 @@
-from typing import Iterable, Optional, Any, Union
+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,
@@ -18,15 +18,14 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 16,
         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,65 +1,72 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Optional, Sequence, Type
 
-
-from fastembed.common.types import NumpyArray
+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,
@@ -104,20 +111,20 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         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]
+        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:
@@ -129,7 +136,7 @@ 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,
@@ -137,22 +144,22 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         )
 
     @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: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 16,
         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.
@@ -178,31 +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,
             **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, Optional, Sequence, Type, Union
+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
 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: Optional[Compose] = 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.
         """
@@ -69,9 +58,8 @@ class OnnxImageModel(OnnxModel[T]):
     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():
@@ -79,11 +67,10 @@ class OnnxImageModel(OnnxModel[T]):
                 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)
 
@@ -91,14 +78,12 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 256,
         parallel: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -115,7 +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()
@@ -125,8 +110,6 @@ class OnnxImageModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -138,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,10 +1,8 @@
-from typing import Union
+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":
@@ -15,9 +13,9 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Union[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:
@@ -42,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
@@ -63,35 +61,38 @@ def center_crop(
 
 
 def normalize(
-    image: NumpyArray,
-    mean: Union[float, list[float]],
-    std: Union[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(
@@ -113,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: Union[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

fastembed/image/transform/operators.py

@@ -1,8 +1,8 @@
 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,
@@ -15,7 +15,7 @@ from fastembed.image.transform.functional import (
 
 
 class Transform:
-    def __call__(self, images: list[Any]) -> Union[list[Image.Image], list[NumpyArray]]:
+    def __call__(self, images: list) -> Union[list[Image.Image], list[np.ndarray]]:
         raise NotImplementedError("Subclasses must implement this method")
 
 
@@ -28,7 +28,7 @@ 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]
 
 
@@ -37,7 +37,7 @@ class Normalize(Transform):
         self.mean = mean
         self.std = std
 
-    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+    def __call__(self, images: list[np.ndarray]) -> list[np.ndarray]:
         return [normalize(image, mean=self.mean, std=self.std) for image in images]
 
 
@@ -58,12 +58,12 @@ class Rescale(Transform):
     def __init__(self, scale: float = 1 / 255):
         self.scale = scale
 
-    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+    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[Union[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]
 
 
@@ -71,7 +71,7 @@ class PadtoSquare(Transform):
     def __init__(
         self,
         size: int,
-        fill_color: Union[str, int, tuple[int, ...]],
+        fill_color: Optional[Union[str, int, tuple[int, ...]]] = None,
     ):
         self.size = size
         self.fill_color = fill_color
@@ -87,8 +87,8 @@ class Compose:
         self.transforms = transforms
 
     def __call__(
-        self, images: Union[list[Image.Image], list[NumpyArray]]
-    ) -> Union[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
@@ -122,7 +122,7 @@ class Compose:
         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)
@@ -139,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:
@@ -202,16 +202,15 @@ class Compose:
     @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

fastembed/late_interaction/colbert.py

@@ -4,7 +4,6 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 import numpy as np
 from tokenizers import Encoding
 
-from fastembed.common.types import NumpyArray
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
@@ -12,42 +11,45 @@ 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="Late interaction model",
-        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), Multilingual (~100 languages), 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
+            return output.model_output.astype(np.float32)
 
         if output.input_ids is None or output.attention_mask is None:
             raise ValueError(
@@ -55,28 +57,22 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             )
 
         for i, token_sequence in enumerate(output.input_ids):
-            for j, token_id in enumerate(token_sequence):  # type: ignore
+            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
 
-        output.model_output *= np.expand_dims(output.attention_mask, 2)
+        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
-
-        for embedding, attention_mask in zip(output.model_output, output.attention_mask):
-            yield embedding[attention_mask == 1]
+        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]:
@@ -87,7 +83,6 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     def _tokenize_query(self, query: str) -> list[Encoding]:
-        assert self.tokenizer is not None
         encoded = self.tokenizer.encode_batch([query])
         # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
         if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
@@ -107,15 +102,15 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         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
 
     @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
 
@@ -163,25 +158,25 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         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]
+        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: Optional[int] = None
-        self.pad_token_id: Optional[int] = None
-        self.skip_list: set[int] = set()
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.skip_list = set()
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -189,13 +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,
         )
-        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 = {
@@ -212,7 +206,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         batch_size: int = 256,
         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.
@@ -237,12 +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,
             **kwargs,
         )
 
-    def query_embed(self, query: Union[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]
 
@@ -255,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,11 +1,11 @@
 from typing import Iterable, Optional, Union, Any
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
+import numpy as np
+
 from fastembed.common.model_management import ModelManagement
 
 
-class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
+class LateInteractionTextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
@@ -17,7 +17,6 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -25,10 +24,10 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         batch_size: int = 256,
         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: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -51,21 +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")

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,8 +1,7 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
+import numpy as np
+
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -39,13 +38,9 @@ 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__(
@@ -61,8 +56,8 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
     ):
         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: Optional[int] = None
-        for description in descriptions:
-            if description.model.lower() == model_name.lower():
-                embedding_size = description.dim
-                break
-        if embedding_size is None:
-            model_names = [description.model for description in descriptions]
-            raise ValueError(
-                f"Embedding size for model {model_name} was None. "
-                f"Available model names: {model_names}"
-            )
-        return embedding_size
-
     def embed(
         self,
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         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: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -146,7 +107,7 @@ 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

fastembed/late_interaction/token_embeddings.py

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

fastembed/late_interaction_multimodal/__init__.py

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

fastembed/late_interaction_multimodal/colpali.py

@@ -1,305 +0,0 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
-
-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
-from fastembed.common.utils import define_cache_dir
-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: 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
-
-        # 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
-
-        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,
-        )
-
-    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 _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: Union[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,
-            **kwargs,
-        )
-
-    def embed_image(
-        self,
-        images: Union[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,
-            **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,164 +0,0 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
-
-from fastembed.common import OnnxProvider, ImageInput
-from fastembed.common.types import NumpyArray
-from fastembed.late_interaction_multimodal.colpali import ColPali
-
-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]
-
-    @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: 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):
-                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: Optional[int] = None
-        for description in descriptions:
-            if description.model.lower() == model_name.lower():
-                embedding_size = description.dim
-                break
-        if embedding_size is None:
-            model_names = [description.model for description in descriptions]
-            raise ValueError(
-                f"Embedding size for model {model_name} was None. "
-                f"Available model names: {model_names}"
-            )
-        return embedding_size
-
-    def embed_text(
-        self,
-        documents: Union[str, Iterable[str]],
-        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.model.embed_text(documents, batch_size, parallel, **kwargs)
-
-    def embed_image(
-        self,
-        images: Union[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 image
-        """
-        yield from self.model.embed_image(images, batch_size, parallel, **kwargs)

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -1,78 +0,0 @@
-from typing import Iterable, Optional, Union, 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: 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: Optional[int] = None
-
-    def embed_text(
-        self,
-        documents: Union[str, Iterable[str]],
-        batch_size: int = 256,
-        parallel: Optional[int] = 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: Union[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 image
-        """
-        raise NotImplementedError()
-
-    @classmethod
-    def get_embedding_size(cls, model_name: str) -> int:
-        """Returns embedding size of the chosen model."""
-        raise NotImplementedError("Subclasses must implement this method")
-
-    @property
-    def embedding_size(self) -> int:
-        """Returns embedding size for the current model"""
-        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -1,279 +0,0 @@
-import contextlib
-import os
-from multiprocessing import get_all_start_methods
-from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
-
-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
-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: Optional[list[str]] = None
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.tokenizer: Optional[Tokenizer] = None
-        self.processor: Optional[Compose] = 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: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-    ) -> None:
-        super()._load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_file,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_id=device_id,
-        )
-        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: Union[str, Iterable[str]],
-        batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
-    ) -> Iterable[T]:
-        is_small = False
-
-        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,
-            }
-
-            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():
-            image_files = [
-                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: Union[Iterable[ImageInput], ImageInput],
-        batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
-    ) -> Iterable[T]:
-        is_small = False
-
-        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,
-            }
-
-            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

@@ -140,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):

fastembed/postprocess/__init__.py

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

fastembed/postprocess/muvera.py

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

fastembed/py.typed

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

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

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

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -10,67 +10,78 @@ from fastembed.rerank.cross_encoder.onnx_text_model import (
     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
 
@@ -123,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: 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]
+        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:
@@ -145,7 +156,7 @@ 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,
@@ -190,8 +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,
             **kwargs,
         )
 
@@ -199,9 +208,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     def _get_worker_class(cls) -> Type[TextRerankerWorker]:
         return TextCrossEncoderWorker
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[float]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
         return (float(elem) for elem in output.model_output)
 
 

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,9 +1,10 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, 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,7 +13,6 @@ from fastembed.common.onnx_model import (
     OnnxOutputContext,
     OnnxProvider,
 )
-from fastembed.common.types import NumpyArray
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
@@ -43,14 +43,15 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             device_id=device_id,
         )
         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:
@@ -71,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(
@@ -94,8 +95,6 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -122,8 +121,6 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -135,49 +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
 
 
-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,21 +1,13 @@
 from typing import Any, Iterable, Optional, Sequence, Type
-from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
-from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
-
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import (
-    ModelSource,
-    BaseModelDescription,
-)
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
     CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
         OnnxTextCrossEncoder,
-        CustomTextCrossEncoder,
     ]
 
     @classmethod
@@ -23,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:
                 ```
@@ -41,13 +33,9 @@ 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__(
@@ -64,8 +52,8 @@ class TextCrossEncoder(TextCrossEncoderBase):
         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,
@@ -130,34 +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: Optional[list[str]] = None,
-    ) -> None:
-        registered_models = cls._list_supported_models()
-        for registered_model in registered_models:
-            if model == registered_model.model:
-                raise ValueError(
-                    f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, "
-                    f"please use another model name"
-                )
-
-        CustomTextCrossEncoder.add_model(
-            BaseModelDescription(
-                model=model,
-                sources=sources,
-                model_file=model_file,
-                description=description,
-                license=license,
-                size_in_GB=size_in_gb,
-                additional_files=additional_files or [],
-            )
-        )

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -1,10 +1,9 @@
 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,

fastembed/sparse/bm25.py

@@ -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,
+    },
 ]
 
 
@@ -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
 
@@ -161,8 +170,6 @@ class Bm25(SparseTextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
     ) -> Iterable[SparseEmbedding]:
         is_small = False
 
@@ -191,8 +198,6 @@ class Bm25(SparseTextEmbeddingBase):
                 "language": self.language,
                 "token_max_length": self.token_max_length,
                 "disable_stemmer": self.disable_stemmer,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
             }
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
@@ -201,7 +206,7 @@ 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,
@@ -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)
@@ -283,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
 
@@ -340,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

@@ -15,20 +15,21 @@ from fastembed.sparse.sparse_embedding_base import (
     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 = {
@@ -101,27 +102,27 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         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]
+        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(MODEL_TO_LANGUAGE[model_name])
@@ -133,21 +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,
         )
-
-        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
@@ -155,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))
@@ -165,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))
@@ -174,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:
@@ -206,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))
@@ -218,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
@@ -243,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)
@@ -253,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
 
@@ -302,13 +300,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             cuda=self.cuda,
             device_ids=self.device_ids,
             alpha=self.alpha,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
         )
 
     @classmethod
     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
@@ -329,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)
@@ -338,11 +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
 
 
-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,356 +0,0 @@
-from pathlib import Path
-
-from typing import Any, Optional, Sequence, Iterable, Union, Type
-
-import numpy as np
-from numpy.typing import NDArray
-from py_rust_stemmers import SnowballStemmer
-from tokenizers import Tokenizer
-
-from fastembed.common.model_description import SparseModelDescription, ModelSource
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common import OnnxProvider
-from fastembed.common.utils import define_cache_dir
-from fastembed.sparse.sparse_embedding_base import (
-    SparseEmbedding,
-    SparseTextEmbeddingBase,
-)
-from fastembed.sparse.utils.minicoil_encoder import Encoder
-from fastembed.sparse.utils.sparse_vectors_converter import SparseVectorConverter, WordEmbedding
-from fastembed.sparse.utils.vocab_resolver import VocabResolver, VocabTokenizer
-from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-
-
-MINICOIL_MODEL_FILE = "minicoil.triplet.model.npy"
-MINICOIL_VOCAB_FILE = "minicoil.triplet.model.vocab"
-STOPWORDS_FILE = "stopwords.txt"
-
-
-supported_minicoil_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="Qdrant/minicoil-v1",
-        vocab_size=19125,
-        description="Sparse embedding model, that resolves semantic meaning of the words, "
-        "while keeping exact keyword match behavior. "
-        "Based on jinaai/jina-embeddings-v2-small-en-tokens",
-        license="apache-2.0",
-        size_in_GB=0.09,
-        sources=ModelSource(hf="Qdrant/minicoil-v1"),
-        model_file="onnx/model.onnx",
-        additional_files=[
-            STOPWORDS_FILE,
-            MINICOIL_MODEL_FILE,
-            MINICOIL_VOCAB_FILE,
-        ],
-        requires_idf=True,
-    ),
-]
-
-MODEL_TO_LANGUAGE = {
-    "Qdrant/minicoil-v1": "english",
-}
-
-
-class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
-    """
-        MiniCOIL is a sparse embedding model, that resolves semantic meaning of the words,
-        while keeping exact keyword match behavior.
-
-        Each vocabulary token is converted into 4d component of a sparse vector, which is then weighted by the token frequency in the corpus.
-        If the token is not found in the corpus, it is treated exactly like in BM25.
-    `
-        The model is based on `jinaai/jina-embeddings-v2-small-en-tokens`
-    """
-
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        k: float = 1.2,
-        b: float = 0.75,
-        avg_len: float = 150.0,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
-        lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
-    ):
-        """
-        Args:
-            model_name (str): The name of the model to use.
-            cache_dir (str, optional): The path to the cache directory.
-                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
-                                       Defaults to `fastembed_cache` in the system's temp directory.
-            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
-            providers (Optional[Sequence[OnnxProvider]], optional): The providers to use for onnxruntime.
-            k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.
-                I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
-            b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
-                Defaults to 0.75.
-            avg_len (float, optional): The average length of the documents in the corpus. Defaults to 150.0.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
-            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
-            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
-                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
-            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
-            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
-
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-
-        super().__init__(model_name, cache_dir, threads, **kwargs)
-        self.providers = providers
-        self.lazy_load = lazy_load
-        self.device_ids = device_ids
-        self.cuda = cuda
-        self.device_id = device_id
-        self.k = k
-        self.b = b
-        self.avg_len = avg_len
-
-        # Initialize class attributes
-        self.tokenizer: Optional[Tokenizer] = None
-        self.invert_vocab: dict[int, str] = {}
-        self.special_tokens: set[str] = set()
-        self.special_tokens_ids: set[int] = set()
-        self.stopwords: set[str] = set()
-        self.vocab_resolver: Optional[VocabResolver] = None
-        self.encoder: Optional[Encoder] = None
-        self.output_dim: Optional[int] = None
-        self.sparse_vector_converter: Optional[SparseVectorConverter] = None
-
-        self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
-        self._specific_model_path = specific_model_path
-        self._model_dir = self.download_model(
-            self.model_description,
-            self.cache_dir,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-        )
-
-        if not self.lazy_load:
-            self.load_onnx_model()
-
-    def load_onnx_model(self) -> None:
-        self._load_onnx_model(
-            model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
-            threads=self.threads,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_id=self.device_id,
-        )
-
-        assert self.tokenizer is not None
-
-        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
-            self.invert_vocab[idx] = token
-        self.special_tokens = set(self.special_token_to_id.keys())
-        self.special_tokens_ids = set(self.special_token_to_id.values())
-        self.stopwords = set(self._load_stopwords(self._model_dir))
-
-        stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[self.model_name])
-
-        self.vocab_resolver = VocabResolver(
-            tokenizer=VocabTokenizer(self.tokenizer),
-            stopwords=self.stopwords,
-            stemmer=stemmer,
-        )
-        self.vocab_resolver.load_json_vocab(str(self._model_dir / MINICOIL_VOCAB_FILE))
-
-        weights = np.load(str(self._model_dir / MINICOIL_MODEL_FILE), mmap_mode="r")
-        self.encoder = Encoder(weights)
-        self.output_dim = self.encoder.output_dim
-
-        self.sparse_vector_converter = SparseVectorConverter(
-            stopwords=self.stopwords,
-            stemmer=stemmer,
-            k=self.k,
-            b=self.b,
-            avg_len=self.avg_len,
-        )
-
-    def embed(
-        self,
-        documents: Union[str, Iterable[str]],
-        batch_size: int = 256,
-        parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[SparseEmbedding]:
-        """
-        Encode a list of documents into list of embeddings.
-        We use mean pooling with attention so that the model can handle variable-length inputs.
-
-        Args:
-            documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per document
-        """
-        yield from self._embed_documents(
-            model_name=self.model_name,
-            cache_dir=str(self.cache_dir),
-            documents=documents,
-            batch_size=batch_size,
-            parallel=parallel,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_ids=self.device_ids,
-            k=self.k,
-            b=self.b,
-            avg_len=self.avg_len,
-            is_query=False,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            **kwargs,
-        )
-
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
-        """
-        Encode a list of queries into list of embeddings.
-        """
-        yield from self._embed_documents(
-            model_name=self.model_name,
-            cache_dir=str(self.cache_dir),
-            documents=query,
-            providers=self.providers,
-            cuda=self.cuda,
-            device_ids=self.device_ids,
-            k=self.k,
-            b=self.b,
-            avg_len=self.avg_len,
-            is_query=True,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
-            **kwargs,
-        )
-
-    @classmethod
-    def _load_stopwords(cls, model_dir: Path) -> list[str]:
-        stopwords_path = model_dir / STOPWORDS_FILE
-        if not stopwords_path.exists():
-            return []
-
-        with open(stopwords_path, "r") as f:
-            return f.read().splitlines()
-
-    @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
-        """Lists the supported models.
-
-        Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
-        """
-        return supported_minicoil_models
-
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, is_query: bool = False, **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
-        if output.input_ids is None:
-            raise ValueError("input_ids must be provided for document post-processing")
-
-        assert self.vocab_resolver is not None
-        assert self.encoder is not None
-        assert self.sparse_vector_converter is not None
-
-        # Size: (batch_size, sequence_length, hidden_size)
-        embeddings = output.model_output
-        # Size: (batch_size, sequence_length)
-        assert output.attention_mask is not None
-        masks = output.attention_mask
-
-        vocab_size = self.vocab_resolver.vocab_size()
-        embedding_size = self.encoder.output_dim
-
-        # For each document we only select those embeddings that are not masked out
-
-        for i in range(embeddings.shape[0]):
-            # Size: (sequence_length, hidden_size)
-            token_embeddings = embeddings[i, masks[i] == 1]
-
-            # Size: (sequence_length)
-            token_ids: NDArray[np.int64] = output.input_ids[i, masks[i] == 1]
-
-            word_ids_array, counts, oov, forms = self.vocab_resolver.resolve_tokens(token_ids)
-
-            # Size: (1, words)
-            word_ids_array_expanded: NDArray[np.int64] = np.expand_dims(word_ids_array, axis=0)
-
-            # Size: (1, words, embedding_size)
-            token_embeddings_array: NDArray[np.float32] = np.expand_dims(token_embeddings, axis=0)
-
-            assert word_ids_array_expanded.shape[1] == token_embeddings_array.shape[1]
-
-            # Size of word_ids_mapping: (unique_words, 2) - [vocab_id, batch_id]
-            # Size of embeddings: (unique_words, embedding_size)
-            ids_mapping, minicoil_embeddings = self.encoder.forward(
-                word_ids_array_expanded, token_embeddings_array
-            )
-
-            # Size of counts: (unique_words)
-            words_ids: list[int] = ids_mapping[:, 0].tolist()  # type: ignore[assignment]
-
-            sentence_result: dict[str, WordEmbedding] = {}
-
-            words = [self.vocab_resolver.lookup_word(word_id) for word_id in words_ids]
-
-            for word, word_id, emb in zip(words, words_ids, minicoil_embeddings.tolist()):  # type: ignore[arg-type]
-                if word_id == 0:
-                    continue
-
-                sentence_result[word] = WordEmbedding(
-                    word=word,
-                    forms=forms[word],
-                    count=int(counts[word_id]),
-                    word_id=int(word_id),
-                    embedding=emb,  # type: ignore[arg-type]
-                )
-
-            for oov_word, count in oov.items():
-                # {
-                #     "word": oov_word,
-                #     "forms": [oov_word],
-                #     "count": int(count),
-                #     "word_id": -1,
-                #     "embedding": [1]
-                # }
-                sentence_result[oov_word] = WordEmbedding(
-                    word=oov_word, forms=[oov_word], count=int(count), word_id=-1, embedding=[1]
-                )
-
-            if not is_query:
-                yield self.sparse_vector_converter.embedding_to_vector(
-                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
-                )
-            else:
-                yield self.sparse_vector_converter.embedding_to_vector_query(
-                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
-                )
-
-    @classmethod
-    def _get_worker_class(cls) -> Type["MiniCoilTextEmbeddingWorker"]:
-        return MiniCoilTextEmbeddingWorker
-
-
-class MiniCoilTextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> MiniCOIL:
-        return MiniCOIL(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=1,
-            **kwargs,
-        )

fastembed/sparse/sparse_embedding_base.py

@@ -2,26 +2,23 @@ from dataclasses import dataclass
 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: Union[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,7 +28,7 @@ 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,
@@ -84,5 +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)

fastembed/sparse/sparse_text_embedding.py

@@ -1,21 +1,18 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.bm42 import Bm42
-from fastembed.sparse.minicoil import MiniCOIL
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
 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]]:
@@ -41,13 +38,9 @@ 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__(
@@ -62,7 +55,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "prithvida/Splade_PP_en_v1".lower():
+        if model_name == "prithvida/Splade_PP_en_v1":
             warnings.warn(
                 "The right spelling is prithivida/Splade_PP_en_v1. "
                 "Support of this name will be removed soon, please fix the model_name",
@@ -72,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,

fastembed/sparse/splade_pp.py

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

fastembed/sparse/utils/vocab_resolver.py

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

fastembed/text/clip_embedding.py

@@ -1,43 +1,47 @@
 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):
+    supported_models = supported_clip_models
+
     @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
+        return cls.supported_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    @classmethod
+    def add_custom_model(cls, model_info: dict[str, Any]):
+        cls.supported_models.append(model_info)
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return output.model_output
 
 

fastembed/text/custom_text_embedding.py

@@ -1,98 +0,0 @@
-from typing import Optional, 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
-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: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
-        lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
-    ):
-        super().__init__(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_ids=device_ids,
-            lazy_load=lazy_load,
-            device_id=device_id,
-            specific_model_path=specific_model_path,
-            **kwargs,
-        )
-        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: Optional[NDArray[np.int64]] = 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

@@ -3,33 +3,30 @@ 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"],
+    },
 ]
 
 
@@ -44,30 +41,28 @@ class Task(int, Enum):
 class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
+    supported_models = supported_multitask_models
 
-    def __init__(self, *args: Any, task_id: Optional[int] = None, **kwargs: Any):
+    def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.default_task_id: Union[Task, int] = (
-            task_id if task_id is not None else self.PASSAGE_TASK
-        )
+        self._current_task_id = 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]:
-        return supported_multitask_models
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        return cls.supported_models
+
+    @classmethod
+    def add_custom_model(cls, model_info: dict[str, Any]):
+        cls.supported_models.append(model_info)
 
     def _preprocess_onnx_input(
-        self,
-        onnx_input: dict[str, NumpyArray],
-        task_id: Optional[Union[int, Task]] = 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(
@@ -75,19 +70,20 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        task_id: Optional[int] = 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)
+        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: Union[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):
@@ -95,17 +91,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,200 +1,193 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+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."""
 
+    supported_models = supported_onnx_models
+
     @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
+        return cls.supported_models
+
+    @classmethod
+    def add_custom_model(cls, model_info: dict[str, Any]):
+        cls.supported_models.append(model_info)
 
     def __init__(
         self,
@@ -239,20 +232,20 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         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]
+        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:
@@ -264,7 +257,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         batch_size: int = 256,
         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.
@@ -289,40 +282,35 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
-            local_files_only=self._local_files_only,
-            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 
     @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,
@@ -330,7 +318,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
 
-class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(
         self,
         model_name: str,

fastembed/text/onnx_text_model.py

@@ -4,10 +4,9 @@ from pathlib import Path
 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
+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
@@ -18,29 +17,20 @@ class OnnxTextModel(OnnxModel[T]):
     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: Optional[Tokenizer] = 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, Union[NumpyArray, NDArray[np.int64]]]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -54,6 +44,7 @@ class OnnxTextModel(OnnxModel[T]):
         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,
@@ -69,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,
@@ -79,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:
@@ -91,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),
@@ -108,8 +99,6 @@ class OnnxTextModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -126,9 +115,7 @@ class OnnxTextModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, **kwargs), **kwargs
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -138,8 +125,6 @@ class OnnxTextModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
-                "local_files_only": local_files_only,
-                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -151,11 +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
+                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,121 @@
 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):
+    supported_models = supported_pooled_models
+
     @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
+        return cls.supported_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    @classmethod
+    def add_custom_model(cls, model_info: dict[str, Any]):
+        cls.supported_models.append(model_info)
+
+    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,119 @@
 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):
+    supported_models = supported_pooled_normalized_models
+
     @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
+        return cls.supported_models
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext, **kwargs: Any
-    ) -> Iterable[NumpyArray]:
+    @classmethod
+    def add_custom_model(cls, model_info: dict[str, Any]):
+        cls.supported_models.append(model_info)
+
+    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, Optional, Sequence, Type, Union
-from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+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,61 +18,69 @@ 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.
-        """
-        return [asdict(model) for model in cls._list_supported_models()]
 
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
+            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",
+                        }
+                    }
+                ]
+                ```
+        """
+        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: Optional[list[str]] = None,
+        cls, model_info: dict[str, Any], mean_pooling: bool = True, normalization: bool = False
     ) -> 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"
-                )
+        """
+        Register a custom model so that TextEmbedding(...) can find it later.
 
-        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,
-        )
+        Args:
+            model_info: Dictionary describing the model, e.g.:
+                {
+                    "model": "alibaba/blablabla",
+                    "dim": 512,
+                    "description": "...",
+                    "license": "apache-2.0",
+                    "size_in_GB": 1.23,
+                    "sources": { ... }  # optional
+                }
+            mean_pooling: apply mean_pooling or not.
+            normalization: apply normalization or not.
+
+        Returns:
+            None
+        """
+        if mean_pooling and not normalization:
+            PooledEmbedding.add_custom_model(model_info)
+        elif mean_pooling and normalization:
+            PooledNormalizedEmbedding.add_custom_model(model_info)
+        elif "clip" in model_info["model"].lower():
+            CLIPOnnxEmbedding.add_custom_model(model_info)
+        else:
+            OnnxTextEmbedding.add_custom_model(model_info)
 
     def __init__(
         self,
@@ -88,29 +94,34 @@ class TextEmbedding(TextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
+        if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
             warnings.warn(
-                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
-                "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
+                "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 +139,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: Optional[int] = None
-        for description in descriptions:
-            if description.model.lower() == model_name.lower():
-                embedding_size = description.dim
-                break
-        if embedding_size is None:
-            model_names = [description.model for description in descriptions]
-            raise ValueError(
-                f"Embedding size for model {model_name} was None. "
-                f"Available model names: {model_names}"
-            )
-        return embedding_size
-
     def embed(
         self,
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         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 +163,7 @@ class TextEmbedding(TextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[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 +171,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.
 

fastembed/text/text_embedding_base.py

@@ -1,11 +1,11 @@
 from typing import Iterable, Optional, Union, Any
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
+import numpy as np
+
 from fastembed.common.model_management import ModelManagement
 
 
-class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
+class TextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
@@ -17,7 +17,6 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -25,10 +24,10 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         batch_size: int = 256,
         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: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+
         """
         Embeds queries
 
@@ -51,21 +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")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.7.2"
+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"
@@ -13,15 +13,15 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 [tool.poetry.dependencies]
 python = ">=3.9.0"
 numpy = [
+    { 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" },
-    { version = ">=1.21,<2.1.0", python = "<3.10" },
+    { version = ">=2.1.0", python = ">=3.13" }
 ]
 onnxruntime = [
     { version = ">=1.17.0,<1.20.0", python = "<3.10" },
-    { version = ">1.20.0", python = ">=3.13" },
     { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
+    { version = ">1.20.0", python = ">=3.13" }
 ]
 tqdm = "^4.66"
 requests = "^2.31"
@@ -29,7 +29,7 @@ tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
 pillow = ">=10.3.0,<12.0.0"
-mmh3 = ">=4.1.0,<6.0.0"
+mmh3 = "^4.1.0"
 py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.test.dependencies]

tests/test_add_custom_model.py

@@ -0,0 +1,115 @@
+import os
+import numpy as np
+import pytest
+
+from fastembed.text.text_embedding import TextEmbedding
+from tests.utils import delete_model_cache
+
+canonical_vectors = [
+    {
+        "model": "intfloat/multilingual-e5-small",
+        "mean_pooling": True,
+        "normalization": True,
+        "canonical_vector": [3.1317e-02, 3.0939e-02, -3.5117e-02, -6.7274e-02, 8.5084e-02],
+    },
+    {
+        "model": "intfloat/multilingual-e5-small",
+        "mean_pooling": True,
+        "normalization": False,
+        "canonical_vector": [1.4604e-01, 1.4428e-01, -1.6376e-01, -3.1372e-01, 3.9677e-01],
+    },
+    {
+        "model": "mixedbread-ai/mxbai-embed-xsmall-v1",
+        "mean_pooling": False,
+        "normalization": False,
+        "canonical_vector": [
+            2.49407589e-02,
+            1.00189969e-02,
+            1.07807154e-02,
+            3.63860987e-02,
+            -2.27128249e-02,
+        ],
+    },
+]
+
+DIMENSIONS = {
+    "intfloat/multilingual-e5-small": 384,
+    "mixedbread-ai/mxbai-embed-xsmall-v1": 384,
+}
+
+SOURCES = {
+    "intfloat/multilingual-e5-small": "intfloat/multilingual-e5-small",
+    "mixedbread-ai/mxbai-embed-xsmall-v1": "mixedbread-ai/mxbai-embed-xsmall-v1",
+}
+
+
+@pytest.mark.parametrize("scenario", canonical_vectors)
+def test_add_custom_model_variations(scenario):
+    is_ci = bool(os.getenv("CI", False))
+
+    base_model_name = scenario["model"]
+    mean_pooling = scenario["mean_pooling"]
+    normalization = scenario["normalization"]
+    cv = np.array(scenario["canonical_vector"], dtype=np.float32)
+
+    backup_supported_models = {}
+    for embedding_cls in TextEmbedding.EMBEDDINGS_REGISTRY:
+        backup_supported_models[embedding_cls] = embedding_cls.list_supported_models().copy()
+
+    suffixes = []
+    suffixes.append("mean" if mean_pooling else "no-mean")
+    suffixes.append("norm" if normalization else "no-norm")
+    suffix_str = "-".join(suffixes)
+
+    custom_model_name = f"{base_model_name}-{suffix_str}"
+
+    dim = DIMENSIONS[base_model_name]
+    hf_source = SOURCES[base_model_name]
+
+    model_info = {
+        "model": custom_model_name,
+        "dim": dim,
+        "description": f"{base_model_name} with {suffix_str}",
+        "license": "mit",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": hf_source,
+        },
+        "model_file": "onnx/model.onnx",
+        "additional_files": [],
+    }
+
+    if is_ci and model_info["size_in_GB"] > 1.0:
+        pytest.skip(
+            f"Skipping {custom_model_name} on CI due to size_in_GB={model_info['size_in_GB']}"
+        )
+
+    try:
+        TextEmbedding.add_custom_model(
+            model_info=model_info, mean_pooling=mean_pooling, normalization=normalization
+        )
+
+        model = TextEmbedding(model_name=custom_model_name)
+
+        docs = ["hello world", "flag embedding"]
+        embeddings = list(model.embed(docs))
+        embeddings = np.stack(embeddings, axis=0)
+
+        assert embeddings.shape == (
+            2,
+            dim,
+        ), f"Expected shape (2, {dim}) for {custom_model_name}, but got {embeddings.shape}"
+
+        num_compare_dims = cv.shape[0]
+        assert np.allclose(
+            embeddings[0, :num_compare_dims], cv, atol=1e-3
+        ), f"Embedding mismatch for {custom_model_name} (first {num_compare_dims} dims)."
+
+        assert not np.allclose(embeddings[0, :], 0.0), "Embedding should not be all zeros."
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+    finally:
+        for embedding_cls, old_list in backup_supported_models.items():
+            embedding_cls.supported_models = old_list

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,229 +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)
-
-
-def test_cross_encoder_custom_model():
-    is_ci = os.getenv("CI")
-    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
-    size_in_gb = 0.08
-    source = ModelSource(hf=custom_model_name)
-    canonical_vector = np.array([-5.7170815, -11.112114], dtype=np.float32)
-
-    TextCrossEncoder.add_custom_model(
-        custom_model_name,
-        model_file="onnx/model.onnx",
-        sources=source,
-        size_in_gb=size_in_gb,
-    )
-
-    assert CustomTextCrossEncoder.SUPPORTED_MODELS[0] == BaseModelDescription(
-        model=custom_model_name,
-        sources=source,
-        model_file="onnx/model.onnx",
-        description="",
-        license="",
-        size_in_GB=size_in_gb,
-    )
-
-    model = TextCrossEncoder(custom_model_name)
-    pairs = [
-        ("What is AI?", "Artificial intelligence is ..."),
-        ("What is ML?", "Machine learning is ..."),
-    ]
-    scores = list(model.rerank_pairs(pairs))
-
-    embeddings = np.stack(scores, axis=0)
-    assert embeddings.shape == (2,)
-    assert np.allclose(embeddings, canonical_vector, atol=1e-3)
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
-def test_mock_add_custom_models():
-    dim = 5
-    size_in_gb = 0.1
-    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)
-
-
-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,
-        )
-
-
-def test_do_not_add_existing_cross_encoder():
-    existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
-    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
-
-    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
-        TextCrossEncoder.add_custom_model(
-            existing_base_model,
-            sources=ModelSource(hf=existing_base_model),
-            size_in_gb=0.08,
-        )
-
-    TextCrossEncoder.add_custom_model(
-        custom_model_name,
-        sources=ModelSource(hf=existing_base_model),
-        size_in_gb=0.08,
-    )
-
-    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
-        TextCrossEncoder.add_custom_model(
-            custom_model_name,
-            sources=ModelSource(hf=custom_model_name),
-            size_in_gb=0.08,
-        )

tests/test_image_onnx_embeddings.py

@@ -8,7 +8,7 @@ from PIL import Image
 
 from fastembed import ImageEmbedding
 from tests.config import TEST_MISC_DIR
-from tests.utils import delete_model_cache, 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,18 +27,16 @@ CANONICAL_VECTOR_VALUES = {
 }
 
 
-@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
-def test_embedding(model_name: str) -> None:
+def test_embedding() -> 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"]
 
-        model = ImageEmbedding(model_name=model_desc.model)
+        model = ImageEmbedding(model_name=model_desc["model"])
 
         images = [
             TEST_MISC_DIR / "image.jpeg",
@@ -50,13 +48,13 @@ def test_embedding(model_name: str) -> None:
         embeddings = np.stack(embeddings, axis=0)
         assert embeddings.shape == (len(images), dim)
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
 
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc.model
+        ), model_desc["model"]
 
-        assert np.allclose(embeddings[1], embeddings[2]), model_desc.model
+        assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
@@ -76,12 +74,8 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
 
     embeddings = list(model.embed(images, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
-    assert np.allclose(embeddings[1], embeddings[2])
-
-    canonical_vector = CANONICAL_VECTOR_VALUES[model_name]
 
     assert embeddings.shape == (len(test_images) * n_images, n_dims)
-    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
@@ -127,21 +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)

tests/test_late_interaction_embeddings.py

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

tests/test_late_interaction_multimodal.py

@@ -1,103 +0,0 @@
-import os
-
-import pytest
-from PIL import Image
-import numpy as np
-
-from fastembed import LateInteractionMultimodalEmbedding
-from tests.config import TEST_MISC_DIR
-
-
-# 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],
-        ]
-    ),
-}
-
-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],
-        ]
-    ),
-}
-
-queries = ["hello world", "flag embedding"]
-images = [
-    TEST_MISC_DIR / "image.jpeg",
-    str(TEST_MISC_DIR / "image.jpeg"),
-    Image.open((TEST_MISC_DIR / "image.jpeg")),
-]
-
-
-def test_batch_embedding():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
-    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = 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():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
-    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = next(iter(model.embed_image(images, batch_size=6)))
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_single_embedding_query():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
-    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = next(iter(model.embed_text(queries)))
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_get_embedding_size():
-    model_name = "Qdrant/colpali-v1.3-fp16"
-    assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
-
-    model_name = "Qdrant/ColPali-v1.3-fp16"
-    assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
-
-
-def test_embedding_size():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-    model_name = "Qdrant/colpali-v1.3-fp16"
-    model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 128
-
-    model_name = "Qdrant/ColPali-v1.3-fp16"
-    model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 128

tests/test_postprocess.py

@@ -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

@@ -5,10 +5,10 @@ import numpy as np
 
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
-from tests.utils import delete_model_cache, 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,
@@ -43,100 +43,46 @@ CANONICAL_COLUMN_VALUES = {
             2.1904349327087402,
             1.0531445741653442,
         ],
-    },
-    "Qdrant/minicoil-v1": {
-        "indices": [80, 81, 82, 83, 6664, 6665, 6666, 6667],
-        "values": [
-            0.52634597,
-            0.8711344,
-            1.2264385,
-            0.52123857,
-            0.974713,
-            -0.97803956,
-            -0.94312465,
-            -0.12508166,
-        ],
-    },
+    }
 }
 
-CANONICAL_QUERY_VALUES = {
-    "Qdrant/minicoil-v1": {
-        "indices": [80, 81, 82, 83, 6664, 6665, 6666, 6667],
-        "values": [
-            0.31389374,
-            0.5195128,
-            0.7314033,
-            0.3108479,
-            0.5812834,
-            -0.5832673,
-            -0.5624452,
-            -0.0745942,
-        ],
-    },
-}
-
-
 docs = ["Hello World"]
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"],
-)
-def test_batch_embedding(model_name: str) -> None:
+def test_batch_embedding() -> None:
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    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)
-
-
-@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1", "Qdrant/minicoil-v1"])
-def test_single_embedding(model_name: str) -> None:
-    is_ci = os.getenv("CI")
-    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
-
-    for model_desc in SparseTextEmbedding._list_supported_models():
-        if (
-            model_desc.model not in CANONICAL_COLUMN_VALUES
-        ):  # attention models and bm25 are also parts of
-            # SparseTextEmbedding, however, they have their own tests
-            continue
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
-
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         model = SparseTextEmbedding(model_name=model_name)
+        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        assert result.indices.tolist() == expected_result["indices"]
 
-        passage_result = next(iter(model.embed(docs, batch_size=6)))
-        query_result = next(iter(model.query_embed(docs)))
-        expected_result = CANONICAL_COLUMN_VALUES[model_name]
-        expected_query_result = CANONICAL_QUERY_VALUES.get(model_name, expected_result)
-        assert passage_result.indices.tolist() == expected_result["indices"]
-        for i, value in enumerate(passage_result.values):
+        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_name: str) -> None:
+def test_single_embedding() -> None:
     is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name=model_name)
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+        model = SparseTextEmbedding(model_name=model_name)
+
+        passage_result = next(iter(model.embed(docs, batch_size=6)))
+        query_result = next(iter(model.query_embed(docs)))
+        for result in [passage_result, query_result]:
+            assert result.indices.tolist() == expected_result["indices"]
+
+            for i, value in enumerate(result.values):
+                assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+def test_parallel_processing() -> None:
+    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))
@@ -226,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)

tests/test_text_cross_encoder.py

@@ -4,7 +4,7 @@ import numpy as np
 import pytest
 
 from fastembed.rerank.cross_encoder import TextCrossEncoder
-from tests.utils import delete_model_cache, 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]),
@@ -15,37 +15,44 @@ CANONICAL_SCORE_VALUES = {
     "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]),
 }
 
+SELECTED_MODELS = {
+    "Xenova": "Xenova/ms-marco-MiniLM-L-6-v2",
+    "BAAI": "BAAI/bge-reranker-base",
+    "jinaai": "jinaai/jina-reranker-v1-tiny-en",
+}
 
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+
+@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")
-    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)
 
-        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."]
+    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."]
-        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}"
 
-        pairs = [(query, doc) for doc in documents]
-        scores2 = np.array(list(model.rerank_pairs(pairs)))
-        assert np.allclose(
-            scores, scores2, atol=1e-5
-        ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
-
-        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
-        assert np.allclose(
-            scores, canonical_scores, atol=1e-3
-        ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+    assert np.allclose(
+        scores, canonical_scores, atol=1e-3
+    ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+@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")
 
@@ -71,7 +78,10 @@ def test_batch_rerank(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+@pytest.mark.parametrize(
+    "model_name",
+    ["Xenova/ms-marco-MiniLM-L-6-v2"],
+)
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextCrossEncoder(model_name=model_name, lazy_load=True)
@@ -85,7 +95,10 @@ 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"])
+@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")
 

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

@@ -5,7 +5,7 @@ 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]),
@@ -52,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,37 +72,38 @@ CANONICAL_VECTOR_VALUES = {
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
 
-@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
-def test_embedding(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"]
 
-        model = TextEmbedding(model_name=model_desc.model)
+        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)
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc.model
+        ), model_desc["model"]
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
+@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)
@@ -111,12 +112,15 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
     embeddings = list(model.embed(docs, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (len(docs), n_dims)
+    assert embeddings.shape == (200, n_dims)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
+@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)
@@ -131,7 +135,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
     embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (len(docs), n_dims)
+    assert embeddings.shape == (200, n_dims)
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
@@ -139,7 +143,10 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+@pytest.mark.parametrize(
+    "model_name",
+    ["BAAI/bge-small-en-v1.5"],
+)
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name, lazy_load=True)
@@ -156,22 +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)

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