new: gpu package (#224)

* fix: stage GCS downloads instead of deleting the caller's cache dir

* fix: verify archive integrity and give each download its own staging dir

.github/dependabot.yml

@@ -0,0 +1,40 @@
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 5
+    groups:
+      security-updates:
+        applies-to: security-updates
+        patterns:
+          - "*"
+      version-updates:
+        applies-to: version-updates
+        update-types:
+          - "minor"
+          - "patch"
+        patterns:
+          - "*"
+    cooldown:
+      default-days: 7
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 5
+    groups:
+      security-updates:
+        applies-to: security-updates
+        patterns:
+          - "*"
+      version-updates:
+        applies-to: version-updates
+        update-types:
+          - "minor"
+          - "patch"
+        patterns:
+          - "*"
+    cooldown:
+      default-days: 7

.github/workflows/ci.yml

@@ -10,12 +10,12 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         with:
           python-version: 3.x
       - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
-      - uses: actions/cache@v3
+      - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
           key: mkdocs-material-${{ env.cache_id }}
           path: .cache

.github/workflows/python-publish.yml

@@ -21,11 +21,11 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
       with:
-        python-version: '3.9.x'
+        python-version: '3.10.x'
     - name: Install dependencies
       run: |
         python -m pip install poetry
@@ -33,7 +33,7 @@ jobs:
     - name: Build package
       run: poetry build
     - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2
       with:
         user: __token__
         password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/python-tests.yml

@@ -16,7 +16,6 @@ jobs:
     strategy:
       matrix:
         python-version:
-          - '3.9.x'
           - '3.10.x'
           - '3.11.x'
           - '3.12.x'
@@ -29,9 +28,9 @@ jobs:
     name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies

.github/workflows/type-checkers.yml

@@ -8,16 +8,16 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
         os: [ubuntu-latest]
 
     name: Python ${{ matrix.python-version }} test
 
     steps:
-      - uses: actions/checkout@v1
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         with:
           python-version: ${{ matrix.python-version }}
 

LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [yyyy] [name of copyright owner]
+   Copyright 2026 Qdrant Solutions GmbH
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

NOTICE

@@ -15,6 +15,8 @@ These models are developed by Jina (https://jina.ai/) and are subject to Jina AI
 This distribution includes the following Google models, each with its respective license:
 - vidore/colpali-v1.3
     - License: gemma
+- google/embeddinggemma-300m
+    - License: gemma
 
 Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms
 

experiments/if_splade_to_onnx.py

@@ -0,0 +1,134 @@
+"""Export an inference-free SPLADE document encoder to ONNX.
+
+Converts `opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte` (an MLM head
+over a GTE backbone) into an onnx model producing token logits, and assembles a model dir
+with everything fastembed's `IfSplade` needs: model.onnx, tokenizer files and idf.json.
+
+Usage:
+    python experiments/if_splade_to_onnx.py --output-dir models/opensearch-neural-sparse-encoding-doc-v3-gte
+"""
+
+import argparse
+import shutil
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+MODEL_ID = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte"
+# revision of the remote modeling code (Alibaba-NLP/new-impl), pinned in the model card
+CODE_REVISION = "40ced75c3017eb27626c9d4ea981bde21a2662f4"
+
+TOKENIZER_FILES = [
+    "config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "vocab.txt",
+    "idf.json",
+]
+
+
+class LogitsOnly(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits
+
+
+def export(model_id: str, output_dir: Path, opset: int = 14) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    model = AutoModelForMaskedLM.from_pretrained(
+        model_id, trust_remote_code=True, code_revision=CODE_REVISION
+    )
+    model.eval()
+    wrapped = LogitsOnly(model)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    dummy = tokenizer(
+        ["fastembed is a library", "onnx export"],
+        padding=True,
+        truncation=True,
+        return_tensors="pt",
+        return_token_type_ids=False,
+    )
+
+    onnx_path = output_dir / "model.onnx"
+    with torch.inference_mode():
+        torch.onnx.export(
+            wrapped,
+            (dummy["input_ids"], dummy["attention_mask"]),
+            f=onnx_path.as_posix(),
+            input_names=["input_ids", "attention_mask"],
+            output_names=["logits"],
+            dynamic_axes={
+                "input_ids": {0: "batch_size", 1: "sequence_length"},
+                "attention_mask": {0: "batch_size", 1: "sequence_length"},
+                "logits": {0: "batch_size", 1: "sequence_length"},
+            },
+            do_constant_folding=True,
+            opset_version=opset,
+            dynamo=False,
+        )
+
+    for file_name in TOKENIZER_FILES:
+        local_path = hf_hub_download(repo_id=model_id, filename=file_name)
+        shutil.copy(local_path, output_dir / file_name)
+
+    return onnx_path
+
+
+def parity_check(model_id: str, output_dir: Path) -> None:
+    import numpy as np
+    import onnxruntime as ort
+
+    model = AutoModelForMaskedLM.from_pretrained(
+        model_id, trust_remote_code=True, code_revision=CODE_REVISION
+    )
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    documents = [
+        "Currently New York is rainy.",
+        "fastembed is a lightweight library for generating embeddings",
+        "hello world",
+    ]
+    features = tokenizer(
+        documents, padding=True, truncation=True, return_tensors="pt", return_token_type_ids=False
+    )
+
+    with torch.inference_mode():
+        torch_logits = model(**features).logits.numpy()
+
+    session = ort.InferenceSession(output_dir / "model.onnx")
+    onnx_logits = session.run(
+        ["logits"],
+        {
+            "input_ids": features["input_ids"].numpy(),
+            "attention_mask": features["attention_mask"].numpy(),
+        },
+    )[0]
+
+    max_diff = np.abs(torch_logits - onnx_logits).max()
+    print(f"max |torch - onnx| logits diff: {max_diff}")
+    assert max_diff < 1e-3, "onnx export does not match the torch model"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-id", default=MODEL_ID)
+    parser.add_argument("--output-dir", default=f"models/{MODEL_ID.replace('/', '_')}", type=Path)
+    parser.add_argument("--opset", default=14, type=int)
+    args = parser.parse_args()
+
+    onnx_path = export(args.model_id, args.output_dir, args.opset)
+    print(f"Exported to {onnx_path}")
+    parity_check(args.model_id, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

fastembed/common/model_description.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Optional, Any
+from typing import Any
 
 
 @dataclass(frozen=True)
 class ModelSource:
-    hf: Optional[str] = None
-    url: Optional[str] = None
+    hf: str | None = None
+    url: str | None = None
     _deprecated_tar_struct: bool = False
 
     @property
@@ -33,8 +33,8 @@ class BaseModelDescription:
 
 @dataclass(frozen=True)
 class DenseModelDescription(BaseModelDescription):
-    dim: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = field(default_factory=dict)
+    dim: int | None = None
+    tasks: dict[str, Any] | None = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         assert self.dim is not None, "dim is required for dense model description"
@@ -42,11 +42,12 @@ class DenseModelDescription(BaseModelDescription):
 
 @dataclass(frozen=True)
 class SparseModelDescription(BaseModelDescription):
-    requires_idf: Optional[bool] = None
-    vocab_size: Optional[int] = None
+    requires_idf: bool | None = None
+    vocab_size: int | None = None
 
 
 class PoolingType(str, Enum):
     CLS = "CLS"
     MEAN = "MEAN"
+    LAST_TOKEN = "LAST_TOKEN"
     DISABLED = "DISABLED"

fastembed/common/model_management.py

@@ -1,11 +1,14 @@
 import os
 import time
+import gzip
 import json
 import shutil
 import tarfile
+import tempfile
+import contextlib
 from copy import deepcopy
-from pathlib import Path
-from typing import Any, Optional, Union, TypeVar, Generic
+from pathlib import Path, PureWindowsPath
+from typing import Any, TypeVar, Generic
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -21,6 +24,8 @@ from fastembed.common.model_description import BaseModelDescription
 
 T = TypeVar("T", bound=BaseModelDescription)
 
+_DOWNLOAD_CHUNK_SIZE = 256 * 1024
+
 
 class ModelManagement(Generic[T]):
     METADATA_FILE = "files_metadata.json"
@@ -97,9 +102,7 @@ class ModelManagement(Generic[T]):
             str: The path to the downloaded file.
         """
 
-        if os.path.exists(output_path):
-            return output_path
-        response = requests.get(url, stream=True)
+        response = requests.get(url, stream=True, timeout=(10, 120))
 
         # Handle HTTP errors
         if response.status_code == 403:
@@ -107,6 +110,8 @@ class ModelManagement(Generic[T]):
                 "Authentication Error: You do not have permission to access this resource. "
                 "Please check your credentials."
             )
+        # Otherwise an error page gets written out as though it were the archive.
+        response.raise_for_status()
 
         # Get the total size of the file
         total_size_in_bytes = int(response.headers.get("content-length", 0))
@@ -124,7 +129,7 @@ class ModelManagement(Generic[T]):
             disable=not show_progress,
         ) as progress_bar:
             with open(output_path, "wb") as file:
-                for chunk in response.iter_content(chunk_size=1024):
+                for chunk in response.iter_content(chunk_size=_DOWNLOAD_CHUNK_SIZE):
                     if chunk:  # Filter out keep-alive new chunks
                         progress_bar.update(len(chunk))
                         file.write(chunk)
@@ -180,8 +185,8 @@ class ModelManagement(Generic[T]):
 
         def _collect_file_metadata(
             model_dir: Path, repo_files: list[RepoFile]
-        ) -> dict[str, dict[str, Union[int, str]]]:
-            meta: dict[str, dict[str, Union[int, str]]] = {}
+        ) -> dict[str, dict[str, int | str]]:
+            meta: dict[str, dict[str, 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:
@@ -193,9 +198,7 @@ class ModelManagement(Generic[T]):
                         }
             return meta
 
-        def _save_file_metadata(
-            model_dir: Path, meta: dict[str, dict[str, Union[int, str]]]
-        ) -> None:
+        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) -> None:
             try:
                 if not model_dir.exists():
                     model_dir.mkdir(parents=True, exist_ok=True)
@@ -288,12 +291,18 @@ class ModelManagement(Generic[T]):
         """
         Decompresses a .tar.gz file to a cache directory.
 
+        Nothing is deleted on failure, since `cache_dir` may hold more than this archive.
+        Cleaning up a partial extraction is the caller's job.
+
         Args:
             targz_path (str): Path to the .tar.gz file.
             cache_dir (str): Path to the cache directory.
 
         Returns:
             cache_dir (str): Path to the cache directory.
+
+        Raises:
+            ValueError: If the archive is missing, corrupt, or holds an unsafe member.
         """
         # Check if targz_path exists and is a file
         if not os.path.isfile(targz_path):
@@ -306,20 +315,50 @@ class ModelManagement(Generic[T]):
         try:
             # Open the tar.gz file
             with tarfile.open(targz_path, "r:gz") as tar:
-                # Extract all files into the cache directory
-                tar.extractall(
-                    path=cache_dir,
-                )
-        except tarfile.TarError as e:
-            # If any error occurs while opening or extracting the tar.gz file,
-            # delete the cache directory (if it was created in this function)
-            # and raise the error again
-            if "tmp" in cache_dir:
-                shutil.rmtree(cache_dir)
-            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
+                if hasattr(tarfile, "data_filter"):
+                    tar.extractall(path=cache_dir, filter="data")
+                else:
+                    # No PEP 706 filter before 3.10.12, so vet the members by hand.
+                    members = tar.getmembers()
+                    for member in members:
+                        cls._validate_tar_member(member)
+                    tar.extractall(path=cache_dir, members=members)
+                # tarfile stops at the end-of-archive marker, short of the gzip trailer, so
+                # the CRC is only checked if the rest of the stream is read.
+                while tar.fileobj.read(1 << 20):
+                    pass
+        except (tarfile.TarError, ValueError, EOFError, gzip.BadGzipFile) as e:
+            # gzip raises EOFError for a truncated stream and BadGzipFile for a corrupted one.
+            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}") from e
 
         return cache_dir
 
+    @staticmethod
+    def _is_unsafe_tar_path(path: str) -> bool:
+        """Checks whether a tar member name or link target may escape the extraction dir.
+
+        Lexical on purpose: resolving against the extraction directory is unsound before
+        extraction, since `link/../escape` only escapes once an earlier member has been
+        written as a symlink. Any `..` component is therefore rejected outright.
+        """
+        # PureWindowsPath splits on both separators, so `root` covers POSIX "/evil" as
+        # well as "\\evil", which escapes on Windows without being absolute.
+        windows_path = PureWindowsPath(path)
+        return bool(windows_path.drive or windows_path.root) or ".." in windows_path.parts
+
+    @classmethod
+    def _validate_tar_member(cls, member: tarfile.TarInfo) -> None:
+        """Raises ValueError if a member could write outside the extraction directory."""
+        if cls._is_unsafe_tar_path(member.name):
+            raise ValueError(f"Unsafe tar member path: {member.name}")
+
+        if member.issym() or member.islnk():
+            if cls._is_unsafe_tar_path(member.linkname):
+                raise ValueError(f"Unsafe tar link target: {member.name} -> {member.linkname}")
+        elif not (member.isfile() or member.isdir()):
+            # Devices, fifos and the like have no place in a model archive.
+            raise ValueError(f"Unsupported tar member type: {member.name}")
+
     @classmethod
     def retrieve_model_gcs(
         cls,
@@ -331,36 +370,13 @@ class ModelManagement(Generic[T]):
     ) -> Path:
         fast_model_name = f"{'fast-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
         cache_tmp_dir = Path(cache_dir) / "tmp"
-        model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
 
         # check if the model_dir and the model files are both present for macOS
         if model_dir.exists() and len(list(model_dir.glob("*"))) > 0:
             return model_dir
 
-        if model_tmp_dir.exists():
-            shutil.rmtree(model_tmp_dir)
-
-        cache_tmp_dir.mkdir(parents=True, exist_ok=True)
-
-        model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
-
-        if model_tar_gz.exists():
-            model_tar_gz.unlink()
-
-        if not local_files_only:
-            cls.download_file_from_gcs(
-                source_url,
-                output_path=str(model_tar_gz),
-            )
-
-            cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
-            assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
-
-            model_tar_gz.unlink()
-            # Rename from tmp to final name is atomic
-            model_tmp_dir.rename(model_dir)
-        else:
+        if local_files_only:
             logger.error(
                 f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
             )
@@ -368,6 +384,45 @@ class ModelManagement(Generic[T]):
                 f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
             )
 
+        if cache_tmp_dir.is_symlink():
+            raise ValueError(
+                f"{cache_tmp_dir} is a symlink, refusing to stage downloads through it"
+            )
+        cache_tmp_dir.mkdir(parents=True, exist_ok=True)
+
+        # The archive and everything extracted from it go in a directory of this attempt's own,
+        # so removing it undoes the attempt without touching any other download of the model.
+        staging_dir = Path(tempfile.mkdtemp(dir=cache_tmp_dir, prefix=f"{fast_model_name}-"))
+        try:
+            model_tar_gz = staging_dir / f"{fast_model_name}.tar.gz"
+            cls.download_file_from_gcs(
+                source_url,
+                output_path=str(model_tar_gz),
+            )
+
+            cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(staging_dir))
+
+            model_tmp_dir = staging_dir / fast_model_name
+            if not model_tmp_dir.is_dir() or model_tmp_dir.is_symlink():
+                raise ValueError(
+                    f"The archive from {source_url} has no {fast_model_name} directory"
+                )
+
+            # Replace a stale empty model_dir, which Windows will not rename onto. rmdir leaves
+            # anything else alone, including one another download has just filled.
+            with contextlib.suppress(OSError):
+                model_dir.rmdir()
+
+            try:
+                # Rename from the staging dir to the final name is atomic
+                model_tmp_dir.rename(model_dir)
+            except OSError:
+                # Another download of the same model finished first, so keep its copy.
+                if not (model_dir.is_dir() and any(model_dir.iterdir())):
+                    raise
+        finally:
+            shutil.rmtree(staging_dir, ignore_errors=True)
+
         return model_dir
 
     @classmethod
@@ -397,7 +452,11 @@ 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)
+        hf_offline = os.environ.get("HF_HUB_OFFLINE", "").strip().upper()
+        if not local_files_only and hf_offline in {"1", "TRUE", "YES", "ON"}:
+            local_files_only = True
+            kwargs["local_files_only"] = True
+        specific_model_path: str | None = kwargs.pop("specific_model_path", None)
         if specific_model_path:
             return Path(specific_model_path)
         retries = 1 if local_files_only else retries
@@ -411,7 +470,7 @@ class ModelManagement(Generic[T]):
             try:
                 cache_kwargs = deepcopy(kwargs)
                 cache_kwargs["local_files_only"] = True
-                return Path(
+                resolved_path = Path(
                     cls.download_files_from_huggingface(
                         hf_source,
                         cache_dir=cache_dir,
@@ -419,6 +478,10 @@ class ModelManagement(Generic[T]):
                         **cache_kwargs,
                     )
                 )
+                if (resolved_path / model.model_file).exists() and all(
+                    (resolved_path / file).exists() for file in extra_patterns
+                ):
+                    return resolved_path
             except Exception:
                 pass
             finally:

fastembed/common/onnx_model.py

@@ -1,7 +1,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
+from typing import Any, Generic, Iterable, Sequence, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort
@@ -9,7 +9,7 @@ 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, NumpyArray, Device
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -19,8 +19,9 @@ T = TypeVar("T")
 @dataclass
 class OnnxOutputContext:
     model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
+    attention_mask: NDArray[np.int64] | None = None
+    input_ids: NDArray[np.int64] | None = None
+    metadata: dict[str, Any] | None = None
 
 
 class OnnxModel(Generic[T]):
@@ -30,6 +31,18 @@ class OnnxModel(Generic[T]):
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
 
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        """Additional kwargs a worker process needs to reconstruct this model.
+
+        Workers are started with `spawn`/`forkserver`, hence they don't inherit class-level state
+        which has been set up in runtime, e.g. models registered via `add_custom_model`.
+        Such state has to be shipped to the workers explicitly.
+
+        Returns:
+            dict[str, Any]: kwargs to pass to `_get_worker_class().init_embedding`.
+        """
+        return {}
+
     def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
         """Post-process the ONNX model output to convert it into a usable format.
 
@@ -43,8 +56,8 @@ class OnnxModel(Generic[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: ort.InferenceSession | None = None
+        self.tokenizer: Tokenizer | None = None
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -58,25 +71,30 @@ class OnnxModel(Generic[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        available_providers = ort.get_available_providers()
+        cuda_available = "CUDAExecutionProvider" in available_providers
+        explicit_cuda = cuda is True or cuda == Device.CUDA
 
-        if cuda and providers is not None:
+        if explicit_cuda and providers is not None:
             warnings.warn(
-                f"`cuda` and `providers` are mutually exclusive parameters, cuda: {cuda}, providers: {providers}",
+                f"`cuda` and `providers` are mutually exclusive parameters, "
+                f"cuda: {cuda}, providers: {providers}. If you'd like to use providers, cuda should be one of "
+                f"[False, Device.CPU, Device.AUTO].",
                 category=UserWarning,
                 stacklevel=6,
             )
 
         if providers is not None:
             onnx_providers = list(providers)
-        elif cuda:
+        elif explicit_cuda or (cuda == Device.AUTO and cuda_available):
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
@@ -84,7 +102,6 @@ class OnnxModel(Generic[T]):
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
-        available_providers = ort.get_available_providers()
         requested_provider_names: list[str] = []
         for provider in onnx_providers:
             # check providers available

fastembed/common/preprocessor_utils.py

@@ -1,5 +1,6 @@
 import json
-from typing import Any
+import sys
+from typing import Any, Iterator
 from pathlib import Path
 
 from tokenizers import AddedToken, Tokenizer
@@ -8,9 +9,10 @@ from fastembed.image.transform.operators import Compose
 
 
 def load_special_tokens(model_dir: Path) -> dict[str, Any]:
+    """Read special_tokens_map.json, treating an absent file as an empty map."""
     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}")
+        return {}
 
     with open(str(tokens_map_path)) as tokens_map_file:
         tokens_map = json.load(tokens_map_file)
@@ -18,11 +20,58 @@ def load_special_tokens(model_dir: Path) -> dict[str, Any]:
     return tokens_map
 
 
-def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
-    config_path = model_dir / "config.json"
-    if not config_path.exists():
-        raise ValueError(f"Could not find config.json in {model_dir}")
+def iter_special_tokens(tokens_map: dict[str, Any]) -> Iterator[str | dict[str, Any]]:
+    """Yield the individual tokens declared in a special tokens map.
 
+    Most keys hold one token, but `additional_special_tokens` holds a list of them.
+    """
+    for value in tokens_map.values():
+        if isinstance(value, list):
+            yield from value
+        else:
+            yield value
+
+
+def _valid_context(value: Any) -> int | None:
+    """Return `value` if it can be used as a truncation limit, `None` otherwise.
+
+    Config files do not always carry a real limit: transformers writes `model_max_length` as
+    1e30 when the value is unknown, and some repos ship a 0 or a null. `enable_truncation`
+    raises an `OverflowError` on the former and silently produces empty encodings on the
+    latter, so both are rejected here rather than passed through.
+    """
+    if isinstance(value, bool) or not isinstance(value, int):
+        return None
+    if not 0 < value <= sys.maxsize:
+        return None
+    return value
+
+
+def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> int:
+    """Pick the truncation limit, preferring the stricter of the two tokenizer config keys.
+
+    `config.json:max_position_embeddings` deliberately is not used as a fallback: it is the size
+    of the position table, not the usable context, and the two differ per architecture, e.g.
+    roberta reports 514 for a usable 512.
+    """
+    candidates = [
+        context
+        for context in (
+            _valid_context(tokenizer_config.get("model_max_length")),
+            _valid_context(tokenizer_config.get("max_length")),
+        )
+        if context is not None
+    ]
+    if not candidates:
+        raise ValueError(
+            f"Could not determine the maximum context length for {model_dir}. Set a positive "
+            "`model_max_length` or `max_length` in tokenizer_config.json."
+        )
+
+    return min(candidates)
+
+
+def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     tokenizer_path = model_dir / "tokenizer.json"
     if not tokenizer_path.exists():
         raise ValueError(f"Could not find tokenizer.json in {model_dir}")
@@ -31,43 +80,62 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     if not tokenizer_config_path.exists():
         raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
 
-    with open(str(config_path)) as config_file:
-        config = json.load(config_file)
+    # config.json is optional: transformers v5 no longer writes it for every model.
+    config_path = model_dir / "config.json"
+    config: dict[str, Any] = {}
+    if config_path.exists():
+        with open(str(config_path)) as config_file:
+            config = json.load(config_file)
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
-        assert "model_max_length" in tokenizer_config or "max_length" in tokenizer_config, (
-            "Models without model_max_length or max_length are not supported."
-        )
-        if "model_max_length" not in tokenizer_config:
-            max_context = tokenizer_config["max_length"]
-        elif "max_length" not in tokenizer_config:
-            max_context = tokenizer_config["model_max_length"]
-        else:
-            max_context = min(tokenizer_config["model_max_length"], tokenizer_config["max_length"])
+
+    max_context = _resolve_max_context(tokenizer_config, model_dir)
 
     tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=max_context)
-    tokenizer.enable_padding(
-        pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
-    )
 
-    for token in tokens_map.values():
+    # Registered before the padding is resolved: the map may name a pad token that
+    # tokenizer.json does not carry, and it only gets an id once it is added.
+    for token in iter_special_tokens(tokens_map):
         if isinstance(token, str):
             tokenizer.add_special_tokens([token])
         elif isinstance(token, dict):
             tokenizer.add_special_tokens([AddedToken(**token)])
 
-    special_token_to_id: dict[str, int] = {}
+    # Padding is always normalized to batch-longest. A serialized fixed length shorter than the
+    # truncation limit leaves longer encodings untouched, which produces ragged batches, and a
+    # fixed length equal to it pads every batch to the maximum. Direction and pad token metadata
+    # are taken from the serialized settings, since some models pad on the left.
+    padding = tokenizer.padding or {}
+    pad_token = padding.get("pad_token") or tokenizer_config.get("pad_token")
+    if pad_token is None:
+        raise ValueError(f"Could not find a pad token for {model_dir}")
 
-    for token in tokens_map.values():
-        if isinstance(token, str):
-            special_token_to_id[token] = tokenizer.token_to_id(token)
-        elif isinstance(token, dict):
-            token_str = token.get("content", "")
-            special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
+    # The vocabulary is the last resort, not a hardcoded 0: that silently disagrees with
+    # `pad_token` for every model whose pad token is not the first entry.
+    pad_id = padding.get("pad_id", config.get("pad_token_id"))
+    if pad_id is None:
+        pad_id = tokenizer.token_to_id(pad_token)
+    if pad_id is None:
+        raise ValueError(f"Could not resolve an id for the pad token {pad_token!r} in {model_dir}")
+
+    tokenizer.enable_padding(
+        direction=padding.get("direction", "right"),
+        pad_id=pad_id,
+        pad_type_id=padding.get("pad_type_id", 0),
+        pad_token=pad_token,
+        pad_to_multiple_of=padding.get("pad_to_multiple_of"),
+        length=None,
+    )
+
+    special_token_to_id = {
+        token.content: token_id
+        for token_id, token in tokenizer.get_added_tokens_decoder().items()
+        if token.special
+    }
 
     return tokenizer, special_token_to_id
 

fastembed/common/types.py

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

fastembed/common/utils.py

@@ -5,7 +5,7 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar
+from typing import Iterable, TypeVar
 
 import numpy as np
 from numpy.typing import NDArray
@@ -32,6 +32,16 @@ def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) ->
     return pooled_embeddings
 
 
+def last_token_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
+    """Take the embedding of the last non-padding token of each sequence.
+
+    Locates the last position the attention mask marks as real, so it holds whichever
+    side the tokenizer pads on.
+    """
+    last_token_indices = attention_mask.shape[1] - 1 - np.argmax(attention_mask[:, ::-1], axis=1)
+    return input_array[np.arange(input_array.shape[0]), last_token_indices]
+
+
 def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))
@@ -45,7 +55,7 @@ def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
         yield b
 
 
-def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
+def define_cache_dir(cache_dir: str | None = None) -> Path:
     """
     Define the cache directory for fastembed
     """

fastembed/embedding.py

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

fastembed/image/image_embedding.py

@@ -1,15 +1,21 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.image.normalized_embedding import NormalizedEmbedding
+from fastembed.image.siglip_embedding import SiglipOnnxImageEmbedding
 from fastembed.common.model_description import DenseModelDescription
 
 
 class ImageEmbedding(ImageEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
+    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [
+        OnnxImageEmbedding,
+        NormalizedEmbedding,
+        SiglipOnnxImageEmbedding,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -48,11 +54,11 @@ class ImageEmbedding(ImageEmbeddingBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -98,7 +104,7 @@ class ImageEmbedding(ImageEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -113,9 +119,9 @@ class ImageEmbedding(ImageEmbeddingBase):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

fastembed/image/image_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Any, Union
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -10,21 +10,21 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

fastembed/image/normalized_embedding.py

@@ -0,0 +1,69 @@
+from typing import Any, Iterable, Type
+
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_normalized_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="nomic-ai/nomic-embed-vision-v1.5",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.37,
+        sources=ModelSource(hf="nomic-ai/nomic-embed-vision-v1.5"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="nomic-ai/nomic-embed-vision-v1.5-Q",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.1,
+        sources=ModelSource(hf="nomic-ai/nomic-embed-vision-v1.5"),
+        model_file="onnx/model_quantized.onnx",
+    ),
+]
+
+
+class NormalizedEmbedding(OnnxImageEmbedding):
+    @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_normalized_models
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
+        return NormalizedEmbeddingWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        # The model emits last_hidden_state, which onnx_embed flattens to (batch, tokens * dim).
+        # Recover the token axis, take the CLS token (index 0) and normalize, matching the reference
+        # F.normalize(last_hidden_state[:, 0], p=2, dim=1).
+        dim = self.model_description.dim
+        assert dim is not None, "Model description is missing the embedding dim"
+        hidden_states = output.model_output.reshape(output.model_output.shape[0], -1, dim)
+        return normalize(hidden_states[:, 0])
+
+
+class NormalizedEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs: Any
+    ) -> NormalizedEmbedding:
+        return NormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/onnx_embedding.py

@@ -1,7 +1,7 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
@@ -63,14 +63,14 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -82,10 +82,11 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -105,7 +106,7 @@ 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
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -150,9 +151,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

fastembed/image/onnx_image_model.py

@@ -2,13 +2,13 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, 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.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
@@ -19,6 +19,8 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxImageModel(OnnxModel[T]):
+    ONNX_OUTPUT_NAMES: list[str] | None = None
+
     @classmethod
     def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
@@ -37,7 +39,7 @@ class OnnxImageModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.processor: Optional[Compose] = None
+        self.processor: Compose | None = None
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -51,11 +53,11 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -76,16 +78,18 @@ class OnnxImageModel(OnnxModel[T]):
         return {input_name: encoded}
 
     def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack():
+        with contextlib.ExitStack() as stack:
             image_files = [
-                Image.open(image) if not isinstance(image, Image.Image) else image
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
                 for image in images
             ]
             assert self.processor is not None, "Processor is not initialized"
             encoded = np.array(self.processor(image_files))
         onnx_input = 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(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -93,15 +97,15 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False

fastembed/image/siglip_embedding.py

@@ -0,0 +1,44 @@
+from typing import Any, Type
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.image.onnx_embedding import OnnxImageEmbedding, OnnxImageEmbeddingWorker
+
+supported_siglip_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="google/siglip2-base-patch16-224",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2025 year",
+        license="apache-2.0",
+        size_in_GB=0.37,
+        sources=ModelSource(hf="onnx-community/siglip2-base-patch16-224-ONNX"),
+        model_file="onnx/vision_model.onnx",
+    ),
+]
+
+
+class SiglipOnnxImageEmbedding(OnnxImageEmbedding):
+    """SigLIP vision tower.
+
+    The exported graph returns both the per-patch `last_hidden_state` and the pooled
+    `pooler_output`; only the latter is the image embedding, so it must be selected explicitly.
+    """
+
+    ONNX_OUTPUT_NAMES = ["pooler_output"]
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["OnnxImageEmbeddingWorker"]:
+        return SiglipImageEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return supported_siglip_models
+
+
+class SiglipImageEmbeddingWorker(OnnxImageEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
+        return SiglipOnnxImageEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/transform/functional.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import numpy as np
 from PIL import Image
 
@@ -15,7 +13,7 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Union[Image.Image, NumpyArray],
+    image: Image.Image | NumpyArray,
     size: tuple[int, int],
 ) -> NumpyArray:
     if isinstance(image, np.ndarray):
@@ -64,10 +62,16 @@ def center_crop(
 
 def normalize(
     image: NumpyArray,
-    mean: Union[float, list[float]],
-    std: Union[float, list[float]],
+    mean: float | list[float],
+    std: float | list[float],
 ) -> NumpyArray:
-    num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
+    if image.ndim < 3:
+        raise ValueError(f"image must be (C, H, W) or (N, C, H, W), got shape {image.shape}")
+
+    # Channels sit on the third axis from the end, which covers (C, H, W) and
+    # (N, C, H, W) alike. Transposing instead reversed every axis, which put the
+    # batch dimension where the channels were meant to be.
+    num_channels = image.shape[-3]
 
     if not np.issubdtype(image.dtype, np.floating):
         image = image.astype(np.float32)
@@ -80,7 +84,9 @@ def normalize(
             f"{len(mean_list)}"
         )
 
-    mean_arr = np.array(mean_list, dtype=np.float32)
+    # (C, 1, 1) lines the channels up with the trailing (C, H, W) axes under numpy
+    # broadcasting, whatever batch dimensions lead them.
+    mean_arr = np.array(mean_list, dtype=np.float32).reshape(-1, 1, 1)
 
     std_list = std if isinstance(std, list) else [std] * num_channels
     if len(std_list) != num_channels:
@@ -88,19 +94,24 @@ def normalize(
             f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std_list)}"
         )
 
-    std_arr = np.array(std_list, dtype=np.float32)
+    std_arr = np.array(std_list, dtype=np.float32).reshape(-1, 1, 1)
 
-    image_upd = ((image.T - mean_arr) / std_arr).T
+    image_upd = (image - mean_arr) / std_arr
     return image_upd
 
 
 def resize(
     image: Image.Image,
-    size: Union[int, tuple[int, int]],
-    resample: Union[int, Image.Resampling] = Image.Resampling.BILINEAR,
+    size: int | tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.BILINEAR,
 ) -> Image.Image:
     if isinstance(size, tuple):
-        return image.resize(size, resample)
+        # fastembed keeps sizes as (height, width) — `Compose.from_config` builds the
+        # tuple as (size["height"], size["width"]) — while Pillow's resize takes
+        # (width, height). The two agree for square sizes, so this only shows up on a
+        # non-square image processor configuration.
+        height, width = size
+        return image.resize((width, height), resample)
 
     height, width = image.height, image.width
     short, long = (width, height) if width <= height else (height, width)
@@ -117,7 +128,7 @@ def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyA
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
+def pil2ndarray(image: Image.Image | NumpyArray) -> NumpyArray:
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image
@@ -126,7 +137,7 @@ def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
 def pad2square(
     image: Image.Image,
     size: int,
-    fill_color: Union[str, int, tuple[int, ...]] = 0,
+    fill_color: str | int | tuple[int, ...] = 0,
 ) -> Image.Image:
     height, width = image.height, image.width
 
@@ -147,3 +158,77 @@ def pad2square(
     new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
     new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
     return new_image
+
+
+def resize_longest_edge(
+    image: Image.Image,
+    max_size: int,
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+) -> Image.Image:
+    height, width = image.height, image.width
+    aspect_ratio = width / height
+
+    if width >= height:
+        # Width is longer
+        new_width = max_size
+        new_height = int(new_width / aspect_ratio)
+    else:
+        # Height is longer
+        new_height = max_size
+        new_width = int(new_height * aspect_ratio)
+
+    # Ensure even dimensions
+    if new_height % 2 != 0:
+        new_height += 1
+    if new_width % 2 != 0:
+        new_width += 1
+
+    return image.resize((new_width, new_height), resample)
+
+
+def crop_ndarray(
+    image: NumpyArray,
+    x1: int,
+    y1: int,
+    x2: int,
+    y2: int,
+    channel_first: bool = True,
+) -> NumpyArray:
+    if channel_first:
+        # (C, H, W) format
+        return image[:, y1:y2, x1:x2]
+    else:
+        # (H, W, C) format
+        return image[y1:y2, x1:x2, :]
+
+
+def resize_ndarray(
+    image: NumpyArray,
+    size: tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+    channel_first: bool = True,
+) -> NumpyArray:
+    # Convert to PIL-friendly format (H, W, C)
+    if channel_first:
+        img_hwc = image.transpose((1, 2, 0))
+    else:
+        img_hwc = image
+
+    # Handle different dtypes
+    if img_hwc.dtype == np.float32 or img_hwc.dtype == np.float64:
+        # Assume normalized, scale to 0-255 for PIL
+        img_hwc_scaled = (img_hwc * 255).astype(np.uint8)
+        pil_img = Image.fromarray(img_hwc_scaled, mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized).astype(np.float32) / 255.0
+    else:
+        # uint8 or similar
+        pil_img = Image.fromarray(img_hwc.astype(np.uint8), mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized)
+
+    # Convert back to original format
+    if channel_first:
+        result = result.transpose((2, 0, 1))
+
+    return result

fastembed/image/transform/operators.py

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

fastembed/late_interaction/colbert.py

@@ -1,11 +1,11 @@
 import string
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, iter_batch
@@ -19,7 +19,7 @@ supported_colbert_models: list[DenseModelDescription] = [
     DenseModelDescription(
         model="colbert-ir/colbertv2.0",
         dim=128,
-        description="Late interaction model",
+        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2023 year",
         license="mit",
         size_in_GB=0.44,
         sources=ModelSource(hf="colbert-ir/colbertv2.0"),
@@ -28,7 +28,7 @@ supported_colbert_models: list[DenseModelDescription] = [
     DenseModelDescription(
         model="answerdotai/answerai-colbert-small-v1",
         dim=96,
-        description="Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2024 year",
         license="apache-2.0",
         size_in_GB=0.13,
         sources=ModelSource(hf="answerdotai/answerai-colbert-small-v1"),
@@ -98,7 +98,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def token_count(
         self,
-        texts: Union[str, Iterable[str]],
+        texts: str | Iterable[str],
         batch_size: int = 1024,
         is_doc: bool = True,
         include_extension: bool = False,
@@ -140,14 +140,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -159,10 +159,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -182,7 +183,7 @@ 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
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -198,11 +199,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
         )
-        self.mask_token_id: Optional[int] = None
-        self.pad_token_id: Optional[int] = None
+        self.mask_token_id: int | None = None
+        self.pad_token_id: int | None = None
         self.skip_list: set[int] = set()
 
-        self.query_tokenizer: Optional[Tokenizer] = None
+        self.query_tokenizer: Tokenizer | None = None
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -238,9 +239,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -273,7 +274,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: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         if isinstance(query, str):
             query = [query]
 

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -9,21 +9,21 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
@@ -43,7 +43,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         # 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: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -72,7 +72,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
 
     def token_count(
         self,
-        texts: Union[str, Iterable[str]],
+        texts: str | Iterable[str],
         batch_size: int = 1024,
         **kwargs: Any,
     ) -> int:

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,8 +1,8 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common.model_description import DenseModelDescription
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -51,11 +51,11 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -101,7 +101,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -116,9 +116,9 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -138,7 +138,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: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -154,7 +154,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
 
     def token_count(
         self,
-        texts: Union[str, Iterable[str]],
+        texts: str | Iterable[str],
         batch_size: int = 1024,
         is_doc: bool = True,
         include_extension: bool = False,

fastembed/late_interaction/token_embeddings.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict
-from typing import Union, Iterable, Optional, Any, Type
+from typing import Iterable, Any, Type
 
 from fastembed.common.model_description import DenseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -63,9 +63,9 @@ class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         yield from super().embed(documents, batch_size=batch_size, parallel=parallel, **kwargs)

fastembed/late_interaction_multimodal/colmodernvbert.py

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

fastembed/late_interaction_multimodal/colpali.py

@@ -1,11 +1,11 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from tokenizers import Encoding
 
 from fastembed.common import OnnxProvider, ImageInput
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common.utils import define_cache_dir, iter_batch
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
@@ -46,14 +46,14 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -65,10 +65,11 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             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.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -87,7 +88,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -174,7 +175,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
 
     def token_count(
         self,
-        texts: Union[str, Iterable[str]],
+        texts: str | Iterable[str],
         batch_size: int = 1024,
         include_extension: bool = False,
         **kwargs: Any,
@@ -227,9 +228,9 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
 
     def embed_text(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -263,9 +264,9 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
 
     def embed_image(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -1,9 +1,10 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common import OnnxProvider, ImageInput
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.late_interaction_multimodal.colpali import ColPali
+from fastembed.late_interaction_multimodal.colmodernvbert import ColModernVBERT
 
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
@@ -12,7 +13,10 @@ from fastembed.common.model_description import DenseModelDescription
 
 
 class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [ColPali]
+    EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [
+        ColPali,
+        ColModernVBERT,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -54,11 +58,11 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -104,7 +108,7 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -119,9 +123,9 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
 
     def embed_text(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -142,9 +146,9 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
 
     def embed_image(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -165,7 +169,7 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
 
     def token_count(
         self,
-        texts: Union[str, Iterable[str]],
+        texts: str | Iterable[str],
         batch_size: int = 1024,
         include_extension: bool = False,
         **kwargs: Any,

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 
 from fastembed.common import ImageInput
@@ -11,21 +11,21 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed_text(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -47,9 +47,9 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
 
     def embed_image(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -79,7 +79,7 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
 
     def token_count(
         self,
-        texts: Union[str, Iterable[str]],
+        texts: str | Iterable[str],
         **kwargs: Any,
     ) -> int:
         """Returns the number of tokens in the texts."""

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -2,7 +2,7 @@ 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, Sequence, Type
 
 import numpy as np
 from PIL import Image
@@ -11,19 +11,19 @@ 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.types import NumpyArray, Device
 from fastembed.common.utils import iter_batch
 from fastembed.image.transform.operators import Compose
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxMultimodalModel(OnnxModel[T]):
-    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+    ONNX_OUTPUT_NAMES: list[str] | None = None
 
     def __init__(self) -> None:
         super().__init__()
-        self.tokenizer: Optional[Tokenizer] = None
-        self.processor: Optional[Compose] = None
+        self.tokenizer: Tokenizer | None = None
+        self.processor: Compose | None = None
         self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_text_input(
@@ -60,11 +60,11 @@ class OnnxMultimodalModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -116,15 +116,15 @@ class OnnxMultimodalModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        documents: Union[str, Iterable[str]],
+        documents: 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,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -170,9 +170,11 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 yield from self._post_process_onnx_text_output(batch)  # type: ignore
 
     def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack():
+        with contextlib.ExitStack() as stack:
             image_files = [
-                Image.open(image) if not isinstance(image, Image.Image) else image
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
                 for image in images
             ]
             assert self.processor is not None, "Processor is not initialized"
@@ -187,15 +189,15 @@ class OnnxMultimodalModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[Iterable[ImageInput], ImageInput],
+        images: 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,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False

fastembed/parallel_processor.py

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

fastembed/postprocess/muvera.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import numpy as np
 
 from fastembed.common.types import NumpyArray
@@ -11,7 +9,7 @@ from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding
 )
 
 
-MultiVectorModel = Union[LateInteractionTextEmbeddingBase, LateInteractionMultimodalEmbeddingBase]
+MultiVectorModel = LateInteractionTextEmbeddingBase | LateInteractionMultimodalEmbeddingBase
 MAX_HAMMING_DISTANCE = 65  # 64 bits + 1
 POPCOUNT_LUT = np.array([bin(x).count("1") for x in range(256)], dtype=np.uint8)
 

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -1,8 +1,10 @@
-from typing import Optional, Sequence, Any
+from typing import Sequence, Any, Type
 
 from fastembed.common import OnnxProvider
 from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.types import Device
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.onnx_text_model import TextRerankerWorker
 
 
 class CustomTextCrossEncoder(OnnxTextCrossEncoder):
@@ -11,14 +13,14 @@ class CustomTextCrossEncoder(OnnxTextCrossEncoder):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(
@@ -38,9 +40,39 @@ class CustomTextCrossEncoder(OnnxTextCrossEncoder):
     def _list_supported_models(cls) -> list[BaseModelDescription]:
         return cls.SUPPORTED_MODELS
 
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextRerankerWorker]:
+        return CustomTextCrossEncoderWorker
+
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        return {"model_description": self.model_description}
+
     @classmethod
     def add_model(
         cls,
         model_description: BaseModelDescription,
     ) -> None:
         cls.SUPPORTED_MODELS.append(model_description)
+
+
+class CustomTextCrossEncoderWorker(TextRerankerWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        model_description: BaseModelDescription | None = None,
+        **kwargs: Any,
+    ) -> CustomTextCrossEncoder:
+        if model_description is None:
+            raise ValueError(
+                "`model_description` is required to initialize a custom model in a worker "
+                "process, it is provided by `CustomTextCrossEncoder._get_worker_init_kwargs`"
+            )
+        # custom models live in a class-level registry, which spawned workers don't inherit
+        CustomTextCrossEncoder.add_model(model_description)
+        return CustomTextCrossEncoder(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -1,9 +1,10 @@
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Sequence, Type
 
 from loguru import logger
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.rerank.cross_encoder.onnx_text_model import (
     OnnxCrossEncoderModel,
@@ -77,14 +78,14 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -96,10 +97,11 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -124,7 +126,7 @@ 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
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -180,7 +182,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         yield from self._rerank_pairs(

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,7 +1,7 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from tokenizers import Encoding
@@ -12,14 +12,14 @@ from fastembed.common.onnx_model import (
     OnnxOutputContext,
     OnnxProvider,
 )
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxCrossEncoderModel(OnnxModel[float]):
-    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+    ONNX_OUTPUT_NAMES: list[str] | None = None
 
     @classmethod
     def _get_worker_class(cls) -> Type["TextRerankerWorker"]:
@@ -29,11 +29,11 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -92,13 +92,13 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         cache_dir: str,
         pairs: Iterable[tuple[str, str]],
         batch_size: int,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -128,6 +128,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 "local_files_only": local_files_only,
                 "specific_model_path": specific_model_path,
                 **kwargs,
+                **self._get_worker_init_kwargs(),
             }
 
             if extra_session_options is not None:

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -1,7 +1,8 @@
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
 from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
 
@@ -53,11 +54,11 @@ class TextCrossEncoder(TextCrossEncoderBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -102,7 +103,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         """
@@ -140,7 +141,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
         description: str = "",
         license: str = "",
         size_in_gb: float = 0.0,
-        additional_files: Optional[list[str]] = None,
+        additional_files: list[str] | None = None,
     ) -> None:
         registered_models = cls._list_supported_models()
         for registered_model in registered_models:

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -1,4 +1,4 @@
-from typing import Any, Iterable, Optional
+from typing import Any, Iterable
 
 from fastembed.common.model_description import BaseModelDescription
 from fastembed.common.model_management import ModelManagement
@@ -8,8 +8,8 @@ class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
@@ -41,7 +41,7 @@ class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
         self,
         pairs: Iterable[tuple[str, str]],
         batch_size: int = 64,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         """Rerank query-document pairs.

fastembed/sparse/bm25.py

@@ -2,7 +2,7 @@ import os
 from collections import defaultdict
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Type, Union
+from typing import Any, Iterable, Type
 
 import mmh3
 import numpy as np
@@ -91,14 +91,14 @@ class Bm25(SparseTextEmbeddingBase):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
+        cache_dir: str | None = None,
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 256.0,
         language: str = "english",
         token_max_length: int = 40,
         disable_stemmer: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
@@ -158,11 +158,11 @@ class Bm25(SparseTextEmbeddingBase):
         self,
         model_name: str,
         cache_dir: str,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
+        specific_model_path: str | None = None,
     ) -> Iterable[SparseEmbedding]:
         is_small = False
 
@@ -205,9 +205,9 @@ class Bm25(SparseTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -268,7 +268,7 @@ class Bm25(SparseTextEmbeddingBase):
             embeddings.append(SparseEmbedding.from_dict(token_id2value))
         return embeddings
 
-    def token_count(self, texts: Union[str, Iterable[str]], **kwargs: Any) -> int:
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
         token_num = 0
         texts = [texts] if isinstance(texts, str) else texts
         for text in texts:
@@ -311,9 +311,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: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """To emulate BM25 behaviour, we don't need to use weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
         """

fastembed/sparse/bm42.py

@@ -1,7 +1,7 @@
 import math
 import string
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import mmh3
 import numpy as np
@@ -9,6 +9,7 @@ from py_rust_stemmers import SnowballStemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -65,15 +66,15 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
         alpha: float = 0.5,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -87,10 +88,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             alpha (float, optional): Parameter, that defines the importance of the token weight in the document
                 versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
                 It is recommended to only change this parameter based on training data for a specific dataset.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -110,7 +112,7 @@ 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
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -282,9 +284,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -325,9 +327,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             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: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         To emulate BM25 behaviour, we don't need to use smart weights in the query, and
         it's enough to just hash the tokens and assign a weight of 1.0 to them.
@@ -353,7 +353,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return Bm42TextEmbeddingWorker
 
     def token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **kwargs: Any
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
         if not hasattr(self, "model") or self.model is None:
             self.load_onnx_model()  # loads the tokenizer as well

fastembed/sparse/if_splade.py

@@ -0,0 +1,246 @@
+import json
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import ModelSource, SparseModelDescription
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+
+IDF_FILE = "idf.json"
+
+supported_if_splade_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte",
+        vocab_size=30522,
+        description="Inference-free SPLADE model. Documents are expanded with an ONNX encoder at index "
+        "time, queries are encoded with a tokenizer and an IDF lookup table only, "
+        "without any model inference.",
+        license="apache-2.0",
+        size_in_GB=0.55,
+        sources=ModelSource(hf="Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte"),
+        model_file="model.onnx",
+        additional_files=[IDF_FILE],
+        requires_idf=None,
+    ),
+]
+
+
+class IfSplade(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    """Inference-free (asymmetric) SPLADE model.
+
+    Documents are encoded with a neural encoder which expands them into a sparse vocabulary-sized
+    vector, while queries are encoded by tokenizing the text and looking up a precomputed IDF
+    weight per token — no neural inference happens at query time.
+
+    Query and document embeddings are compared with a dot product.
+    Special tokens are excluded from both document and query embeddings.
+    """
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        # Max-pool token logits over the sequence, masking out the padding
+        pooled = np.max(
+            output.model_output * np.expand_dims(output.attention_mask, axis=-1), axis=1
+        )
+        # v3 models of the opensearch-neural-sparse family use a double log activation,
+        # log(1 + log(1 + relu(x))), to increase sparsity of document embeddings
+        scores = np.log1p(np.log1p(np.maximum(pooled, 0.0)))
+
+        if self.special_tokens_ids:
+            scores[:, list(self.special_tokens_ids)] = 0.0
+
+        for row_scores in scores:
+            indices = row_scores.nonzero()[0]
+            yield SparseEmbedding(values=row_scores[indices], indices=indices)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        # unlike `OnnxTextModel._token_count`, does not require the onnx model to be loaded
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        for batch in iter_batch(texts, batch_size):
+            for tokens in self.tokenizer.encode_batch(batch):  # type: ignore[union-attr]
+                token_num += sum(tokens.attention_mask)
+        return token_num
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_if_splade_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        # The tokenizer and the idf table are lightweight and are required for query embedding,
+        # which does not involve any model inference, so they are loaded eagerly, while
+        # `lazy_load` only defers the initialization of the onnx model
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir)
+        self.special_tokens_ids: set[int] = set(self.special_token_to_id.values())
+        self._token_id_to_idf = self._load_idf()
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    def _load_idf(self) -> dict[int, float]:
+        with open(self._model_dir / IDF_FILE) as f:
+            token_to_idf: dict[str, float] = json.load(f)
+
+        vocab: dict[str, int] = self.tokenizer.get_vocab()  # type: ignore[union-attr]
+        return {vocab[token]: idf for token, idf in token_to_idf.items() if token in vocab}
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of queries into list of sparse embeddings without any model inference.
+
+        A query is tokenized, and each unique token is assigned its IDF weight from
+        a precomputed lookup table shipped with the model. Special tokens are ignored.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        for text in query:
+            token_ids = set(self.tokenizer.encode(text).ids) - self.special_tokens_ids  # type: ignore[union-attr]
+            embedding = {
+                token_id: self._token_id_to_idf[token_id]
+                for token_id in sorted(token_ids)
+                if token_id in self._token_id_to_idf
+            }
+            yield SparseEmbedding.from_dict(embedding)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+        return IfSpladeEmbeddingWorker
+
+
+class IfSpladeEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> IfSplade:
+        return IfSplade(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/minicoil.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from typing import Any, Optional, Sequence, Iterable, Union, Type
+from typing import Any, Sequence, Iterable, Type
 
 import numpy as np
 from numpy.typing import NDArray
@@ -10,6 +10,7 @@ from tokenizers import Tokenizer
 from fastembed.common.model_description import SparseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -72,17 +73,17 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 150.0,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -98,10 +99,11 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
                 Defaults to 0.75.
             avg_len (float, optional): The average length of the documents in the corpus. Defaults to 150.0.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -124,15 +126,15 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.avg_len = avg_len
 
         # Initialize class attributes
-        self.tokenizer: Optional[Tokenizer] = None
+        self.tokenizer: Tokenizer | None = None
         self.invert_vocab: dict[int, str] = {}
         self.special_tokens: set[str] = set()
         self.special_tokens_ids: set[int] = set()
         self.stopwords: set[str] = set()
-        self.vocab_resolver: Optional[VocabResolver] = None
-        self.encoder: Optional[Encoder] = None
-        self.output_dim: Optional[int] = None
-        self.sparse_vector_converter: Optional[SparseVectorConverter] = None
+        self.vocab_resolver: VocabResolver | None = None
+        self.encoder: Encoder | None = None
+        self.output_dim: int | None = None
+        self.sparse_vector_converter: SparseVectorConverter | None = None
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
@@ -188,15 +190,15 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         )
 
     def token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **kwargs: Any
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
         return self._token_count(texts, batch_size=batch_size, **kwargs)
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -233,9 +235,7 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             **kwargs,
         )
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs: Any
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Encode a list of queries into list of embeddings.
         """

fastembed/sparse/sparse_embedding_base.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 import numpy as np
 from numpy.typing import NDArray
@@ -12,7 +12,7 @@ from fastembed.common.model_management import ModelManagement
 @dataclass
 class SparseEmbedding:
     values: NumpyArray
-    indices: Union[NDArray[np.int64], NDArray[np.int32]]
+    indices: NDArray[np.int64] | NDArray[np.int32]
 
     def as_object(self) -> dict[str, NumpyArray]:
         return {
@@ -35,8 +35,8 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
@@ -46,9 +46,9 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         raise NotImplementedError()
@@ -68,9 +68,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: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -87,6 +85,6 @@ class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
         else:
             yield from self.embed(query, **kwargs)
 
-    def token_count(self, texts: Union[str, Iterable[str]], **kwargs: Any) -> int:
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
         """Returns the number of tokens in the texts."""
         raise NotImplementedError("Subclasses must implement this method")

fastembed/sparse/sparse_text_embedding.py

@@ -1,9 +1,11 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.bm42 import Bm42
+from fastembed.sparse.if_splade import IfSplade
 from fastembed.sparse.minicoil import MiniCOIL
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -15,7 +17,13 @@ from fastembed.common.model_description import SparseModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25, MiniCOIL]
+    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [
+        SpladePP,
+        Bm42,
+        Bm25,
+        MiniCOIL,
+        IfSplade,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -53,11 +61,11 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -93,9 +101,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """
@@ -115,9 +123,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: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 
@@ -130,7 +136,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         yield from self.model.query_embed(query, **kwargs)
 
     def token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **kwargs: Any
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
         """Returns the number of tokens in the texts.
 

fastembed/sparse/splade_pp.py

@@ -1,8 +1,9 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
@@ -54,7 +55,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             yield SparseEmbedding(values=scores, indices=indices)
 
     def token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **kwargs: Any
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
         return self._token_count(texts, batch_size=batch_size, **kwargs)
 
@@ -70,14 +71,14 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -89,10 +90,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.
             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.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -111,7 +113,7 @@ 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
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -144,9 +146,9 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[SparseEmbedding]:
         """

fastembed/sparse/utils/sparse_vectors_converter.py

@@ -1,12 +1,11 @@
-from typing import Dict, List, Set
-from py_rust_stemmers import SnowballStemmer
-from fastembed.common.utils import get_all_punctuation, remove_non_alphanumeric
-import mmh3
 import copy
 from dataclasses import dataclass
 
+import mmh3
 import numpy as np
+from py_rust_stemmers import SnowballStemmer
 
+from fastembed.common.utils import get_all_punctuation, remove_non_alphanumeric
 from fastembed.sparse.sparse_embedding_base import SparseEmbedding
 
 GAP = 32000
@@ -16,16 +15,16 @@ INT32_MAX = 2**31 - 1
 @dataclass
 class WordEmbedding:
     word: str
-    forms: List[str]
+    forms: list[str]
     count: int
     word_id: int
-    embedding: List[float]
+    embedding: list[float]
 
 
 class SparseVectorConverter:
     def __init__(
         self,
-        stopwords: Set[str],
+        stopwords: set[str],
         stemmer: SnowballStemmer,
         k: float = 1.2,
         b: float = 0.75,
@@ -58,15 +57,15 @@ class SparseVectorConverter:
         return res
 
     @classmethod
-    def normalize_vector(cls, vector: List[float]) -> List[float]:
+    def normalize_vector(cls, vector: list[float]) -> list[float]:
         norm = sum([x**2 for x in vector]) ** 0.5
         if norm < 1e-8:
             return vector
         return [x / norm for x in vector]
 
     def clean_words(
-        self, sentence_embedding: Dict[str, WordEmbedding], token_max_length: int = 40
-    ) -> Dict[str, WordEmbedding]:
+        self, sentence_embedding: dict[str, WordEmbedding], token_max_length: int = 40
+    ) -> dict[str, WordEmbedding]:
         """
         Clean miniCOIL-produced sentence_embedding, as unknown to the miniCOIL's stemmer tokens should fully resemble
         our BM25 token representation.
@@ -85,7 +84,7 @@ class SparseVectorConverter:
                 }
         """
 
-        new_sentence_embedding: Dict[str, WordEmbedding] = {}
+        new_sentence_embedding: dict[str, WordEmbedding] = {}
 
         for word, embedding in sentence_embedding.items():
             # embedding = {
@@ -127,7 +126,7 @@ class SparseVectorConverter:
 
     def embedding_to_vector(
         self,
-        sentence_embedding: Dict[str, WordEmbedding],
+        sentence_embedding: dict[str, WordEmbedding],
         embedding_size: int,
         vocab_size: int,
     ) -> SparseEmbedding:
@@ -156,14 +155,14 @@ class SparseVectorConverter:
 
         """
 
-        indices: List[int] = []
-        values: List[float] = []
-        
+        indices: list[int] = []
+        values: list[float] = []
+
         # Example:
         # vocab_size = 10000
         # embedding_size = 4
         # GAP = 32000
-        # 
+        #
         # We want to start random words section from the bucket, that is guaranteed to not
         # include any vocab words.
         # We need (vocab_size * embedding_size) slots for vocab words.
@@ -171,9 +170,7 @@ class SparseVectorConverter:
         # Therefore, we can start random words from bucket (vocab_size * embedding_size) // GAP + 1 + 1
 
         # ID at which the scope of OOV words starts
-        unknown_words_shift = (
-            (vocab_size * embedding_size) // GAP + 2
-        ) * GAP
+        unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP
         sentence_embedding_cleaned = self.clean_words(sentence_embedding)
 
         # Calculate sentence length after cleaning
@@ -208,7 +205,7 @@ class SparseVectorConverter:
 
     def embedding_to_vector_query(
         self,
-        sentence_embedding: Dict[str, WordEmbedding],
+        sentence_embedding: dict[str, WordEmbedding],
         embedding_size: int,
         vocab_size: int,
     ) -> SparseEmbedding:
@@ -216,8 +213,8 @@ class SparseVectorConverter:
         Same as `embedding_to_vector`, but no TF
         """
 
-        indices: List[int] = []
-        values: List[float] = []
+        indices: list[int] = []
+        values: list[float] = []
 
         # ID at which the scope of OOV words starts
         unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP

fastembed/text/builtin_sentence_embedding.py

@@ -0,0 +1,69 @@
+from typing import Any, Iterable, Type
+
+
+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
+
+
+supported_builtin_sentence_embedding_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="google/embeddinggemma-300m",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), multilingual, 2048 input tokens truncation, "
+            "Prefixes for queries/documents: `task: search result | query: {content}` for query, "
+            "`title: {title | 'none'} | text: {content}` for documents, 2025 year."
+        ),
+        license="gemma",
+        size_in_GB=1.24,
+        sources=ModelSource(
+            hf="onnx-community/embeddinggemma-300m-ONNX",
+        ),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    ),
+]
+
+
+class BuiltinSentenceEmbedding(OnnxTextEmbedding):
+    """Builtin Sentence Embedding uses built-in pooling and normalization of underlying onnx models"""
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return BuiltinSentenceEmbeddingWorker
+
+    @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_builtin_sentence_embedding_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return output.model_output
+
+    def _run_model(
+        self, onnx_input: dict[str, Any], onnx_output_names: list[str] | None = None
+    ) -> NumpyArray:
+        return self.model.run(onnx_output_names, onnx_input)[1]  # type: ignore[union-attr]
+
+
+class BuiltinSentenceEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return BuiltinSentenceEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/custom_text_embedding.py

@@ -1,5 +1,4 @@
-from typing import Optional, Sequence, Any, Iterable
-
+from typing import Sequence, Any, Iterable, Type
 from dataclasses import dataclass
 
 import numpy as np
@@ -11,9 +10,10 @@ from fastembed.common.model_description import (
     DenseModelDescription,
 )
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray
-from fastembed.common.utils import normalize, mean_pooling
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.utils import normalize, mean_pooling, last_token_pooling
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
 
 
 @dataclass(frozen=True)
@@ -29,14 +29,14 @@ class CustomTextEmbedding(OnnxTextEmbedding):
     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,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(
@@ -51,20 +51,31 @@ class CustomTextEmbedding(OnnxTextEmbedding):
             specific_model_path=specific_model_path,
             **kwargs,
         )
-        self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
-        self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization
+        postprocessing_config = self.POSTPROCESSING_MAPPING[self.model_description.model]
+        self._pooling = postprocessing_config.pooling
+        self._normalization = postprocessing_config.normalization
 
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         return cls.SUPPORTED_MODELS
 
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
+        return CustomTextEmbeddingWorker
+
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        return {
+            "model_description": self.model_description,
+            "postprocessing_config": self.POSTPROCESSING_MAPPING[self.model_description.model],
+        }
+
     def _post_process_onnx_output(
         self, output: OnnxOutputContext, **kwargs: Any
     ) -> Iterable[NumpyArray]:
         return self._normalize(self._pool(output.model_output, output.attention_mask))
 
     def _pool(
-        self, embeddings: NumpyArray, attention_mask: Optional[NDArray[np.int64]] = None
+        self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None
     ) -> NumpyArray:
         if self._pooling == PoolingType.CLS:
             return embeddings[:, 0]
@@ -74,12 +85,18 @@ class CustomTextEmbedding(OnnxTextEmbedding):
                 raise ValueError("attention_mask must be provided for mean pooling")
             return mean_pooling(embeddings, attention_mask)
 
+        if self._pooling == PoolingType.LAST_TOKEN:
+            if attention_mask is None:
+                raise ValueError("attention_mask must be provided for last token pooling")
+            return last_token_pooling(embeddings, attention_mask)
+
         if self._pooling == PoolingType.DISABLED:
             return embeddings
 
         raise ValueError(
             f"Unsupported pooling type {self._pooling}. "
-            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, {PoolingType.DISABLED}."
+            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, "
+            f"{PoolingType.LAST_TOKEN}, {PoolingType.DISABLED}."
         )
 
     def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
@@ -96,3 +113,32 @@ class CustomTextEmbedding(OnnxTextEmbedding):
         cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
             pooling=pooling, normalization=normalization
         )
+
+
+class CustomTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        model_description: DenseModelDescription | None = None,
+        postprocessing_config: PostprocessingConfig | None = None,
+        **kwargs: Any,
+    ) -> CustomTextEmbedding:
+        if model_description is None or postprocessing_config is None:
+            raise ValueError(
+                "`model_description` and `postprocessing_config` are required to initialize a "
+                "custom model in a worker process, they are provided by "
+                "`CustomTextEmbedding._get_worker_init_kwargs`"
+            )
+        # custom models live in a class-level registry, which spawned workers don't inherit
+        CustomTextEmbedding.add_model(
+            model_description,
+            pooling=postprocessing_config.pooling,
+            normalization=postprocessing_config.normalization,
+        )
+        return CustomTextEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/last_token_normalized_embedding.py

@@ -0,0 +1,94 @@
+from typing import Any, Iterable, Type
+
+import onnxruntime as ort
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import last_token_pooling, normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_last_token_normalized_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qwen/Qwen3-Embedding-0.6B",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), multilingual, 32768 input tokens truncation, "
+            "Prefixes for queries/documents: `Instruct: {task_description}\\nQuery:{query}` "
+            "for queries, none for documents, 2025 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=2.38,
+        sources=ModelSource(hf="Qdrant/Qwen3-Embedding-0.6B-onnx"),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    ),
+    DenseModelDescription(
+        model="Qwen/Qwen3-Embedding-0.6B-Q",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), multilingual, 32768 input tokens truncation, "
+            "Prefixes for queries/documents: `Instruct: {task_description}\\nQuery:{query}` "
+            "for queries, none for documents, int8 weights, requires onnxruntime>=1.23, "
+            "2025 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=1.12,
+        sources=ModelSource(hf="Qdrant/Qwen3-Embedding-0.6B-onnx"),
+        model_file="onnx/model_quantized.onnx",
+    ),
+]
+
+
+class LastTokenNormalizedEmbedding(OnnxTextEmbedding):
+    """Decoder-based embedding models, which pool the last non-padding token and normalize it"""
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return LastTokenNormalizedEmbeddingWorker
+
+    def load_onnx_model(self) -> None:
+        try:
+            super().load_onnx_model()
+        except Exception as e:
+            # int8 weights are stored as 8-bit MatMulNBits, which onnxruntime only
+            # implements since 1.23; older versions fail with "nbits_ == 4 was false"
+            if "nbits" not in str(e).lower():
+                raise
+            raise RuntimeError(
+                f"Could not load {self.model_name}: its int8 weights require "
+                f"onnxruntime>=1.23, but onnxruntime {ort.__version__} is installed. "
+                f"Either upgrade onnxruntime or use a non-quantized model."
+            ) from e
+
+    @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_last_token_normalized_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for last token pooling")
+
+        return normalize(last_token_pooling(output.model_output, output.attention_mask))
+
+
+class LastTokenNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return LastTokenNormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/multitask_embedding.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Type, Iterable, Union, Optional
+from typing import Any, Type, Iterable
 
 import numpy as np
 
@@ -45,11 +45,9 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
 
-    def __init__(self, *args: Any, task_id: Optional[int] = None, **kwargs: Any):
+    def __init__(self, *args: Any, task_id: int | None = None, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.default_task_id: Union[Task, int] = (
-            task_id if task_id is not None else self.PASSAGE_TASK
-        )
+        self.default_task_id: Task | int = task_id if task_id is not None else self.PASSAGE_TASK
 
     @classmethod
     def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
@@ -62,7 +60,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     def _preprocess_onnx_input(
         self,
         onnx_input: dict[str, NumpyArray],
-        task_id: Optional[Union[int, Task]] = None,
+        task_id: int | Task | None = None,
         **kwargs: Any,
     ) -> dict[str, NumpyArray]:
         if task_id is None:
@@ -72,10 +70,10 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        task_id: Optional[int] = None,
+        parallel: int | None = None,
+        task_id: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         task_id = (
@@ -83,7 +81,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         )  # required for multiprocessing
         yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
 
     def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:

fastembed/text/onnx_embedding.py

@@ -1,11 +1,11 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device, NumpyArray, OnnxProvider
 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(
@@ -34,7 +34,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         license="mit",
         size_in_GB=0.21,
         sources=ModelSource(
-            hf="qdrant/bge-base-en-v1.5-onnx-q",
+            hf="Qdrant/bge-base-en-v1.5-onnx-Q",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
             _deprecated_tar_struct=True,
         ),
@@ -77,7 +77,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         ),
         license="mit",
         size_in_GB=0.067,
-        sources=ModelSource(hf="qdrant/bge-small-en-v1.5-onnx-q"),
+        sources=ModelSource(hf="Qdrant/bge-small-en-v1.5-onnx-Q"),
         model_file="model_optimized.onnx",
     ),
     DenseModelDescription(
@@ -180,6 +180,42 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(hf="jinaai/jina-clip-v1"),
         model_file="onnx/text_model.onnx",
     ),
+    DenseModelDescription(
+        model="minishlab/potion-base-8M",
+        dim=256,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2024 year."
+        ),
+        license="mit",
+        size_in_GB=0.030,
+        sources=ModelSource(hf="minishlab/potion-base-8m-onnx"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="minishlab/potion-retrieval-32M",
+        dim=512,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2025 year."
+        ),
+        license="mit",
+        size_in_GB=0.129,
+        sources=ModelSource(hf="minishlab/potion-retrieval-32m-onnx"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="minishlab/potion-multilingual-128M",
+        dim=256,
+        description=(
+            "Text embeddings, Unimodal (text), Multilingual, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2025 year."
+        ),
+        license="mit",
+        size_in_GB=0.512,
+        sources=ModelSource(hf="minishlab/potion-multilingual-128m-onnx"),
+        model_file="model.onnx",
+    ),
 ]
 
 
@@ -199,14 +235,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -218,10 +254,11 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -239,7 +276,7 @@ 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
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -260,9 +297,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -332,7 +369,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     def token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **kwargs: Any
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
         return self._token_count(texts, batch_size=batch_size, **kwargs)
 

fastembed/text/onnx_text_model.py

@@ -1,13 +1,13 @@
 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, Sequence, Type
 
 import numpy as np
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider, Device
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
@@ -15,7 +15,7 @@ from fastembed.parallel_processor import ParallelWorkerPool
 
 
 class OnnxTextModel(OnnxModel[T]):
-    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+    ONNX_OUTPUT_NAMES: list[str] | None = None
 
     @classmethod
     def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
@@ -35,12 +35,12 @@ class OnnxTextModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.tokenizer: Optional[Tokenizer] = None
+        self.tokenizer: Tokenizer | None = None
         self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
-    ) -> dict[str, Union[NumpyArray, NDArray[np.int64]]]:
+    ) -> dict[str, NumpyArray | NDArray[np.int64]]:
         """
         Preprocess the onnx input.
         """
@@ -50,11 +50,11 @@ class OnnxTextModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -92,27 +92,34 @@ class OnnxTextModel(OnnxModel[T]):
                 [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
+        model_output = self._run_model(
+            onnx_input=onnx_input, onnx_output_names=self.ONNX_OUTPUT_NAMES
+        )
 
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         return OnnxOutputContext(
-            model_output=model_output[0],
+            model_output=model_output,
             attention_mask=onnx_input.get("attention_mask", attention_mask),
             input_ids=onnx_input.get("input_ids", input_ids),
         )
 
+    def _run_model(
+        self, onnx_input: dict[str, Any], onnx_output_names: list[str] | None = None
+    ) -> NumpyArray:
+        return self.model.run(onnx_output_names, onnx_input)[0]  # type: ignore[union-attr]
+
     def _embed_documents(
         self,
         model_name: str,
         cache_dir: str,
-        documents: Union[str, Iterable[str]],
+        documents: 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,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -144,6 +151,7 @@ class OnnxTextModel(OnnxModel[T]):
                 "local_files_only": local_files_only,
                 "specific_model_path": specific_model_path,
                 **kwargs,
+                **self._get_worker_init_kwargs(),
             }
 
             if extra_session_options is not None:
@@ -159,9 +167,7 @@ class OnnxTextModel(OnnxModel[T]):
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
 
-    def _token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **_: Any
-    ) -> int:
+    def _token_count(self, texts: str | Iterable[str], batch_size: int = 1024, **_: Any) -> int:
         if not hasattr(self, "model") or self.model is None:
             self.load_onnx_model()  # loads the tokenizer as well
 

fastembed/text/pooled_normalized_embedding.py

@@ -57,9 +57,9 @@ supported_pooled_normalized_models: list[DenseModelDescription] = [
             "Prefixes for queries/documents: not necessary, 2024 year."
         ),
         license="apache-2.0",
-        size_in_GB=0.32,
+        size_in_GB=0.64,
         sources=ModelSource(hf="jinaai/jina-embeddings-v2-base-de"),
-        model_file="onnx/model_fp16.onnx",
+        model_file="onnx/model.onnx",
     ),
     DenseModelDescription(
         model="jinaai/jina-embeddings-v2-base-code",

fastembed/text/siglip_embedding.py

@@ -0,0 +1,71 @@
+from typing import Any, Type
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+
+supported_siglip_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="google/siglip2-base-patch16-224",
+        dim=768,
+        description=(
+            "Text embeddings, Multimodal (text&image), multilingual, 64 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2025 year"
+        ),
+        license="apache-2.0",
+        size_in_GB=1.13,
+        sources=ModelSource(hf="onnx-community/siglip2-base-patch16-224-ONNX"),
+        model_file="onnx/text_model.onnx",
+    ),
+]
+
+
+class SiglipOnnxTextEmbedding(OnnxTextEmbedding):
+    """SigLIP text tower.
+
+    SigLIP always pools the hidden state at the last sequence position, whether or not it is
+    padding, so every batch must be padded to the model's fixed training length (rather than to
+    the longest sequence in the batch) or the resulting embeddings become dependent on what else
+    is in the batch.
+
+    The exported graph also returns both `last_hidden_state` and `pooler_output`; only the
+    latter is the text embedding, so it must be selected explicitly.
+    """
+
+    ONNX_OUTPUT_NAMES = ["pooler_output"]
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return SiglipTextEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return supported_siglip_models
+
+    def load_onnx_model(self) -> None:
+        super().load_onnx_model()
+        if self.tokenizer is not None:
+            truncation = self.tokenizer.truncation
+            padding = self.tokenizer.padding
+            if truncation and padding and padding.get("length") is None:
+                self.tokenizer.enable_padding(
+                    direction=padding["direction"],
+                    pad_id=padding["pad_id"],
+                    pad_type_id=padding["pad_type_id"],
+                    pad_token=padding["pad_token"],
+                    length=truncation["max_length"],
+                )
+
+
+class SiglipTextEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return SiglipOnnxTextEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/text_embedding.py

@@ -1,14 +1,17 @@
 import warnings
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.types import NumpyArray, OnnxProvider, Device
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 from fastembed.text.custom_text_embedding import CustomTextEmbedding
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.multitask_embedding import JinaEmbeddingV3
+from fastembed.text.builtin_sentence_embedding import BuiltinSentenceEmbedding
+from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.siglip_embedding import SiglipOnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
 
@@ -17,9 +20,12 @@ class TextEmbedding(TextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [
         OnnxTextEmbedding,
         CLIPOnnxEmbedding,
+        SiglipOnnxTextEmbedding,
         PooledNormalizedEmbedding,
         PooledEmbedding,
         JinaEmbeddingV3,
+        BuiltinSentenceEmbedding,
+        LastTokenNormalizedEmbedding,
         CustomTextEmbedding,
     ]
 
@@ -51,7 +57,7 @@ class TextEmbedding(TextEmbeddingBase):
         description: str = "",
         license: str = "",
         size_in_gb: float = 0.0,
-        additional_files: Optional[list[str]] = None,
+        additional_files: list[str] | None = None,
     ) -> None:
         registered_models = cls._list_supported_models()
         for registered_model in registered_models:
@@ -79,32 +85,18 @@ class TextEmbedding(TextEmbeddingBase):
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
+        if model_name.lower() == "jinaai/jina-embeddings-v2-base-de":
             warnings.warn(
-                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
-                "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
-                UserWarning,
-                stacklevel=2,
-            )
-        if model_name.lower() in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2".lower(),
-            "thenlper/gte-large".lower(),
-            "intfloat/multilingual-e5-large".lower(),
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2".lower(),
-        }:
-            warnings.warn(
-                f"The model {model_name} now uses mean pooling instead of CLS embedding. "
-                f"In order to preserve the previous behaviour, consider either pinning fastembed version to 0.5.1 or "
-                "using `add_custom_model` functionality.",
+                "The model 'jinaai/jina-embeddings-v2-base-de' used to run with fp16 model, but due to onnxruntime updates, now it runs with the original fp32 model.",
                 UserWarning,
                 stacklevel=2,
             )
@@ -149,7 +141,7 @@ class TextEmbedding(TextEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -164,9 +156,9 @@ class TextEmbedding(TextEmbeddingBase):
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """
@@ -186,7 +178,7 @@ class TextEmbedding(TextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -214,7 +206,7 @@ class TextEmbedding(TextEmbeddingBase):
         yield from self.model.passage_embed(texts, **kwargs)
 
     def token_count(
-        self, texts: Union[str, Iterable[str]], batch_size: int = 1024, **kwargs: Any
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
         """Returns the number of tokens in the texts.
 

fastembed/text/text_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Union, Any
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -9,21 +9,21 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        documents: Union[str, Iterable[str]],
+        documents: str | Iterable[str],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
@@ -43,7 +43,7 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         # 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: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
@@ -70,6 +70,6 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         """Returns embedding size for the current model"""
         raise NotImplementedError("Subclasses must implement this method")
 
-    def token_count(self, texts: Union[str, Iterable[str]], **kwargs: Any) -> int:
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
         """Returns the number of tokens in the texts."""
         raise NotImplementedError("Subclasses must implement this method")

mkdocs.yml

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

pyproject.toml

@@ -1,9 +1,9 @@
 [tool.poetry]
 name = "fastembed-gpu"
-version = "0.7.4"
+version = "0.8.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
-license = "Apache License"
+license = "Apache-2.0"
 readme = "README.md"
 packages = [{include = "fastembed"}]
 homepage = "https://github.com/qdrant/fastembed"
@@ -11,19 +11,19 @@ repository = "https://github.com/qdrant/fastembed"
 keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
 
 [tool.poetry.dependencies]
-python = ">=3.9.0"
+python = ">=3.10.0"
 numpy = [
-    { version = ">=1.21,<2.1.0", python = "<3.10" },
-    { version = ">=1.21,<2.3.0", python = ">=3.10,<3.11" },
-    { version = ">=1.21", python = ">=3.11,<3.12" },
-    { version = ">=1.26", python = ">=3.12,<3.13" },
-    { version = ">=2.1.0", python = ">=3.13,<3.14" },
+    { version = ">=1.21,<2.3.0", python = "3.10" },
+    { version = ">=1.21", python = "3.11" },
+    { version = ">=1.26", python = "3.12" },
+    { version = ">=2.1.0", python = "3.13" },
     { version = ">=2.3.0", python = ">=3.14" },
 ]
 onnxruntime-gpu = [
-    { version = ">=1.17.0,<1.20.0", python = "<3.10" },
-    { version = ">1.20.0", python = ">=3.13" },
-    { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
+    { version = ">=1.17.0,!=1.20.0,<1.24", python = "3.10" },
+    { version = ">=1.17.0,!=1.20.0,!=1.24.0,!=1.24.1", python = ">=3.11,<3.13" },
+    { version = ">1.21.0,!=1.24.0,!=1.24.1", python = "3.13" },
+    { version = ">=1.24.2", python = ">=3.14" },
 ]
 tqdm = "^4.66"
 requests = "^2.31"
@@ -31,35 +31,36 @@ tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<2.0"
 loguru = "^0.7.2"
 pillow = [
-    { version = ">=10.3.0,<11.0", python = "<3.10" },
-    { version = ">=10.3.0,<12.0", python = ">=3.10,<3.13" },
-    { version = ">=11.0.0,<12.0", python = ">=3.13" },
+    { version = ">=10.3.0,<13.0", python = ">=3.10,<3.13" },
+    { version = ">=11.0.0,<13.0", python = "3.13" },
+    { version = ">=12.0.0,<13.0", python = ">=3.14" },
 ]
 mmh3 = ">=4.1.0,<6.0.0"
 py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.test.dependencies]
-pytest = "^7.4.2"
+pytest = ">=7.4.2,<10.0.0"
 ruff = ">=0.3.1,<1.0"
 
 [tool.poetry.group.dev.dependencies]
 notebook = ">=7.0.2"
-pre-commit = "^3.6.2"
+pre-commit = ">=3.6.2,<5.0.0"
 onnx = [
-    { version = ">=1.15.0", python = "<3.13" },
-    { version = ">=1.18.0", python = ">=3.13" },
+    { version = ">=1.15.0", python = ">=3.10,<3.13" },
+    { version = ">=1.18.0", python = "3.13" },
+    { version = ">=1.20.0", python = ">=3.14" },
 ]
 
 [tool.poetry.group.docs.dependencies]
 mkdocs-material = "^9.5.10"
-mkdocstrings = "^0.24.0"
+mkdocstrings = ">=0.24,<1.1"
 pillow = ">=10.3.0,<13.0.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 
 [tool.poetry.group.types.dependencies]
 pyright = ">=1.1.293"
-mypy = "^1.0.0"
+mypy = ">=1,<3"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_common.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from fastembed import (
     TextEmbedding,
     SparseTextEmbedding,
@@ -5,6 +7,7 @@ from fastembed import (
     LateInteractionMultimodalEmbedding,
     LateInteractionTextEmbedding,
 )
+from fastembed.common.utils import last_token_pooling
 
 
 def test_text_list_supported_models():
@@ -28,3 +31,31 @@ def test_text_list_supported_models():
         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"]
+
+
+def test_last_token_pooling():
+    token_embeddings = np.array(
+        [
+            [[1.0, 1.0], [2.0, 2.0], [9.0, 9.0], [9.0, 9.0]],  # 2 real tokens, then padding
+            [[3.0, 3.0], [4.0, 4.0], [5.0, 5.0], [6.0, 6.0]],  # no padding
+        ]
+    )
+    attention_mask = np.array([[1, 1, 0, 0], [1, 1, 1, 1]], dtype=np.int64)
+
+    pooled = last_token_pooling(token_embeddings, attention_mask)
+
+    assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])
+
+
+def test_last_token_pooling_with_left_padding():
+    token_embeddings = np.array(
+        [
+            [[9.0, 9.0], [9.0, 9.0], [1.0, 1.0], [2.0, 2.0]],  # padding, then 2 real tokens
+            [[3.0, 3.0], [4.0, 4.0], [5.0, 5.0], [6.0, 6.0]],  # no padding
+        ]
+    )
+    attention_mask = np.array([[0, 0, 1, 1], [1, 1, 1, 1]], dtype=np.int64)
+
+    pooled = last_token_pooling(token_embeddings, attention_mask)
+
+    assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])

tests/test_custom_models.py

@@ -10,7 +10,7 @@ from fastembed.common.model_description import (
     BaseModelDescription,
 )
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, mean_pooling
+from fastembed.common.utils import normalize, mean_pooling, last_token_pooling
 from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
 from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
 from fastembed.rerank.cross_encoder import TextCrossEncoder
@@ -21,9 +21,11 @@ from tests.utils import delete_model_cache
 @pytest.fixture(autouse=True)
 def restore_custom_models_fixture():
     CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextEmbedding.POSTPROCESSING_MAPPING = {}
     CustomTextCrossEncoder.SUPPORTED_MODELS = []
     yield
     CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextEmbedding.POSTPROCESSING_MAPPING = {}
     CustomTextCrossEncoder.SUPPORTED_MODELS = []
 
 
@@ -74,8 +76,29 @@ def test_text_custom_model():
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
-    CustomTextEmbedding.SUPPORTED_MODELS.clear()
-    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
+
+def test_text_custom_model_parallel_processing():
+    is_ci = os.getenv("CI")
+    custom_model_name = "intfloat/multilingual-e5-small"
+    dim = 384
+
+    TextEmbedding.add_custom_model(
+        custom_model_name,
+        pooling=PoolingType.MEAN,
+        normalization=True,
+        sources=ModelSource(hf=custom_model_name),
+        dim=dim,
+        size_in_gb=0.47,
+    )
+
+    model = TextEmbedding(custom_model_name)
+    docs = ["hello world", "flag embedding"] * 50
+    embeddings = np.stack(list(model.embed(docs, batch_size=10, parallel=2)), axis=0)
+
+    assert embeddings.shape == (len(docs), dim)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 def test_cross_encoder_custom_model():
@@ -114,7 +137,26 @@ def test_cross_encoder_custom_model():
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
-    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()
+
+def test_cross_encoder_custom_model_parallel_processing():
+    is_ci = os.getenv("CI")
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        model_file="onnx/model.onnx",
+        sources=ModelSource(hf=custom_model_name),
+        size_in_gb=0.08,
+    )
+
+    model = TextCrossEncoder(custom_model_name)
+    pairs = [("What is AI?", "Artificial intelligence is ...")] * 50
+    scores = np.stack(list(model.rerank_pairs(pairs, batch_size=10, parallel=2)), axis=0)
+
+    assert scores.shape == (len(pairs),)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 def test_mock_add_custom_models():
@@ -136,6 +178,8 @@ def test_mock_add_custom_models():
         f"{PoolingType.MEAN.lower()}": dummy_token_output,
         f"{PoolingType.CLS.lower()}-normalized": dummy_token_output,
         f"{PoolingType.CLS.lower()}": dummy_token_output,
