Bump version

.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.5.1
+      placeholder: v0.4.2
     validations:
       required: true
   - type: dropdown

.github/workflows/ci.yml

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

.github/workflows/python-tests.yml

@@ -38,7 +38,7 @@ jobs:
         run: |
           python -m pip install poetry
           poetry config virtualenvs.create false
-          poetry install --no-interaction --no-ansi --without dev,docs
+          poetry install --no-interaction --no-ansi --without docs
 
       - name: Run pytest
         run: |

.github/workflows/type-checkers.yml

@@ -1,40 +0,0 @@
-name: type-checkers
-
-on: [push]
-
-jobs:
-  build:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      fail-fast: true
-      matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-        os: [ubuntu-latest]
-
-    name: Python ${{ matrix.python-version }} test
-
-    steps:
-      - uses: actions/checkout@v1
-
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip poetry
-          poetry install --no-interaction --no-ansi --without dev,docs,test
-
-          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,16 +7,8 @@ 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.

docs/examples/ColBERT_with_FastEmbed.ipynb

@@ -54,14 +54,7 @@
     },
     {
      "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": {},
@@ -219,9 +212,7 @@
    "outputs": [
     {
      "data": {
-      "text/plain": [
-       "((26, 128), (32, 128))"
-      ]
+      "text/plain": "((26, 128), (32, 128))"
      },
      "execution_count": 18,
      "metadata": {},
@@ -280,9 +271,7 @@
     "import numpy as np\n",
     "\n",
     "\n",
-    "def compute_relevance_scores(\n",
-    "    query_embedding: np.array, document_embeddings: np.array, k: int\n",
-    ") -> list[int]:\n",
+    "def compute_relevance_scores(query_embedding: np.array, document_embeddings: np.array, k: 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) -> None:\n",
+    "    def __init__(self, model_id: str):\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]) -> list[SparseEmbedding]:\n",
+    "def make_sparse_embedding(texts: list[str]):\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) -> dict[str, float]:\n",
+    "def get_tokens_and_weights(sparse_embedding, model_name):\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: dict[str, float] = {}\n",
+    "    token_weight_dict = {}\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,7 +2,6 @@ import importlib.metadata
 
 from fastembed.image import ImageEmbedding
 from fastembed.late_interaction import LateInteractionTextEmbedding
-from fastembed.late_interaction_multimodal import LateInteractionMultimodalEmbedding
 from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
 from fastembed.text import TextEmbedding
 
@@ -18,5 +17,4 @@ __all__ = [
     "SparseEmbedding",
     "ImageEmbedding",
     "LateInteractionTextEmbedding",
-    "LateInteractionMultimodalEmbedding",
 ]

fastembed/common/__init__.py

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

fastembed/common/model_description.py

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

fastembed/common/model_management.py

@@ -1,14 +1,12 @@
 import os
 import time
-import json
 import shutil
 import tarfile
 from pathlib import Path
-from typing import Any, Optional, Union, TypeVar, Generic
+from typing import Any, Optional
 
 import requests
-from huggingface_hub import snapshot_download, model_info, list_repo_tree
-from huggingface_hub.hf_api import RepoFile
+from huggingface_hub import snapshot_download
 from huggingface_hub.utils import (
     RepositoryNotFoundError,
     disable_progress_bars,
@@ -16,29 +14,20 @@ from huggingface_hub.utils import (
 )
 from loguru import logger
 from tqdm import tqdm
-from fastembed.common.model_description import BaseModelDescription
-
-T = TypeVar("T", bound=BaseModelDescription)
 
 
-class ModelManagement(Generic[T]):
-    METADATA_FILE = "files_metadata.json"
-
+class ModelManagement:
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[T]: A list of dictionaries containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         raise NotImplementedError()
 
     @classmethod
-    def _list_supported_models(cls) -> list[T]:
-        raise NotImplementedError()
-
-    @classmethod
-    def _get_model_description(cls, model_name: str) -> T:
+    def _get_model_description(cls, model_name: str) -> dict[str, Any]:
         """
         Gets the model description from the model_name.
 
@@ -49,10 +38,10 @@ class ModelManagement(Generic[T]):
             ValueError: If the model_name is not supported.
 
         Returns:
-            T: The model description.
+            dict[str, Any]: The model description.
         """
-        for model in cls._list_supported_models():
-            if model_name.lower() == model.model.lower():
+        for model in cls.list_supported_models():
+            if model_name.lower() == model["model"].lower():
                 return model
 
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
@@ -109,74 +98,21 @@ class ModelManagement(Generic[T]):
         cls,
         hf_source_repo: str,
         cache_dir: str,
-        extra_patterns: list[str],
+        extra_patterns: Optional[list[str]] = None,
         local_files_only: bool = False,
-        **kwargs: Any,
+        **kwargs,
     ) -> 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 (list[str]): extra patterns to allow in the snapshot download, typically
+            extra_patterns (Optional[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",
@@ -184,59 +120,16 @@ class ModelManagement(Generic[T]):
             "special_tokens_map.json",
             "preprocessor_config.json",
         ]
-
-        allow_patterns.extend(extra_patterns)
+        if extra_patterns is not None:
+            allow_patterns.extend(extra_patterns)
 
         snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
-        metadata_file = snapshot_dir / cls.METADATA_FILE
+        is_cached = snapshot_dir.exists()
 
-        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:
+        if is_cached:
             disable_progress_bars()
 
-        result = snapshot_download(
+        return snapshot_download(
             repo_id=hf_source_repo,
             allow_patterns=allow_patterns,
             cache_dir=cache_dir,
@@ -244,26 +137,8 @@ class ModelManagement(Generic[T]):
             **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) -> str:
+    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
         """
         Decompresses a .tar.gz file to a cache directory.
 
@@ -301,11 +176,7 @@ class ModelManagement(Generic[T]):
 
     @classmethod
     def retrieve_model_gcs(
-        cls,
-        model_name: str,
-        source_url: str,
-        cache_dir: str,
-        local_files_only: bool = False,
+        cls, model_name: str, source_url: str, cache_dir: str, local_files_only: bool = False
     ) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
         cache_tmp_dir = Path(cache_dir) / "tmp"
@@ -349,12 +220,14 @@ class ModelManagement(Generic[T]):
         return model_dir
 
     @classmethod
-    def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: Any) -> Path:
+    def download_model(
+        cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
+    ) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
         Args:
-            model (T): The model description.
+            model (dict[str, Any]): The model description.
                 Example:
                 ```
                 {
@@ -375,26 +248,23 @@ class ModelManagement(Generic[T]):
             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.sources.hf
-        url_source = model.sources.url
+        hf_source = model.get("sources", {}).get("hf")
+        url_source = model.get("sources", {}).get("url")
 
         sleep = 3.0
         while retries > 0:
             retries -= 1
 
             if hf_source:
-                extra_patterns = [model.model_file]
-                extra_patterns.extend(model.additional_files)
+                extra_patterns = [model["model_file"]]
+                extra_patterns.extend(model.get("additional_files", []))
 
                 try:
                     return Path(
                         cls.download_files_from_huggingface(
                             hf_source,
-                            cache_dir=cache_dir,
+                            cache_dir=str(cache_dir),
                             extra_patterns=extra_patterns,
                             **kwargs,
                         )
@@ -410,8 +280,8 @@ class ModelManagement(Generic[T]):
             if url_source or local_files_only:
                 try:
                     return cls.retrieve_model_gcs(
-                        model.model,
-                        str(url_source),
+                        model["model"],
+                        url_source,
                         str(cache_dir),
                         local_files_only=local_files_only,
                     )
@@ -428,4 +298,4 @@ class ModelManagement(Generic[T]):
             time.sleep(sleep)
             sleep *= 3
 
-        raise ValueError(f"Could not load model {model.model} from any source.")
+        raise ValueError(f"Could not load model {model['model']} from any source.")

fastembed/common/onnx_model.py

@@ -6,10 +6,7 @@ from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
 import numpy as np
 import onnxruntime as ort
 
-from numpy.typing import NDArray
-from tokenizers import Tokenizer
-
-from fastembed.common.types import OnnxProvider, NumpyArray
+from fastembed.common.types import OnnxProvider
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -18,26 +15,26 @@ T = TypeVar("T")
 
 @dataclass
 class OnnxOutputContext:
-    model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
+    model_output: np.ndarray
+    attention_mask: Optional[np.ndarray] = None
+    input_ids: Optional[np.ndarray] = None
 
 
 class OnnxModel(Generic[T]):
     @classmethod
-    def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
-        self.model: Optional[ort.InferenceSession] = None
-        self.tokenizer: Optional[Tokenizer] = None
+        self.model = None
+        self.tokenizer = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -73,7 +70,7 @@ class OnnxModel(Generic[T]):
             onnx_providers = ["CPUExecutionProvider"]
 
         available_providers = ort.get_available_providers()
-        requested_provider_names: list[str] = []
+        requested_provider_names = []
         for provider in onnx_providers:
             # check providers available
             provider_name = provider if isinstance(provider, str) else provider[0]
@@ -94,7 +91,6 @@ class OnnxModel(Generic[T]):
             str(model_path), providers=onnx_providers, sess_options=so
         )
         if "CUDAExecutionProvider" in requested_provider_names:
-            assert self.model is not None
             current_providers = self.model.get_providers()
             if "CUDAExecutionProvider" not in current_providers:
                 warnings.warn(
@@ -107,29 +103,29 @@ class OnnxModel(Generic[T]):
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def onnx_embed(self, *args: Any, **kwargs: Any) -> OnnxOutputContext:
+    def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
         raise NotImplementedError("Subclasses must implement this method")
 
 
-class EmbeddingWorker(Worker, Generic[T]):
+class EmbeddingWorker(Worker):
     def init_embedding(
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
-    ) -> OnnxModel[T]:
+        **kwargs,
+    ) -> OnnxModel:
         raise NotImplementedError()
 
     def __init__(
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
+        **kwargs,
     ):
         self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
     @classmethod
-    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:

fastembed/common/preprocessor_utils.py

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

fastembed/common/types.py

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

fastembed/common/utils.py

@@ -5,16 +5,12 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar
+from typing import Generator, Iterable, Optional, Union
 
 import numpy as np
 
-from fastembed.common.types import NumpyArray
 
-T = TypeVar("T")
-
-
-def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
+def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
     # Calculate the Lp norm along the specified dimension
     norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
     norm = np.maximum(norm, eps)  # Avoid division by zero
@@ -22,7 +18,7 @@ def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e
     return normalized_array
 
 
-def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
+def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))
     [[1, 2, 3], [4, 5]]

fastembed/config/dense_models.json

@@ -1,16 +0,0 @@
-{
-  "models": [
-    {
-      "model": "BAAI/bge-base-en",
-      "dim": 768,
-      "description": "Text embeddings, Unimodal (text), English...",
-      "license": "mit",
-      "size_in_GB": 0.42,
-      "sources": {
-        "hf": "Qdrant/fast-bge-base-en",
-        "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz"
-      },
-      "model_file": "model_optimized.onnx"
-    }
-  ]
-}

fastembed/embedding.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any
+from typing import Optional
 
 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: Any,
+        **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)

fastembed/image/image_embedding.py

@@ -1,11 +1,10 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
 
-from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
-from fastembed.common.model_description import DenseModelDescription
 
 
 class ImageEmbedding(ImageEmbeddingBase):
@@ -36,13 +35,9 @@ class ImageEmbedding(ImageEmbeddingBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
     def __init__(
@@ -54,12 +49,12 @@ class ImageEmbedding(ImageEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs: Any,
+        **kwargs,
     ):
         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,
@@ -79,13 +74,14 @@ class ImageEmbedding(ImageEmbeddingBase):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 16,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         """
-        Encode a list of images into list of embeddings.
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
 
         Args:
             images: Iterator of image paths or single image path to embed

fastembed/image/image_embedding_base.py

@@ -1,18 +1,18 @@
-from typing import Iterable, Optional, Any, Union
+from typing import Iterable, Optional
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 from fastembed.common.types import ImageInput
 
 
-class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
+class ImageEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -21,11 +21,11 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 16,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         """
         Embeds a list of images into a list of embeddings.
 
@@ -39,6 +39,6 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
         raise NotImplementedError()

fastembed/image/onnx_embedding.py

@@ -1,66 +1,73 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
 
-from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
 
-from fastembed.common.model_description import DenseModelDescription, ModelSource
-
-supported_onnx_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="Qdrant/clip-ViT-B-32-vision",
-        dim=512,
-        description="Image embeddings, Multimodal (text&image), 2021 year",
-        license="mit",
-        size_in_GB=0.34,
-        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-vision"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="Qdrant/resnet50-onnx",
-        dim=2048,
-        description="Image embeddings, Unimodal (image), 2016 year",
-        license="apache-2.0",
-        size_in_GB=0.1,
-        sources=ModelSource(hf="Qdrant/resnet50-onnx"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="Qdrant/Unicom-ViT-B-16",
-        dim=768,
-        description="Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
-        license="apache-2.0",
-        size_in_GB=0.82,
-        sources=ModelSource(hf="Qdrant/Unicom-ViT-B-16"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="Qdrant/Unicom-ViT-B-32",
-        dim=512,
-        description="Image embeddings, Multimodal (text&image), 2023 year",
-        license="apache-2.0",
-        size_in_GB=0.48,
-        sources=ModelSource(hf="Qdrant/Unicom-ViT-B-32"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-clip-v1",
-        dim=768,
-        description="Image embeddings, Multimodal (text&image), 2024 year",
-        license="apache-2.0",
-        size_in_GB=0.34,
-        sources=ModelSource(hf="jinaai/jina-clip-v1"),
-        model_file="onnx/vision_model.onnx",
-    ),
+supported_onnx_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-vision",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2021 year",
+        "license": "mit",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-vision",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/resnet50-onnx",
+        "dim": 2048,
+        "description": "Image embeddings, Unimodal (image), 2016 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.1,
+        "sources": {
+            "hf": "Qdrant/resnet50-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-16",
+        "dim": 768,
+        "description": "Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.82,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-16",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-32",
+        "dim": 512,
+        "description": "Image embeddings, Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.48,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-32",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "jinaai/jina-clip-v1",
+        "dim": 768,
+        "description": "Image embeddings, Multimodal (text&image), 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "jinaai/jina-clip-v1",
+        },
+        "model_file": "onnx/vision_model.onnx",
+    },
 ]
 
 
-class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
+class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
     def __init__(
         self,
         model_name: str,
@@ -71,8 +78,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         """
         Args:
@@ -90,7 +96,6 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             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.
@@ -105,19 +110,17 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
         self._model_dir = self.download_model(
-            self.model_description,
-            self.cache_dir,
-            local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         if not self.lazy_load:
@@ -129,7 +132,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         """
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
@@ -137,22 +140,22 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         )
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[Dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 16,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of images into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -182,24 +185,24 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
         return OnnxImageEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
 
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return normalize(output.model_output).astype(np.float32)
 
 
-class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
+class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
         return OnnxImageEmbedding(
             model_name=model_name,
             cache_dir=cache_dir,

fastembed/image/onnx_image_model.py

@@ -2,13 +2,11 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
 from PIL import Image
 
-from fastembed.image.transform.operators import Compose
-from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
@@ -20,7 +18,7 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 class OnnxImageModel(OnnxModel[T]):
     @classmethod
-    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
@@ -28,11 +26,11 @@ class OnnxImageModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.processor: Optional[Compose] = None
+        self.processor = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -60,21 +58,19 @@ class OnnxImageModel(OnnxModel[T]):
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _build_onnx_input(self, encoded: NumpyArray) -> dict[str, NumpyArray]:
-        input_name = self.model.get_inputs()[0].name  # type: ignore[union-attr]
-        return {input_name: encoded}
+    def _build_onnx_input(self, encoded: np.ndarray) -> dict[str, np.ndarray]:
+        return {node.name: encoded for node in self.model.get_inputs()}
 
-    def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+    def onnx_embed(self, images: list[ImageInput], **kwargs) -> 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))
+            encoded = self.processor(image_files)
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
-        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+        model_output = self.model.run(None, onnx_input)
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -82,13 +78,13 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput,
         batch_size: int = 256,
         parallel: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[T]:
         is_small = False
 
@@ -125,10 +121,10 @@ class OnnxImageModel(OnnxModel[T]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
-                yield from self._post_process_onnx_output(batch)  # type: ignore
+                yield from self._post_process_onnx_output(batch)
 
 
-class ImageEmbeddingWorker(EmbeddingWorker[T]):
+class ImageEmbeddingWorker(EmbeddingWorker):
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             embeddings = self.model.onnx_embed(batch)

fastembed/image/transform/functional.py

@@ -1,10 +1,8 @@
-from typing import Union
+from typing import Sized, Union
 
 import numpy as np
 from PIL import Image
 
-from fastembed.common.types import NumpyArray
-
 
 def convert_to_rgb(image: Image.Image) -> Image.Image:
     if image.mode == "RGB":
@@ -15,9 +13,9 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Union[Image.Image, NumpyArray],
+    image: Union[Image.Image, np.ndarray],
     size: tuple[int, int],
-) -> NumpyArray:
+) -> np.ndarray:
     if isinstance(image, np.ndarray):
         _, orig_height, orig_width = image.shape
     else:
@@ -42,7 +40,7 @@ def center_crop(
     new_height = max(crop_height, orig_height)
     new_width = max(crop_width, orig_width)
     new_shape = image.shape[:-2] + (new_height, new_width)
-    new_image = np.zeros_like(image, shape=new_shape, dtype=np.float32)
+    new_image = np.zeros_like(image, shape=new_shape)
 
     top_pad = (new_height - orig_height) // 2
     bottom_pad = top_pad + orig_height
@@ -63,34 +61,37 @@ def center_crop(
 
 
 def normalize(
-    image: NumpyArray,
-    mean: Union[float, list[float]],
-    std: Union[float, list[float]],
-) -> NumpyArray:
+    image: np.ndarray,
+    mean: Union[float, np.ndarray],
+    std: Union[float, np.ndarray],
+) -> np.ndarray:
+    if not isinstance(image, np.ndarray):
+        raise ValueError("image must be a numpy array")
+
     num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
 
     if not np.issubdtype(image.dtype, np.floating):
         image = image.astype(np.float32)
 
-    mean = mean if isinstance(mean, list) else [mean] * num_channels
+    if isinstance(mean, Sized):
+        if len(mean) != num_channels:
+            raise ValueError(
+                f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}"
+            )
+    else:
+        mean = [mean] * num_channels
+    mean = np.array(mean, dtype=image.dtype)
 
-    if len(mean) != 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)}"
-        )
+    if isinstance(std, Sized):
+        if len(std) != num_channels:
+            raise ValueError(
+                f"std must have {num_channels} elements if it is an iterable, got {len(std)}"
+            )
+    else:
+        std = [std] * num_channels
+    std = np.array(std, dtype=image.dtype)
 
-    mean_arr = np.array(mean, 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
+    image = ((image.T - mean) / std).T
     return image
 
 
@@ -113,11 +114,11 @@ def resize(
     return image.resize(new_size, resample)
 
 
-def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyArray:
+def rescale(image: np.ndarray, scale: float, dtype=np.float32) -> np.ndarray:
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
+def pil2ndarray(image: Union[Image.Image, np.ndarray]):
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image

fastembed/image/transform/operators.py

@@ -1,8 +1,8 @@
 from typing import Any, Union, Optional
 
+import numpy as np
 from PIL import Image
 
-from fastembed.common.types import NumpyArray
 from fastembed.image.transform.functional import (
     center_crop,
     convert_to_rgb,
@@ -15,7 +15,7 @@ from fastembed.image.transform.functional import (
 
 
 class Transform:
-    def __call__(self, images: list[Any]) -> Union[list[Image.Image], list[NumpyArray]]:
+    def __call__(self, images: list) -> Union[list[Image.Image], list[np.ndarray]]:
         raise NotImplementedError("Subclasses must implement this method")
 
 
@@ -28,7 +28,7 @@ class CenterCrop(Transform):
     def __init__(self, size: tuple[int, int]):
         self.size = size
 
-    def __call__(self, images: list[Image.Image]) -> list[NumpyArray]:
+    def __call__(self, images: list[Image.Image]) -> list[np.ndarray]:
         return [center_crop(image=image, size=self.size) for image in images]
 
 
@@ -37,7 +37,7 @@ class Normalize(Transform):
         self.mean = mean
         self.std = std
 
-    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+    def __call__(self, images: list[np.ndarray]) -> list[np.ndarray]:
         return [normalize(image, mean=self.mean, std=self.std) for image in images]
 
 
@@ -58,12 +58,12 @@ class Rescale(Transform):
     def __init__(self, scale: float = 1 / 255):
         self.scale = scale
 
-    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+    def __call__(self, images: list[np.ndarray]) -> list[np.ndarray]:
         return [rescale(image, scale=self.scale) for image in images]
 
 
 class PILtoNDarray(Transform):
-    def __call__(self, images: list[Union[Image.Image, NumpyArray]]) -> list[NumpyArray]:
+    def __call__(self, images: list[Union[Image.Image, np.ndarray]]) -> list[np.ndarray]:
         return [pil2ndarray(image) for image in images]
 
 
@@ -71,7 +71,7 @@ class PadtoSquare(Transform):
     def __init__(
         self,
         size: int,
-        fill_color: Union[str, int, tuple[int, ...]],
+        fill_color: Optional[Union[str, int, tuple[int, ...]]] = None,
     ):
         self.size = size
         self.fill_color = fill_color
@@ -87,8 +87,8 @@ class Compose:
         self.transforms = transforms
 
     def __call__(
-        self, images: Union[list[Image.Image], list[NumpyArray]]
-    ) -> Union[list[NumpyArray], list[Image.Image]]:
+        self, images: Union[list[Image.Image], list[np.ndarray]]
+    ) -> Union[list[np.ndarray], list[Image.Image]]:
         for transform in self.transforms:
             images = transform(images)
         return images
@@ -122,7 +122,7 @@ class Compose:
         Returns:
             Compose: Image processor.
         """
-        transforms: list[Transform] = []
+        transforms = []
         cls._get_convert_to_rgb(transforms, config)
         cls._get_resize(transforms, config)
         cls._get_pad2square(transforms, config)
@@ -133,13 +133,13 @@ class Compose:
         return cls(transforms=transforms)
 
     @staticmethod
-    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> None:
+    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]):
         transforms.append(ConvertToRGB())
 
     @classmethod
-    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
+    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]):
         mode = config.get("image_processor_type", "CLIPImageProcessor")
-        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
+        if mode == "CLIPImageProcessor":
             if config.get("do_resize", False):
                 size = config["size"]
                 if "shortest_edge" in size:
@@ -200,18 +200,17 @@ class Compose:
             raise ValueError(f"Preprocessor {mode} is not supported")
 
     @staticmethod
-    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]) -> None:
+    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]):
         mode = config.get("image_processor_type", "CLIPImageProcessor")
