draft

* Migration of models to dataclasses

* Model description file

* Test fix

* kw_only support

* Multitask embeddings test fix

* list_supported_models type fix

* Dim fix for sparsemodels

* Dim fix for sparsemodels

* Dim fix for sparsemodels (x2)

* Model management type fix

* Interface docstring fixes

* Mypy fixes

* Typing fix again

* Typing fix again

* Special cast to SparseModelDescription

* Special cast to SparseModelDescription

* Special cast to SparseModelDescription

* typing fix for colpali

* typing fix for colpali

* typing fix for colpali

* typing fix for colpali

* Let's try generic typing for ModelManagment

* wip: dataclass idea, small fixes (#475)

* wip: dataclass idea, small fixes

* fix: fix exception message in base model description

* remove custom model descriptions

* make license, description and size in gb mandatory in model description

* fix: introduce _list_supported_models which returns model description objects

* test: add test for list supported models

* fix: fix list supported models usage in tests

---------

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

.github/workflows/ci.yml

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

.github/workflows/type-checkers.yml

@@ -0,0 +1,40 @@
+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
+
+          poetry run pip install "numpy<2.0.0"  # https://github.com/python/mypy/issues/17396
+
+      - 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,5 +12,11 @@ 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.

fastembed/__init__.py

@@ -2,6 +2,7 @@ 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
 
@@ -17,4 +18,5 @@ __all__ = [
     "SparseEmbedding",
     "ImageEmbedding",
     "LateInteractionTextEmbedding",
+    "LateInteractionMultimodalEmbedding",
 ]

fastembed/common/__init__.py

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

fastembed/common/model_description.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass, field
+from typing import Optional, Any
+
+
+@dataclass(frozen=True)
+class ModelSource:
+    hf: Optional[str] = None
+    url: Optional[str] = None
+
+    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]] = None
+
+    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

fastembed/common/model_management.py

@@ -4,7 +4,7 @@ import json
 import shutil
 import tarfile
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, Union, TypeVar, Generic
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -16,9 +16,12 @@ 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:
+class ModelManagement(Generic[T]):
     METADATA_FILE = "files_metadata.json"
 
     @classmethod
@@ -26,12 +29,16 @@ class ModelManagement:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[T]: A list of dictionaries containing the model information.
         """
         raise NotImplementedError()
 
     @classmethod
-    def _get_model_description(cls, model_name: str) -> dict[str, Any]:
+    def _list_supported_models(cls) -> list[T]:
+        raise NotImplementedError()
+
+    @classmethod
+    def _get_model_description(cls, model_name: str) -> T:
         """
         Gets the model description from the model_name.
 
@@ -42,10 +49,10 @@ class ModelManagement:
             ValueError: If the model_name is not supported.
 
         Returns:
-            dict[str, Any]: The model description.
+            T: 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__}.")
@@ -114,7 +121,6 @@ class ModelManagement:
             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.
         """
@@ -148,8 +154,8 @@ class ModelManagement:
 
         def _collect_file_metadata(
             model_dir: Path, repo_files: list[RepoFile]
-        ) -> dict[str, dict[str, int]]:
-            meta = {}
+        ) -> dict[str, dict[str, Union[int, str]]]:
+            meta: dict[str, dict[str, Union[int, str]]] = {}
             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:
@@ -161,7 +167,9 @@ class ModelManagement:
                         }
             return meta
 
-        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int]]) -> None:
+        def _save_file_metadata(
+            model_dir: Path, meta: dict[str, dict[str, Union[int, str]]]
+        ) -> None:
             try:
                 if not model_dir.exists():
                     model_dir.mkdir(parents=True, exist_ok=True)
@@ -293,7 +301,11 @@ class ModelManagement:
 
     @classmethod
     def retrieve_model_gcs(
-        cls, model_name: str, source_url: str, cache_dir: str, 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-{model_name.split('/')[-1]}"
         cache_tmp_dir = Path(cache_dir) / "tmp"
@@ -337,14 +349,12 @@ class ModelManagement:
         return model_dir
 
     @classmethod
-    def download_model(
-        cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs: Any
-    ) -> Path:
+    def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: Any) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
         Args:
-            model (dict[str, Any]): The model description.
+            model (T): The model description.
                 Example:
                 ```
                 {
@@ -369,22 +379,22 @@ class ModelManagement:
         if specific_model_path:
             return Path(specific_model_path)
         retries = 1 if local_files_only else retries
-        hf_source = model.get("sources", {}).get("hf")
-        url_source = model.get("sources", {}).get("url")
+        hf_source = model.sources.hf
+        url_source = model.sources.url
 
         sleep = 3.0
         while retries > 0:
             retries -= 1
 
             if hf_source:
-                extra_patterns = [model["model_file"]]
-                extra_patterns.extend(model.get("additional_files", []))
+                extra_patterns = [model.model_file]
+                extra_patterns.extend(model.additional_files)
 
                 try:
                     return Path(
                         cls.download_files_from_huggingface(
                             hf_source,
-                            cache_dir=str(cache_dir),
+                            cache_dir=cache_dir,
                             extra_patterns=extra_patterns,
                             **kwargs,
                         )
@@ -400,8 +410,8 @@ class ModelManagement:
             if url_source or local_files_only:
                 try:
                     return cls.retrieve_model_gcs(
-                        model["model"],
-                        url_source,
+                        model.model,
+                        str(url_source),
                         str(cache_dir),
                         local_files_only=local_files_only,
                     )
@@ -418,4 +428,4 @@ class ModelManagement:
             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,7 +6,10 @@ from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
 import numpy as np
 import onnxruntime as ort
 
-from fastembed.common.types import OnnxProvider
+from numpy.typing import NDArray
+from tokenizers import Tokenizer
+
+from fastembed.common.types import OnnxProvider, NumpyArray
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -15,26 +18,26 @@ T = TypeVar("T")
 
 @dataclass
 class OnnxOutputContext:
-    model_output: np.ndarray
-    attention_mask: Optional[np.ndarray] = None
-    input_ids: Optional[np.ndarray] = None
+    model_output: NumpyArray
+    attention_mask: Optional[NDArray[np.int64]] = None
+    input_ids: Optional[NDArray[np.int64]] = None
 
 
 class OnnxModel(Generic[T]):
     @classmethod
-    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+    def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
-        self.model = None
-        self.tokenizer = None
+        self.model: Optional[ort.InferenceSession] = None
+        self.tokenizer: Optional[Tokenizer] = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
@@ -70,7 +73,7 @@ class OnnxModel(Generic[T]):
             onnx_providers = ["CPUExecutionProvider"]
 
         available_providers = ort.get_available_providers()
-        requested_provider_names = []
+        requested_provider_names: list[str] = []
         for provider in onnx_providers:
             # check providers available
             provider_name = provider if isinstance(provider, str) else provider[0]
@@ -91,6 +94,7 @@ 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(
@@ -107,13 +111,13 @@ class OnnxModel(Generic[T]):
         raise NotImplementedError("Subclasses must implement this method")
 
 
-class EmbeddingWorker(Worker):
+class EmbeddingWorker(Worker, Generic[T]):
     def init_embedding(
         self,
         model_name: str,
         cache_dir: str,
         **kwargs: Any,
-    ) -> OnnxModel:
+    ) -> OnnxModel[T]:
         raise NotImplementedError()
 
     def __init__(
@@ -125,7 +129,7 @@ class EmbeddingWorker(Worker):
         self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
     @classmethod
-    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:

fastembed/common/preprocessor_utils.py

@@ -1,4 +1,5 @@
 import json
+from typing import Any
 from pathlib import Path
 
 from tokenizers import AddedToken, Tokenizer
@@ -6,7 +7,7 @@ from tokenizers import AddedToken, Tokenizer
 from fastembed.image.transform.operators import Compose
 
 
-def load_special_tokens(model_dir: Path) -> dict:
+def load_special_tokens(model_dir: Path) -> dict[str, Any]:
     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}")
@@ -17,7 +18,7 @@ def load_special_tokens(model_dir: Path) -> dict:
     return tokens_map
 
 
-def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict]:
+def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     config_path = model_dir / "config.json"
     if not config_path.exists():
         raise ValueError(f"Could not find config.json in {model_dir}")
@@ -59,7 +60,7 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict]:
         elif isinstance(token, dict):
             tokenizer.add_special_tokens([AddedToken(**token)])
 
-    special_token_to_id = {}
+    special_token_to_id: dict[str, int] = {}
 
     for token in tokens_map.values():
         if isinstance(token, str):

fastembed/common/types.py

@@ -1,7 +1,9 @@
 from pathlib import Path
 import sys
 from PIL import Image
-from typing import Any, Iterable, Union
+from typing import Any, Union
+import numpy as np
+from numpy.typing import NDArray
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -10,7 +12,13 @@ else:
 
 
 PathInput: TypeAlias = Union[str, Path]
-PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
-ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
+ImageInput: TypeAlias = Union[PathInput, Image.Image]
 
 OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]
+NumpyArray = Union[
+    NDArray[np.float32],
+    NDArray[np.float16],
+    NDArray[np.int8],
+    NDArray[np.int64],
+    NDArray[np.int32],
+]

fastembed/common/utils.py

@@ -9,10 +9,12 @@ from typing import Iterable, Optional, TypeVar
 
 import numpy as np
 
+from fastembed.common.types import NumpyArray
+
 T = TypeVar("T")
 
 
-def normalize(input_array: np.ndarray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> np.ndarray:
+def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
     # 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

fastembed/config/dense_models.json

@@ -0,0 +1,16 @@
+{
+  "models": [
+    {
+      "model": "BAAI/bge-base-en",
+      "dim": 768,
+      "description": "Text embeddings, Unimodal (text), English...",
+      "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"
+    }
+  ]
+}

fastembed/image/image_embedding.py

@@ -1,10 +1,11 @@
-from typing import Any, Iterable, Optional, Sequence, Type
-
-import numpy as np
+from typing import Any, Iterable, Optional, Sequence, Type, Union
+from dataclasses import asdict
 
+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):
@@ -35,9 +36,13 @@ class ImageEmbedding(ImageEmbeddingBase):
                 ]
                 ```
         """