+        f"{PoolingType.LAST_TOKEN.lower()}-normalized": dummy_token_output,
+        f"{PoolingType.LAST_TOKEN.lower()}": dummy_token_output,
         f"{PoolingType.DISABLED.lower()}-normalized": dummy_pooled_output,
         f"{PoolingType.DISABLED.lower()}": dummy_pooled_output,
     }
@@ -147,12 +191,19 @@ def test_mock_add_custom_models():
         f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
         f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]),
         f"{PoolingType.CLS.lower()}": dummy_token_embedding[:, 0],
+        f"{PoolingType.LAST_TOKEN.lower()}-normalized": normalize(
+            last_token_pooling(dummy_token_embedding, dummy_attention_mask)
+        ),
+        f"{PoolingType.LAST_TOKEN.lower()}": last_token_pooling(
+            dummy_token_embedding, dummy_attention_mask
+        ),
         f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding),
         f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
     }
 
     for pooling, normalization in itertools.product(
-        (PoolingType.MEAN, PoolingType.CLS, PoolingType.DISABLED), (True, False)
+        (PoolingType.MEAN, PoolingType.CLS, PoolingType.LAST_TOKEN, PoolingType.DISABLED),
+        (True, False),
     ):
         model_name = f"{pooling.name.lower()}{'-normalized' if normalization else ''}"
         TextEmbedding.add_custom_model(
@@ -175,8 +226,24 @@ def test_mock_add_custom_models():
         )
         assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3)
 