-        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
+        if mode == "CLIPImageProcessor":
             if config.get("do_center_crop", False):
-                crop_size_raw = config["crop_size"]
-                crop_size: tuple[int, int]
-                if isinstance(crop_size_raw, int):
-                    crop_size = (crop_size_raw, crop_size_raw)
-                elif isinstance(crop_size_raw, dict):
-                    crop_size = (crop_size_raw["height"], crop_size_raw["width"])
+                crop_size = config["crop_size"]
+                if isinstance(crop_size, int):
+                    crop_size = (crop_size, crop_size)
+                elif isinstance(crop_size, dict):
+                    crop_size = (crop_size["height"], crop_size["width"])
                 else:
-                    raise ValueError(f"Invalid crop size: {crop_size_raw}")
+                    raise ValueError(f"Invalid crop size: {crop_size}")
                 transforms.append(CenterCrop(size=crop_size))
         elif mode == "ConvNextFeatureExtractor":
             pass
@@ -221,24 +220,24 @@ class Compose:
             raise ValueError(f"Preprocessor {mode} is not supported")
 
     @staticmethod
-    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
+    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]):
         transforms.append(PILtoNDarray())
 
     @staticmethod
-    def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
+    def _get_rescale(transforms: list[Transform], config: dict[str, Any]):
         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]) -> None:
