improve: Added pool of shared memory and reuse mechanisim

* Custom rerankers support

* Test for reranker_custom_model

* test fix

* Model description type fix

* Test fix

* fix: fix naming

* fix: remove redundant arg from tests

* new: update readme

---------

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

.github/ISSUE_TEMPLATE/bug-report.yml

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

.github/workflows/ci.yml

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

.github/workflows/python-tests.yml

@@ -1,9 +1,10 @@
 name: Tests
 
 on:
-  push:
-    branches: [ master, main, gpu ]
   pull_request:
+    branches: [ master, main, gpu ]
+  workflow_dispatch:
+  
 
 env:
   CARGO_TERM_COLOR: always
@@ -38,8 +39,8 @@ jobs:
         run: |
           python -m pip install poetry
           poetry config virtualenvs.create false
-          poetry install --no-interaction --no-ansi --without docs
+          poetry install --no-interaction --no-ansi --without dev,docs
 
       - name: Run pytest
         run: |
-          poetry run pytest 
+          poetry run pytest 

.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

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

README.md

@@ -63,6 +63,23 @@ embeddings = list(model.embed(documents))
 
 ```
 
+Dense text embedding can also be extended with models which are not in the list of supported models.
+
+```python
+from fastembed import TextEmbedding
+from fastembed.common.model_description import PoolingType, ModelSource
+
+TextEmbedding.add_custom_model(
+    model="intfloat/multilingual-e5-small",
+    pooling=PoolingType.MEAN,
+    normalization=True,
+    sources=ModelSource(hf="intfloat/multilingual-e5-small"),  # can be used with an `url` to load files from a private storage
+    dim=384,
+    model_file="onnx/model.onnx",  # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
+)
+model = TextEmbedding(model_name="intfloat/multilingual-e5-small")
+embeddings = list(model.embed(documents))
+```
 
 
 ### 🔱 Sparse text embeddings
@@ -137,6 +154,27 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
+### Late interaction multimodal models (ColPali)
+
+```python
+from fastembed import LateInteractionMultimodalEmbedding
+
+doc_images = [
+    "./path/to/qdrant_pdf_doc_1_screenshot.jpg",
+    "./path/to/colpali_pdf_doc_2_screenshot.jpg",
+]
+
+query = "What is Qdrant?"
+
+model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
+doc_images_embeddings = list(model.embed_image(doc_images))
+# shape (2, 1030, 128)
+# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
+query_embedding = model.embed_text(query)
+# shape (1, 20, 128)
+# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
+```
+
 ### 🔄 Rerankers
 ```python
 from fastembed.rerank.cross_encoder import TextCrossEncoder