-    CustomTextEmbedding.SUPPORTED_MODELS.clear()
-    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
+
+def test_custom_text_model_lookup_is_case_insensitive():
+    model_name = "Org/Model"
+
+    TextEmbedding.add_custom_model(
+        model_name,
+        pooling=PoolingType.MEAN,
+        normalization=True,
+        sources=ModelSource(hf="artificial"),
+        dim=5,
+        size_in_gb=0.1,
+    )
+
+    model = TextEmbedding("org/model", lazy_load=True, specific_model_path="./")
+
+    assert isinstance(model.model, CustomTextEmbedding)
+    assert model.model._pooling == PoolingType.MEAN
+    assert model.model._normalization is True
 
 
 def test_do_not_add_existing_model():
@@ -212,9 +279,6 @@ def test_do_not_add_existing_model():
             size_in_gb=0.47,
         )
 
-    CustomTextEmbedding.SUPPORTED_MODELS.clear()
-    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()
-
 
 def test_do_not_add_existing_cross_encoder():
     existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
@@ -239,5 +303,3 @@ def test_do_not_add_existing_cross_encoder():
             sources=ModelSource(hf=custom_model_name),
             size_in_gb=0.08,
         )
-
-    CustomTextCrossEncoder.SUPPORTED_MODELS.clear()