+    def _get_normalize(transforms: list[Transform], config: dict[str, Any]):
         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]) -> None:
+    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]):
         mode = config.get("image_processor_type", "CLIPImageProcessor")
         if mode == "CLIPImageProcessor":
             pass

fastembed/late_interaction/colbert.py

@@ -4,7 +4,6 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 import numpy as np
 from tokenizers import Encoding
 
-from fastembed.common.types import NumpyArray
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
@@ -12,31 +11,35 @@ from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_colbert_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="colbert-ir/colbertv2.0",
-        dim=128,
-        description="Late interaction model",
-        license="mit",
-        size_in_GB=0.44,
-        sources=ModelSource(hf="colbert-ir/colbertv2.0"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="answerdotai/answerai-colbert-small-v1",
-        dim=96,
-        description="Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="answerdotai/answerai-colbert-small-v1"),
-        model_file="vespa_colbert.onnx",
-    ),
+
+supported_colbert_models = [
+    {
+        "model": "colbert-ir/colbertv2.0",
+        "dim": 128,
+        "description": "Late interaction model",
+        "license": "mit",
+        "size_in_GB": 0.44,
+        "sources": {
+            "hf": "colbert-ir/colbertv2.0",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "answerdotai/answerai-colbert-small-v1",
+        "dim": 96,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "answerdotai/answerai-colbert-small-v1",
+        },
+        "model_file": "vespa_colbert.onnx",
+    },
 ]
 
 
-class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
+class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
     QUERY_MARKER_TOKEN_ID = 1
     DOCUMENT_MARKER_TOKEN_ID = 2
     MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
@@ -44,7 +47,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def _post_process_onnx_output(
         self, output: OnnxOutputContext, is_doc: bool = True
-    ) -> Iterable[NumpyArray]:
+    ) -> Iterable[np.ndarray]:
         if not is_doc:
             return output.model_output.astype(np.float32)
 
@@ -54,7 +57,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             )
 
         for i, token_sequence in enumerate(output.input_ids):
-            for j, token_id in enumerate(token_sequence):  # type: ignore
+            for j, token_id in enumerate(token_sequence):
                 if token_id in self.skip_list or token_id == self.pad_token_id:
                     output.attention_mask[i, j] = 0
 
@@ -65,15 +68,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         return output.model_output.astype(np.float32)
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
-        onnx_input["input_ids"] = np.insert(
-            onnx_input["input_ids"].astype(np.int64), 1, marker_token, axis=1
-        )
-        onnx_input["attention_mask"] = np.insert(
-            onnx_input["attention_mask"].astype(np.int64), 1, 1, axis=1
-        )
+        onnx_input["input_ids"] = np.insert(onnx_input["input_ids"], 1, marker_token, axis=1)
+        onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
         return onnx_input
 
     def tokenize(self, documents: list[str], is_doc: bool = True, **kwargs: Any) -> list[Encoding]:
@@ -84,7 +83,6 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     def _tokenize_query(self, query: str) -> list[Encoding]:
-        assert self.tokenizer is not None
         encoded = self.tokenizer.encode_batch([query])
         # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
         if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
@@ -104,15 +102,15 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         return encoded
 
     def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
-        encoded = self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+        encoded = self.tokenizer.encode_batch(documents)
         return encoded
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_colbert_models
 
@@ -126,8 +124,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         """
         Args:
@@ -145,7 +142,6 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             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.
@@ -160,24 +156,22 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
         self._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.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
-        self.mask_token_id: Optional[int] = None
-        self.pad_token_id: Optional[int] = None
-        self.skip_list: set[int] = set()
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.skip_list = set()
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -185,13 +179,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
         )
-        assert self.tokenizer is not None
         self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
         self.skip_list = {
@@ -207,8 +200,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -236,7 +229,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
         if isinstance(query, str):
             query = [query]
 
@@ -249,12 +242,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return ColbertEmbeddingWorker
 
 
-class ColbertEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
+class ColbertEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
         return Colbert(
             model_name=model_name,
             cache_dir=cache_dir,

fastembed/late_interaction/jina_colbert.py

@@ -1,20 +1,24 @@
 from typing import Any, Type
 
-from fastembed.common.types import NumpyArray
-from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
+import numpy as np
 
-supported_jina_colbert_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="jinaai/jina-colbert-v2",
-        dim=128,
-        description="New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
-        license="cc-by-nc-4.0",
-        size_in_GB=2.24,
-        sources=ModelSource(hf="jinaai/jina-colbert-v2"),
-        model_file="onnx/model.onnx",
-        additional_files=["onnx/model.onnx_data"],
-    )
+from fastembed.late_interaction.colbert import Colbert
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+supported_jina_colbert_models = [
+    {
+        "model": "jinaai/jina-colbert-v2",
+        "dim": 128,
+        "description": "New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
+        "license": "cc-by-nc-4.0",
+        "size_in_GB": 2.24,
+        "sources": {
+            "hf": "jinaai/jina-colbert-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "additional_files": ["onnx/model.onnx_data"],
+    },
 ]
 
 
@@ -25,21 +29,21 @@ class JinaColbert(Colbert):
     MASK_TOKEN = "<mask>"
 
     @classmethod
-    def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return JinaColbertEmbeddingWorker
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_jina_colbert_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
 
         # the attention mask for jina-colbert-v2 is always 1 in queries
@@ -48,8 +52,8 @@ class JinaColbert(Colbert):
         return onnx_input
 
 
-class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
+class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> 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, Any
+from typing import Iterable, Optional, Union
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
+class LateInteractionTextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -23,11 +23,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -36,13 +36,15 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -50,11 +52,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[NdArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        else:
+        if isinstance(query, Iterable):
             yield from self.embed(query, **kwargs)

fastembed/late_interaction/late_interaction_text_embedding.py

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

fastembed/late_interaction_multimodal/__init__.py

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

fastembed/late_interaction_multimodal/colpali.py

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

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

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

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -1,67 +0,0 @@
-from typing import Iterable, Optional, Union, Any
-
-
-from fastembed.common import ImageInput
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.model_management import ModelManagement
-from fastembed.common.types import NumpyArray
-
-
-class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescription]):
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        **kwargs: Any,
-    ):
-        self.model_name = model_name
-        self.cache_dir = cache_dir
-        self.threads = threads
-        self._local_files_only = kwargs.pop("local_files_only", False)
-
-    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

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

fastembed/model_loader.py

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

fastembed/parallel_processor.py

@@ -140,7 +140,7 @@ class ParallelWorkerPool:
             self.processes.append(process)
 
     def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]:
-        buffer: defaultdict[int, Any] = defaultdict(Any)  # type: ignore
+        buffer = defaultdict(Any)
         next_expected = 0
 
         for idx, item in self.semi_ordered_map(stream, *args, **kwargs):

fastembed/py.typed

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

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -10,67 +10,78 @@ from fastembed.rerank.cross_encoder.onnx_text_model import (
     TextRerankerWorker,
 )
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import BaseModelDescription, ModelSource
 
-supported_onnx_models: list[BaseModelDescription] = [
-    BaseModelDescription(
-        model="Xenova/ms-marco-MiniLM-L-6-v2",
-        description="MiniLM-L-6-v2 model optimized for re-ranking tasks.",
-        license="apache-2.0",
-        size_in_GB=0.08,
-        sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-6-v2"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="Xenova/ms-marco-MiniLM-L-12-v2",
-        description="MiniLM-L-12-v2 model optimized for re-ranking tasks.",
-        license="apache-2.0",
-        size_in_GB=0.12,
-        sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-12-v2"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="BAAI/bge-reranker-base",
-        description="BGE reranker base model for cross-encoder re-ranking.",
-        license="mit",
-        size_in_GB=1.04,
-        sources=ModelSource(hf="BAAI/bge-reranker-base"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="jinaai/jina-reranker-v1-tiny-en",
-        description="Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="jinaai/jina-reranker-v1-tiny-en"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="jinaai/jina-reranker-v1-turbo-en",
-        description="Designed for blazing-fast re-ranking with 8K context length.",
-        license="apache-2.0",
-        size_in_GB=0.15,
-        sources=ModelSource(hf="jinaai/jina-reranker-v1-turbo-en"),
-        model_file="onnx/model.onnx",
-    ),
-    BaseModelDescription(
-        model="jinaai/jina-reranker-v2-base-multilingual",
-        description="A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
-        license="cc-by-nc-4.0",
-        size_in_GB=1.11,
-        sources=ModelSource(hf="jinaai/jina-reranker-v2-base-multilingual"),
-        model_file="onnx/model.onnx",
-    ),
+supported_onnx_models = [
+    {
+        "model": "Xenova/ms-marco-MiniLM-L-6-v2",
+        "size_in_GB": 0.08,
+        "sources": {
+            "hf": "Xenova/ms-marco-MiniLM-L-6-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "Xenova/ms-marco-MiniLM-L-12-v2",
+        "size_in_GB": 0.12,
+        "sources": {
+            "hf": "Xenova/ms-marco-MiniLM-L-12-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "MiniLM-L-12-v2 model optimized for re-ranking tasks.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "BAAI/bge-reranker-base",
+        "size_in_GB": 1.04,
+        "sources": {
+            "hf": "BAAI/bge-reranker-base",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "BGE reranker base model for cross-encoder re-ranking.",
+        "license": "mit",
+    },
+    {
+        "model": "jinaai/jina-reranker-v1-tiny-en",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-tiny-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v1-turbo-en",
+        "size_in_GB": 0.15,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-turbo-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v2-base-multilingual",
+        "size_in_GB": 1.11,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v2-base-multilingual",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
+        "license": "cc-by-nc-4.0",
+    },
 ]
 
 
 class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     @classmethod
-    def _list_supported_models(cls) -> list[BaseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[BaseModelDescription]: A list of BaseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
@@ -84,7 +95,6 @@ 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,
     ):
         """
@@ -103,7 +113,6 @@ 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.
@@ -123,19 +132,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 = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
         )
 
         if not self.lazy_load:
@@ -144,7 +153,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,
@@ -205,7 +214,7 @@ class TextCrossEncoderWorker(TextRerankerWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
+        **kwargs,
     ) -> OnnxTextCrossEncoder:
         return OnnxTextCrossEncoder(
             model_name=model_name,

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -12,7 +12,6 @@ from fastembed.common.onnx_model import (
     OnnxOutputContext,
     OnnxProvider,
 )
-from fastembed.common.types import NumpyArray
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
@@ -43,14 +42,13 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             device_id=device_id,
         )
         self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
-        assert self.tokenizer is not None
 
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
-        return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
+        return self.tokenizer.encode_batch(pairs)
 
-    def _build_onnx_input(self, tokenized_input: list[Encoding]) -> dict[str, NumpyArray]:
-        input_names: set[str] = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
-        inputs: dict[str, NumpyArray] = {
+    def _build_onnx_input(self, tokenized_input):
+        input_names = {node.name for node in self.model.get_inputs()}
+        inputs = {
             "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
         }
         if "token_type_ids" in input_names:
@@ -71,9 +69,9 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         tokenized_input = self.tokenize(pairs, **kwargs)
         inputs = self._build_onnx_input(tokenized_input)
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
-        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
         relevant_output = outputs[0]
-        scores: NumpyArray = relevant_output[:, 0]
+        scores = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
 
     def _rerank_documents(
@@ -131,38 +129,21 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(pairs, batch_size), **params):
-                yield from self._post_process_onnx_output(batch)  # type: ignore
+                yield from self._post_process_onnx_output(batch)
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
 
-class TextRerankerWorker(EmbeddingWorker[float]):
-    def __init__(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ):
-        self.model: OnnxCrossEncoderModel
-        super().__init__(model_name, cache_dir, **kwargs)
-
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-        **kwargs: Any,
-    ) -> OnnxCrossEncoderModel:
-        raise NotImplementedError()
-
+class TextRerankerWorker(EmbeddingWorker):
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             onnx_output = self.model.onnx_embed_pairs(batch)

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -1,10 +1,8 @@
 from typing import Any, Iterable, Optional, Sequence, Type