-        result = []
+        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())
+            result.extend(embedding._list_supported_models())
         return result
 
     def __init__(
@@ -53,8 +58,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,
@@ -74,14 +79,13 @@ class ImageEmbedding(ImageEmbeddingBase):
 
     def embed(
         self,
-        images: ImageInput,
+        images: Union[ImageInput, Iterable[ImageInput]],
         batch_size: int = 16,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
-        Encode a list of documents into list of embeddings.
-        We use mean pooling with attention so that the model can handle variable-length inputs.
+        Encode a list of images into list of embeddings.
 
         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
-
-import numpy as np
+from typing import Iterable, Optional, Any, Union
 
+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):
+class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
@@ -21,11 +21,11 @@ class ImageEmbeddingBase(ModelManagement):
 
     def embed(
         self,
-        images: ImageInput,
+        images: Union[ImageInput, Iterable[ImageInput]],
         batch_size: int = 16,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Embeds a list of images into a list of embeddings.
 
@@ -39,6 +39,6 @@ class ImageEmbeddingBase(ModelManagement):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NdArray]: The embeddings.
         """
         raise NotImplementedError()

fastembed/image/onnx_embedding.py

@@ -1,72 +1,66 @@
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
+
+from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 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
 
-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",
-    },
+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",
+    ),
 ]
 
 
-class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
+class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
     def __init__(
         self,
         model_name: str,
@@ -111,15 +105,14 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         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 = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
@@ -136,7 +129,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         """
         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,
@@ -144,22 +137,22 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         )
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """
         Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_onnx_models
 
     def embed(
         self,
-        images: ImageInput,
+        images: Union[ImageInput, Iterable[ImageInput]],
         batch_size: int = 16,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of images into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -189,23 +182,23 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
         return OnnxImageEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
 
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         return normalize(output.model_output).astype(np.float32)
 
 
-class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
+class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
     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,11 +2,13 @@ import contextlib
 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 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
@@ -18,19 +20,19 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 class OnnxImageModel(OnnxModel[T]):
     @classmethod
-    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
-        self.processor = None
+        self.processor: Optional[Compose] = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
@@ -58,8 +60,9 @@ class OnnxImageModel(OnnxModel[T]):
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _build_onnx_input(self, encoded: np.ndarray) -> dict[str, np.ndarray]:
-        return {node.name: encoded for node in self.model.get_inputs()}
+    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 onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
         with contextlib.ExitStack():
@@ -67,10 +70,11 @@ class OnnxImageModel(OnnxModel[T]):
                 Image.open(image) if not isinstance(image, Image.Image) else image
                 for image in images
             ]
-            encoded = self.processor(image_files)
+            assert self.processor is not None, "Processor is not initialized"
+            encoded = np.array(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)
+        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)
 
@@ -78,7 +82,7 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: ImageInput,
+        images: Union[ImageInput, Iterable[ImageInput]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
@@ -121,10 +125,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)
+                yield from self._post_process_onnx_output(batch)  # type: ignore
 
 
-class ImageEmbeddingWorker(EmbeddingWorker):
+class ImageEmbeddingWorker(EmbeddingWorker[T]):
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             embeddings = self.model.onnx_embed(batch)

fastembed/image/transform/functional.py

@@ -1,8 +1,10 @@
-from typing import Sized, Union
+from typing import Union
 
 import numpy as np
 from PIL import Image
 
+from fastembed.common.types import NumpyArray
+
 
 def convert_to_rgb(image: Image.Image) -> Image.Image:
     if image.mode == "RGB":
@@ -13,9 +15,9 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Union[Image.Image, np.ndarray],
+    image: Union[Image.Image, NumpyArray],
     size: tuple[int, int],
-) -> np.ndarray:
+) -> NumpyArray:
     if isinstance(image, np.ndarray):
         _, orig_height, orig_width = image.shape
     else:
@@ -40,7 +42,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)
+    new_image = np.zeros_like(image, shape=new_shape, dtype=np.float32)
 
     top_pad = (new_height - orig_height) // 2
     bottom_pad = top_pad + orig_height
@@ -61,37 +63,34 @@ def center_crop(
 
 
 def normalize(
-    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")
-
+    image: NumpyArray,
+    mean: Union[float, list[float]],
+    std: Union[float, list[float]],
+) -> NumpyArray:
     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)
 
-    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)
+    mean = mean if isinstance(mean, list) else [mean] * num_channels
 
-    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)
+    if len(mean) != num_channels:
+        raise ValueError(
+            f"mean must have the same number of channels as the image, image has {num_channels} channels, got "
+            f"{len(mean)}"
+        )
 
-    image = ((image.T - mean) / std).T
+    mean_arr = np.array(mean, dtype=np.float32)
+
+    std = std if isinstance(std, list) else [std] * num_channels
+    if len(std) != num_channels:
+        raise ValueError(
+            f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std)}"
+        )
+
+    std_arr = np.array(std, dtype=np.float32)
+
+    image = ((image.T - mean_arr) / std_arr).T
     return image
 
 
@@ -114,11 +113,11 @@ def resize(
     return image.resize(new_size, resample)
 
 
-def rescale(image: np.ndarray, scale: float, dtype: type = np.float32) -> np.ndarray:
+def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyArray:
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Union[Image.Image, np.ndarray]) -> np.ndarray:
+def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
     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) -> Union[list[Image.Image], list[np.ndarray]]:
+    def __call__(self, images: list[Any]) -> Union[list[Image.Image], list[NumpyArray]]:
         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[np.ndarray]:
+    def __call__(self, images: list[Image.Image]) -> list[NumpyArray]:
         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[np.ndarray]) -> list[np.ndarray]:
+    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
         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[np.ndarray]) -> list[np.ndarray]:
+    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
         return [rescale(image, scale=self.scale) for image in images]
 
 
 class PILtoNDarray(Transform):
-    def __call__(self, images: list[Union[Image.Image, np.ndarray]]) -> list[np.ndarray]:
+    def __call__(self, images: list[Union[Image.Image, NumpyArray]]) -> list[NumpyArray]:
         return [pil2ndarray(image) for image in images]
 
 
@@ -71,7 +71,7 @@ class PadtoSquare(Transform):
     def __init__(
         self,
         size: int,
-        fill_color: Optional[Union[str, int, tuple[int, ...]]] = None,
+        fill_color: Union[str, int, tuple[int, ...]],
     ):
         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[np.ndarray]]
-    ) -> Union[list[np.ndarray], list[Image.Image]]:
+        self, images: Union[list[Image.Image], list[NumpyArray]]
+    ) -> Union[list[NumpyArray], list[Image.Image]]:
         for transform in self.transforms:
             images = transform(images)
         return images
@@ -122,7 +122,7 @@ class Compose:
         Returns:
             Compose: Image processor.
         """
-        transforms = []
+        transforms: list[Transform] = []
         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 == "CLIPImageProcessor":