tests/test_image_onnx_embeddings.py

@@ -1,4 +1,5 @@
 import os
+import platform
 from contextlib import contextmanager
 from io import BytesIO
 
@@ -25,6 +26,15 @@ CANONICAL_VECTOR_VALUES = {
     "jinaai/jina-clip-v1": np.array(
         [-0.029, 0.0216, 0.0396, 0.0283, -0.0023, 0.0151, 0.011, -0.0235, 0.0251, -0.0343]
     ),
+    "nomic-ai/nomic-embed-vision-v1.5": np.array(
+        [0.0048, -0.0254, 0.0067, -0.0296, -0.0435, -0.0123, 0.0024, -0.0361, -0.0703, -0.0186]
+    ),
+    "nomic-ai/nomic-embed-vision-v1.5-Q": np.array(
+        [-0.0011, -0.0477, 0.0024, -0.049, -0.0458, -0.0314, 0.017, -0.0383, -0.0537, -0.021]
+    ),
+    "google/siglip2-base-patch16-224": np.array(
+        [-0.02095927, -0.0075177, -0.00144479, -0.0080948, 0.05031789]
+    ),
 }
 
 _MODELS_TO_CACHE = ("Qdrant/clip-ViT-B-32-vision",)
@@ -59,9 +69,13 @@ def model_cache():
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
 def test_embedding(model_cache, model_name: str) -> None:
     is_ci = os.getenv("CI")
