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

{fastembed_gpu-0.8.1 → fastembed_gpu-0.9.0}/NOTICE

@@ -12,8 +12,8 @@ This distribution includes the following Jina AI models, each with its respectiv
 
 These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
 
-This distribution includes the following Google models, each with its respective license:
-- vidore/colpali-v1.3
+This distribution includes the following Gemma-based models, each with its respective license:
+- Qdrant/colpali-v1.3-fp16
     - License: gemma
 - google/embeddinggemma-300m
     - License: gemma

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fastembed-gpu
-Version: 0.8.1
+Version: 0.9.0
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache-2.0
 License-File: LICENSE
@@ -211,10 +211,10 @@ query = "What is Qdrant?"
 model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
 doc_images_embeddings = list(model.embed_image(doc_images))
 # shape (2, 1030, 128)
-# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
+# [array([[-0.03354, -0.0209, ..., -0.1558, -0.0768]], dtype=float16)]
 query_embedding = model.embed_text(query)
 # shape (1, 20, 128)
-# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
+# [array([[-0.002182, 0.1476, ..., -0.02208, 0.1683]], dtype=float16)]
 ```
 
 ### 🔄 Rerankers

{fastembed_gpu-0.8.1 → fastembed_gpu-0.9.0}/README.md

@@ -169,10 +169,10 @@ query = "What is Qdrant?"
 model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
 doc_images_embeddings = list(model.embed_image(doc_images))
 # shape (2, 1030, 128)
-# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
+# [array([[-0.03354, -0.0209, ..., -0.1558, -0.0768]], dtype=float16)]
 query_embedding = model.embed_text(query)
 # shape (1, 20, 128)
-# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
+# [array([[-0.002182, 0.1476, ..., -0.02208, 0.1683]], dtype=float16)]
 ```
 
 ### 🔄 Rerankers

{fastembed_gpu-0.8.1 → fastembed_gpu-0.9.0}/fastembed/common/model_management.py

@@ -5,15 +5,18 @@ import json
 import shutil
 import tarfile
 import tempfile
+import warnings
 import contextlib
 from copy import deepcopy
 from pathlib import Path, PureWindowsPath
 from typing import Any, TypeVar, Generic
 
 import requests