+        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
             if config.get("do_resize", False):
                 size = config["size"]
                 if "shortest_edge" in size:
@@ -202,15 +202,16 @@ class Compose:
     @staticmethod
     def _get_center_crop(transforms: list[Transform], config: dict[str, Any]) -> None:
         mode = config.get("image_processor_type", "CLIPImageProcessor")
-        if mode == "CLIPImageProcessor":
+        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
             if config.get("do_center_crop", False):
-                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"])
+                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"])
                 else:
-                    raise ValueError(f"Invalid crop size: {crop_size}")
+                    raise ValueError(f"Invalid crop size: {crop_size_raw}")
                 transforms.append(CenterCrop(size=crop_size))
         elif mode == "ConvNextFeatureExtractor":
             pass

fastembed/late_interaction/colbert.py

@@ -4,6 +4,7 @@ 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
@@ -11,35 +12,31 @@ 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 = [
-    {
-        "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",
-    },
+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",
+    ),
 ]
 
 
-class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
+class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     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
@@ -47,7 +44,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
 
     def _post_process_onnx_output(
         self, output: OnnxOutputContext, is_doc: bool = True
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         if not is_doc:
             return output.model_output.astype(np.float32)
 
@@ -57,7 +54,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             )
 
         for i, token_sequence in enumerate(output.input_ids):
-            for j, token_id in enumerate(token_sequence):
+            for j, token_id in enumerate(token_sequence):  # type: ignore
                 if token_id in self.skip_list or token_id == self.pad_token_id:
                     output.attention_mask[i, j] = 0
 
@@ -68,11 +65,15 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         return output.model_output.astype(np.float32)
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         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"], 1, marker_token, axis=1)
-        onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
+        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
+        )
         return onnx_input
 
     def tokenize(self, documents: list[str], is_doc: bool = True, **kwargs: Any) -> list[Encoding]:
@@ -83,6 +84,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         )
 
     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:
@@ -102,15 +104,15 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         return encoded
 
     def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
-        encoded = self.tokenizer.encode_batch(documents)
+        encoded = self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
         return encoded
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_colbert_models
 
@@ -158,15 +160,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         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 = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
 
         self._model_dir = self.download_model(
             self.model_description,
@@ -174,9 +175,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             local_files_only=self._local_files_only,
             specific_model_path=specific_model_path,
         )
-        self.mask_token_id = None
-        self.pad_token_id = None
-        self.skip_list = set()
+        self.mask_token_id: Optional[int] = None
+        self.pad_token_id: Optional[int] = None
+        self.skip_list: set[int] = set()
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -184,12 +185,13 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
     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 = {
@@ -206,7 +208,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -234,7 +236,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         if isinstance(query, str):
             query = [query]
 
@@ -247,11 +249,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
         return ColbertEmbeddingWorker
 
 
-class ColbertEmbeddingWorker(TextEmbeddingWorker):
+class ColbertEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
     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,24 +1,20 @@
 from typing import Any, Type
 
-import numpy as np
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-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"],
-    },
+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"],
+    )
 ]
 
 
@@ -29,21 +25,21 @@ class JinaColbert(Colbert):
     MASK_TOKEN = "<mask>"
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
         return JinaColbertEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_jina_colbert_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
 
         # the attention mask for jina-colbert-v2 is always 1 in queries
@@ -52,7 +48,7 @@ class JinaColbert(Colbert):
         return onnx_input
 
 
-class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
+class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):
     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
 
-import numpy as np
-
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class LateInteractionTextEmbeddingBase(ModelManagement):
+class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
@@ -24,10 +24,10 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -36,13 +36,13 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[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[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -50,11 +50,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[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)
-        if isinstance(query, Iterable):
+        else:
             yield from self.embed(query, **kwargs)

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,7 +1,8 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
+from dataclasses import asdict
 
-import numpy as np
-
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -38,9 +39,13 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
                 ]
                 ```
         """
-        result = []
+        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())
+            result.extend(embedding._list_supported_models())
         return result
 
     def __init__(
@@ -56,8 +61,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,
@@ -81,7 +86,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -99,7 +104,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[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -107,7 +112,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NdArray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations

fastembed/late_interaction_multimodal/__init__.py

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

fastembed/late_interaction_multimodal/colpali.py

@@ -0,0 +1,301 @@
+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._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=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
+        ).astype(np.float32)
+
+    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.astype(np.float32)
+
+    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()
+                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["input_ids"]]
+        )
+        onnx_input["attention_mask"] = np.array(
+            [self.EVEN_ATTENTION_MASK for _ in onnx_input["input_ids"]]
+        )
+        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,
+            **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,
+            **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

@@ -0,0 +1,130 @@
+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()`"
+        )
+
+    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