-from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import BaseModelDescription
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
@@ -17,7 +15,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         """Lists the supported models.
 
         Returns:
-            list[BaseModelDescription]: A list of dictionaries containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
 
             Example:
                 ```
@@ -35,13 +33,9 @@ class TextCrossEncoder(TextCrossEncoderBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[BaseModelDescription]:
-        result: list[BaseModelDescription] = []
+        result = []
         for encoder in cls.CROSS_ENCODER_REGISTRY:
-            result.extend(encoder._list_supported_models())
+            result.extend(encoder.list_supported_models())
         return result
 
     def __init__(
@@ -58,8 +52,8 @@ class TextCrossEncoder(TextCrossEncoderBase):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         for CROSS_ENCODER_TYPE in self.CROSS_ENCODER_REGISTRY:
-            supported_models = CROSS_ENCODER_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = CROSS_ENCODER_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = CROSS_ENCODER_TYPE(
                     model_name=model_name,
                     cache_dir=cache_dir,

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -1,16 +1,15 @@
 from typing import Any, Iterable, Optional
 
-from fastembed.common.model_description import BaseModelDescription
 from fastembed.common.model_management import ModelManagement
 
 
-class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
+class TextCrossEncoderBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -22,7 +21,7 @@ class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
         query: str,
         documents: Iterable[str],
         batch_size: int = 64,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[float]:
         """Rerank a list of documents given a query.
 

fastembed/sparse/bm25.py

@@ -19,7 +19,6 @@ 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",
@@ -53,18 +52,19 @@ supported_languages = [
     "turkish",
 ]
 
-supported_bm25_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="Qdrant/bm25",
-        vocab_size=0,
-        description="BM25 as sparse embeddings meant to be used with Qdrant",
-        license="apache-2.0",
-        size_in_GB=0.01,
-        sources=ModelSource(hf="Qdrant/bm25"),
-        additional_files=[f"{lang}.txt" for lang in supported_languages],
-        requires_idf=True,
-        model_file="mock.file",
-    ),
+supported_bm25_models = [
+    {
+        "model": "Qdrant/bm25",
+        "description": "BM25 as sparse embeddings meant to be used with Qdrant",
+        "license": "apache-2.0",
+        "size_in_GB": 0.01,
+        "sources": {
+            "hf": "Qdrant/bm25",
+        },
+        "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
+        "additional_files": [f"{lang}.txt" for lang in supported_languages],
+        "requires_idf": True,
+    },
 ]
 
 
@@ -108,8 +108,7 @@ class Bm25(SparseTextEmbeddingBase):
         language: str = "english",
         token_max_length: int = 40,
         disable_stemmer: bool = False,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
 
@@ -123,13 +122,10 @@ class Bm25(SparseTextEmbeddingBase):
         self.avg_len = avg_len
 
         model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
         self._model_dir = self.download_model(
-            model_description,
-            self.cache_dir,
-            local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.token_max_length = token_max_length
@@ -137,7 +133,7 @@ class Bm25(SparseTextEmbeddingBase):
         self.disable_stemmer = disable_stemmer
 
         if disable_stemmer:
-            self.stopwords: set[str] = set()
+            self.stopwords = set()
             self.stemmer = None
         else:
             self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
@@ -146,11 +142,11 @@ class Bm25(SparseTextEmbeddingBase):
         self.tokenizer = SimpleTokenizer
 
     @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_bm25_models
 
@@ -206,14 +202,14 @@ class Bm25(SparseTextEmbeddingBase):
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 for record in batch:
-                    yield record  # type: ignore
+                    yield record
 
     def embed(
         self,
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -239,7 +235,7 @@ class Bm25(SparseTextEmbeddingBase):
         )
 
     def _stem(self, tokens: list[str]) -> list[str]:
-        stemmed_tokens: list[str] = []
+        stemmed_tokens = []
         for token in tokens:
             lower_token = token.lower()
 
@@ -262,7 +258,7 @@ class Bm25(SparseTextEmbeddingBase):
         self,
         documents: list[str],
     ) -> list[SparseEmbedding]:
-        embeddings: list[SparseEmbedding] = []
+        embeddings = []
         for document in documents:
             document = remove_non_alphanumeric(document)
             tokens = self.tokenizer.tokenize(document)
@@ -286,8 +282,8 @@ class Bm25(SparseTextEmbeddingBase):
         Returns:
             dict[int, float]: The token_id to term frequency mapping.
         """
-        tf_map: dict[int, float] = {}
-        counter: defaultdict[str, int] = defaultdict(int)
+        tf_map = {}
+        counter = defaultdict(int)
         for stemmed_token in tokens:
             counter[stemmed_token] += 1
 
@@ -305,9 +301,7 @@ 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: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> 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.
         """
@@ -335,7 +329,7 @@ class Bm25Worker(Worker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
+        **kwargs,
     ):
         self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
@@ -343,13 +337,11 @@ class Bm25Worker(Worker):
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
-    def process(
-        self, items: Iterable[tuple[int, Any]]
-    ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             onnx_output = self.model.raw_embed(batch)
             yield idx, onnx_output
 
     @staticmethod
-    def init_embedding(model_name: str, cache_dir: str, **kwargs: Any) -> Bm25:
+    def init_embedding(model_name: str, cache_dir: str, **kwargs) -> Bm25:
         return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)

fastembed/sparse/bm42.py

