fastembed-gpu 0.7.4__tar.gz → 0.8.1__tar.gz

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.1}/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.

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.1}/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
 

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.1}/PKG-INFO

@@ -1,37 +1,37 @@
 Metadata-Version: 2.4
 Name: fastembed-gpu
-Version: 0.7.4
+Version: 0.8.1
 Summary: Fast, light, accurate library built for retrieval embedding generation
-License: Apache License
+License: Apache-2.0
 License-File: LICENSE
 License-File: NOTICE
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
 Author: Qdrant Team
 Author-email: info@qdrant.tech
-Requires-Python: >=3.9.0
-Classifier: License :: Other/Proprietary License
+Requires-Python: >=3.10.0
+Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
 Requires-Dist: huggingface-hub (>=0.20,<2.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: mmh3 (>=4.1.0,<6.0.0)
 Requires-Dist: numpy (>=1.21) ; python_version == "3.11"
-Requires-Dist: numpy (>=1.21,<2.1.0) ; python_version < "3.10"
 Requires-Dist: numpy (>=1.21,<2.3.0) ; python_version == "3.10"
 Requires-Dist: numpy (>=1.26) ; python_version == "3.12"
 Requires-Dist: numpy (>=2.1.0) ; python_version == "3.13"
 Requires-Dist: numpy (>=2.3.0) ; python_version >= "3.14"
-Requires-Dist: onnxruntime-gpu (>1.20.0) ; python_version >= "3.13"
-Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0) ; python_version >= "3.10" and python_version < "3.13"
-Requires-Dist: onnxruntime-gpu (>=1.17.0,<1.20.0) ; python_version < "3.10"
-Requires-Dist: pillow (>=10.3.0,<11.0) ; python_version < "3.10"
-Requires-Dist: pillow (>=10.3.0,<12.0) ; python_version >= "3.10" and python_version < "3.13"
-Requires-Dist: pillow (>=11.0.0,<12.0) ; python_version >= "3.13"
+Requires-Dist: onnxruntime-gpu (>1.21.0,!=1.24.0,!=1.24.1) ; python_version == "3.13"
+Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0,!=1.24.0,!=1.24.1) ; python_version >= "3.11" and python_version < "3.13"
+Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0,<1.24) ; python_version == "3.10"
+Requires-Dist: onnxruntime-gpu (>=1.24.2) ; python_version >= "3.14"
+Requires-Dist: pillow (>=10.3.0,<13.0) ; python_version >= "3.10" and python_version < "3.13"
+Requires-Dist: pillow (>=11.0.0,<13.0) ; python_version == "3.13"
+Requires-Dist: pillow (>=12.0.0,<13.0) ; python_version >= "3.14"
 Requires-Dist: py-rust-stemmers (>=0.1.0,<0.2.0)
 Requires-Dist: requests (>=2.31,<3.0)
 Requires-Dist: tokenizers (>=0.15,<1.0)

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.1}/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_gpu-0.7.4 → fastembed_gpu-0.8.1}/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,42 +370,58 @@ 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)
+        if local_files_only:
+            logger.error(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
+            raise ValueError(
+                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)
 
-        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:
+        # 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(cache_tmp_dir))
-            assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+            cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(staging_dir))
 
-            model_tar_gz.unlink()
-            # Rename from tmp to final name is atomic
-            model_tmp_dir.rename(model_dir)
-        else:
-            logger.error(
-                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
-            )
-            raise ValueError(
-                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
-            )
+            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
 
@@ -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_gpu-0.7.4 → fastembed_gpu-0.8.1}/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_gpu-0.8.1/fastembed/common/preprocessor_utils.py

@@ -0,0 +1,151 @@
+import json
+import sys
+from typing import Any, Iterator
+from pathlib import Path
+
+from tokenizers import AddedToken, Tokenizer
+
+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():
+        return {}
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    return tokens_map
+
+
+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}")
+
+    tokenizer_config_path = model_dir / "tokenizer_config.json"
+    if not tokenizer_config_path.exists():
+        raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
+
+    # 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)
+
+    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)
+
+    # 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)])
+
+    # 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}")
+
+    # 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
+
+
+def load_preprocessor(model_dir: Path) -> Compose:
+    preprocessor_config_path = model_dir / "preprocessor_config.json"
+    if not preprocessor_config_path.exists():
+        raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
+
+    with open(str(preprocessor_config_path)) as preprocessor_config_file:
+        preprocessor_config = json.load(preprocessor_config_file)
+        transforms = Compose.from_config(preprocessor_config)
+    return transforms

fastembed_gpu-0.8.1/fastembed/common/types.py

@@ -0,0 +1,27 @@
+from enum import Enum
+from pathlib import Path
+from typing import Any, TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
+from PIL import Image
+
+
+class Device(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+    AUTO = "auto"
+
+
+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]
+)