@@ -0,0 +1,67 @@
+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)
+
+    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()

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -0,0 +1,275 @@
+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,
+        )
+        assert self.tokenizer is not None
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        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,
+        **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,
+                **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 _build_onnx_image_input(self, encoded: NumpyArray) -> dict[str, NumpyArray]:
+        input_name = self.model.get_inputs()[0].name  # type: ignore[union-attr]
+        return {input_name: encoded}
+
+    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 = self._build_onnx_image_input(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,
+        **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,
+                **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/model_loader.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+import json
+from typing import Dict, List
+
+
+class ModelLoader:
+    def __init__(self):
+        self.config_dir = Path(__file__).parent / "configs"
+        self._models: Dict[str, List[Dict]] = {}
+
+    def load_models(self, model_type: str) -> List[Dict]:
+        if model_type not in self._models:
+            config_path = self.config_dir / f"{model_type}_models.json"
+            with open(config_path) as f:
+                self._models[model_type] = json.load(f)["models"]
+        return self._models[model_type]

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(Any)
+        buffer: defaultdict[int, Any] = defaultdict(Any)  # type: ignore
         next_expected = 0
 
         for idx, item in self.semi_ordered_map(stream, *args, **kwargs):

fastembed/py.typed

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

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -10,78 +10,67 @@ 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 = [
-    {
-        "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",
-    },
+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",
+    ),
 ]
 
 
 class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[BaseModelDescription]: A list of BaseModelDescription objects containing the model information.
         """
         return supported_onnx_models
 
@@ -134,15 +123,14 @@ 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 = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
@@ -156,7 +144,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,

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,10 +1,9 @@
 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 numpy.typing import NDArray
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -13,6 +12,7 @@ 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,15 +43,14 @@ 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)
+        return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
 
-    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 = {
+    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] = {
             "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
         }
         if "token_type_ids" in input_names:
@@ -72,9 +71,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)
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         relevant_output = outputs[0]
-        scores = relevant_output[:, 0]
+        scores: NumpyArray = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
 
     def _rerank_documents(
@@ -132,21 +131,38 @@ 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)
+                yield from self._post_process_onnx_output(batch)  # type: ignore
 
     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, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
 
-class TextRerankerWorker(EmbeddingWorker):
+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()
+
     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,8 +1,10 @@
 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.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.common.model_description import BaseModelDescription
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
@@ -15,7 +17,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[BaseModelDescription]: A list of dictionaries containing the model information.
 
             Example:
                 ```
@@ -33,9 +35,13 @@ class TextCrossEncoder(TextCrossEncoderBase):
                 ]
                 ```
         """
-        result = []
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        result: list[BaseModelDescription] = []
         for encoder in cls.CROSS_ENCODER_REGISTRY:
-            result.extend(encoder.list_supported_models())
+            result.extend(encoder._list_supported_models())
         return result
 
     def __init__(
@@ -52,8 +58,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,

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -1,9 +1,10 @@
 from typing import Any, Iterable, Optional
 
+from fastembed.common.model_description import BaseModelDescription
 from fastembed.common.model_management import ModelManagement
 
 
-class TextCrossEncoderBase(ModelManagement):
+class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
     def __init__(
         self,
         model_name: str,

fastembed/sparse/bm25.py

@@ -19,6 +19,7 @@ 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",
@@ -52,19 +53,18 @@ supported_languages = [
     "turkish",
 ]
 
-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,
-    },
+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",
+    ),
 ]
 
 
@@ -123,7 +123,7 @@ class Bm25(SparseTextEmbeddingBase):
         self.avg_len = avg_len
 
         model_description = self._get_model_description(model_name)
-        self.cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
 
         self._model_dir = self.download_model(
             model_description,
@@ -137,7 +137,7 @@ class Bm25(SparseTextEmbeddingBase):
         self.disable_stemmer = disable_stemmer
 
         if disable_stemmer:
-            self.stopwords = set()
+            self.stopwords: set[str] = set()
             self.stemmer = None
         else:
             self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
@@ -146,11 +146,11 @@ class Bm25(SparseTextEmbeddingBase):
         self.tokenizer = SimpleTokenizer
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
         """
         return supported_bm25_models
 
@@ -206,7 +206,7 @@ class Bm25(SparseTextEmbeddingBase):
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 for record in batch:
-                    yield record
+                    yield record  # type: ignore
 
     def embed(
         self,
@@ -239,7 +239,7 @@ class Bm25(SparseTextEmbeddingBase):
         )
 
     def _stem(self, tokens: list[str]) -> list[str]:
-        stemmed_tokens = []
+        stemmed_tokens: list[str] = []
         for token in tokens:
             lower_token = token.lower()
 
@@ -262,7 +262,7 @@ class Bm25(SparseTextEmbeddingBase):
         self,
         documents: list[str],
     ) -> list[SparseEmbedding]:
-        embeddings = []
+        embeddings: list[SparseEmbedding] = []
         for document in documents:
             document = remove_non_alphanumeric(document)
             tokens = self.tokenizer.tokenize(document)
@@ -286,8 +286,8 @@ class Bm25(SparseTextEmbeddingBase):
         Returns:
             dict[int, float]: The token_id to term frequency mapping.
         """
-        tf_map = {}
-        counter = defaultdict(int)
+        tf_map: dict[int, float] = {}
+        counter: defaultdict[str, int] = defaultdict(int)
         for stemmed_token in tokens:
             counter[stemmed_token] += 1
 
@@ -343,7 +343,9 @@ 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, Any]]:
+    def process(
+        self, items: Iterable[tuple[int, Any]]
+    ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
         for idx, batch in items:
             onnx_output = self.model.raw_embed(batch)
             yield idx, onnx_output

fastembed/sparse/bm42.py

@@ -15,21 +15,20 @@ 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 = [
-    {
-        "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,
-    },
+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,
+    ),
 ]
 
 MODEL_TO_LANGUAGE = {
@@ -102,15 +101,14 @@ 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 = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
 
         self._model_dir = self.download_model(
             self.model_description,
@@ -119,10 +117,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             specific_model_path=specific_model_path,
         )
 
-        self.invert_vocab = {}
+        self.invert_vocab: dict[int, str] = {}
 
-        self.special_tokens = set()
-        self.special_tokens_ids = set()
+        self.special_tokens: set[str] = set()
+        self.special_tokens_ids: set[int] = set()
         self.punctuation = set(string.punctuation)
         self.stopwords = set(self._load_stopwords(self._model_dir))
         self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
@@ -134,20 +132,21 @@ 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():
+
+        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))
 
     def _filter_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
-        result = []
+        result: list[tuple[str, Any]] = []
         for token, value in tokens:
             if token in self.stopwords or token in self.punctuation:
                 continue