@@ -15,20 +15,21 @@ from fastembed.sparse.sparse_embedding_base import (
     SparseTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import SparseModelDescription, ModelSource
 
-supported_bm42_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="Qdrant/bm42-all-minilm-l6-v2-attentions",
-        vocab_size=30522,
-        description="Light sparse embedding model, which assigns an importance score to each token in the text",
-        license="apache-2.0",
-        size_in_GB=0.09,
-        sources=ModelSource(hf="Qdrant/all_miniLM_L6_v2_with_attentions"),
-        model_file="model.onnx",
-        additional_files=["stopwords.txt"],
-        requires_idf=True,
-    ),
+supported_bm42_models = [
+    {
+        "model": "Qdrant/bm42-all-minilm-l6-v2-attentions",
+        "vocab_size": 30522,
+        "description": "Light sparse embedding model, which assigns an importance score to each token in the text",
+        "license": "apache-2.0",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "Qdrant/all_miniLM_L6_v2_with_attentions",
+        },
+        "model_file": "model.onnx",
+        "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
+    },
 ]
 
 MODEL_TO_LANGUAGE = {
@@ -65,8 +66,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         """
         Args:
@@ -86,7 +86,6 @@ 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.
@@ -101,26 +100,24 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
         self._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.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
-        self.invert_vocab: dict[int, str] = {}
+        self.invert_vocab = {}
 
-        self.special_tokens: set[str] = set()
-        self.special_tokens_ids: set[int] = set()
+        self.special_tokens = set()
+        self.special_tokens_ids = set()
         self.punctuation = set(string.punctuation)
         self.stopwords = set(self._load_stopwords(self._model_dir))
         self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
@@ -132,21 +129,20 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
         )
-
-        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
+        for token, idx in self.tokenizer.get_vocab().items():
             self.invert_vocab[idx] = token
         self.special_tokens = set(self.special_token_to_id.keys())
         self.special_tokens_ids = set(self.special_token_to_id.values())
         self.stopwords = set(self._load_stopwords(self._model_dir))
 
     def _filter_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
-        result: list[tuple[str, Any]] = []
+        result = []
         for token, value in tokens:
             if token in self.stopwords or token in self.punctuation:
                 continue
@@ -154,7 +150,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return result
 
     def _stem_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
-        result: list[tuple[str, Any]] = []
+        result = []
         for token, value in tokens:
             processed_token = self.stemmer.stem_word(token)
             result.append((processed_token, value))
@@ -164,7 +160,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _aggregate_weights(
         cls, tokens: list[tuple[str, list[int]]], weights: list[float]
     ) -> list[tuple[str, float]]:
-        result: list[tuple[str, float]] = []
+        result = []
         for token, idxs in tokens:
             sum_weight = sum(weights[idx] for idx in idxs)
             result.append((token, sum_weight))
@@ -173,11 +169,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _reconstruct_bpe(
         self, bpe_tokens: Iterable[tuple[int, str]]
     ) -> list[tuple[str, list[int]]]:
-        result: list[tuple[str, list[int]]] = []
-        acc: str = ""
-        acc_idx: list[int] = []
+        result = []
+        acc = ""
+        acc_idx = []
 
-        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix  # type: ignore[union-attr]
+        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix
         continuing_subword_prefix_len = len(continuing_subword_prefix)
 
         for idx, token in bpe_tokens:
@@ -205,7 +201,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         So that the scoring doesn't depend on absolute values assigned by the model, but on the relative importance.
         """
 
-        new_vector: dict[int, float] = {}
+        new_vector = {}
 
         for token, value in vector.items():
             token_id = abs(mmh3.hash(token))
@@ -221,7 +217,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.astype(int)
+        token_ids_batch = output.input_ids
 
         # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
         pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask
@@ -240,7 +236,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
             weighted = self._aggregate_weights(stemmed, attention_value)
 
-            max_token_weight: dict[str, float] = {}
+            max_token_weight = {}
 
             for token, weight in weighted:
                 max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
@@ -250,11 +246,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding.from_dict(rescored)
 
     @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_bm42_models
 
@@ -272,7 +268,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -303,15 +299,13 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     @classmethod
     def _query_rehash(cls, tokens: Iterable[str]) -> dict[int, float]:
-        result: dict[int, float] = {}
+        result = {}
         for token in tokens:
             token_id = abs(mmh3.hash(token))
             result[token_id] = 1.0
         return result
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> 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.
@@ -324,7 +318,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             self.load_onnx_model()
 
         for text in query:
-            encoded = self.tokenizer.encode(text)  # type: ignore[union-attr]
+            encoded = self.tokenizer.encode(text)
             document_tokens_with_ids = enumerate(encoded.tokens)
             reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
             filtered = self._filter_pair_tokens(reconstructed)
@@ -333,12 +327,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[SparseEmbedding]]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return Bm42TextEmbeddingWorker
 
 
-class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
+class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
         return Bm42(
             model_name=model_name,
             cache_dir=cache_dir,

fastembed/sparse/sparse_embedding_base.py

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

fastembed/sparse/sparse_text_embedding.py

@@ -1,5 +1,4 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.sparse.bm25 import Bm25
@@ -10,7 +9,6 @@ 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):
@@ -40,13 +38,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
                 ]
                 ```
         """
-        return [asdict(model) for model in cls._list_supported_models()]
-
-    @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
-        result: list[SparseModelDescription] = []
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
     def __init__(
@@ -58,7 +52,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs: Any,
+        **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
         if model_name == "prithvida/Splade_PP_en_v1":
@@ -71,8 +65,8 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
             model_name = "prithivida/Splade_PP_en_v1"
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -95,7 +89,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -114,9 +108,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 

fastembed/sparse/splade_pp.py

@@ -9,27 +9,30 @@ from fastembed.sparse.sparse_embedding_base import (
     SparseTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import SparseModelDescription, ModelSource
 
-supported_splade_models: list[SparseModelDescription] = [
-    SparseModelDescription(
-        model="prithivida/Splade_PP_en_v1",
-        vocab_size=30522,
-        description="Independent Implementation of SPLADE++ Model for English.",
-        license="apache-2.0",
-        size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
-        model_file="model.onnx",
-    ),
-    SparseModelDescription(
-        model="prithvida/Splade_PP_en_v1",
-        vocab_size=30522,
-        description="Independent Implementation of SPLADE++ Model for English.",
-        license="apache-2.0",
-        size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
-        model_file="model.onnx",
-    ),
+supported_splade_models = [
+    {
+        "model": "prithivida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "prithvida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
 
@@ -52,11 +55,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding(values=scores, indices=indices)
 
     @classmethod
-    def _list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_splade_models
 
@@ -70,8 +73,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         """
         Args:
@@ -89,7 +91,6 @@ 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.
@@ -103,20 +104,18 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
 
         self._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.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         if not self.lazy_load:
@@ -125,7 +124,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,
@@ -137,7 +136,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -167,12 +166,12 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return SpladePPEmbeddingWorker
 
 
-class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
+class SpladePPEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> 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
 
-from fastembed.common.types import NumpyArray
+import numpy as np
+
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
 
-supported_clip_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="Qdrant/clip-ViT-B-32-text",
-        dim=512,
-        description=(
-            "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2021 year"
-        ),
-        license="mit",
-        size_in_GB=0.25,
-        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-text"),
-        model_file="model.onnx",
-    ),
+supported_clip_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-text",
+        "dim": 512,
+        "description": "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
+        "license": "mit",
+        "size_in_GB": 0.25,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-text",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
 
 class CLIPOnnxEmbedding(OnnxTextEmbedding):
     @classmethod
-    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return CLIPEmbeddingWorker
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_clip_models
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return output.model_output
 
 
@@ -44,7 +44,7 @@ class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
+        **kwargs,
     ) -> OnnxTextEmbedding:
         return CLIPOnnxEmbedding(
             model_name=model_name,

fastembed/text/e5_onnx_embedding.py

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

@@ -1,100 +0,0 @@
-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,207 +1,196 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from fastembed.common.types import NumpyArray, OnnxProvider
+
+from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 from fastembed.text.text_embedding_base import TextEmbeddingBase
-from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_onnx_models: list[DenseModelDescription] = [
-    DenseModelDescription(
-        model="BAAI/bge-base-en",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.42,
-        sources=ModelSource(
-            hf="Qdrant/fast-bge-base-en",
-            url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
-        ),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-base-en-v1.5",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.21,
-        sources=ModelSource(
-            hf="qdrant/bge-base-en-v1.5-onnx-q",
-            url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
-        ),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-large-en-v1.5",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=1.20,
-        sources=ModelSource(hf="qdrant/bge-large-en-v1.5-onnx"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-small-en",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.13,
-        sources=ModelSource(
-            hf="Qdrant/bge-small-en",
-            url="https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
-        ),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-small-en-v1.5",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.067,
-        sources=ModelSource(hf="qdrant/bge-small-en-v1.5-onnx-q"),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="BAAI/bge-small-zh-v1.5",
-        dim=512,
-        description=(
-            "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not so necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=0.09,
-        sources=ModelSource(
-            hf="Qdrant/bge-small-zh-v1.5",
-            url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
-        ),
-        model_file="model_optimized.onnx",
-    ),
-    DenseModelDescription(
-        model="thenlper/gte-large",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=1.20,
-        sources=ModelSource(hf="qdrant/gte-large-onnx"),
-        model_file="model.onnx",
-    ),
-    DenseModelDescription(
-        model="mixedbread-ai/mxbai-embed-large-v1",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.64,
-        sources=ModelSource(hf="mixedbread-ai/mxbai-embed-large-v1"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-xs",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.09,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-xs"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-s",
-        dim=384,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.13,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-s"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-m",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.43,
-        sources=ModelSource(hf="Snowflake/snowflake-arctic-embed-m"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-m-long",
-        dim=768,
-        description=(
-            "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=0.54,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-m-long"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="snowflake/snowflake-arctic-embed-l",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: necessary, 2024 year."
-        ),
-        license="apache-2.0",
-        size_in_GB=1.02,
-        sources=ModelSource(hf="snowflake/snowflake-arctic-embed-l"),
-        model_file="onnx/model.onnx",
-    ),
-    DenseModelDescription(
-        model="jinaai/jina-clip-v1",
-        dim=768,
-        description=(
-            "Text embeddings, Multimodal (text&image), English, Prefixes for queries/documents: "
-            "not necessary, 2024 year"
-        ),
-        license="apache-2.0",
-        size_in_GB=0.55,
-        sources=ModelSource(hf="jinaai/jina-clip-v1"),
-        model_file="onnx/text_model.onnx",
-    ),
+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",
+    },
 ]
 
 
-class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
+class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
     """Implementation of the Flag Embedding model."""
 
     @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         Lists the supported models.
 
         Returns:
-            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
@@ -215,8 +204,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         """
         Args:
@@ -234,7 +222,6 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             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.
@@ -248,19 +235,17 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = str(define_cache_dir(cache_dir))
+        self.cache_dir = define_cache_dir(cache_dir)
         self._model_dir = self.download_model(
-            self.model_description,
-            self.cache_dir,
-            local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         if not self.lazy_load:
@@ -271,8 +256,8 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -301,18 +286,18 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
         return OnnxTextEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, NumpyArray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         embeddings = output.model_output
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
             processed_embeddings = embeddings[:, 0]
@@ -325,7 +310,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
-            model_file=self.model_description.model_file,
+            model_file=self.model_description["model_file"],
             threads=self.threads,
             providers=self.providers,
             cuda=self.cuda,
@@ -333,12 +318,12 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
 
-class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
     def init_embedding(
         self,
         model_name: str,
         cache_dir: str,
-        **kwargs: Any,
+        **kwargs,
     ) -> OnnxTextEmbedding:
         return OnnxTextEmbedding(
             model_name=model_name,

fastembed/text/onnx_text_model.py

@@ -4,10 +4,9 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from numpy.typing import NDArray
-from tokenizers import Encoding, Tokenizer
+from tokenizers import Encoding
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
@@ -18,7 +17,7 @@ class OnnxTextModel(OnnxModel[T]):
     ONNX_OUTPUT_NAMES: Optional[list[str]] = None
 
     @classmethod
-    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
@@ -26,12 +25,12 @@ class OnnxTextModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.tokenizer: Optional[Tokenizer] = None
-        self.special_token_to_id: dict[str, int] = {}
+        self.tokenizer = None
+        self.special_token_to_id = {}
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, Union[NumpyArray, NDArray[np.int64]]]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -59,19 +58,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: Any) -> list[Encoding]:
-        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+    def tokenize(self, documents: list[str], **kwargs) -> list[Encoding]:
+        return self.tokenizer.encode_batch(documents)
 
     def onnx_embed(
         self,
         documents: list[str],
-        **kwargs: Any,
+        **kwargs,
     ) -> 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()}  # type: ignore[union-attr]
-        onnx_input: dict[str, NumpyArray] = {
+        input_names = {node.name for node in self.model.get_inputs()}
+        onnx_input = {
             "input_ids": np.array(input_ids, dtype=np.int64),
         }
         if "attention_mask" in input_names:
@@ -82,7 +81,7 @@ class OnnxTextModel(OnnxModel[T]):
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -99,7 +98,7 @@ class OnnxTextModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
-        **kwargs: Any,
+        **kwargs,
     ) -> Iterable[T]:
         is_small = False
 
@@ -136,11 +135,11 @@ class OnnxTextModel(OnnxModel[T]):
                 start_method=start_method,
             )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                yield from self._post_process_onnx_output(batch)  # type: ignore
+                yield from self._post_process_onnx_output(batch)
 
 
-class TextEmbeddingWorker(EmbeddingWorker[T]):
-    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
+class TextEmbeddingWorker(EmbeddingWorker):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             onnx_output = self.model.onnx_embed(batch)
             yield idx, onnx_output

fastembed/text/pooled_embedding.py

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

fastembed/text/pooled_normalized_embedding.py

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

fastembed/text/text_embedding.py

@@ -1,40 +1,53 @@
-import warnings
 from typing import Any, Iterable, Optional, Sequence, Type, Union
-from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+import numpy as np
+
+from fastembed.common import OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
+from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.pooled_embedding import PooledEmbedding
-from fastembed.text.multitask_embedding import JinaEmbeddingV3
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
-from fastembed.common.model_description import DenseModelDescription
 
 
 class TextEmbedding(TextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [
         OnnxTextEmbedding,
+        E5OnnxEmbedding,
         CLIPOnnxEmbedding,
         PooledNormalizedEmbedding,
         PooledEmbedding,
-        JinaEmbeddingV3,
     ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
-        """Lists the supported models.
+        """
+        Lists the supported models.
 
         Returns:
             list[dict[str, Any]]: A list of dictionaries containing the model information.
-        """
-        return [asdict(model) for model in cls._list_supported_models()]
 
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        result: list[DenseModelDescription] = []
+            Example:
+                ```
+                [
+                    {
+                        "model": "intfloat/multilingual-e5-large",
+                        "dim": 1024,
+                        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+                        "license": "mit",
+                        "size_in_GB": 2.24,
+                        "sources": {
+                            "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+                            "hf": "qdrant/multilingual-e5-large-onnx",
+                        }
+                    }
+                ]
+                ```
+        """
+        result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding._list_supported_models())
+            result.extend(embedding.list_supported_models())
         return result
 
     def __init__(
@@ -46,37 +59,12 @@ class TextEmbedding(TextEmbeddingBase):
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
-        **kwargs: Any,
+        **kwargs,
     ):
         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 and release notes to ensure compatibility with your workflow. ",
-                UserWarning,
-                stacklevel=2,
-            )
-        if model_name == "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":
-            warnings.warn(
-                "The model 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2' has been updated to "
-                "include a mean pooling layer. Please ensure your usage aligns with the new functionality. "
-                "Support for the previous version without mean pooling will be removed as of version 0.5.2.",
-                UserWarning,
-                stacklevel=2,
-            )
-        if model_name in {
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
-            "intfloat/multilingual-e5-large",
-        }:
-            warnings.warn(
-                f"{model_name} has been updated as of fastembed 0.5.2, outputs are now average pooled.",
-                UserWarning,
-                stacklevel=2,
-            )
-
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
-            supported_models = EMBEDDING_MODEL_TYPE._list_supported_models()
-            if any(model_name.lower() == model.model.lower() for model in supported_models):
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name=model_name,
                     cache_dir=cache_dir,
@@ -99,8 +87,8 @@ class TextEmbedding(TextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -117,30 +105,3 @@ 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, Any
+from typing import Iterable, Optional, Union
+
+import numpy as np
 
-from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
+class TextEmbeddingBase(ModelManagement):
     def __init__(
         self,
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -23,11 +23,11 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Iterable[NumpyArray]:
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -36,13 +36,15 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[NumpyArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[np.ndarray]:
         """
         Embeds queries
 
@@ -50,11 +52,11 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[NumpyArray]: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        else:
+        if isinstance(query, Iterable):
             yield from self.embed(query, **kwargs)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.5.1"
+version = "0.5.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,52 +12,44 @@ 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 = ">=1.21,<2.1.0", python = "<3.10" },
+    { version = ">=2.1.0", python = ">=3.13" }
 ]
 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,<12.0.0"
+pillow = "^10.3.0"
 mmh3 = "^4.1.0"
 py-rust-stemmers = "^0.1.0"
 
-[tool.poetry.group.test.dependencies]
+[tool.poetry.group.dev.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.3.0,<12.0.0"
+pillow = "^10.2.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: str) -> None:
+def test_attention_embeddings(model_name):
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name)
 
@@ -71,7 +71,7 @@ def test_attention_embeddings(model_name: str) -> None:
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
-def test_parallel_processing(model_name: str) -> None:
+def test_parallel_processing(model_name):
     is_ci = os.getenv("CI")
 
     model = SparseTextEmbedding(model_name=model_name)
@@ -96,7 +96,7 @@ def test_parallel_processing(model_name: str) -> None:
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
-def test_multilanguage(model_name: str) -> None:
+def test_multilanguage(model_name):
     is_ci = os.getenv("CI")
 
     docs = ["Mangez-vous vraiment des grenouilles?", "Je suis au lit"]
@@ -122,7 +122,7 @@ def test_multilanguage(model_name: str) -> None:
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
-def test_special_characters(model_name: str) -> None:
+def test_special_characters(model_name):
     is_ci = os.getenv("CI")
 
     docs = [
@@ -145,7 +145,7 @@ def test_special_characters(model_name: str) -> None:
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])
-def test_lazy_load(model_name: str) -> None:
+def test_lazy_load(model_name):
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
     docs = ["hello world", "flag embedding"]

tests/test_common.py

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

tests/test_image_onnx_embeddings.py

@@ -27,16 +27,16 @@ CANONICAL_VECTOR_VALUES = {
 }
 
 
-def test_embedding() -> None:
+def test_embedding():
     is_ci = os.getenv("CI")
 
-    for model_desc in ImageEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
+    for model_desc in ImageEmbedding.list_supported_models():
+        if not is_ci and model_desc["size_in_GB"] > 1:
             continue
 
-        dim = model_desc.dim
+        dim = model_desc["dim"]
 
-        model = ImageEmbedding(model_name=model_desc.model)
+        model = ImageEmbedding(model_name=model_desc["model"])
 
         images = [
             TEST_MISC_DIR / "image.jpeg",
@@ -48,20 +48,20 @@ def test_embedding() -> None:
         embeddings = np.stack(embeddings, axis=0)
         assert embeddings.shape == (len(images), dim)
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
 
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc.model
+        ), model_desc["model"]
 
-        assert np.allclose(embeddings[1], embeddings[2]), model_desc.model
+        assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
 
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
-def test_batch_embedding(n_dims: int, model_name: str) -> None:
+def test_batch_embedding(n_dims, model_name):
     is_ci = os.getenv("CI")
     model = ImageEmbedding(model_name=model_name)
     n_images = 32
@@ -81,7 +81,7 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
-def test_parallel_processing(n_dims: int, model_name: str) -> None:
+def test_parallel_processing(n_dims, model_name):
     is_ci = os.getenv("CI")
     model = ImageEmbedding(model_name=model_name)
 
@@ -109,7 +109,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
 
 
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
-def test_lazy_load(model_name: str) -> None:
+def test_lazy_load(model_name):
     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

@@ -226,7 +226,7 @@ def test_parallel_processing():
     "model_name",
     ["colbert-ir/colbertv2.0"],
 )
-def test_lazy_load(model_name: str):
+def test_lazy_load(model_name):
     is_ci = os.getenv("CI")
 
     model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)

tests/test_late_interaction_multimodal.py

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

tests/test_multi_gpu.py

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

tests/test_sparse_embeddings.py

@@ -49,7 +49,7 @@ CANONICAL_COLUMN_VALUES = {
 docs = ["Hello World"]
 
 
-def test_batch_embedding() -> None:
+def test_batch_embedding():
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
@@ -64,7 +64,7 @@ def test_batch_embedding() -> None:
             delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding() -> None:
+def test_single_embedding():
     is_ci = os.getenv("CI")
     for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         model = SparseTextEmbedding(model_name=model_name)
@@ -80,7 +80,7 @@ def test_single_embedding() -> None:
             delete_model_cache(model.model._model_dir)
 
 
-def test_parallel_processing() -> None:
+def test_parallel_processing():
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
     docs = ["hello world", "flag embedding"] * 30
@@ -111,7 +111,7 @@ def test_parallel_processing() -> None:
 
 
 @pytest.fixture
-def bm25_instance() -> None:
+def bm25_instance():
     ci = os.getenv("CI", True)
     model = Bm25("Qdrant/bm25", language="english")
     yield model
@@ -119,7 +119,7 @@ def bm25_instance() -> None:
         delete_model_cache(model._model_dir)
 
 
-def test_stem_with_stopwords_and_punctuation(bm25_instance: Bm25) -> None:
+def test_stem_with_stopwords_and_punctuation(bm25_instance):
     # Setup
     bm25_instance.stopwords = {"the", "is", "a"}
     bm25_instance.punctuation = {".", ",", "!"}
@@ -135,7 +135,7 @@ def test_stem_with_stopwords_and_punctuation(bm25_instance: Bm25) -> None:
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
-def test_stem_case_insensitive_stopwords(bm25_instance: Bm25) -> None:
+def test_stem_case_insensitive_stopwords(bm25_instance):
     # Setup
     bm25_instance.stopwords = {"the", "is", "a"}
     bm25_instance.punctuation = {".", ",", "!"}
@@ -152,7 +152,7 @@ def test_stem_case_insensitive_stopwords(bm25_instance: Bm25) -> None:
 
 
 @pytest.mark.parametrize("disable_stemmer", [True, False])
-def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
+def test_disable_stemmer_behavior(disable_stemmer):
     # Setup
     model = Bm25("Qdrant/bm25", language="english", disable_stemmer=disable_stemmer)
     model.stopwords = {"the", "is", "a"}
@@ -176,7 +176,7 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
     "model_name",
     ["prithivida/Splade_PP_en_v1"],
 )
-def test_lazy_load(model_name: str) -> None:
+def test_lazy_load(model_name):
     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

@@ -26,7 +26,7 @@ SELECTED_MODELS = {
     "model_name",
     [model_name for model_name in CANONICAL_SCORE_VALUES],
 )
-def test_rerank(model_name: str) -> None:
+def test_rerank(model_name):
     is_ci = os.getenv("CI")
 
     model = TextCrossEncoder(model_name=model_name)
@@ -53,7 +53,7 @@ def test_rerank(model_name: str) -> None:
     "model_name",
     [model_name for model_name in SELECTED_MODELS.values()],
 )
-def test_batch_rerank(model_name: str) -> None:
+def test_batch_rerank(model_name):
     is_ci = os.getenv("CI")
 
     model = TextCrossEncoder(model_name=model_name)
@@ -82,7 +82,7 @@ def test_batch_rerank(model_name: str) -> None:
     "model_name",
     ["Xenova/ms-marco-MiniLM-L-6-v2"],
 )
-def test_lazy_load(model_name: str) -> None:
+def test_lazy_load(model_name):
     is_ci = os.getenv("CI")
     model = TextCrossEncoder(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
@@ -99,7 +99,7 @@ def test_lazy_load(model_name: str) -> None:
     "model_name",
     [model_name for model_name in SELECTED_MODELS.values()],
 )
-def test_rerank_pairs_parallel(model_name: str) -> None:
+def test_rerank_pairs_parallel(model_name):
     is_ci = os.getenv("CI")
 
     model = TextCrossEncoder(model_name=model_name)

tests/test_text_multitask_embeddings.py

@@ -1,253 +0,0 @@
-import os
-
-import numpy as np
-import pytest
-
-from fastembed import TextEmbedding
-from fastembed.text.multitask_embedding import 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."]
-
-
-def test_batch_embedding():
-    is_ci = os.getenv("CI")
-    docs_to_embed = docs * 10
-    default_task = Task.RETRIEVAL_PASSAGE
-
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
-        model_name = model_desc.model
-        dim = model_desc.dim
-
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
-        model = TextEmbedding(model_name=model_name)
-
-        print(f"evaluating {model_name} default task")
-
-        embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
-        embeddings = np.stack(embeddings, axis=0)
-
-        assert embeddings.shape == (len(docs_to_embed), dim)
-
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_name][default_task]["vectors"]
-        assert np.allclose(
-            embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
-        ), model_desc.model
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-
-def test_single_embedding():
-    is_ci = os.getenv("CI")
-
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
-        model_name = model_desc.model
-        dim = model_desc.dim
-
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
-        model = TextEmbedding(model_name=model_name)
-
-        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")
-    task_id = Task.RETRIEVAL_QUERY
-
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
-        model_name = model_desc.model
-        dim = model_desc.dim
-
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
-        model = TextEmbedding(model_name=model_name)
-
-        print(f"evaluating {model_name} query_embed task_id: {task_id}")
-
-        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")
-    task_id = Task.RETRIEVAL_PASSAGE
-
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
-        model_name = model_desc.model
-        dim = model_desc.dim
-
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
-        model = TextEmbedding(model_name=model_name)
-
-        print(f"evaluating {model_name} passage_embed task_id: {task_id}")
-
-        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)
-
-
-def test_parallel_processing():
-    is_ci = os.getenv("CI")
-
-    docs = ["Hello World", "Follow the white rabbit."] * 10
-
-    model_name = "jinaai/jina-embeddings-v3"
-    dim = 1024
-
-    model = TextEmbedding(model_name=model_name)
-
-    task_id = Task.SEPARATION
-    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")
-
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
-        model_name = model_desc.model
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
-
-        model = TextEmbedding(model_name=model_name)
-
-        for i, task_id in enumerate(Task):
-            _ = list(model.embed(documents=docs, batch_size=1, task_id=i))
-            assert model.model.current_task_id == task_id
-
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
-
-
-@pytest.mark.parametrize(
-    "model_name",
-    ["jinaai/jina-embeddings-v3"],
-)
-def test_lazy_load(model_name: str):
-    is_ci = os.getenv("CI")
-    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,5 +1,4 @@
 import os
-import platform
 
 import numpy as np
 import pytest
@@ -32,11 +31,11 @@ CANONICAL_VECTOR_VALUES = {
         [-0.034478, 0.03102, 0.00673, 0.02611, -0.039362]
     ),
     "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array(
-        [0.0361, 0.1862, 0.2776, 0.2461, -0.1904]
+        [0.0094, 0.0184, 0.0328, 0.0072, -0.0351]
     ),
-    "intfloat/multilingual-e5-large": np.array([0.4544, -0.0968, 0.1054, -1.3753, 0.1500]),
+    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
     "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
-        [0.0047, 0.1334, -0.0102, 0.0714, 0.1930]
+        [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
     ),
     "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]),
@@ -49,7 +48,7 @@ CANONICAL_VECTOR_VALUES = {
         [-0.15407836, -0.03053198, -3.9138033, 0.1910364, 0.13224715]
     ),
     "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
-        [0.0802303, 0.3700881, -4.3053818, 0.4431803, -0.271572]
+        [-0.12525563, 0.38030425, -3.961622, 0.04176439, -0.0758301]
     ),
     "thenlper/gte-large": np.array(
         [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
@@ -69,30 +68,23 @@ 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() -> None:
+def test_embedding():
     is_ci = os.getenv("CI")
-    is_mac = platform.system() == "Darwin"
 
-    for model_desc in TextEmbedding._list_supported_models():
-        if (
-            (not is_ci and model_desc.size_in_GB > 1)
-            or model_desc.model in MULTI_TASK_MODELS
-            or (is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q")
-        ):
+    for model_desc in TextEmbedding.list_supported_models():
+        if not is_ci and model_desc["size_in_GB"] > 1:
             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"]
@@ -104,7 +96,7 @@ def test_embedding() -> None:
     "n_dims,model_name",
     [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
 )
-def test_batch_embedding(n_dims: int, model_name: str) -> None:
+def test_batch_embedding(n_dims, model_name):
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
 
@@ -121,7 +113,7 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
     "n_dims,model_name",
     [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
 )
-def test_parallel_processing(n_dims: int, model_name: str) -> None:
+def test_parallel_processing(n_dims, model_name):
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
 
@@ -147,7 +139,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
     "model_name",
     ["BAAI/bge-small-en-v1.5"],
 )
-def test_lazy_load(model_name: str) -> None:
+def test_lazy_load(model_name):
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")

tests/type_stub.py

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

tests/utils.py

@@ -2,8 +2,7 @@ import shutil
 import traceback
 
 from pathlib import Path
-from types import TracebackType
-from typing import Union, Callable, Any, Type
+from typing import Union
 
 
 def delete_model_cache(model_dir: Union[str, Path]) -> None:
@@ -17,11 +16,7 @@ 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: Callable[..., Any],
-        path: str,
-        exc_info: tuple[Type[BaseException], BaseException, TracebackType],
-    ) -> None:
+    def on_error(func, path, exc_info):
         print("Failed to remove: ", path)
         print("Exception: ", exc_info)
         traceback.print_exception(*exc_info)