@@ -152,6 +190,23 @@ scores = list(encoder.rerank(query, documents))
 # [-11.48061752319336, 5.472434997558594]
 ```
 
+Text cross encoders can also be extended with models which are not in the list of supported models.
+
+```python
+from fastembed.rerank.cross_encoder import TextCrossEncoder 
+from fastembed.common.model_description import ModelSource
+
+TextCrossEncoder.add_custom_model(
+    model="Xenova/ms-marco-MiniLM-L-4-v2",
+    model_file="onnx/model.onnx",
+    sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-4-v2"),
+)
+model = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-4-v2")
+scores = list(model.rerank_pairs(
+    [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ..."),]
+))
+```
+
 ## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.

docs/examples/ColBERT_with_FastEmbed.ipynb

@@ -54,7 +54,14 @@
     },
     {
      "data": {
-      "text/plain": "[{'model': 'colbert-ir/colbertv2.0',\n  'dim': 128,\n  'description': 'Late interaction model',\n  'size_in_GB': 0.44,\n  'sources': {'hf': 'colbert-ir/colbertv2.0'},\n  'model_file': 'model.onnx'}]"
+      "text/plain": [
+       "[{'model': 'colbert-ir/colbertv2.0',\n",
+       "  'dim': 128,\n",
+       "  'description': 'Late interaction model',\n",
+       "  'size_in_GB': 0.44,\n",
+       "  'sources': {'hf': 'colbert-ir/colbertv2.0'},\n",
+       "  'model_file': 'model.onnx'}]"
+      ]
      },
      "execution_count": 1,
      "metadata": {},
@@ -212,7 +219,9 @@
    "outputs": [
     {
      "data": {
-      "text/plain": "((26, 128), (32, 128))"
+      "text/plain": [
+       "((26, 128), (32, 128))"
+      ]
      },
      "execution_count": 18,
      "metadata": {},
@@ -271,7 +280,9 @@
     "import numpy as np\n",
     "\n",
     "\n",
-    "def compute_relevance_scores(query_embedding: np.array, document_embeddings: np.array, k: int):\n",
+    "def compute_relevance_scores(\n",
+    "    query_embedding: np.array, document_embeddings: np.array, k: int\n",
+    ") -> list[int]:\n",
     "    \"\"\"\n",
     "    Compute relevance scores for top-k documents given a query.\n",
     "\n",

docs/examples/FastEmbed_vs_HF_Comparison.ipynb

@@ -152,7 +152,7 @@
     "    HuggingFace Transformer implementation of FlagEmbedding\n",
     "    \"\"\"\n",
     "\n",
-    "    def __init__(self, model_id: str):\n",
+    "    def __init__(self, model_id: str) -> None:\n",
     "        self.model = AutoModel.from_pretrained(model_id)\n",
     "        self.tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
     "\n",

docs/examples/Hybrid_Search.ipynb

@@ -488,7 +488,7 @@
     }
    ],
    "source": [
-    "def make_sparse_embedding(texts: list[str]):\n",
+    "def make_sparse_embedding(texts: list[str]) -> list[SparseEmbedding]:\n",
     "    return list(sparse_model.embed(texts, batch_size=32))\n",
     "\n",
     "\n",
@@ -615,7 +615,7 @@
     }
    ],
    "source": [
-    "def get_tokens_and_weights(sparse_embedding, model_name):\n",
+    "def get_tokens_and_weights(sparse_embedding, model_name) -> dict[str, float]:\n",
     "    # Find the tokenizer for the model\n",
     "    tokenizer_source = None\n",
     "    for model_info in SparseTextEmbedding.list_supported_models():\n",
@@ -626,7 +626,7 @@
     "            raise ValueError(f\"Model {model_name} not found in the supported models.\")\n",
     "\n",
     "    tokenizer = AutoTokenizer.from_pretrained(tokenizer_source)\n",
-    "    token_weight_dict = {}\n",
+    "    token_weight_dict: dict[str, float] = {}\n",
     "    for i in range(len(sparse_embedding.indices)):\n",
     "        token = tokenizer.decode([sparse_embedding.indices[i]])\n",
     "        weight = sparse_embedding.values[i]\n",

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

fastembed/common/model_management.py

@@ -1,12 +1,14 @@
 import os
 import time
+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
+from huggingface_hub import snapshot_download, model_info, list_repo_tree
+from huggingface_hub.hf_api import RepoFile
 from huggingface_hub.utils import (
     RepositoryNotFoundError,
     disable_progress_bars,
@@ -14,20 +16,54 @@ 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
     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.
+            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 add_custom_model(
+        cls,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Add a custom model to the existing embedding classes based on the passed model descriptions
+
+        Model description dict should contain the fields same as in one of the model descriptions presented
+         in fastembed.common.model_description
+
+         E.g. for BaseModelDescription:
+              model: str
+              sources: ModelSource
+              model_file: str
+              description: str
+              license: str
+              size_in_GB: float
+              additional_files: list[str]
+
+        Returns:
+            None
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def _list_supported_models(cls) -> list[T]:
+        raise NotImplementedError()
+
+    @classmethod
+    def _get_model_description(cls, model_name: str) -> T:
         """
         Gets the model description from the model_name.
 
@@ -38,10 +74,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__}.")
@@ -98,21 +134,74 @@ class ModelManagement:
         cls,
         hf_source_repo: str,
         cache_dir: str,
-        extra_patterns: Optional[list[str]] = None,
+        extra_patterns: list[str],
         local_files_only: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> str:
         """
         Downloads a model from HuggingFace Hub.
         Args:
             hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
             cache_dir (Optional[str]): The path to the cache directory.
-            extra_patterns (Optional[list[str]]): extra patterns to allow in the snapshot download, typically
+            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.
         Returns:
             Path: The path to the model directory.
         """
+
+        def _verify_files_from_metadata(
+            model_dir: Path, stored_metadata: dict[str, Any], repo_files: list[RepoFile]
+        ) -> bool:
+            try:
+                for rel_path, meta in stored_metadata.items():
+                    file_path = model_dir / rel_path
+
+                    if not file_path.exists():
+                        return False
+
+                    if repo_files:  # online verification
+                        file_info = next((f for f in repo_files if f.path == file_path.name), None)
+                        if (
+                            not file_info
+                            or file_info.size != meta["size"]
+                            or file_info.blob_id != meta["blob_id"]
+                        ):
+                            return False
+
+                    else:  # offline verification
+                        if file_path.stat().st_size != meta["size"]:
+                            return False
+                return True
+            except (OSError, KeyError) as e:
+                logger.error(f"Error verifying files: {str(e)}")
+                return False
+
+        def _collect_file_metadata(
+            model_dir: Path, repo_files: list[RepoFile]
+        ) -> dict[str, dict[str, Union[int, str]]]:
+            meta: dict[str, dict[str, Union[int, str]]] = {}
+            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:
+                    repo_file = file_info_map.get(file_path.name)
+                    if repo_file:
+                        meta[str(file_path.relative_to(model_dir))] = {
+                            "size": repo_file.size,
+                            "blob_id": repo_file.blob_id,
+                        }
+            return meta
+
+        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)
+                (model_dir / cls.METADATA_FILE).write_text(json.dumps(meta))
+            except (OSError, ValueError) as e:
+                logger.warning(f"Error saving metadata: {str(e)}")
+
         allow_patterns = [
             "config.json",
             "tokenizer.json",
@@ -120,16 +209,59 @@ class ModelManagement:
             "special_tokens_map.json",
             "preprocessor_config.json",
         ]
-        if extra_patterns is not None:
-            allow_patterns.extend(extra_patterns)
+
+        allow_patterns.extend(extra_patterns)
 
         snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
-        is_cached = snapshot_dir.exists()
+        metadata_file = snapshot_dir / cls.METADATA_FILE
 
-        if is_cached:
+        if local_files_only:
+            disable_progress_bars()
+            if metadata_file.exists():
+                metadata = json.loads(metadata_file.read_text())
+                verified = _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[])
+                if not verified:
+                    logger.warning(
+                        "Local file sizes do not match the metadata."
+                    )  # do not raise, still make an attempt to load the model
+            else:
+                logger.warning(
+                    "Metadata file not found. Proceeding without checking local files."
+                )  # if users have downloaded models from hf manually, or they're updating from previous versions of
+                # fastembed
+            result = snapshot_download(
+                repo_id=hf_source_repo,
+                allow_patterns=allow_patterns,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+                **kwargs,
+            )
+            return result
+
+        repo_revision = model_info(hf_source_repo).sha
+        repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model"))
+
+        allowed_extensions = {".json", ".onnx", ".txt"}
+        repo_files = (
+            [
+                f
+                for f in repo_tree
+                if isinstance(f, RepoFile) and Path(f.path).suffix in allowed_extensions
+            ]
+            if repo_tree
+            else []
+        )
+
+        verified_metadata = False
+
+        if snapshot_dir.exists() and metadata_file.exists():
+            metadata = json.loads(metadata_file.read_text())
+            verified_metadata = _verify_files_from_metadata(snapshot_dir, metadata, repo_files)
+
+        if verified_metadata:
             disable_progress_bars()
 
-        return snapshot_download(
+        result = snapshot_download(
             repo_id=hf_source_repo,
             allow_patterns=allow_patterns,
             cache_dir=cache_dir,
@@ -137,8 +269,26 @@ class ModelManagement:
             **kwargs,
         )
 
+        if (
+            not verified_metadata
+        ):  # metadata is not up-to-date, update it and check whether the files have been
+            # downloaded correctly
+            metadata = _collect_file_metadata(snapshot_dir, repo_files)
+
+            download_successful = _verify_files_from_metadata(
+                snapshot_dir, metadata, repo_files=[]
+            )  # offline verification
+            if not download_successful:
+                raise ValueError(
+                    "Files have been corrupted during downloading process. "
+                    "Please check your internet connection and try again."
+                )
+            _save_file_metadata(snapshot_dir, metadata)
+
+        return result
+
     @classmethod
-    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
+    def decompress_to_cache(cls, targz_path: str, cache_dir: str) -> str:
         """
         Decompresses a .tar.gz file to a cache directory.
 
@@ -176,9 +326,14 @@ 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,
+        deprecated_tar_struct: bool = False,
+        local_files_only: bool = False,
     ) -> Path:
-        fast_model_name = f"fast-{model_name.split('/')[-1]}"
+        fast_model_name = f"{'fast-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
         cache_tmp_dir = Path(cache_dir) / "tmp"
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
@@ -220,14 +375,12 @@ class ModelManagement:
         return model_dir
 
     @classmethod
-    def download_model(
-        cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
-    ) -> 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:
                 ```
                 {
@@ -248,23 +401,26 @@ class ModelManagement:
             Path: The path to the downloaded model directory.
         """
         local_files_only = kwargs.get("local_files_only", False)
+        specific_model_path: Optional[str] = kwargs.pop("specific_model_path", None)
+        if specific_model_path:
+            return Path(specific_model_path)
         retries = 1 if local_files_only else retries
-        hf_source = model.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,
                         )
@@ -280,9 +436,10 @@ 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),
+                        deprecated_tar_struct=model.sources.deprecated_tar_struct,
                         local_files_only=local_files_only,
                     )
                 except Exception:
@@ -298,4 +455,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
-    ) -> 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(
@@ -103,29 +107,29 @@ class OnnxModel(Generic[T]):
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
+    def onnx_embed(self, *args: Any, **kwargs: Any) -> OnnxOutputContext:
         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,
-    ) -> OnnxModel:
+        **kwargs: Any,
+    ) -> OnnxModel[T]:
         raise NotImplementedError()
 
     def __init__(
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ):
         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 @@
-import os
+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
@@ -9,8 +11,14 @@ else:
     from typing_extensions import TypeAlias
 
 
-PathInput: TypeAlias = Union[str, os.PathLike]
-PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
-ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
+PathInput: TypeAlias = Union[str, Path]
+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

@@ -5,12 +5,17 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Generator, Iterable, Optional, Union
+from typing import Iterable, Optional, TypeVar
 
 import numpy as np
+from numpy.typing import NDArray
+
+from fastembed.common.types import NumpyArray
+
+T = TypeVar("T")
 
 
-def normalize(input_array, p=2, dim=1, eps=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
@@ -18,7 +23,16 @@ def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
     return normalized_array
 
 
-def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
+def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
+    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
+    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
+    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
+    sum_mask = np.sum(input_mask_expanded, axis=1)
+    pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+    return pooled_embeddings
+
+
+def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))
     [[1, 2, 3], [4, 5]]

fastembed/embedding.py

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

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__(
@@ -49,12 +54,12 @@ class ImageEmbedding(ImageEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
+            if any(model_name.lower() == model.model.lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -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,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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,18 +1,18 @@
-from typing import Iterable, Optional
-
-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,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -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,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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,73 +1,65 @@
-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,
@@ -78,7 +70,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        **kwargs,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -96,6 +89,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -110,17 +104,19 @@ 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, local_files_only=self._local_files_only
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -132,7 +128,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,
@@ -140,22 +136,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,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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.
@@ -185,24 +181,24 @@ 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
-    ) -> 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]:
-        return normalize(output.model_output).astype(np.float32)
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+        return normalize(output.model_output)
 
 
-class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
+class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
         return OnnxImageEmbedding(
             model_name=model_name,
             cache_dir=cache_dir,

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,7 +20,7 @@ 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]:
@@ -26,11 +28,11 @@ class OnnxImageModel(OnnxModel[T]):
 
     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
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
@@ -58,19 +60,21 @@ 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) -> OnnxOutputContext:
+    def onnx_embed(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
             ]
-            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,13 +82,13 @@ 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,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
 
@@ -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=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]):
+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)
@@ -133,13 +133,13 @@ class Compose:
         return cls(transforms=transforms)
 
     @staticmethod
-    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]):
+    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> None:
         transforms.append(ConvertToRGB())
 
     @classmethod
-    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]):
+    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:
@@ -200,17 +200,18 @@ class Compose:
             raise ValueError(f"Preprocessor {mode} is not supported")
 
     @staticmethod
-    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]):
+    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
@@ -220,24 +221,24 @@ class Compose:
             raise ValueError(f"Preprocessor {mode} is not supported")
 
     @staticmethod
-    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]):
+    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
         transforms.append(PILtoNDarray())
 
     @staticmethod
-    def _get_rescale(transforms: list[Transform], config: dict[str, Any]):
+    def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
         if config.get("do_rescale", True):
             rescale_factor = config.get("rescale_factor", 1 / 255)
             transforms.append(Rescale(scale=rescale_factor))
 
     @staticmethod
-    def _get_normalize(transforms: list[Transform], config: dict[str, Any]):
+    def _get_normalize(transforms: list[Transform], config: dict[str, Any]) -> None:
         if config.get("do_normalize", False):
             transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
         elif "mean" in config and "std" in config:
             transforms.append(Normalize(mean=config["mean"], std=config["std"]))
 
     @staticmethod
-    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]):
+    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]) -> None:
         mode = config.get("image_processor_type", "CLIPImageProcessor")
         if mode == "CLIPImageProcessor":
             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,9 +44,9 @@ 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)
+            return output.model_output
 
         if output.input_ids is None or output.attention_mask is None:
             raise ValueError(
@@ -57,22 +54,26 @@ 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
 
-        output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32)
+        output.model_output *= np.expand_dims(output.attention_mask, 2)
         norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
         norm_clamped = np.maximum(norm, 1e-12)
         output.model_output /= norm_clamped
-        return output.model_output.astype(np.float32)
+        return output.model_output
 
     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
 
@@ -124,7 +126,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        **kwargs,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -142,6 +145,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -156,22 +160,24 @@ 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, self.cache_dir, local_files_only=self._local_files_only
+            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
-        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()
@@ -179,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 = {
@@ -200,8 +207,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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.
@@ -229,7 +236,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         if isinstance(query, str):
             query = [query]
 
@@ -242,12 +249,12 @@ 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):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
+class ColbertEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
         return Colbert(
             model_name=model_name,
             cache_dir=cache_dir,

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,8 +48,8 @@ class JinaColbert(Colbert):
         return onnx_input
 
 
-class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> JinaColbert:
+class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
         return JinaColbert(
             model_name=model_name,
             cache_dir=cache_dir,

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -1,17 +1,17 @@
-from typing import Iterable, Optional, Union
-
-import numpy as np
+from typing import Iterable, Optional, Union, Any
 
+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,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -23,11 +23,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs) -> 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,15 +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
-    ) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -52,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__(
@@ -52,12 +57,12 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
+            if any(model_name.lower() == model.model.lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -80,8 +85,8 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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) -> 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,300 @@
+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
+        )
+
+    def _post_process_onnx_text_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        return output.model_output
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        texts_query: list[str] = []
+        for query in documents:
+            query = self.BOS_TOKEN + self.QUERY_PREFIX + query + self.PAD_TOKEN * 10
+            query += "\n"
+
+            texts_query.append(query)
+        encoded = self.tokenizer.encode_batch(texts_query)  # type: ignore[union-attr]
+        return encoded
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        onnx_input["input_ids"] = np.array(
+            [
+                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()
+                for input_ids in onnx_input["input_ids"]
+            ]
+        )
+        empty_image_placeholder: NumpyArray = np.zeros(
+            self.IMAGE_PLACEHOLDER_SIZE, dtype=np.float32
+        )
+        onnx_input["pixel_values"] = np.array(
+            [empty_image_placeholder for _ in onnx_input["input_ids"]],
+        )
+        return onnx_input
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Add placeholders for text input when processing image data for ONNX.
+        Args:
+            onnx_input (Dict[str, NumpyArray]): Preprocessed image inputs.
+            **kwargs: Additional arguments.
+        Returns:
+            Dict[str, NumpyArray]: ONNX input with text placeholders.
+        """
+        onnx_input["input_ids"] = np.array(
+            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["pixel_values"]]
+        )
+        onnx_input["attention_mask"] = np.array(
+            [self.EVEN_ATTENTION_MASK for _ in onnx_input["pixel_values"]]
+        )
+        return onnx_input
+
+    def embed_text(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            **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,271 @@
+import contextlib
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Optional, Sequence, Type, Union
+
+import numpy as np
+from PIL import Image
+from tokenizers import Encoding, Tokenizer
+
+from fastembed.common import OnnxProvider, ImageInput
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
+from fastembed.common.preprocessor_utils import load_tokenizer, load_preprocessor
+from fastembed.common.types import NumpyArray
+from fastembed.common.utils import iter_batch
+from fastembed.image.transform.operators import Compose
+from fastembed.parallel_processor import ParallelWorkerPool
+
+
+class OnnxMultimodalModel(OnnxModel[T]):
+    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.tokenizer: Optional[Tokenizer] = None
+        self.processor: Optional[Compose] = None
+        self.special_token_to_id: dict[str, int] = {}
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    @classmethod
+    def _get_text_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @classmethod
+    def _get_image_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_image_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_text_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+        )
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        assert self.tokenizer is not None
+        self.processor = load_preprocessor(model_dir=model_dir)
+
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+
+    def onnx_embed_text(
+        self,
+        documents: list[str],
+        **kwargs: Any,
+    ) -> OnnxOutputContext:
+        encoded = self.tokenize(documents, **kwargs)
+        input_ids = np.array([e.ids for e in encoded])
+        attention_mask = np.array([e.attention_mask for e in encoded])  # type: ignore[union-attr]
+        input_names = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
+        onnx_input: dict[str, NumpyArray] = {
+            "input_ids": np.array(input_ids, dtype=np.int64),
+        }
+        if "attention_mask" in input_names:
+            onnx_input["attention_mask"] = np.array(attention_mask, dtype=np.int64)
+        if "token_type_ids" in input_names:
+            onnx_input["token_type_ids"] = np.array(
+                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
+            )
+
+        onnx_input = self._preprocess_onnx_text_input(onnx_input, **kwargs)
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        return OnnxOutputContext(
+            model_output=model_output[0],
+            attention_mask=onnx_input.get("attention_mask", attention_mask),
+            input_ids=onnx_input.get("input_ids", input_ids),
+        )
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        **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 onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+        with contextlib.ExitStack():
+            image_files = [
+                Image.open(image) if not isinstance(image, Image.Image) else image
+                for image in images
+            ]
+            assert self.processor is not None, "Processor is not initialized"
+            encoded = np.array(self.processor(image_files))
+        onnx_input = {"pixel_values": encoded}
+        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+        embeddings = model_output[0].reshape(len(images), -1)
+        return OnnxOutputContext(model_output=embeddings)
+
+    def _embed_images(
+        self,
+        model_name: str,
+        cache_dir: str,
+        images: Union[Iterable[ImageInput], ImageInput],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        **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/parallel_processor.py

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

fastembed/py.typed

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

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

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

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -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
 
@@ -95,6 +84,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ):
         """
@@ -113,6 +103,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. Xenova/ms-marco-MiniLM-L-6-v2.
@@ -132,19 +123,19 @@ 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,
             local_files_only=self._local_files_only,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -153,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,
@@ -214,7 +205,7 @@ class TextCrossEncoderWorker(TextRerankerWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ) -> OnnxTextCrossEncoder:
         return OnnxTextCrossEncoder(
             model_name=model_name,

fastembed/rerank/cross_encoder/onnx_text_model.py

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

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -1,15 +1,16 @@
 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,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -21,7 +22,7 @@ class TextCrossEncoderBase(ModelManagement):
         query: str,
         documents: Iterable[str],
         batch_size: int = 64,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[float]:
         """Rerank a list of documents given a query.
 

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",
+    ),
 ]
 
 