@@ -155,7 +154,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return result
 
     def _stem_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
-        result = []
+        result: list[tuple[str, Any]] = []
         for token, value in tokens:
             processed_token = self.stemmer.stem_word(token)
             result.append((processed_token, value))
@@ -165,7 +164,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _aggregate_weights(
         cls, tokens: list[tuple[str, list[int]]], weights: list[float]
     ) -> list[tuple[str, float]]:
-        result = []
+        result: list[tuple[str, float]] = []
         for token, idxs in tokens:
             sum_weight = sum(weights[idx] for idx in idxs)
             result.append((token, sum_weight))
@@ -174,11 +173,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _reconstruct_bpe(
         self, bpe_tokens: Iterable[tuple[int, str]]
     ) -> list[tuple[str, list[int]]]:
-        result = []
-        acc = ""
-        acc_idx = []
+        result: list[tuple[str, list[int]]] = []
+        acc: str = ""
+        acc_idx: list[int] = []
 
-        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix
+        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix  # type: ignore[union-attr]
         continuing_subword_prefix_len = len(continuing_subword_prefix)
 
         for idx, token in bpe_tokens:
@@ -206,7 +205,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 = {}
+        new_vector: dict[int, float] = {}
 
         for token, value in vector.items():
             token_id = abs(mmh3.hash(token))
@@ -222,7 +221,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         if output.input_ids is None:
             raise ValueError("input_ids must be provided for document post-processing")
 
-        token_ids_batch = output.input_ids
+        token_ids_batch = output.input_ids.astype(int)
 
         # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
         pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask
@@ -241,7 +240,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
             weighted = self._aggregate_weights(stemmed, attention_value)
 
-            max_token_weight = {}
+            max_token_weight: dict[str, float] = {}
 
             for token, weight in weighted:
                 max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
@@ -251,11 +250,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding.from_dict(rescored)
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
         """
         return supported_bm42_models
 
@@ -304,7 +303,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     @classmethod
     def _query_rehash(cls, tokens: Iterable[str]) -> dict[int, float]:
-        result = {}
+        result: dict[int, float] = {}
         for token in tokens:
             token_id = abs(mmh3.hash(token))
             result[token_id] = 1.0
@@ -325,7 +324,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             self.load_onnx_model()
 
         for text in query:
-            encoded = self.tokenizer.encode(text)
+            encoded = self.tokenizer.encode(text)  # type: ignore[union-attr]
             document_tokens_with_ids = enumerate(encoded.tokens)
             reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
             filtered = self._filter_pair_tokens(reconstructed)
@@ -334,11 +333,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]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
         return Bm42TextEmbeddingWorker
 
 
-class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
+class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
         return Bm42(
             model_name=model_name,

fastembed/sparse/sparse_embedding_base.py

@@ -2,23 +2,26 @@ 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: np.ndarray
-    indices: np.ndarray
+    values: NumpyArray
+    indices: Union[NDArray[np.int64], NDArray[np.int32]]
 
-    def as_object(self) -> dict[str, np.ndarray]:
+    def as_object(self) -> dict[str, NumpyArray]:
         return {
             "values": self.values,
             "indices": self.indices,
         }
 
     def as_dict(self) -> dict[int, float]:
-        return {i: v for i, v in zip(self.indices, self.values)}
+        return {int(i): float(v) for i, v in zip(self.indices, self.values)}  # type: ignore
 
     @classmethod
     def from_dict(cls, data: dict[int, float]) -> "SparseEmbedding":
@@ -28,7 +31,7 @@ class SparseEmbedding:
         return cls(values=np.array(values), indices=np.array(indices))
 
 
-class SparseTextEmbeddingBase(ModelManagement):
+class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
     def __init__(
         self,
         model_name: str,
@@ -81,5 +84,5 @@ class SparseTextEmbeddingBase(ModelManagement):
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        if isinstance(query, Iterable):
+        else:
             yield from self.embed(query, **kwargs)

fastembed/sparse/sparse_text_embedding.py

@@ -1,4 +1,5 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
+from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.sparse.bm25 import Bm25
@@ -9,6 +10,7 @@ from fastembed.sparse.sparse_embedding_base import (
 )
 from fastembed.sparse.splade_pp import SpladePP
 import warnings
+from fastembed.common.model_description import SparseModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
@@ -38,9 +40,13 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
                 ]
                 ```
         """
-        result = []
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        result: list[SparseModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding.list_supported_models())
+            result.extend(embedding._list_supported_models())
         return result
 
     def __init__(
@@ -65,8 +71,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,30 +9,27 @@ 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 = [
-    {
-        "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",
-    },
+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",
+    ),
 ]
 
 