+    is_mac = platform.system() == "Darwin"
     is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
     for model_desc in ImageEmbedding._list_supported_models():
+        # quantized int8 ops diverge on macOS; canonical vector is generated on linux/amd64 (CI)
+        if is_mac and model_desc.model == "nomic-ai/nomic-embed-vision-v1.5-Q":
+            continue
         if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue
 

tests/test_image_transform.py

@@ -0,0 +1,90 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from fastembed.image.transform.functional import normalize, resize
+
+
+@pytest.mark.parametrize(
+    ("size", "expected"),
+    [
+        ((100, 200), (200, 100)),  # the bug: a non-square size came back transposed
+        ((224, 224), (224, 224)),  # the square path every shipped model takes
+    ],
+)
+def test_resize_tuple_is_height_width(size: tuple[int, int], expected: tuple[int, int]) -> None:
+    """A ``(height, width)`` size must reach Pillow as ``(width, height)``."""
+    resized = resize(Image.new("RGB", (300, 300)), size=size)
+
+    assert resized.size == expected  # PIL reports (width, height)
+
+
+def test_resize_int_keeps_shortest_edge_behaviour() -> None:
+    """The int branch already emitted Pillow order; it must not be disturbed."""
+    landscape = Image.new("RGB", (400, 200))
+    portrait = Image.new("RGB", (200, 400))
+
+    # size sets the shortest edge, and the aspect ratio is preserved.
+    assert resize(landscape, size=100).size == (200, 100)
+    assert resize(portrait, size=100).size == (100, 200)
+
+
+@pytest.mark.parametrize(
+    ("mean", "std"),
+    [
+        ([0.1, 0.2, 0.3], [0.5, 0.6, 0.7]),  # per-channel, as every model config gives it
+        (0.5, 0.25),  # scalar, expanded to one value per channel
+    ],
+)
+def test_normalize_chw_is_channel_wise(
+    mean: list[float] | float, std: list[float] | float
+) -> None:
+    """Each channel must be normalized by its own mean/std, not by any other axis."""
+    rng = np.random.default_rng(0)
+    image = rng.random((3, 5, 7)).astype(np.float32)
+    means = mean if isinstance(mean, list) else [mean] * 3
+    stds = std if isinstance(std, list) else [std] * 3
+
+    result = normalize(image, mean=mean, std=std)
+
+    for c in range(3):
+        assert np.allclose(result[c], (image[c] - means[c]) / stds[c], atol=1e-6)
+
+
+@pytest.mark.parametrize("batch_size", [2, 3])
+def test_normalize_batched_matches_per_image(batch_size: int) -> None:
+    """A batch must give exactly what the (C, H, W) path gives image by image.
+
+    batch_size 2 used to raise, since transposing reversed every axis; batch_size 3
+    matched the channel count and silently normalized along the batch axis instead.
+    """
+    rng = np.random.default_rng(2)
+    batch = rng.random((batch_size, 3, 4, 4)).astype(np.float32)
+    mean, std = [0.1, 0.2, 0.3], [0.5, 0.6, 0.7]
+
+    result = normalize(batch, mean=mean, std=std)
+
+    per_image = np.stack([normalize(image, mean=mean, std=std) for image in batch])
+    assert result.shape == batch.shape
+    assert np.array_equal(result, per_image)
+
+
+def test_normalize_rejects_input_without_a_channel_axis() -> None:
+    """Every pipeline runs ConvertToRGB first, so normalize only ever sees (C, H, W)."""
+    with pytest.raises(ValueError, match=r"must be \(C, H, W\)"):
+        normalize(np.zeros((4, 6), dtype=np.float32), mean=0.5, std=0.25)
+
+
+@pytest.mark.parametrize(
+    ("mean", "std", "expected"),
+    [
+        ([0.1, 0.2], [1.0, 1.0, 1.0], "mean must"),
+        ([0.1, 0.2, 0.3], [1.0, 1.0], "std must"),
+    ],
+)
+def test_normalize_channel_count_mismatch_raises(
+    mean: list[float], std: list[float], expected: str
+) -> None:
+    image = np.zeros((3, 4, 4), dtype=np.float32)
+    with pytest.raises(ValueError, match=expected):
+        normalize(image, mean=mean, std=std)