@@ -108,7 +108,8 @@ class Bm25(SparseTextEmbeddingBase):
         language: str = "english",
         token_max_length: int = 40,
         disable_stemmer: bool = False,
-        **kwargs,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
 
@@ -122,10 +123,13 @@ 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, self.cache_dir, local_files_only=self._local_files_only
+            model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=specific_model_path,
         )
 
         self.token_max_length = token_max_length
@@ -133,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))
@@ -142,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
 
@@ -202,14 +206,14 @@ 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,
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -235,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()
 
@@ -258,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)
@@ -282,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
 
@@ -301,7 +305,9 @@ class Bm25(SparseTextEmbeddingBase):
     def compute_token_id(cls, token: str) -> int:
         return abs(mmh3.hash(token))
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """To emulate BM25 behaviour, we don't need to use weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
         """
@@ -329,7 +335,7 @@ class Bm25Worker(Worker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
@@ -337,11 +343,13 @@ 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
 
     @staticmethod
-    def init_embedding(model_name: str, cache_dir: str, **kwargs) -> Bm25:
+    def init_embedding(model_name: str, cache_dir: str, **kwargs: Any) -> Bm25:
         return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)

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 = {
@@ -66,7 +65,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        **kwargs,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -86,6 +86,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -100,24 +101,26 @@ 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, self.cache_dir, local_files_only=self._local_files_only
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            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])
@@ -129,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
@@ -150,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))
@@ -160,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))
@@ -169,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:
@@ -201,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))
@@ -217,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
@@ -236,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)
@@ -246,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
 