@@ -55,11 +52,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding(values=scores, indices=indices)
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
         """
         return supported_splade_models
 
@@ -106,15 +103,14 @@ 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 = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
 
         self._model_dir = self.download_model(
             self.model_description,
@@ -129,7 +125,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,
@@ -171,11 +167,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
         return SpladePPEmbeddingWorker
 
 
-class SpladePPEmbeddingWorker(TextEmbeddingWorker):
+class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
         return SpladePP(
             model_name=model_name,

fastembed/text/clip_embedding.py

@@ -1,41 +1,41 @@
 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.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-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",
-    },
+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",
+    ),
 ]
 
 
 class CLIPOnnxEmbedding(OnnxTextEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return CLIPEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_clip_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         return output.model_output
 
 

fastembed/text/multitask_embedding.py

@@ -3,30 +3,32 @@ from typing import Any, Type, Iterable, Union, Optional
 
 import numpy as np
 
+from fastembed.common.types import NumpyArray
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_multitask_models = [
-    {
-        "model": "jinaai/jina-embeddings-v3",
-        "dim": 1024,
-        "tasks": {
+supported_multitask_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        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": {
-            "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=ModelSource(hf="jinaai/jina-embeddings-v3"),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    ),
 ]
 
 
@@ -44,20 +46,20 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self._current_task_id = self.PASSAGE_TASK
+        self.current_task_id: Union[Task, int] = self.PASSAGE_TASK
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return JinaEmbeddingV3Worker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         return supported_multitask_models
 
     def _preprocess_onnx_input(
-        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)
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        onnx_input["task_id"] = np.array(self.current_task_id, dtype=np.int64)
         return onnx_input
 
     def embed(
@@ -66,18 +68,18 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         task_id: int = PASSAGE_TASK,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
-        self._current_task_id = task_id
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        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) -> Iterable[np.ndarray]:
-        self._current_task_id = self.QUERY_TASK
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+        self.current_task_id = self.QUERY_TASK
         yield from super().embed(query, **kwargs)
 
-    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
-        self._current_task_id = self.PASSAGE_TASK
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        self.current_task_id = self.PASSAGE_TASK
         yield from super().embed(texts, **kwargs)
 
 
@@ -86,7 +88,7 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ) -> JinaEmbeddingV3:
         model = JinaEmbeddingV3(
             model_name=model_name,
@@ -94,5 +96,5 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
             threads=1,
             **kwargs,
         )
-        model._current_task_id = kwargs["task_id"]
+        model.current_task_id = kwargs["task_id"]
         return model

fastembed/text/onnx_embedding.py

@@ -1,185 +1,207 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-
-from fastembed.common import OnnxProvider
+from fastembed.common.types import NumpyArray, 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 = [
-    {
-        "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",
-    },
+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",
+        ),
+        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",
+        ),
+        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",
+        ),
+        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",
+        ),
+        model_file="model_optimized.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",
+    ),
+    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",
+    ),
 ]
 
 
-class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
+class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     """Implementation of the Flag Embedding model."""
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """
         Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_onnx_models
 
@@ -226,15 +248,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         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 = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
@@ -251,7 +272,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -280,18 +301,18 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
         return OnnxTextEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         embeddings = output.model_output
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
             processed_embeddings = embeddings[:, 0]
@@ -304,7 +325,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
     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,
@@ -312,7 +333,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         )
 
 
-class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
+class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
     def init_embedding(
         self,
         model_name: str,

fastembed/text/onnx_text_model.py

@@ -4,9 +4,10 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from tokenizers import Encoding
+from numpy.typing import NDArray
+from tokenizers import Encoding, Tokenizer
 
-from fastembed.common import OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
@@ -17,20 +18,20 @@ class OnnxTextModel(OnnxModel[T]):
     ONNX_OUTPUT_NAMES: Optional[list[str]] = None
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
-        self.tokenizer = None
-        self.special_token_to_id = {}
+        self.tokenizer: Optional[Tokenizer] = None
+        self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, Union[NumpyArray, NDArray[np.int64]]]:
         """
         Preprocess the onnx input.
         """
@@ -44,7 +45,6 @@ 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,
@@ -60,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)
+        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
 
     def onnx_embed(
         self,
@@ -70,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()}
-        onnx_input = {
+        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:
@@ -82,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)
+        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),
@@ -136,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)
+                yield from self._post_process_onnx_output(batch)  # type: ignore
 
 
-class TextEmbeddingWorker(EmbeddingWorker):
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+class TextEmbeddingWorker(EmbeddingWorker[T]):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
         for idx, batch in items:
             onnx_output = self.model.onnx_embed(batch)
             yield idx, onnx_output

fastembed/text/pooled_embedding.py

@@ -2,108 +2,118 @@ 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.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-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"],
-    },
+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",
+        ),
+        model_file="model.onnx",
+        additional_files=["model.onnx_data"],
+    ),
 ]
 
 
 class PooledEmbedding(OnnxTextEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return PooledEmbeddingWorker
 
     @classmethod
-    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
-        token_embeddings = model_output
+    def mean_pooling(cls, model_output: NumpyArray, attention_mask: NumpyArray) -> NumpyArray:
+        token_embeddings = model_output.astype(np.float32)
+        attention_mask = attention_mask.astype(np.float32)
         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)
+        input_mask_expanded = input_mask_expanded.astype(np.float32)
         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[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_pooled_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 

fastembed/text/pooled_normalized_embedding.py

@@ -2,106 +2,131 @@ 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 = [
-    {
-        "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",
-    },
+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",
+        ),
+        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",
+    ),
 ]
 
 
 class PooledNormalizedEmbedding(PooledEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return PooledNormalizedEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
         """
         return supported_pooled_normalized_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")
 

fastembed/text/text_embedding.py

@@ -1,14 +1,15 @@
 import warnings
 from typing import Any, Iterable, Optional, Sequence, Type, Union
+from dataclasses import asdict
 
-import numpy as np
-from fastembed.common import OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 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
 
 
 class TextEmbedding(TextEmbeddingBase):
@@ -22,32 +23,18 @@ class TextEmbedding(TextEmbeddingBase):
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
-        """
-        Lists the supported models.
+        """Lists the supported models.
 
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
-
-            Example:
-                ```
-                [
-                    {
-                        "model": "intfloat/multilingual-e5-large",
-                        "dim": 1024,
-                        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
-                        "license": "mit",
-                        "size_in_GB": 2.24,
-                        "sources": {
-                            "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
-                            "hf": "qdrant/multilingual-e5-large-onnx",
-                        }
-                    }
-                ]
-                ```
         """
-        result = []
+        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())
+            result.extend(embedding._list_supported_models())
         return result
 
     def __init__(
@@ -88,8 +75,8 @@ class TextEmbedding(TextEmbeddingBase):
             )
 
         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,