tests/test_late_interaction_multimodal.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 
 import pytest
 from PIL import Image
@@ -6,7 +7,7 @@ import numpy as np
 
 from fastembed import LateInteractionMultimodalEmbedding
 from tests.config import TEST_MISC_DIR
-
+from tests.utils import delete_model_cache
 
 # vectors are abridged and rounded for brevity
 CANONICAL_IMAGE_VALUES = {
@@ -21,6 +22,17 @@ CANONICAL_IMAGE_VALUES = {
             [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
         ]
     ),
+    "Qdrant/colmodernvbert": np.array(
+        [
+            [0.11614, -0.15793, -0.11194, 0.0688, 0.08001, 0.10575, -0.07871],
+            [0.10094, -0.13301, -0.12069, 0.10932, 0.04645, 0.09884, 0.04048],
+            [0.13106, -0.18613, -0.13469, 0.10566, 0.03659, 0.07712, -0.03916],
+            [0.09754, -0.09596, -0.04839, 0.14991, 0.05692, 0.10569, -0.08349],
+            [0.02576, -0.15651, -0.09977, 0.09707, 0.13412, 0.09994, -0.09931],
+            [-0.06741, -0.1787, -0.19677, -0.07618, 0.13102, -0.02131, -0.02437],
+            [-0.02776, -0.10187, -0.13793, 0.03835, 0.04766, 0.04701, -0.15635],
+        ]
+    ),
 }
 
 CANONICAL_QUERY_VALUES = {
@@ -35,6 +47,17 @@ CANONICAL_QUERY_VALUES = {
             [-0.0165, -0.0106, 0.1672, -0.0768, 0.0389, -0.0038, 0.1137],
         ]
     ),