@@ -268,7 +272,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -299,13 +303,15 @@ 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
         return result
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """
         To emulate BM25 behaviour, we don't need to use smart weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
@@ -318,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)
@@ -327,12 +333,12 @@ 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):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
+class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
         return Bm42(
             model_name=model_name,
             cache_dir=cache_dir,

fastembed/sparse/sparse_embedding_base.py

@@ -1,24 +1,27 @@
 from dataclasses import dataclass
-from typing import Iterable, Optional, Union
+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,13 +31,13 @@ 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,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -46,11 +49,11 @@ class SparseTextEmbeddingBase(ModelManagement):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -65,7 +68,9 @@ class SparseTextEmbeddingBase(ModelManagement):
         # 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) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -79,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__(
@@ -52,7 +58,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         if model_name == "prithvida/Splade_PP_en_v1":
@@ -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,
@@ -89,7 +95,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -108,7 +114,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 

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
 
@@ -73,7 +70,8 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        **kwargs,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -91,6 +89,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -104,18 +103,20 @@ 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, self.cache_dir, local_files_only=self._local_files_only
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -124,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,
@@ -136,7 +137,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -166,12 +167,12 @@ 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):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
+class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
         return SpladePP(
             model_name=model_name,
             cache_dir=cache_dir,

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
 
 
@@ -44,7 +44,7 @@ class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ) -> OnnxTextEmbedding:
         return CLIPOnnxEmbedding(
             model_name=model_name,

fastembed/text/custom_text_embedding.py

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

fastembed/text/e5_onnx_embedding.py

@@ -1,72 +0,0 @@
-from typing import Any, Type
-
-import numpy as np
-
-from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
-
-supported_multilingual_e5_models = [
-    {
-        "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"],
-    },
-    {
-        "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",
-    },
-]
-
-
-class E5OnnxEmbedding(OnnxTextEmbedding):
-    @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
-        return E5OnnxEmbeddingWorker
-
-    @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
-        """Lists the supported models.
-
-        Returns:
-            list[dict[str, Any]]: A list of dictionaries containing the model information.
-        """
-        return supported_multilingual_e5_models
-
-    def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs
-    ) -> dict[str, np.ndarray]:
-        """
-        Preprocess the onnx input.
-        """
-        onnx_input.pop("token_type_ids", None)
-        return onnx_input
-
-
-class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs,
-    ) -> E5OnnxEmbedding:
-        return E5OnnxEmbedding(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=1,
-            **kwargs,
-        )

fastembed/text/multitask_embedding.py

@@ -0,0 +1,100 @@
+from enum import Enum
+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.common.model_description import DenseModelDescription, ModelSource
+
+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=ModelSource(hf="jinaai/jina-embeddings-v3"),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    ),
+]
+
+
+class Task(int, Enum):
+    RETRIEVAL_QUERY = 0
+    RETRIEVAL_PASSAGE = 1
+    SEPARATION = 2
+    CLASSIFICATION = 3
+    TEXT_MATCHING = 4
+
+
+class JinaEmbeddingV3(PooledNormalizedEmbedding):
+    PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
+    QUERY_TASK = Task.RETRIEVAL_QUERY
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.current_task_id: Union[Task, int] = self.PASSAGE_TASK
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return JinaEmbeddingV3Worker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return supported_multitask_models
+
+    def _preprocess_onnx_input(
+        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(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        task_id: int = PASSAGE_TASK,
+        **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: Any) -> Iterable[NumpyArray]:
+        self.current_task_id = self.QUERY_TASK
+        yield from super().embed(query, **kwargs)
+
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        self.current_task_id = self.PASSAGE_TASK
+        yield from super().embed(texts, **kwargs)
+
+
+class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> JinaEmbeddingV3:
+        model = JinaEmbeddingV3(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+        model.current_task_id = kwargs["task_id"]
+        return model

fastembed/text/onnx_embedding.py

@@ -1,196 +1,198 @@
 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": "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": "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",
+            _deprecated_tar_struct=True,
+        ),
+        model_file="model_optimized.onnx",
+    ),
+    DenseModelDescription(
+        model="BAAI/bge-base-en-v1.5",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not so necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=0.21,
+        sources=ModelSource(
+            hf="qdrant/bge-base-en-v1.5-onnx-q",
+            url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            _deprecated_tar_struct=True,
+        ),
+        model_file="model_optimized.onnx",
+    ),
+    DenseModelDescription(
+        model="BAAI/bge-large-en-v1.5",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not so necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=1.20,
+        sources=ModelSource(hf="qdrant/bge-large-en-v1.5-onnx"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="BAAI/bge-small-en",
+        dim=384,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=0.13,
+        sources=ModelSource(
+            hf="Qdrant/bge-small-en",
+            url="https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+            _deprecated_tar_struct=True,
+        ),
+        model_file="model_optimized.onnx",
+    ),
+    DenseModelDescription(
+        model="BAAI/bge-small-en-v1.5",
+        dim=384,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not so necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=0.067,
+        sources=ModelSource(hf="qdrant/bge-small-en-v1.5-onnx-q"),
+        model_file="model_optimized.onnx",
+    ),
+    DenseModelDescription(
+        model="BAAI/bge-small-zh-v1.5",
+        dim=512,
+        description=(
+            "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not so necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=0.09,
+        sources=ModelSource(
+            hf="Qdrant/bge-small-zh-v1.5",
+            url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+            _deprecated_tar_struct=True,
+        ),
+        model_file="model_optimized.onnx",
+    ),
+    DenseModelDescription(
+        model="mixedbread-ai/mxbai-embed-large-v1",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="mixedbread-ai/mxbai-embed-large-v1"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="snowflake/snowflake-arctic-embed-xs",
+        dim=384,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.09,
+        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-xs"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="snowflake/snowflake-arctic-embed-s",
+        dim=384,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.13,
+        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-s"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="snowflake/snowflake-arctic-embed-m",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.43,
+        sources=ModelSource(hf="Snowflake/snowflake-arctic-embed-m"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="snowflake/snowflake-arctic-embed-m-long",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.54,
+        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-m-long"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="snowflake/snowflake-arctic-embed-l",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=1.02,
+        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-l"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-clip-v1",
+        dim=768,
+        description=(
+            "Text embeddings, Multimodal (text&image), English, Prefixes for queries/documents: "
+            "not necessary, 2024 year"
+        ),
+        license="apache-2.0",
+        size_in_GB=0.55,
+        sources=ModelSource(hf="jinaai/jina-clip-v1"),
+        model_file="onnx/text_model.onnx",
+    ),
 ]
 
 
-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
 
@@ -204,7 +206,8 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        **kwargs,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -222,6 +225,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -235,17 +239,19 @@ 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, local_files_only=self._local_files_only
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -256,8 +262,8 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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.
@@ -286,31 +292,32 @@ 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
-    ) -> 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]
         elif embeddings.ndim == 2:  # (batch_size, embedding_dim)
             processed_embeddings = embeddings
         else:
             raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
-        return normalize(processed_embeddings).astype(np.float32)
+        return normalize(processed_embeddings)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             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,
@@ -318,12 +325,12 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         )
 
 
-class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
+class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
     def init_embedding(
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ) -> OnnxTextEmbedding:
         return OnnxTextEmbedding(
             model_name=model_name,

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,7 +18,7 @@ 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]:
@@ -25,12 +26,12 @@ class OnnxTextModel(OnnxModel[T]):
 
     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
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, Union[NumpyArray, NDArray[np.int64]]]:
         """
         Preprocess the onnx input.
         """
@@ -58,19 +59,19 @@ class OnnxTextModel(OnnxModel[T]):
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def tokenize(self, documents: list[str], **kwargs) -> list[Encoding]:
-        return self.tokenizer.encode_batch(documents)
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
 
     def onnx_embed(
         self,
         documents: list[str],
-        **kwargs,
+        **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])
-        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:
@@ -81,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),
@@ -98,7 +99,7 @@ class OnnxTextModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
 
@@ -135,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

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

fastembed/text/pooled_normalized_embedding.py

@@ -1,113 +1,150 @@
 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",
+            _deprecated_tar_struct=True,
+        ),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-en",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.52,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-base-en"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-small-en",
+        dim=512,
+        description=(
+            "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.12,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-small-en"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-de",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.32,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-de"),
+        model_file="onnx/model_fp16.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-code",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), "
+            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-code"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-zh",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), supports mixed Chinese-English input text, "
+            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-zh"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-base-es",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), supports mixed Spanish-English input text, "
+            "8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=0.64,
+        sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-es"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="thenlper/gte-base",
+        dim=768,
+        description=(
+            "General text embeddings, Unimodal (text), supports English only input text, "
+            "512 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="mit",
+        size_in_GB=0.44,
+        sources=ModelSource(hf="thenlper/gte-base"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="thenlper/gte-large",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=1.20,
+        sources=ModelSource(hf="qdrant/gte-large-onnx"),
+        model_file="model.onnx",
+    ),
 ]
 
 
 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")
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+        return normalize(self.mean_pooling(embeddings, attn_mask))
 
 
 class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
@@ -115,7 +152,7 @@ class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs,
+        **kwargs: Any,
     ) -> OnnxTextEmbedding:
         return PooledNormalizedEmbedding(
             model_name=model_name,

fastembed/text/text_embedding.py

@@ -1,55 +1,81 @@
+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.e5_onnx_embedding import E5OnnxEmbedding
+from fastembed.text.custom_text_embedding import CustomTextEmbedding
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.pooled_embedding import PooledEmbedding
+from fastembed.text.multitask_embedding import JinaEmbeddingV3
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
+from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
 
 
 class TextEmbedding(TextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [
         OnnxTextEmbedding,
-        E5OnnxEmbedding,
         CLIPOnnxEmbedding,
         PooledNormalizedEmbedding,
         PooledEmbedding,
+        JinaEmbeddingV3,
+        CustomTextEmbedding,
     ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
-        """
-        Lists the supported models.
+        """Lists the supported models.
 
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
-
-            Example:
-                ```
-                [
-                    {
-                        "model": "intfloat/multilingual-e5-large",
-                        "dim": 1024,
-                        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
-                        "license": "mit",
-                        "size_in_GB": 2.24,
-                        "sources": {
-                            "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
-                            "hf": "qdrant/multilingual-e5-large-onnx",
-                        }
-                    }
-                ]
-                ```
         """