@@ -113,7 +100,7 @@ class TextEmbedding(TextEmbeddingBase):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -131,7 +118,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[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -139,12 +126,12 @@ class TextEmbedding(TextEmbeddingBase):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NumpyArray]: 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[np.ndarray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         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
 
-import numpy as np
-
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class TextEmbeddingBase(ModelManagement):
+class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
@@ -24,10 +24,10 @@ class TextEmbeddingBase(ModelManagement):
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -36,14 +36,13 @@ class TextEmbeddingBase(ModelManagement):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NumpyArray]: 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[np.ndarray]:
-
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -51,11 +50,11 @@ class TextEmbeddingBase(ModelManagement):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NumpyArray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        if isinstance(query, Iterable):
+        else:
             yield from self.embed(query, **kwargs)

pyproject.toml

@@ -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 = ">=2.1.0", python = ">=3.13" },
+    { version = ">=1.21,<2.1.0", python = "<3.10" },
 ]
 onnxruntime = [
+    { version = ">1.20.0", python = ">=3.13" },
     { version = ">=1.17.0,<1.20.0", python = "<3.10" },
     { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
-    { version = ">1.20.0", python = ">=3.13" }
 ]
 tqdm = "^4.66"
 requests = "^2.31"

tests/test_common.py

@@ -0,0 +1,30 @@
+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_image_onnx_embeddings.py

@@ -30,13 +30,13 @@ CANONICAL_VECTOR_VALUES = {
 def test_embedding() -> None:
     is_ci = os.getenv("CI")
 
-    for model_desc in ImageEmbedding.list_supported_models():
-        if not is_ci and model_desc["size_in_GB"] > 1:
+    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",
@@ -48,13 +48,13 @@ def test_embedding() -> 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)

tests/test_late_interaction_multimodal.py

@@ -0,0 +1,84 @@
+import os
+
+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():
+    is_ci = os.getenv("CI")
+
+    if not is_ci:
+        for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+            print("evaluating", model_name)
+            model = LateInteractionMultimodalEmbedding(model_name=model_name)
+            result = list(model.embed_image(images, batch_size=2))
+
+            for value in result:
+                batch_size, token_num, abridged_dim = expected_result.shape
+                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=1e-3)
+
+
+def test_single_embedding():
+    is_ci = os.getenv("CI")
+    if not is_ci:
+        for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+            print("evaluating", model_name)
+            model = LateInteractionMultimodalEmbedding(model_name=model_name)
+            result = next(iter(model.embed_image(images, batch_size=6)))
+            batch_size, token_num, abridged_dim = expected_result.shape
+            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+
+
+def test_single_embedding_query():
+    is_ci = os.getenv("CI")
+    if not is_ci:
+        queries_to_embed = queries
+
+        for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+            print("evaluating", model_name)
+            model = LateInteractionMultimodalEmbedding(model_name=model_name)
+            result = next(iter(model.embed_text(queries_to_embed)))
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)

tests/test_text_multitask_embeddings.py

@@ -65,12 +65,12 @@ def test_batch_embedding():
     docs_to_embed = docs * 10
     default_task = Task.RETRIEVAL_PASSAGE
 
-    for model_desc in TextEmbedding.list_supported_models():
-        if not is_ci and model_desc["size_in_GB"] > 1:
+    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
@@ -87,7 +87,7 @@ def test_batch_embedding():
         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"]
+        ), model_desc.model
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
@@ -96,12 +96,12 @@ def test_batch_embedding():
 def test_single_embedding():
     is_ci = os.getenv("CI")
 
-    for model_desc in TextEmbedding.list_supported_models():
-        if not is_ci and model_desc["size_in_GB"] > 1:
+    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
@@ -119,7 +119,7 @@ def test_single_embedding():
             canonical_vector = task["vectors"]
             assert np.allclose(
                 embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-            ), model_desc["model"]
+            ), model_desc.model
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
@@ -129,12 +129,12 @@ def test_single_embedding_query():
     is_ci = os.getenv("CI")
     task_id = Task.RETRIEVAL_QUERY
 
-    for model_desc in TextEmbedding.list_supported_models():
-        if not is_ci and model_desc["size_in_GB"] > 1:
+    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
@@ -151,7 +151,7 @@ def test_single_embedding_query():
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
             embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-        ), model_desc["model"]
+        ), model_desc.model
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
@@ -161,12 +161,12 @@ def test_single_embedding_passage():
     is_ci = os.getenv("CI")
     task_id = Task.RETRIEVAL_PASSAGE
 
-    for model_desc in TextEmbedding.list_supported_models():
-        if not is_ci and model_desc["size_in_GB"] > 1:
+    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
@@ -183,7 +183,7 @@ def test_single_embedding_passage():
         canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
         assert np.allclose(
             embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-        ), model_desc["model"]
+        ), model_desc.model
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
@@ -219,11 +219,11 @@ def test_parallel_processing():
 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:
+    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"]
+        model_name = model_desc.model
         if model_name not in CANONICAL_VECTOR_VALUES.keys():
             continue
 
@@ -231,7 +231,7 @@ def test_task_assignment():
 
         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
+            assert model.model.current_task_id == task_id
 
         if is_ci:
             delete_model_cache(model.model._model_dir)

tests/test_text_onnx_embeddings.py

@@ -76,23 +76,23 @@ def test_embedding() -> None:
     is_ci = os.getenv("CI")
     is_mac = platform.system() == "Darwin"
 
-    for model_desc in TextEmbedding.list_supported_models():
+    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")
+            (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
 
-        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"]

tests/type_stub.py

@@ -0,0 +1,56 @@
+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=[""])