+    "Qdrant/colmodernvbert": np.array(
+        [
+            [0.05, 0.06557, 0.04026, 0.14981, 0.1842, 0.0263, -0.18706],
+            [-0.05664, -0.14028, 0.00649, -0.02849, 0.09034, -0.01494, 0.10693],
+            [-0.10147, -0.00716, 0.09084, -0.08236, -0.01849, -0.00972, -0.00461],
+            [-0.1233, -0.10814, -0.02337, -0.00329, 0.05984, 0.09934, 0.09846],
+            [-0.07053, -0.13119, -0.06487, 0.01508, 0.07459, 0.07655, 0.14821],
+            [0.00526, -0.13842, -0.05837, -0.02721, 0.13009, 0.05076, 0.17962],
+            [0.00924, -0.14383, -0.03057, -0.03691, 0.11718, 0.037, 0.13344],
+        ]
+    ),
 }
 
 queries = ["hello world", "flag embedding"]
@@ -44,43 +67,69 @@ images = [
     Image.open((TEST_MISC_DIR / "image.jpeg")),
 ]
 
+_MODELS_TO_CACHE = ("Qdrant/colmodernvbert",)
+MODELS_TO_CACHE = tuple(model_name.lower() for model_name in _MODELS_TO_CACHE)
 
-def test_batch_embedding():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
 
+@pytest.fixture(scope="module")
+def model_cache():
+    is_ci = os.getenv("CI")
+    cache = {}
+
+    @contextmanager
+    def get_model(model_name: str):
+        lowercase_model_name = model_name.lower()
+        if lowercase_model_name not in cache:
+            cache[lowercase_model_name] = LateInteractionMultimodalEmbedding(lowercase_model_name)
+        yield cache[lowercase_model_name]
+        if lowercase_model_name not in MODELS_TO_CACHE:
+            model_inst = cache.pop(lowercase_model_name)
+            if is_ci:
+                delete_model_cache(model_inst.model._model_dir)
+            del model_inst
+
+    yield get_model
+
+    if is_ci:
+        for _, model in cache.items():
+            delete_model_cache(model.model._model_dir)
+    cache.clear()
+
+
+def test_batch_embedding(model_cache):
     for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = list(model.embed_image(images, batch_size=2))
+        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
+            continue  # colpali is too large for ci
 
-        for value in result:
+        print("evaluating", model_name)
+        with model_cache(model_name) as model:
+            result = list(model.embed_image(images, batch_size=2))
+
+            for value in result:
+                token_num, abridged_dim = expected_result.shape
+                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+
+
+def test_single_embedding(model_cache):
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
+            continue  # colpali is too large for ci
+        print("evaluating", model_name)
+        with model_cache(model_name) as model:
+            result = next(iter(model.embed_image(images, batch_size=6)))
             token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
-def test_single_embedding():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
-    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = next(iter(model.embed_image(images, batch_size=6)))
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
-
-
-def test_single_embedding_query():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-
+def test_single_embedding_query(model_cache):
     for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+        if model_name.lower() == "Qdrant/colpali-v1.3-fp16".lower() and os.getenv("CI"):
+            continue  # colpali is too large for ci
         print("evaluating", model_name)
-        model = LateInteractionMultimodalEmbedding(model_name=model_name)
-        result = next(iter(model.embed_text(queries)))
-        token_num, abridged_dim = expected_result.shape
-        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+        with model_cache(model_name) as model:
+            result = next(iter(model.embed_text(queries)))
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_get_embedding_size():
@@ -90,33 +139,27 @@ def test_get_embedding_size():
     model_name = "Qdrant/ColPali-v1.3-fp16"
     assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
 
+    model_name = "Qdrant/colmodernvbert"
+    assert LateInteractionMultimodalEmbedding.get_embedding_size(model_name) == 128
+
 
 def test_embedding_size():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-    model_name = "Qdrant/colpali-v1.3-fp16"
-    model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
-    assert model.embedding_size == 128
-
-    model_name = "Qdrant/ColPali-v1.3-fp16"
+    model_name = "Qdrant/colmodernvbert"
     model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
     assert model.embedding_size == 128
 
 
-def test_token_count() -> None:
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
-    model_name = "Qdrant/colpali-v1.3-fp16"
-    model = LateInteractionMultimodalEmbedding(model_name=model_name, lazy_load=True)
-
-    documents = ["short doc", "it is a long document to check attention mask for paddings"]
-    short_doc_token_count = model.token_count(documents[0])
-    long_doc_token_count = model.token_count(documents[1])
-    documents_token_count = model.token_count(documents)
-    assert short_doc_token_count + long_doc_token_count == documents_token_count
-    assert short_doc_token_count + long_doc_token_count == model.token_count(
-        documents, batch_size=1
-    )
-    assert short_doc_token_count + long_doc_token_count < model.token_count(
-        documents, include_extension=True
-    )
+def test_token_count(model_cache) -> None:
+    model_name = "Qdrant/colmodernvbert"
+    with model_cache(model_name) as model:
+        documents = ["short doc", "it is a long document to check attention mask for paddings"]
+        short_doc_token_count = model.token_count(documents[0])
+        long_doc_token_count = model.token_count(documents[1])
+        documents_token_count = model.token_count(documents)
+        assert short_doc_token_count + long_doc_token_count == documents_token_count
+        assert short_doc_token_count + long_doc_token_count == model.token_count(
+            documents, batch_size=1
+        )
+        assert short_doc_token_count + long_doc_token_count < model.token_count(
+            documents, include_extension=True
+        )

tests/test_multi_gpu.py

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

tests/test_preprocessor_utils.py