-        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
 
+    @classmethod
+    def add_custom_model(
+        cls,
+        model: str,
+        pooling: PoolingType,
+        normalization: bool,
+        sources: ModelSource,
+        dim: int,
+        model_file: str = "onnx/model.onnx",
+        description: str = "",
+        license: str = "",
+        size_in_gb: float = 0.0,
+        additional_files: Optional[list[str]] = None,
+    ) -> None:
+        registered_models = cls._list_supported_models()
+        for registered_model in registered_models:
+            if model == registered_model.model:
+                raise ValueError(
+                    f"Model {model} is already registered in TextEmbedding, if you still want to add this model, "
+                    f"please use another model name"
+                )
+
+        CustomTextEmbedding.add_model(
+            DenseModelDescription(
+                model=model,
+                sources=sources,
+                dim=dim,
+                model_file=model_file,
+                description=description,
+                license=license,
+                size_in_GB=size_in_gb,
+                additional_files=additional_files or [],
+            ),
+            pooling=pooling,
+            normalization=normalization,
+        )
+
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
@@ -59,12 +85,32 @@ class TextEmbedding(TextEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
+            warnings.warn(
+                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
+                "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
+                UserWarning,
+                stacklevel=2,
+            )
+        if model_name in {
+            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+            "thenlper/gte-large",
+            "intfloat/multilingual-e5-large",
+            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+        }:
+            warnings.warn(
+                f"The model {model_name} now uses mean pooling instead of CLS embedding. "
+                f"In order to preserve the previous behaviour, consider either pinning fastembed version to 0.5.1 or "
+                "using `add_custom_model` functionality.",
+                UserWarning,
+                stacklevel=2,
+            )
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
+            if any(model_name.lower() == model.model.lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name=model_name,
                     cache_dir=cache_dir,
@@ -87,8 +133,8 @@ class TextEmbedding(TextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> 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.
@@ -105,3 +151,30 @@ class TextEmbedding(TextEmbeddingBase):
             List of embeddings, one per document
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
+
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            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[NumpyArray]:
+        """
+        Embeds a list of text passages into a list of embeddings.
+
+        Args:
+            texts (Iterable[str]): The list of texts to embed.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[SparseEmbedding]: The sparse embeddings.
+        """
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.model.passage_embed(texts, **kwargs)

fastembed/text/text_embedding_base.py

@@ -1,17 +1,17 @@
-from typing import Iterable, Optional, Union
-
-import numpy as np
+from typing import Iterable, Optional, Union, Any
 
+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,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -23,11 +23,11 @@ class TextEmbeddingBase(ModelManagement):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs,
-    ) -> Iterable[np.ndarray]:
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs) -> 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,15 +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
-    ) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -52,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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.5.0"
+version = "0.6.0"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -12,44 +12,52 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 
 [tool.poetry.dependencies]
 python = ">=3.9.0"
-onnx = ">=1.15.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"
 tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
-pillow = "^10.3.0"
-mmh3 = "^4.1.0"
+pillow = ">=10.3.0,<12.0.0"
+mmh3 = ">=4.1.0,<6.0.0"
 py-rust-stemmers = "^0.1.0"
 
-[tool.poetry.group.dev.dependencies]
+[tool.poetry.group.test.dependencies]
 pytest = "^7.4.2"
 ruff = ">=0.3.1,<1.0"
+
+[tool.poetry.group.dev.dependencies]
 notebook = ">=7.0.2"
 pre-commit = "^3.6.2"
+onnx = ">=1.15.0"
 
 [tool.poetry.group.docs.dependencies]
 mkdocs-material = "^9.5.10"
 mkdocstrings = "^0.24.0"
-pillow = "^10.2.0"
+pillow = ">=10.3.0,<12.0.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 
+[tool.poetry.group.types.dependencies]
+pyright = ">=1.1.293"
+mypy = "^1.0.0"
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
+[tool.pyright]
+typeCheckingMode = "strict"
+
 [tool.ruff]
 line-length = 99

tests/test_attention_embeddings.py

@@ -8,7 +8,7 @@ from tests.utils import delete_model_cache
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
-def test_attention_embeddings(model_name):
+def test_attention_embeddings(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name)
 
@@ -71,7 +71,7 @@ def test_attention_embeddings(model_name):
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
-def test_parallel_processing(model_name):
+def test_parallel_processing(model_name: str) -> None:
     is_ci = os.getenv("CI")
 
     model = SparseTextEmbedding(model_name=model_name)
@@ -96,7 +96,7 @@ def test_parallel_processing(model_name):
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
-def test_multilanguage(model_name):
+def test_multilanguage(model_name: str) -> None:
     is_ci = os.getenv("CI")
 
     docs = ["Mangez-vous vraiment des grenouilles?", "Je suis au lit"]
@@ -122,7 +122,7 @@ def test_multilanguage(model_name):
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
-def test_special_characters(model_name):
+def test_special_characters(model_name: str) -> None:
     is_ci = os.getenv("CI")
 
     docs = [
@@ -145,7 +145,7 @@ def test_special_characters(model_name):
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])
-def test_lazy_load(model_name):
+def test_lazy_load(model_name: str) -> None:
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
     docs = ["hello world", "flag embedding"]

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_custom_models.py