-from huggingface_hub import snapshot_download, model_info, list_repo_tree
+from huggingface_hub import constants, snapshot_download, model_info, list_repo_tree
+from huggingface_hub.file_download import repo_folder_name
 from huggingface_hub.hf_api import RepoFile
 from huggingface_hub.utils import (
+    HFValidationError,
     RepositoryNotFoundError,
     disable_progress_bars,
     enable_progress_bars,
@@ -21,12 +24,44 @@ from huggingface_hub.utils import (
 from loguru import logger
 from tqdm import tqdm
 from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.onnx_external_data import link_external_data
 
 T = TypeVar("T", bound=BaseModelDescription)
 
 _DOWNLOAD_CHUNK_SIZE = 256 * 1024
 
 
+def _hf_transport_errors() -> tuple[type[Exception], ...]:
+    """Network errors of huggingface_hub's HTTP library that aren't OSError.
+
+    A refused connection, a DNS failure or a timeout raises an OSError in huggingface_hub 0.x,
+    which is built on requests, but a TransportError in 1.x (httpx) and 2.x (httpx2).
+    """
+    try:
+        # huggingface_hub>=1.30 re-exports whichever of httpx and httpx2 it's built on.
+        from huggingface_hub.utils import httpx
+    except ImportError:
+        try:
+            import httpx  # huggingface_hub 1.0 to 1.29
+        except ImportError:  # huggingface_hub 0.x
+            return ()
+    return (httpx.TransportError,)
+
+
+# Errors from an HF download that download_model handles by falling back to url and retrying.
+_HF_DOWNLOAD_ERRORS = (OSError, RepositoryNotFoundError, ValueError) + _hf_transport_errors()
+
+
+class CorruptedCacheError(ValueError):
+    """A cached model file is present but does not match the size recorded for it.
+
+    Distinct from the download failures above so that callers can tell "this snapshot is
+    unusable, look elsewhere" apart from "the hub could not be reached": only the former
+    is worth re-fetching a file the hub's existence-only cache check would otherwise keep.
+    Subclasses ValueError to stay inside _HF_DOWNLOAD_ERRORS for callers that don't care.
+    """
+
+
 class ModelManagement(Generic[T]):
     METADATA_FILE = "files_metadata.json"
 
@@ -135,6 +170,41 @@ class ModelManagement(Generic[T]):
                         file.write(chunk)
         return output_path
 
+    @classmethod
+    def _find_legacy_cased_source(cls, cache_dir: str, hf_source_repo: str) -> str | None:
+        """Looks for a cached snapshot of the same repo spelled with a different casing.
+
+        Built-in sources used to be lowercase and were canonicalized once it turned out that
+        relying on the hub's normalizing redirect breaks proxies. Both the hub and fastembed
+        derive the cache directory from the source verbatim, so on a case-sensitive filesystem
+        an offline load would otherwise miss a model an older version had already cached.
+
+        Args:
+            cache_dir (str): The path to the cache directory.
+            hf_source_repo (str): Name of the model on HuggingFace Hub.
+
+        Returns:
+            str | None: The differently cased source found in the cache, None if there is none.
+        """
+        separator = constants.REPO_ID_SEPARATOR
+        try:
+            expected = repo_folder_name(repo_id=hf_source_repo, repo_type="model")
+            entries = list(Path(cache_dir).iterdir())
+        except (HFValidationError, OSError):
+            return None
+
+        # the repo id sits at the tail of the folder name, with every "/" replaced
+        offset = len(expected) - len(hf_source_repo.replace("/", separator))
+
+        for entry in entries:
+            if entry.name == expected or entry.name.lower() != expected.lower():
+                continue
+            if not entry.is_dir():
+                continue
+            # the hub forbids the separator inside a repo id, so it only marks the split
+            return entry.name[offset:].replace(separator, "/")
+        return None
+
     @classmethod
     def download_files_from_huggingface(
         cls,
@@ -147,7 +217,7 @@ class ModelManagement(Generic[T]):
         """
         Downloads a model from HuggingFace Hub.
         Args:
-            hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
+            hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "Qdrant/all-MiniLM-L6-v2-onnx".
             cache_dir (Optional[str]): The path to the cache directory.
             extra_patterns (list[str]): extra patterns to allow in the snapshot download, typically
                 includes the required model files.
@@ -156,6 +226,15 @@ class ModelManagement(Generic[T]):
             Path: The path to the model directory.
         """
 
+        def _repo_relative_path(relative_path: Path) -> str:
+            # Cached files are stored under `snapshots/<revision>/<repo_path>`. Strip that
+            # prefix so the remainder matches a repo file's path, which may itself be nested
+            # (e.g. `onnx/model.onnx`).
+            parts = relative_path.parts
+            if len(parts) > 2 and parts[0] == "snapshots":
+                return "/".join(parts[2:])
+            return relative_path.as_posix()
+
         def _verify_files_from_metadata(
             model_dir: Path, stored_metadata: dict[str, Any], repo_files: list[RepoFile]
         ) -> bool:
@@ -167,7 +246,8 @@ class ModelManagement(Generic[T]):
                         return False
 
                     if repo_files:  # online verification
-                        file_info = next((f for f in repo_files if f.path == file_path.name), None)
+                        repo_path = _repo_relative_path(Path(rel_path))
+                        file_info = next((f for f in repo_files if f.path == repo_path), None)
                         if (
                             not file_info
                             or file_info.size != meta["size"]
@@ -184,15 +264,19 @@ class ModelManagement(Generic[T]):
                 return False
 
         def _collect_file_metadata(
-            model_dir: Path, repo_files: list[RepoFile]
+            model_dir: Path, revision: str, repo_files: list[RepoFile]
         ) -> 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("*"):
+            # Only look at the downloaded revision's folder. The cache may also hold folders of
+            # older revisions, whose files differ from the ones on the hub now, so checking
+            # them would report a good download as corrupted.
+            for file_path in (model_dir / "snapshots" / revision).rglob("*"):
                 if file_path.is_file() and file_path.name != cls.METADATA_FILE:
-                    repo_file = file_info_map.get(file_path.name)
+                    relative_path = file_path.relative_to(model_dir)
+                    repo_file = file_info_map.get(_repo_relative_path(relative_path))
                     if repo_file:
-                        meta[str(file_path.relative_to(model_dir))] = {
+                        meta[str(relative_path)] = {
                             "size": repo_file.size,
                             "blob_id": repo_file.blob_id,
                         }
@@ -216,31 +300,96 @@ class ModelManagement(Generic[T]):
 
         allow_patterns.extend(extra_patterns)
 
-        snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
+        snapshot_dir = Path(cache_dir) / repo_folder_name(
+            repo_id=hf_source_repo, repo_type="model"
+        )
         metadata_file = snapshot_dir / cls.METADATA_FILE
 
         if local_files_only:
             disable_progress_bars()
-            if metadata_file.exists():
-                metadata = json.loads(metadata_file.read_text())
-                verified = _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[])
-                if not verified:
-                    logger.warning(
-                        "Local file sizes do not match the metadata."
-                    )  # do not raise, still make an attempt to load the model
-            result = snapshot_download(
-                repo_id=hf_source_repo,
-                allow_patterns=allow_patterns,
-                cache_dir=cache_dir,
-                local_files_only=local_files_only,
-                **kwargs,
-            )
-            return result
+            # a canonical directory can exist yet hold no usable snapshot, e.g. when an
+            # earlier download was interrupted, so fall back on failure rather than on
+            # the directory being absent
+            sources = [hf_source_repo]
+            legacy_source = cls._find_legacy_cased_source(cache_dir, hf_source_repo)
+            if legacy_source is not None:
+                sources.append(legacy_source)
+
+            # a corrupt candidate is remembered rather than raised on the spot: the next
+            # candidate may still be loadable, and if none is, this is the error worth
+            # reporting, since it is the only one download_model can act on.
+            corruption: CorruptedCacheError | None = None
+
+            for index, source in enumerate(sources):
+                last_source = index == len(sources) - 1
+                snapshot_dir = Path(cache_dir) / repo_folder_name(
+                    repo_id=source, repo_type="model"
+                )
+                metadata_file = snapshot_dir / cls.METADATA_FILE
+                if metadata_file.exists():
+                    metadata = json.loads(metadata_file.read_text())
+                    if not _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[]):
+                        # A size mismatch on one of this model's own files (its weights, or
+                        # whatever else its description lists) means loading would later fail
+                        # with a cryptic error, so the snapshot is not offered at all. A
+                        # mismatch confined to a file nothing asked for is left to best-effort
+                        # loading, preserving the previous behavior.
+                        model_files = set(extra_patterns)
+                        model_file_corrupt = any(
+                            _repo_relative_path(Path(rel_path)) in model_files
+                            and (snapshot_dir / rel_path).is_file()
+                            and (snapshot_dir / rel_path).stat().st_size != meta["size"]
+                            for rel_path, meta in metadata.items()
+                        )
+                        if model_file_corrupt:
+                            corruption = corruption or CorruptedCacheError(
+                                f"Cached model files for {source} in {cache_dir} do not match "
+                                "the stored metadata; the cache appears corrupted."
+                            )
+                            if last_source:
+                                raise corruption
+                            logger.info(
+                                f"{source} is corrupted in {cache_dir}, loading "
+                                f"{sources[index + 1]}, cached from the same repo by an older "
+                                "fastembed version."
+                            )
+                            continue
+                        logger.warning("Local file sizes do not match the metadata.")
+                try:
+                    # a legacy source is only ever resolved against the cache: sending it
+                    # to the hub would ask for the very redirect this casing avoids
+                    return snapshot_download(
+                        repo_id=source,
+                        allow_patterns=allow_patterns,
+                        cache_dir=cache_dir,
+                        local_files_only=True,
+                        **kwargs,
+                    )
+                except _HF_DOWNLOAD_ERRORS:
+                    if last_source:
+                        # an earlier candidate was corrupt and this one is merely absent:
+                        # report the corruption, the only failure a re-download can fix
+                        if corruption is not None:
+                            raise corruption
+                        raise
+                    logger.info(
+                        f"{source} is not usable in {cache_dir}, loading "
+                        f"{sources[index + 1]}, cached from the same repo by an older "
+                        "fastembed version."
+                    )
 
-        repo_revision = model_info(hf_source_repo).sha
-        repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model"))
+        # hub sends this request with no timeout unless given one, so an endpoint that accepts
+        # the connection but never answers would block here for good, before download_model can
+        # retry or fall back to another source. list_repo_tree and snapshot_download's own
+        # metadata requests can't be given one, but a silent endpoint now fails here first.
+        repo_revision = model_info(hf_source_repo, timeout=constants.HF_HUB_ETAG_TIMEOUT).sha
+        repo_tree = list(
+            list_repo_tree(
+                hf_source_repo, revision=repo_revision, repo_type="model", recursive=True
+            )
+        )
 
-        allowed_extensions = {".json", ".onnx", ".txt"}
+        allowed_extensions = {Path(pattern).suffix for pattern in allow_patterns}
         repo_files = (
             [
                 f
@@ -253,9 +402,18 @@ class ModelManagement(Generic[T]):
 
         verified_metadata = False
 
-        if snapshot_dir.exists() and metadata_file.exists():
+        # a forced download re-fetches every file, so show progress and re-collect the metadata
+        if not kwargs.get("force_download") and snapshot_dir.exists() and metadata_file.exists():
             metadata = json.loads(metadata_file.read_text())
-            verified_metadata = _verify_files_from_metadata(snapshot_dir, metadata, repo_files)
+            # metadata lacking a requested file, e.g. a new model_file, can't vouch for it
+            requested_files = {f.path for f in repo_files if f.path in allow_patterns}
+            stored_files = {_repo_relative_path(Path(rel_path)) for rel_path in metadata}
+            # empty metadata lists no files, so it can't vouch for the cached ones
+            verified_metadata = (
+                bool(metadata)
+                and requested_files.issubset(stored_files)
+                and _verify_files_from_metadata(snapshot_dir, metadata, repo_files)
+            )
 
         if verified_metadata:
             disable_progress_bars()
@@ -272,7 +430,10 @@ class ModelManagement(Generic[T]):
             not verified_metadata
         ):  # metadata is not up-to-date, update it and check whether the files have been
             # downloaded correctly
-            metadata = _collect_file_metadata(snapshot_dir, repo_files)
+            # result is <cache_dir>/.../snapshots/<revision>. Only the revision is taken from
+            # it: the hub resolves cache_dir, so the rest of the path may not match ours.
+            downloaded_revision = Path(result).name
+            metadata = _collect_file_metadata(snapshot_dir, downloaded_revision, repo_files)
 
             download_successful = _verify_files_from_metadata(
                 snapshot_dir, metadata, repo_files=[]
@@ -282,7 +443,10 @@ class ModelManagement(Generic[T]):
                     "Files have been corrupted during downloading process. "
                     "Please check your internet connection and try again."
                 )
-            _save_file_metadata(snapshot_dir, metadata)
+            # Empty means no file was checked, e.g. the files went to a local_dir instead of
+            # the cache. Saving it would make later downloads skip this check.
+            if metadata:
+                _save_file_metadata(snapshot_dir, metadata)
 
         return result
 
@@ -374,6 +538,20 @@ class ModelManagement(Generic[T]):
 
         # 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:
+            if deprecated_tar_struct and not source_url:
+                # No built-in model is served from the bucket anymore, only copies of it remain.
+                # A model with a url of its own (a custom one) got its copy from there instead.
+                # stacklevel points at the caller of TextEmbedding(...) via download_model.
+                warnings.warn(
+                    f"Loading {model_name} from {model_dir}, a copy downloaded from Google Cloud "
+                    "Storage by an older fastembed version. Support for such copies is deprecated "
+                    "and will be removed in a future release. To switch to Hugging Face, load the "
+                    "model once while huggingface.co (or an `HF_ENDPOINT` mirror) is reachable, "
+                    "without `local_files_only=True` or `HF_HUB_OFFLINE=1`, then delete "
+                    f"{model_dir}.",
+                    FutureWarning,
+                    stacklevel=5,
+                )
             return model_dir
 
         if local_files_only:
@@ -425,6 +603,19 @@ class ModelManagement(Generic[T]):
 
         return model_dir
 
+    @staticmethod
+    def _link_onnx_external_data(model: BaseModelDescription, model_dir: Path) -> Path:
+        """Links the external data of `model` right after its download, see link_external_data.
+
+        Loading the model links it as well, but linking it here also covers lazy_load=True, e.g.
+        in a Docker build step, where links made by a later step would copy the files into a new
+        image layer. A failure, e.g. in a read-only cache, is reported when the model is loaded.
+        """
+        if model.model_file.endswith(".onnx"):
+            with contextlib.suppress(OSError):
+                link_external_data(model_dir, model.model_file, model.additional_files)
+        return model_dir
+
     @classmethod
     def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: Any) -> Path:
         """
@@ -440,8 +631,7 @@ class ModelManagement(Generic[T]):
                     "description": "Base English model, v1.5",
                     "size_in_GB": 0.44,
                     "sources": {
-                        "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
-                        "hf": "qdrant/bge-base-en-v1.5-onnx-q",
+                        "hf": "Qdrant/bge-base-en-v1.5-onnx-Q",
                     }
                 }
                 ```
@@ -466,6 +656,12 @@ class ModelManagement(Generic[T]):
         extra_patterns = [model.model_file]
         extra_patterns.extend(model.additional_files)
 
+        # Only a cached file that is present but does not match its recorded size needs
+        # force_download: the hub's cache check is existence-only, so a plain snapshot_download
+        # would keep the truncated blob. Every other probe failure (a missing file, an absent
+        # cache) is served by an ordinary download, which re-uses what is already on disk.
+        force_download = False
+
         if hf_source:
             try:
                 cache_kwargs = deepcopy(kwargs)
@@ -481,7 +677,9 @@ class ModelManagement(Generic[T]):
                 if (resolved_path / model.model_file).exists() and all(
                     (resolved_path / file).exists() for file in extra_patterns
                 ):
-                    return resolved_path
+                    return cls._link_onnx_external_data(model, resolved_path)
+            except CorruptedCacheError:
+                force_download = True
             except Exception:
                 pass
             finally:
@@ -493,28 +691,35 @@ class ModelManagement(Generic[T]):
 
             if hf_source and not local_files_only:
                 # we have already tried loading with `local_files_only=True` via hf and we failed
+
+                # Merged into kwargs so a caller-supplied force_download can't raise a
+                # duplicate keyword error. Cleared straight away: force_download deletes the
+                # partially fetched `.incomplete` file, so leaving it on would make every
+                # retry restart a large download from zero instead of resuming it.
+                attempt_kwargs = {**kwargs, "force_download": True} if force_download else kwargs
+                force_download = False
                 try:
-                    return Path(
+                    model_dir = Path(
                         cls.download_files_from_huggingface(
                             hf_source,
                             cache_dir=cache_dir,
                             extra_patterns=extra_patterns,
-                            **kwargs,
+                            **attempt_kwargs,
                         )
                     )
-                except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
-                    if not local_files_only:
-                        logger.error(
-                            f"Could not download model from HuggingFace: {e} "
-                            "Falling back to other sources."
-                        )
+                    return cls._link_onnx_external_data(model, model_dir)
+                except _HF_DOWNLOAD_ERRORS as e:
+                    logger.error(
+                        f"Could not download model from HuggingFace: {e} "
+                        "Falling back to other sources."
+                    )
                 finally:
                     enable_progress_bars()
             if url_source or local_files_only:
                 try:
                     return cls.retrieve_model_gcs(
                         model.model,
-                        str(url_source),
+                        url_source or "",
                         str(cache_dir),
                         deprecated_tar_struct=model.sources.deprecated_tar_struct,
                         local_files_only=local_files_only,
@@ -522,6 +727,17 @@ class ModelManagement(Generic[T]):
                 except Exception:
                     if not local_files_only:
                         logger.error(f"Could not download model from url: {url_source}")
+            elif model.sources.deprecated_tar_struct:
+                # Nothing is downloaded from the bucket anymore, but an old copy may still be cached.
+                legacy_dir = Path(cache_dir) / f"fast-{model.model.split('/')[-1]}"
+                if legacy_dir.is_dir() and any(legacy_dir.iterdir()):
+                    return cls.retrieve_model_gcs(
+                        model.model,
+                        "",
+                        str(cache_dir),
+                        deprecated_tar_struct=True,
+                        local_files_only=True,
+                    )
 
             if local_files_only:
                 logger.error("Could not find model in cache_dir")

fastembed_gpu-0.9.0/fastembed/common/onnx_external_data.py

@@ -0,0 +1,91 @@
+"""Loading ONNX models with external data from a huggingface_hub cache.
+
+onnxruntime>=1.24 refuses external data that, once symlinks are resolved, is located outside of
+the directory of the model file, and of the directory the model file resolves to (a fallback of
+1.24.2, pyproject.toml excludes 1.24.0 and 1.24.1). Since huggingface_hub 1.32, a snapshot is made
+of symlinks into a blob store shared by the whole cache and sharded by hash, where a model and its
+data may resolve into different directories. Then the model and its data are hardlinked into
+`ONNX_SNAPSHOTS_DIR` of the repo cache, laid out as in the snapshot, which takes no extra space.
+"""
+
+import contextlib
+import os
+import shutil
+from pathlib import Path
+
+# Next to `snapshots` in the cache of a huggingface_hub repo: `onnx_snapshots/<revision>/...`
+ONNX_SNAPSHOTS_DIR = "onnx_snapshots"
+
+
+def link_external_data(model_dir: Path, model_file: str, additional_files: list[str]) -> Path:
+    """Returns a path of `model_file` from which onnxruntime can load its external data.
+
+    Args:
+        model_dir (Path): The directory with the model files, e.g. a huggingface_hub snapshot.
+        model_file (str): The path of the ONNX file, relative to `model_dir`.
+        additional_files (list[str]): Other files of the model, relative to `model_dir`,
+            among which its external data.
+
+    Returns:
+        Path: The ONNX file to load, `model_dir / model_file` unless it had to be linked.
+
+    Raises:
+        OSError: If the files had to be linked, but couldn't be, e.g. in a read-only cache or
+            on a filesystem without hardlinks.
+    """
+    model_path = model_dir / model_file
+    snapshots_dir = model_dir.parent
+    repo_dir = snapshots_dir.parent
+    if snapshots_dir.name != "snapshots" or not repo_dir.name.startswith("models--"):
+        return model_path  # not a huggingface_hub cache, there is no place of ours to link into
+    if not model_path.exists():
+        return model_path  # onnxruntime reports it more clearly than a failed hardlink would
+
+    # external data is located relative to the model file, so it can't be anywhere else
+    data_paths = [
+        path
+        for path in (model_dir / file for file in additional_files)
+        if path.parent.is_relative_to(model_path.parent) and path.exists()
+    ]
+    # onnxruntime accepts data within these, as in the caches of older huggingface_hub versions
+    real_model_dirs = (
+        os.path.realpath(model_path.parent),
+        os.path.dirname(os.path.realpath(model_path)),
+    )
+    if all(
+        any(Path(os.path.realpath(path)).is_relative_to(d) for d in real_model_dirs)
+        for path in data_paths
+    ):
+        return model_path
+
+    links_dir = repo_dir / ONNX_SNAPSHOTS_DIR
+    for path in (model_path, *data_paths):
+        _link_file(path, links_dir / model_dir.name / path.relative_to(model_dir))
+
+    # the links keep the blobs on disk, so drop those of revisions deleted from the cache
+    for revision_dir in links_dir.iterdir():
+        if not (snapshots_dir / revision_dir.name).is_dir():
+            shutil.rmtree(revision_dir, ignore_errors=True)
+
+    return links_dir / model_dir.name / model_file
+
+
+def _link_file(source: Path, link: Path) -> None:
+    """Makes `link` a hardlink to the file `source` resolves to, unless it already is one."""
+    target = os.path.realpath(source)
+    try:
+        if os.path.samefile(link, target):
+            return
+        # a link to a blob downloaded again since, or a copy of the cache without its hardlinks
+        link.unlink()
+    except FileNotFoundError:
+        pass
+    except OSError:
+        # e.g. such a copy on a read-only filesystem, which works all the same
+        if os.path.getsize(link) != os.path.getsize(target):
+            raise
+        return
+    link.parent.mkdir(parents=True, exist_ok=True)
+    # os.link is atomic, so if the link exists, another process loading the model just made it
+    with contextlib.suppress(FileExistsError):
+        os.link(target, link)

{fastembed_gpu-0.8.1 → fastembed_gpu-0.9.0}/fastembed/common/onnx_model.py

@@ -6,9 +6,11 @@ from typing import Any, Generic, Iterable, Sequence, Type, TypeVar
 import numpy as np
 import onnxruntime as ort
 
+from loguru import logger
 from numpy.typing import NDArray
 from tokenizers import Tokenizer
 
+from fastembed.common.onnx_external_data import ONNX_SNAPSHOTS_DIR, link_external_data
 from fastembed.common.types import OnnxProvider, NumpyArray, Device
 from fastembed.parallel_processor import Worker
 
@@ -26,6 +28,8 @@ class OnnxOutputContext:
 
 class OnnxModel(Generic[T]):
     EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)
+    # set by the concrete model classes in their __init__
+    _model_dir: Path
 
     @classmethod
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
@@ -76,8 +80,8 @@ class OnnxModel(Generic[T]):
         cuda: bool | Device = Device.AUTO,
         device_id: int | None = None,
         extra_session_options: dict[str, Any] | None = None,
+        additional_files: list[str] | 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
@@ -122,9 +126,29 @@ class OnnxModel(Generic[T]):
         if extra_session_options is not None:
             self.add_extra_session_options(so, extra_session_options)
 
-        self.model = ort.InferenceSession(
-            str(model_path), providers=onnx_providers, sess_options=so
-        )
+        model_path = model_dir / model_file
+        link_error: OSError | None = None
+        try:
+            model_path = link_external_data(model_dir, model_file, additional_files or [])
+        except OSError as e:
+            # e.g. a read-only cache, which onnxruntime<1.24 loads from all the same
+            link_error = e
+
+        try:
+            self.model = ort.InferenceSession(
+                str(model_path), providers=onnx_providers, sess_options=so
+            )
+        except Exception:
+            if link_error is not None:
+                logger.warning(
+                    f"Could not link the files of {model_path} into {ONNX_SNAPSHOTS_DIR}: "
+                    f"{link_error}. onnxruntime>=1.24 refuses external data that resolves "
+                    "outside of the model directory, as in a huggingface_hub>=1.32 cache. "
+                    "Download the model with fastembed into a writable cache_dir on a "
+                    "filesystem with hardlinks, or delete it from the cache and download it "
+                    "again with HF_HUB_DISABLE_SHARED_BLOBS=1."
+                )
+            raise
         if "CUDAExecutionProvider" in requested_provider_names:
             assert self.model is not None
             current_providers = self.model.get_providers()