@@ -0,0 +1,333 @@
+import itertools
+import json
+import os
+import shutil
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+
+from tokenizers import Tokenizer
+
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.text.text_embedding import TextEmbedding
+from tests.utils import delete_model_cache
+
+# transformers writes its VERY_LARGE_INTEGER in place of `model_max_length` when the real
+# value is unknown, which is more than `enable_truncation` can accept
+HF_SENTINEL = int(1e30)
+
+# a lightweight model whose config files serve as a realistic starting point for the cases below
+BASE_MODEL = "BAAI/bge-small-en-v1.5"
+TOKENIZER_FILES = (
+    "config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+)
+
+
+def _patch_json(path: Path, overrides: dict[str, Any], drop: tuple[str, ...] = ()) -> None:
+    with open(path) as source:
+        content = json.load(source)
+
+    content.update(overrides)
+    for key in drop:
+        content.pop(key, None)
+
+    with open(path, "w") as target:
+        json.dump(content, target)
+
+
+def _set_serialized_padding(path: Path, padding: dict[str, Any] | None) -> None:
+    """Rewrite tokenizer.json through the tokenizers library, so the format stays authoritative."""
+    tokenizer = Tokenizer.from_file(str(path))
+    if padding is None:
+        tokenizer.no_padding()
+    else:
+        tokenizer.enable_padding(**padding)
+    tokenizer.save(str(path))
+
+
+@pytest.fixture(scope="module")
+def make_model_dir(tmp_path_factory):
+    """Build model directories from a real model's config files, with targeted overrides.
+
+    `load_tokenizer` reads only the four files in `TOKENIZER_FILES`, so the onnx weights are
+    never copied.
+    """
+    is_ci = os.getenv("CI")
+    base_model = TextEmbedding(BASE_MODEL)
+    source_dir = Path(base_model.model._model_dir)
+    counter = itertools.count()
+
+    def factory(
+        tokenizer_config: dict[str, Any] | None = None,
+        config: dict[str, Any] | None = None,
+        padding: dict[str, Any] | None = None,
+        drop_from_tokenizer_config: tuple[str, ...] = (),
+        drop_from_config: tuple[str, ...] = (),
+        drop_files: tuple[str, ...] = (),
+    ) -> Path:
+        model_dir = tmp_path_factory.mktemp(f"model_dir_{next(counter)}")
+        for file_name in TOKENIZER_FILES:
+            if file_name in drop_files:
+                continue
+            shutil.copy(source_dir / file_name, model_dir / file_name)
+
+        if "tokenizer_config.json" not in drop_files:
+            _patch_json(
+                model_dir / "tokenizer_config.json",
+                tokenizer_config or {},
+                drop_from_tokenizer_config,
+            )
+        if "config.json" not in drop_files:
+            _patch_json(model_dir / "config.json", config or {}, drop_from_config)
+        if padding is not None:
+            _set_serialized_padding(model_dir / "tokenizer.json", padding)
+
+        return model_dir
+
+    yield factory
+
+    if is_ci:
+        delete_model_cache(base_model.model._model_dir)
+
+
+def test_fixed_padding_is_relaxed_to_batch_longest(make_model_dir) -> None:
+    """Fixed padding shorter than the truncation limit leaves longer encodings ragged."""
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": 512, "max_length": None},
+        padding={"length": 128, "pad_id": 0, "pad_token": "[PAD]", "direction": "right"},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["length"] is None
+    assert tokenizer.truncation["max_length"] == 512
+
+    encoded = tokenizer.encode_batch(["hello world", "retrieval " * 200])
+    # ragged encodings make this raise, the same way onnx_embed does
+    input_ids = np.array([encoding.ids for encoding in encoded])
+
+    assert input_ids.shape == (2, len(encoded[0].ids))
+    assert input_ids.shape[1] > 128
+
+
+def test_batch_longest_padding_does_not_pad_to_the_truncation_limit(make_model_dir) -> None:
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": 512, "max_length": None},
+        padding={"length": 128, "pad_id": 0, "pad_token": "[PAD]", "direction": "right"},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir)
+    encoded = tokenizer.encode_batch(["hello world", "hello"])
+
+    assert len(encoded[0].ids) == len(encoded[1].ids) < 128
+
+
+def test_serialized_left_padding_is_preserved(make_model_dir) -> None:
+    """ColModernVBERT pads on the left, normalizing the length must not reset the direction."""
+    model_dir = make_model_dir(
+        padding={"length": None, "pad_id": 0, "pad_token": "[PAD]", "direction": "left"},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["direction"] == "left"
+    assert tokenizer.padding["length"] is None
+
+    encoded = tokenizer.encode_batch(["hello world and then some", "hello"])
+    assert encoded[1].ids[0] == 0
+    assert encoded[1].attention_mask[0] == 0
+
+
+def test_serialized_pad_to_multiple_of_is_preserved(make_model_dir) -> None:
+    """Everything the tokenizer declared is kept; only the fixed length is overridden."""
+    model_dir = make_model_dir(
+        padding={"length": 128, "pad_id": 0, "pad_token": "[PAD]", "pad_to_multiple_of": 8},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["length"] is None
+    assert tokenizer.padding["pad_to_multiple_of"] == 8
+
+    encoded = tokenizer.encode_batch(["hello world", "hello"])
+    assert len(encoded[0].ids) % 8 == 0
+
+
+def test_serialized_pad_id_takes_precedence_over_config(make_model_dir) -> None:
+    mask_pad_id = 103  # [MASK] in the bert-base vocab, any id other than the config's works
+    model_dir = make_model_dir(
+        config={"pad_token_id": 7},
+        padding={"length": None, "pad_id": mask_pad_id, "pad_token": "[MASK]"},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["pad_id"] == mask_pad_id
+    assert tokenizer.padding["pad_token"] == "[MASK]"
+
+
+def test_pad_token_falls_back_to_tokenizer_config(make_model_dir) -> None:
+    model_dir = make_model_dir(config={"pad_token_id": 3}, padding=None)
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["pad_token"] == "[PAD]"
+    assert tokenizer.padding["pad_id"] == 3
+
+
+def test_missing_pad_token_raises(make_model_dir) -> None:
+    model_dir = make_model_dir(drop_from_tokenizer_config=("pad_token",))
+
+    with pytest.raises(ValueError, match="Could not find a pad token"):
+        load_tokenizer(model_dir)
+
+
+@pytest.mark.parametrize(
+    "model_max_length,max_length,expected",
+    [
+        (512, 128, 128),  # both usable, the stricter one wins
+        (128, 512, 128),
+        (512, None, 512),
+        (None, 256, 256),
+        (HF_SENTINEL, 128, 128),  # qdrant/gte-large-onnx
+        (0, 256, 256),
+        (512, 0, 512),
+    ],
+)
+def test_max_context_resolution(make_model_dir, model_max_length, max_length, expected) -> None:
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": model_max_length, "max_length": max_length},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.truncation["max_length"] == expected
+
+
+@pytest.mark.parametrize(
+    "model_max_length,max_length",
+    [
+        (HF_SENTINEL, None),  # transformers' placeholder is not a limit
+        (0, None),  # a zero would truncate everything away
+        (None, 0),
+        (None, None),
+        ("512", None),  # not an integer
+    ],
+)
+def test_unusable_max_context_raises(make_model_dir, model_max_length, max_length) -> None:
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": model_max_length, "max_length": max_length},
+    )
+
+    with pytest.raises(ValueError, match="Could not determine the maximum context length"):
+        load_tokenizer(model_dir)
+
+
+def test_absent_max_context_keys_raise(make_model_dir) -> None:
+    model_dir = make_model_dir(
+        drop_from_tokenizer_config=("model_max_length", "max_length"),
+    )
+
+    with pytest.raises(ValueError, match="Could not determine the maximum context length"):
+        load_tokenizer(model_dir)
+
+
+@pytest.fixture(scope="module")
+def token_id(make_model_dir):
+    """Resolve vocabulary ids by name, so the cases below carry no magic numbers."""
+    return Tokenizer.from_file(str(make_model_dir() / "tokenizer.json")).token_to_id
+
+
+@pytest.mark.parametrize(
+    "dropped",
+    [("config.json",), ("special_tokens_map.json",), ("config.json", "special_tokens_map.json")],
+    ids=["no-config", "no-special-tokens-map", "neither"],
+)
+def test_optional_files_do_not_change_what_is_loaded(make_model_dir, dropped) -> None:
+    """Both files are redundant: everything they carry is already in the tokenizer."""
+    baseline, baseline_specials = load_tokenizer(make_model_dir())
+
+    tokenizer, specials = load_tokenizer(make_model_dir(drop_files=dropped))
+
+    assert specials == baseline_specials
+    assert tokenizer.padding == baseline.padding
+    assert tokenizer.encode("hello world").ids == baseline.encode("hello world").ids
+
+
+@pytest.mark.parametrize("missing", ("tokenizer.json", "tokenizer_config.json"))
+def test_the_remaining_files_are_still_required(make_model_dir, missing) -> None:
+    """Relaxing the optional two must not relax the two that carry irreplaceable data."""
+    model_dir = make_model_dir(drop_files=(missing,))
+
+    with pytest.raises(ValueError, match=f"Could not find {missing}"):
+        load_tokenizer(model_dir)
+
+
+@pytest.mark.parametrize(
+    "model_files",
+    [
+        pytest.param({"drop_from_config": ("pad_token_id",)}, id="config-omits-pad-token-id"),
+        pytest.param({"drop_files": ("config.json",)}, id="config-is-absent"),
+    ],
+)
+def test_pad_id_falls_back_to_the_vocabulary(make_model_dir, token_id, model_files) -> None:
+    """Last link of the chain; a hardcoded 0 would silently disagree with `pad_token`."""
+    expected = token_id("[SEP]")
+    assert expected != 0, "a pad token whose id is 0 would pass even without a lookup"
+    model_dir = make_model_dir(tokenizer_config={"pad_token": "[SEP]"}, **model_files)
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["pad_id"] == expected
+
+
+def test_pad_token_that_resolves_nowhere_raises(make_model_dir) -> None:
+    """Without config.json a pad token outside the vocabulary has no id left to fall back on."""
+    model_dir = make_model_dir(
+        tokenizer_config={"pad_token": "[NOT_IN_VOCAB]"},
+        drop_files=("config.json",),
+    )
+
+    with pytest.raises(ValueError, match="Could not resolve an id for the pad token"):
+        load_tokenizer(model_dir)
+
+
+def test_pad_token_named_only_in_the_map_resolves(make_model_dir) -> None:
+    """The map is read first, so it can name a pad token tokenizer.json does not carry."""
+    model_dir = make_model_dir(
+        tokenizer_config={"pad_token": "<|mypad|>"},
+        drop_files=("config.json",),
+    )
+    _patch_json(model_dir / "special_tokens_map.json", {"pad_token": "<|mypad|>"})
+
+    tokenizer, specials = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["pad_token"] == "<|mypad|>"
+    assert tokenizer.padding["pad_id"] == specials["<|mypad|>"]
+
+
+@pytest.mark.parametrize(
+    "additional",
+    [
+        pytest.param(["<|list_str|>"], id="list-of-strings"),
+        pytest.param([{"content": "<|list_str|>"}], id="list-of-added-token-dicts"),
+    ],
+)
+def test_list_valued_map_entries_are_registered(make_model_dir, additional) -> None:
+    """`additional_special_tokens` holds a list, which the str/dict dispatch alone drops.
+
+    Real repos ship both spellings, and their tokens are in tokenizer.json already, so
+    only a token living nowhere else shows the drop.
+    """
+    model_dir = make_model_dir()
+    _patch_json(model_dir / "special_tokens_map.json", {"additional_special_tokens": additional})
+
+    _, specials = load_tokenizer(model_dir)
+
+    assert "<|list_str|>" in specials

tests/test_sparse_embeddings.py

@@ -8,6 +8,7 @@ from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
 from tests.utils import delete_model_cache, should_test_model
 
+
 CANONICAL_COLUMN_VALUES = {
     "prithivida/Splade_PP_en_v1": {
         "indices": [
@@ -58,9 +59,50 @@ CANONICAL_COLUMN_VALUES = {
             -0.12508166,
         ],
     },
+    # first 15 non-zero dimensions of the embedding
+    "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte": {
+        "indices": [
+            999,
+            1010,
+            1011,
+            1024,
+            1028,
+            1029,
+            1045,
+            1074,
+            1993,
+            2017,
+            2033,
+            2054,
+            2073,
+            2080,
+            2088,
+        ],
+        "values": [
+            0.16544909,
+            0.00529129,
+            0.0392109,
+            0.12337475,
+            0.09640586,
+            0.05325737,
+            0.09611791,
+            0.03159865,
+            0.01349991,
+            0.09392473,
+            0.01928805,
+            0.05238346,
+            0.05515401,
+            0.03156782,
+            0.98263124,
+        ],
+    },
 }
 
 CANONICAL_QUERY_VALUES = {
+    "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte": {
+        "indices": [2088, 7592],
+        "values": [3.42086864, 6.93775654],
+    },
     "Qdrant/minicoil-v1": {
         "indices": [80, 81, 82, 83, 6664, 6665, 6666, 6667],
         "values": [
@@ -82,6 +124,7 @@ _MODELS_TO_CACHE = (
     "Qdrant/minicoil-v1",
     "Qdrant/bm25",
     "Qdrant/bm42-all-minilm-l6-v2-attentions",
+    "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte",
 )
 MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE])
 
@@ -142,18 +185,23 @@ def test_single_embedding(model_cache) -> None:
             continue
         if not should_test_model(model_desc, model_desc.model, is_ci, is_manual):
             continue
-
         with model_cache(model_desc.model) as model:
             passage_result = next(iter(model.embed(docs, batch_size=6)))
             query_result = next(iter(model.query_embed(docs)))
             expected_result = CANONICAL_COLUMN_VALUES[model_desc.model]
             expected_query_result = CANONICAL_QUERY_VALUES.get(model_desc.model, expected_result)
-            assert passage_result.indices.tolist() == expected_result["indices"]
-            for i, value in enumerate(passage_result.values):
+
+            # canonical values might contain only a prefix of the non-zero dimensions
+            num_dims = len(expected_result["indices"])
+            assert passage_result.indices.tolist()[:num_dims] == expected_result["indices"]
+            for i, value in enumerate(passage_result.values[:num_dims]):
                 assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
 
-            assert query_result.indices.tolist() == expected_query_result["indices"]
-            for i, value in enumerate(query_result.values):
+            num_query_dims = len(expected_query_result["indices"])
+            assert (
+                query_result.indices.tolist()[:num_query_dims] == expected_query_result["indices"]
+            )
+            for i, value in enumerate(query_result.values[:num_query_dims]):
                 assert pytest.approx(value, abs=0.001) == expected_query_result["values"][i]
 
 
@@ -263,6 +311,22 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
+def test_if_splade_query_embed_is_inference_free() -> None:
+    is_ci = os.getenv("CI")
+    model = SparseTextEmbedding(
+        model_name="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte",
+        lazy_load=True,
+    )
+    embeddings = list(model.query_embed(["hello world", "flag embedding"]))
+    # queries are embedded with a tokenizer and an idf lookup table only,
+    # the onnx model must stay unloaded
+    assert not hasattr(model.model, "model")
+    assert all(len(embedding.indices) > 0 for embedding in embeddings)
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
 @pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")

tests/test_text_onnx_embeddings.py

@@ -5,6 +5,8 @@ from contextlib import contextmanager
 import numpy as np
 import pytest
 
+from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding import TextEmbedding
 from tests.utils import delete_model_cache, should_test_model
 
@@ -68,6 +70,51 @@ CANONICAL_VECTOR_VALUES = {
     "Qdrant/clip-ViT-B-32-text": np.array([0.0083, 0.0103, -0.0138, 0.0199, -0.0069]),
     "thenlper/gte-base": np.array([0.0038, 0.0355, 0.0181, 0.0092, 0.0654]),
     "jinaai/jina-clip-v1": np.array([-0.0862, -0.0101, -0.0056, 0.0375, -0.0472]),
+    "google/embeddinggemma-300m": np.array(
+        [-0.08181356, 0.0214127, 0.05120273, -0.03690156, -0.0254504]
+    ),
+    "Qwen/Qwen3-Embedding-0.6B": np.array(
+        [-0.01476084, 0.01723184, -0.01195498, -0.07275258, 0.00281229]
+    ),
+    "Qwen/Qwen3-Embedding-0.6B-Q": np.array(
+        [-0.01599521, 0.01676456, -0.01195119, -0.07132675, 0.00346729]
+    ),
+    "google/siglip2-base-patch16-224": np.array(
+        [-0.01181389, 0.00737596, 0.01118064, 0.0103095, 0.3451049]
+    ),
+    "minishlab/potion-base-8M": np.array(
+        [-0.03432461, -0.08020256, -0.14396408, 0.08480079, 0.01958815]
+    ),
+    "minishlab/potion-retrieval-32M": np.array(
+        [0.019733, -0.01530093, -0.08678473, 0.0229059, 0.04700558]
+    ),
+    "minishlab/potion-multilingual-128M": np.array(
+        [0.02366836, 0.02973341, 0.05140258, -0.00745248, -0.06740689]
+    ),
+}
+
+QWEN3_INSTRUCT_PREFIX = (
+    "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:"
+)
+
+DOC_PREFIXES = {
+    "google/embeddinggemma-300m": "title: none | text: ",
+}
+QUERY_PREFIXES = {
+    "google/embeddinggemma-300m": "task: search result | query: ",
+    "Qwen/Qwen3-Embedding-0.6B": QWEN3_INSTRUCT_PREFIX,
+    "Qwen/Qwen3-Embedding-0.6B-Q": QWEN3_INSTRUCT_PREFIX,
+}
+CANONICAL_QUERY_VECTOR_VALUES = {
+    "google/embeddinggemma-300m": np.array(
+        [-0.22990295, 0.03311195, 0.04290345, -0.03558498, -0.01399477]
+    ),
+    "Qwen/Qwen3-Embedding-0.6B": np.array(
+        [-0.01908712, 0.01635596, -0.00356586, -0.03947155, -0.01387356]
+    ),
+    "Qwen/Qwen3-Embedding-0.6B-Q": np.array(
+        [-0.02221339, 0.01932909, -0.00361797, -0.03888897, -0.01362813]
+    ),
 }
 
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
@@ -119,6 +166,9 @@ def test_embedding(model_cache, model_name: str) -> None:
 
         with model_cache(model_desc.model) as model:
             docs = ["hello world", "flag embedding"]
+            if model_desc.model in DOC_PREFIXES:
+                docs = [DOC_PREFIXES[model_desc.model] + doc for doc in docs]
+
             embeddings = list(model.embed(docs))
             embeddings = np.stack(embeddings, axis=0)
             assert embeddings.shape == (2, dim)
@@ -129,6 +179,56 @@ def test_embedding(model_cache, model_name: str) -> None:
             ), model_desc.model
 
 
+def test_query_embedding(model_cache) -> None:
+    is_ci = os.getenv("CI")
+    is_mac = platform.system() == "Darwin"
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    for model_desc in TextEmbedding._list_supported_models():
+        if model_desc.model in MULTI_TASK_MODELS or (
+            is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q"
+        ):
+            continue
+
+        if model_desc.model not in CANONICAL_QUERY_VECTOR_VALUES:
+            continue
+
+        if not should_test_model(model_desc, "", is_ci, is_manual):
+            continue
+
+        dim = model_desc.dim
+        with model_cache(model_desc.model) as model:
+            queries = ["hello world", "flag embedding"]
+            if model_desc.model in QUERY_PREFIXES:
+                queries = [QUERY_PREFIXES[model_desc.model] + query for query in queries]
+
+            embeddings = list(model.query_embed(queries))
+            embeddings = np.stack(embeddings, axis=0)
+            assert embeddings.shape == (2, dim)
+
+            canonical_vector = CANONICAL_QUERY_VECTOR_VALUES[model_desc.model]
+            assert np.allclose(
+                embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
+            ), model_desc.model
+
+
+def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> None:
+    """Old onnxruntime only implements 4-bit MatMulNBits, the error should say so."""
+    monkeypatch.setattr(
+        OnnxTextEmbedding,
+        "load_onnx_model",
+        lambda self: (_ for _ in ()).throw(RuntimeError("nbits_ == 4 was false")),
+    )
+    model = LastTokenNormalizedEmbedding(
+        "Qwen/Qwen3-Embedding-0.6B-Q",
+        lazy_load=True,
+        specific_model_path="./",  # disable model downloading and loading
+    )
+
+    with pytest.raises(RuntimeError, match="onnxruntime>=1.23"):
+        model.load_onnx_model()
+
+
 @pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
     with model_cache(model_name) as model:
@@ -217,3 +317,17 @@ def test_token_count(model_cache, model_name) -> None:
         doc_token_count = model.token_count(documents)
         assert first_doc_token_count + second_doc_token_count == doc_token_count
         assert doc_token_count == model.token_count(documents, batch_size=1)
+
+
+@pytest.mark.parametrize(
+    "model_name,dim",
+    [("sentence-transformers/all-MiniLM-L6-v2", 384), ("thenlper/gte-base", 768)],
+)
+def test_mixed_length_batch_with_fixed_padding(model_cache, model_name: str, dim: int) -> None:
+    # both models serialize a fixed padding length of 128 in tokenizer.json; gte-base truncates
+    # at 512, so a document longer than 128 makes the batch ragged unless the padding is relaxed
+    with model_cache(model_name) as model:
+        assert model.model.tokenizer.padding["length"] is None
+
+        embeddings = np.stack(list(model.embed(["hello world", "retrieval " * 200])), axis=0)
+        assert embeddings.shape == (2, dim)

tests/utils.py

@@ -3,12 +3,12 @@ import traceback
 
 from pathlib import Path
 from types import TracebackType
-from typing import Union, Callable, Any, Type, Optional
+from typing import Callable, Any, Type
 
 from fastembed.common.model_description import BaseModelDescription
 
 
-def delete_model_cache(model_dir: Union[str, Path]) -> None:
+def delete_model_cache(model_dir: str | Path) -> None:
     """Delete the model cache directory.
 
     If a model was downloaded from the HuggingFace model hub, then _model_dir is the dir to snapshots, removing
@@ -42,7 +42,7 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
 def should_test_model(
     model_desc: BaseModelDescription,
     autotest_model_name: str,
-    is_ci: Optional[str],
+    is_ci: str | None,
     is_manual: bool,
 ):
     """Determine if a model should be tested based on environment