@@ -0,0 +1,229 @@
+import itertools
+import os
+import numpy as np
+import pytest
+
+from fastembed.common.model_description import (
+    PoolingType,
+    ModelSource,
+    DenseModelDescription,
+    BaseModelDescription,
+)
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize, mean_pooling
+from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+from fastembed.text.text_embedding import TextEmbedding
+from tests.utils import delete_model_cache
+
+
+@pytest.fixture(autouse=True)
+def restore_custom_models_fixture():
+    CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextCrossEncoder.SUPPORTED_MODELS = []
+    yield
+    CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextCrossEncoder.SUPPORTED_MODELS = []
+
+
+def test_text_custom_model():
+    is_ci = os.getenv("CI")
+    custom_model_name = "intfloat/multilingual-e5-small"
+    canonical_vector = np.array(
+        [3.1317e-02, 3.0939e-02, -3.5117e-02, -6.7274e-02, 8.5084e-02], dtype=np.float32
+    )
+    pooling = PoolingType.MEAN
+    normalization = True
+    dim = 384
+    size_in_gb = 0.47
+    source = ModelSource(hf=custom_model_name)
+
+    TextEmbedding.add_custom_model(
+        custom_model_name,
+        pooling=pooling,
+        normalization=normalization,
+        sources=source,
+        dim=dim,
+        size_in_gb=size_in_gb,
+    )
+
+    assert CustomTextEmbedding.SUPPORTED_MODELS[0] == DenseModelDescription(
+        model=custom_model_name,
+        sources=source,
+        model_file="onnx/model.onnx",
+        description="",
+        license="",
+        size_in_GB=size_in_gb,
+        additional_files=[],
+        dim=dim,
+        tasks={},
+    )
+    assert CustomTextEmbedding.POSTPROCESSING_MAPPING[custom_model_name] == PostprocessingConfig(
+        pooling=pooling, normalization=normalization
+    )
+
+    model = TextEmbedding(custom_model_name)
+    docs = ["hello world", "flag embedding"]
+    embeddings = list(model.embed(docs))
+    embeddings = np.stack(embeddings, axis=0)
+    assert embeddings.shape == (2, dim)
+
+    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+def test_cross_encoder_custom_model():
+    is_ci = os.getenv("CI")
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+    size_in_gb = 0.08
+    source = ModelSource(hf=custom_model_name)
+    canonical_vector = np.array([-5.7170815, -11.112114], dtype=np.float32)
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        model_file="onnx/model.onnx",
+        sources=source,
+        size_in_gb=size_in_gb,
+    )
+
+    assert CustomTextCrossEncoder.SUPPORTED_MODELS[0] == BaseModelDescription(
+        model=custom_model_name,
+        sources=source,
+        model_file="onnx/model.onnx",
+        description="",
+        license="",
+        size_in_GB=size_in_gb,
+    )
+
+    model = TextCrossEncoder(custom_model_name)
+    pairs = [
+        ("What is AI?", "Artificial intelligence is ..."),
+        ("What is ML?", "Machine learning is ..."),
+    ]
+    scores = list(model.rerank_pairs(pairs))
+
+    embeddings = np.stack(scores, axis=0)
+    assert embeddings.shape == (2,)
+    assert np.allclose(embeddings, canonical_vector, atol=1e-3)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+def test_mock_add_custom_models():
+    dim = 5
+    size_in_gb = 0.1
+    source = ModelSource(hf="artificial")
+
+    num_tokens = 10
+    dummy_pooled_embedding = np.random.random((1, dim)).astype(np.float32)
+    dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(np.float32)
+    dummy_attention_mask = np.ones((1, num_tokens)).astype(np.int64)
+
+    dummy_token_output = OnnxOutputContext(
+        model_output=dummy_token_embedding, attention_mask=dummy_attention_mask
+    )
+    dummy_pooled_output = OnnxOutputContext(model_output=dummy_pooled_embedding)
+    input_data = {
+        f"{PoolingType.MEAN.lower()}-normalized": dummy_token_output,
+        f"{PoolingType.MEAN.lower()}": dummy_token_output,
+        f"{PoolingType.CLS.lower()}-normalized": dummy_token_output,
+        f"{PoolingType.CLS.lower()}": dummy_token_output,
+        f"{PoolingType.DISABLED.lower()}-normalized": dummy_pooled_output,
+        f"{PoolingType.DISABLED.lower()}": dummy_pooled_output,
+    }
+
+    expected_output = {
+        f"{PoolingType.MEAN.lower()}-normalized": normalize(
+            mean_pooling(dummy_token_embedding, dummy_attention_mask)
+        ),
+        f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
+        f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]),
+        f"{PoolingType.CLS.lower()}": dummy_token_embedding[:, 0],
+        f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding),
+        f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
+    }
+
+    for pooling, normalization in itertools.product(
+        (PoolingType.MEAN, PoolingType.CLS, PoolingType.DISABLED), (True, False)
+    ):
+        model_name = f"{pooling.name.lower()}{'-normalized' if normalization else ''}"
+        TextEmbedding.add_custom_model(
+            model_name,
+            pooling=pooling,
+            normalization=normalization,
+            sources=source,
+            dim=dim,
+            size_in_gb=size_in_gb,
+        )
+
+        custom_text_embedding = CustomTextEmbedding(
+            model_name,
+            lazy_load=True,
+            specific_model_path="./",  # disable model downloading and loading
+        )
+
+        post_processed_output = next(
+            iter(custom_text_embedding._post_process_onnx_output(input_data[model_name]))
+        )
+        assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3)
+
+
+def test_do_not_add_existing_model():
+    existing_base_model = "sentence-transformers/all-MiniLM-L6-v2"
+    custom_model_name = "intfloat/multilingual-e5-small"
+
+    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
+        TextEmbedding.add_custom_model(
+            existing_base_model,
+            pooling=PoolingType.MEAN,
+            normalization=True,
+            sources=ModelSource(hf=existing_base_model),
+            dim=384,
+            size_in_gb=0.47,
+        )
+
+    TextEmbedding.add_custom_model(
+        custom_model_name,
+        pooling=PoolingType.MEAN,
+        normalization=False,
+        sources=ModelSource(hf=existing_base_model),
+        dim=384,
+        size_in_gb=0.47,
+    )
+
+    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
+        TextEmbedding.add_custom_model(
+            custom_model_name,
+            pooling=PoolingType.MEAN,
+            normalization=True,
+            sources=ModelSource(hf=custom_model_name),
+            dim=384,
+            size_in_gb=0.47,
+        )
+
+
+def test_do_not_add_existing_cross_encoder():
+    existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+
+    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
+        TextCrossEncoder.add_custom_model(
+            existing_base_model,
+            sources=ModelSource(hf=existing_base_model),
+            size_in_gb=0.08,
+        )
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        sources=ModelSource(hf=existing_base_model),
+        size_in_gb=0.08,
+    )
+
+    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
+        TextCrossEncoder.add_custom_model(
+            custom_model_name,
+            sources=ModelSource(hf=custom_model_name),
+            size_in_gb=0.08,
+        )

tests/test_image_onnx_embeddings.py

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

tests/test_late_interaction_embeddings.py

@@ -6,7 +6,7 @@ import numpy as np
 from fastembed.late_interaction.late_interaction_text_embedding import (
     LateInteractionTextEmbedding,
 )
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -153,31 +153,37 @@ CANONICAL_QUERY_VALUES = {
 docs = ["Hello World"]
 
 
-def test_batch_embedding():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_batch_embedding(model_name: str):
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        result = list(model.embed(docs_to_embed, batch_size=6))
+    model = LateInteractionTextEmbedding(model_name=model_name)
+    result = list(model.embed(docs_to_embed, batch_size=6))
+    expected_result = CANONICAL_COLUMN_VALUES[model_name]
 
-        for value in result:
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
+    for value in result:
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
 
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_single_embedding(model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     docs_to_embed = docs
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+    for model_desc in LateInteractionTextEmbedding._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         print("evaluating", model_name)
         model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
         token_num, abridged_dim = expected_result.shape
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
@@ -185,14 +191,20 @@ def test_single_embedding():
             delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding_query():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_single_embedding_query(model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     queries_to_embed = docs
 
-    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+    for model_desc in LateInteractionTextEmbedding._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         print("evaluating", model_name)
         model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.query_embed(queries_to_embed)))
+        expected_result = CANONICAL_QUERY_VALUES[model_name]
         token_num, abridged_dim = expected_result.shape
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
@@ -200,10 +212,11 @@ def test_single_embedding_query():
             delete_model_cache(model.model._model_dir)
 
 
-def test_parallel_processing():
+@pytest.mark.parametrize("token_dim,model_name", [(96, "answerdotai/answerai-colbert-small-v1")])
+def test_parallel_processing(token_dim: int, model_name: str):
     is_ci = os.getenv("CI")
-    model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
-    token_dim = 128
+    model = LateInteractionTextEmbedding(model_name=model_name)
+
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
     embeddings = np.stack(embeddings, axis=0)
@@ -222,11 +235,8 @@ def test_parallel_processing():
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["colbert-ir/colbertv2.0"],
-)
-def test_lazy_load(model_name):
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")
 
     model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)

tests/test_late_interaction_multimodal.py

@@ -0,0 +1,83 @@
+import os
+
+import pytest
+from PIL import Image
+import numpy as np
+
+from fastembed import LateInteractionMultimodalEmbedding
+from tests.config import TEST_MISC_DIR
+
+
+# vectors are abridged and rounded for brevity
+CANONICAL_IMAGE_VALUES = {
+    "Qdrant/colpali-v1.3-fp16": np.array(
+        [
+            [-0.0345, -0.022, 0.0567, -0.0518, -0.0782, 0.1714, -0.1738],
+            [-0.1181, -0.099, 0.0268, 0.0774, 0.0228, 0.0563, -0.1021],
+            [-0.117, -0.0683, 0.0371, 0.0921, 0.0107, 0.0659, -0.0666],
+            [-0.1393, -0.0948, 0.037, 0.0951, -0.0126, 0.0678, -0.087],
+            [-0.0957, -0.081, 0.0404, 0.052, 0.0409, 0.0335, -0.064],
+            [-0.0626, -0.0445, 0.056, 0.0592, -0.0229, 0.0409, -0.0301],
+            [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
+        ]
+    ),
+}
+
+CANONICAL_QUERY_VALUES = {
+    "Qdrant/colpali-v1.3-fp16": np.array(
+        [
+            [-0.0023, 0.1477, 0.1594, 0.046, -0.0196, 0.0554, 0.1567],
+            [-0.0139, -0.0057, 0.0932, 0.0052, -0.0678, 0.0131, 0.0537],
+            [0.0054, 0.0364, 0.2078, -0.074, 0.0355, 0.061, 0.1593],
+            [-0.0076, -0.0154, 0.2266, 0.0103, 0.0089, -0.024, 0.098],
+            [-0.0274, 0.0098, 0.2106, -0.0634, 0.0616, -0.0021, 0.0708],
+            [0.0074, 0.0025, 0.1631, -0.0802, 0.0418, -0.0219, 0.1022],
+            [-0.0165, -0.0106, 0.1672, -0.0768, 0.0389, -0.0038, 0.1137],
+        ]
+    ),
+}
+
+queries = ["hello world", "flag embedding"]
+images = [
+    TEST_MISC_DIR / "image.jpeg",
+    str(TEST_MISC_DIR / "image.jpeg"),
+    Image.open((TEST_MISC_DIR / "image.jpeg")),
+]
+
+
+def test_batch_embedding():
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
+
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = list(model.embed_image(images, batch_size=2))
+
+        for value in result:
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+
+
+def test_single_embedding():
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
+
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = next(iter(model.embed_image(images, batch_size=6)))
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+
+
+def test_single_embedding_query():
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
+
+    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = next(iter(model.embed_text(queries)))
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)

tests/test_multi_gpu.py

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

tests/test_sparse_embeddings.py

@@ -5,10 +5,10 @@ import numpy as np
 
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_COLUMN_VALUES = {
-    "prithvida/Splade_PP_en_v1": {
+    "prithivida/Splade_PP_en_v1": {
         "indices": [
             2040,
             2047,
@@ -49,28 +49,41 @@ CANONICAL_COLUMN_VALUES = {
 docs = ["Hello World"]
 
 
-def test_batch_embedding():
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_batch_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        model = SparseTextEmbedding(model_name=model_name)
-        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
-        assert result.indices.tolist() == expected_result["indices"]
+    model = SparseTextEmbedding(model_name=model_name)
+    result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+    expected_result = CANONICAL_COLUMN_VALUES[model_name]
+    assert result.indices.tolist() == expected_result["indices"]
 
-        for i, value in enumerate(result.values):
-            assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    for i, value in enumerate(result.values):
+        assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding():
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_single_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    for model_desc in SparseTextEmbedding._list_supported_models():
+        if (
+            model_desc.model not in CANONICAL_COLUMN_VALUES
+        ):  # attention models and bm25 are also parts of
+            # SparseTextEmbedding, however, they have their own tests
+            continue
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         model = SparseTextEmbedding(model_name=model_name)
 
         passage_result = next(iter(model.embed(docs, batch_size=6)))
         query_result = next(iter(model.query_embed(docs)))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
         for result in [passage_result, query_result]:
             assert result.indices.tolist() == expected_result["indices"]
 
@@ -80,9 +93,10 @@ def test_single_embedding():
             delete_model_cache(model.model._model_dir)
 
 
-def test_parallel_processing():
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_parallel_processing(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+    model = SparseTextEmbedding(model_name=model_name)
     docs = ["hello world", "flag embedding"] * 30
     sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
     sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))
@@ -111,7 +125,7 @@ def test_parallel_processing():
 
 
 @pytest.fixture
-def bm25_instance():
+def bm25_instance() -> None:
     ci = os.getenv("CI", True)
     model = Bm25("Qdrant/bm25", language="english")
     yield model
@@ -119,7 +133,7 @@ def bm25_instance():
         delete_model_cache(model._model_dir)
 
 
-def test_stem_with_stopwords_and_punctuation(bm25_instance):
+def test_stem_with_stopwords_and_punctuation(bm25_instance: Bm25) -> None:
     # Setup
     bm25_instance.stopwords = {"the", "is", "a"}
     bm25_instance.punctuation = {".", ",", "!"}
@@ -135,7 +149,7 @@ def test_stem_with_stopwords_and_punctuation(bm25_instance):
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
-def test_stem_case_insensitive_stopwords(bm25_instance):
+def test_stem_case_insensitive_stopwords(bm25_instance: Bm25) -> None:
     # Setup
     bm25_instance.stopwords = {"the", "is", "a"}
     bm25_instance.punctuation = {".", ",", "!"}
@@ -152,7 +166,7 @@ def test_stem_case_insensitive_stopwords(bm25_instance):
 
 
 @pytest.mark.parametrize("disable_stemmer", [True, False])
-def test_disable_stemmer_behavior(disable_stemmer):
+def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
     # Setup
     model = Bm25("Qdrant/bm25", language="english", disable_stemmer=disable_stemmer)
     model.stopwords = {"the", "is", "a"}
@@ -172,11 +186,8 @@ def test_disable_stemmer_behavior(disable_stemmer):
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["prithivida/Splade_PP_en_v1"],
-)
-def test_lazy_load(model_name):
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")

tests/test_text_cross_encoder.py

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

tests/test_text_multitask_embeddings.py

@@ -0,0 +1,246 @@
+import os
+
+import numpy as np
+import pytest
+
+from fastembed import TextEmbedding
+from fastembed.text.multitask_embedding import JinaEmbeddingV3, Task
+from tests.utils import delete_model_cache
+
+
+CANONICAL_VECTOR_VALUES = {
+    "jinaai/jina-embeddings-v3": [
+        {
+            "task_id": Task.RETRIEVAL_QUERY,
+            "vectors": np.array(
+                [
+                    [0.0623, -0.0402, 0.1706, -0.0143, 0.0617],
+                    [-0.1064, -0.0733, 0.0353, 0.0096, 0.0667],
+                ]
+            ),
+        },
+        {
+            "task_id": Task.RETRIEVAL_PASSAGE,
+            "vectors": np.array(
+                [
+                    [0.0513, -0.0247, 0.1751, -0.0075, 0.0679],
+                    [-0.0987, -0.0786, 0.09, 0.0087, 0.0577],
+                ]
+            ),
+        },
+        {
+            "task_id": Task.SEPARATION,
+            "vectors": np.array(
+                [
+                    [0.094, -0.1065, 0.1305, 0.0547, 0.0556],
+                    [0.0315, -0.1468, 0.065, 0.0568, 0.0546],
+                ]
+            ),
+        },
+        {
+            "task_id": Task.CLASSIFICATION,
+            "vectors": np.array(
+                [
+                    [0.0606, -0.0877, 0.1384, 0.0065, 0.0722],
+                    [-0.0502, -0.119, 0.032, 0.0514, 0.0689],
+                ]
+            ),
+        },
+        {
+            "task_id": Task.TEXT_MATCHING,
+            "vectors": np.array(
+                [
+                    [0.0911, -0.0341, 0.1305, -0.026, 0.0576],
+                    [-0.1432, -0.05, 0.0133, 0.0464, 0.0789],
+                ]
+            ),
+        },
+    ]
+}
+docs = ["Hello World", "Follow the white rabbit."]
+
+
+@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
+def test_batch_embedding(dim: int, model_name: str):
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
+    docs_to_embed = docs * 10
+    default_task = Task.RETRIEVAL_PASSAGE
+
+    model = TextEmbedding(model_name=model_name)
+
+    embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
+    embeddings = np.stack(embeddings, axis=0)
+
+    assert embeddings.shape == (len(docs_to_embed), dim)
+
+    canonical_vector = CANONICAL_VECTOR_VALUES[model_name][default_task]["vectors"]
+    assert np.allclose(
+        embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+    ), model_name
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+def test_single_embedding():
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
+        model_name = model_desc.model
+        dim = model_desc.dim
+
+        model = TextEmbedding(model_name=model_name)
+
+        for task in CANONICAL_VECTOR_VALUES[model_name]:
+            print(f"evaluating {model_name} task_id: {task['task_id']}")
+
+            embeddings = list(model.embed(documents=docs, task_id=task["task_id"]))
+            embeddings = np.stack(embeddings, axis=0)
+
+            assert embeddings.shape == (len(docs), dim)
+
+            canonical_vector = task["vectors"]
+            assert np.allclose(
+                embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+            ), model_desc.model
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+def test_single_embedding_query():
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
+    task_id = Task.RETRIEVAL_QUERY
+
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
+        model_name = model_desc.model
+        dim = model_desc.dim
+
+        model = TextEmbedding(model_name=model_name)
+
+        print(f"evaluating {model_name} query_embed task_id: {task_id}")
+
+        embeddings = list(model.query_embed(query=docs))
+        embeddings = np.stack(embeddings, axis=0)
+
+        assert embeddings.shape == (len(docs), dim)
+
+        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
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+def test_single_embedding_passage():
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
+    task_id = Task.RETRIEVAL_PASSAGE
+
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
+
+        model_name = model_desc.model
+        dim = model_desc.dim
+
+        model = TextEmbedding(model_name=model_name)
+
+        print(f"evaluating {model_name} passage_embed task_id: {task_id}")
+
+        embeddings = list(model.passage_embed(texts=docs))
+        embeddings = np.stack(embeddings, axis=0)
+
+        assert embeddings.shape == (len(docs), dim)
+
+        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
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
+def test_parallel_processing(dim: int, model_name: str):
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping in CI non-manual mode")
+
+    docs = ["Hello World", "Follow the white rabbit."] * 10
+
+    model = TextEmbedding(model_name=model_name)
+
+    task_id = Task.SEPARATION
+    embeddings_1 = list(model.embed(docs, batch_size=10, parallel=None, task_id=task_id))
+    embeddings_1 = np.stack(embeddings_1, axis=0)
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=1, task_id=task_id))
+    embeddings_2 = np.stack(embeddings_2, axis=0)
+
+    assert embeddings_1.shape[0] == len(docs) and embeddings_1.shape[-1] == dim
+    assert np.allclose(embeddings_1, embeddings_2, atol=1e-4)
+
+    canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
+    assert np.allclose(embeddings_2[:2, : canonical_vector.shape[1]], canonical_vector, atol=1e-4)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
+def test_task_assignment():
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    if is_ci and not is_manual:
+        pytest.skip("Skipping in CI non-manual mode")
+
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
+        model_name = model_desc.model
+
+        model = TextEmbedding(model_name=model_name)
+
+        for i, task_id in enumerate(Task):
+            _ = list(model.embed(documents=docs, batch_size=1, task_id=i))
+            assert model.model.current_task_id == task_id
+
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
+
+
+@pytest.mark.parametrize("model_name", ["jinaai/jina-embeddings-v3"])
+def test_lazy_load(model_name: str):
+    is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    if is_ci and not is_manual:
+        pytest.skip("Skipping in CI non-manual mode")
+
+    model = TextEmbedding(model_name=model_name, lazy_load=True)
+    assert not hasattr(model.model, "model")
+
+    list(model.embed(docs))
+    assert hasattr(model.model, "model")
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)

tests/test_text_onnx_embeddings.py

@@ -1,10 +1,11 @@
 import os
+import platform
 
 import numpy as np
 import pytest
 
 from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
@@ -31,11 +32,11 @@ CANONICAL_VECTOR_VALUES = {
         [-0.034478, 0.03102, 0.00673, 0.02611, -0.039362]
     ),
     "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array(
-        [0.0094, 0.0184, 0.0328, 0.0072, -0.0351]
+        [0.0361, 0.1862, 0.2776, 0.2461, -0.1904]
     ),
-    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
+    "intfloat/multilingual-e5-large": np.array([0.4544, -0.0968, 0.1054, -1.3753, 0.1500]),
     "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
-        [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
+        [0.0047, 0.1334, -0.0102, 0.0714, 0.1930]
     ),
     "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
     "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
@@ -48,10 +49,10 @@ CANONICAL_VECTOR_VALUES = {
         [-0.15407836, -0.03053198, -3.9138033, 0.1910364, 0.13224715]
     ),
     "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
-        [-0.12525563, 0.38030425, -3.961622, 0.04176439, -0.0758301]
+        [0.0802303, 0.3700881, -4.3053818, 0.4431803, -0.271572]
     ),
     "thenlper/gte-large": np.array(
-        [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
+        [-0.00986551, -0.00018734, 0.00605892, -0.03289612, -0.0387564],
     ),
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]
@@ -68,35 +69,41 @@ CANONICAL_VECTOR_VALUES = {
     "jinaai/jina-clip-v1": np.array([-0.0862, -0.0101, -0.0056, 0.0375, -0.0472]),
 }
 
+MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
-def test_embedding():
+
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+def test_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
+    is_mac = platform.system() == "Darwin"
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
-    for model_desc in TextEmbedding.list_supported_models():
-        if not is_ci and model_desc["size_in_GB"] > 1:
+    for model_desc in TextEmbedding._list_supported_models():
+        if model_desc.model in MULTI_TASK_MODELS or (
+            is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q"
+        ):
+            continue
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue
 
-        dim = model_desc["dim"]
+        dim = model_desc.dim
 
-        model = TextEmbedding(model_name=model_desc["model"])
+        model = TextEmbedding(model_name=model_desc.model)
         docs = ["hello world", "flag embedding"]
         embeddings = list(model.embed(docs))
         embeddings = np.stack(embeddings, axis=0)
         assert embeddings.shape == (2, dim)
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc["model"]
+        ), model_desc.model
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "n_dims,model_name",
-    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
-)
-def test_batch_embedding(n_dims, model_name):
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
+def test_batch_embedding(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
 
@@ -104,16 +111,13 @@ def test_batch_embedding(n_dims, model_name):
     embeddings = list(model.embed(docs, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (200, n_dims)
+    assert embeddings.shape == (len(docs), n_dims)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "n_dims,model_name",
-    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
-)
-def test_parallel_processing(n_dims, model_name):
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
+def test_parallel_processing(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
 
@@ -127,7 +131,7 @@ def test_parallel_processing(n_dims, model_name):
     embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (200, n_dims)
+    assert embeddings.shape == (len(docs), n_dims)
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
@@ -135,11 +139,8 @@ def test_parallel_processing(n_dims, model_name):
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["BAAI/bge-small-en-v1.5"],
-)
-def test_lazy_load(model_name):
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "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=[""])

tests/utils.py

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