fix: re-download model when cached files fail verification

fastembed/common/model_management.py

@@ -151,6 +151,15 @@ def download_files_from_huggingface(
             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:
@@ -162,7 +171,8 @@ def _verify_files_from_metadata(
                         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"]
@@ -185,9 +195,10 @@ def _collect_file_metadata(
             file_info_map = {f.path: f for f in repo_files}
             for file_path in model_dir.rglob("*"):
                 if file_path.is_file() and file_path.name != cls.METADATA_FILE:
-                    repo_file = file_info_map.get(file_path.name)
+                    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,
                         }
@@ -218,11 +229,26 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]])
             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
+                if not _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[]):
+                    # A size mismatch on a model file means ONNX Runtime would later fail
+                    # with a cryptic protobuf error. Raise so download_model's offline probe
+                    # (which wraps this call in `except Exception`) falls through to a forced
+                    # online re-download instead of returning a corrupt cached path. A
+                    # mismatch limited to an auxiliary file (e.g. a config) 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:
+                        raise ValueError(
+                            "Cached model files do not match the stored metadata; "
+                            "the cache appears corrupted."
+                        )
+                    logger.warning("Local file sizes do not match the metadata.")
             result = snapshot_download(
                 repo_id=hf_source_repo,
                 allow_patterns=allow_patterns,
@@ -233,9 +259,13 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]])
             return result
 
         repo_revision = model_info(hf_source_repo).sha
-        repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model"))
+        repo_tree = list(
+            list_repo_tree(
+                hf_source_repo, revision=repo_revision, repo_type="model", recursive=True
+            )
+        )
 
-        allowed_extensions = {".json", ".onnx", ".txt"}
+        allowed_extensions = {".json", ".onnx", ".txt", ".onnx_data", ".npy", ".vocab"}
         repo_files = (
             [
                 f
@@ -442,7 +472,12 @@ def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: An
                             hf_source,
                             cache_dir=cache_dir,
                             extra_patterns=extra_patterns,
-                            **kwargs,
+                            # The local probe failed (cache missing or corrupt). force_download so
+                            # huggingface_hub re-fetches a present-but-truncated blob instead of
+                            # trusting it (its cache short-circuit is existence-only). Merge into
+                            # kwargs so a caller-supplied force_download can't raise a duplicate
+                            # keyword error.
+                            **{**kwargs, "force_download": True},
                         )
                     )
                 except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:

tests/test_model_management.py

@@ -0,0 +1,92 @@
+"""Offline cache-verification tests for ModelManagement.
+
+The offline probe must distinguish a corrupt *model* file (which would make ONNX Runtime
+fail with a cryptic protobuf error) from a benign size drift on an auxiliary file. A corrupt
+model file must raise so ``download_model`` falls through to a forced re-download; an auxiliary
+mismatch must be tolerated so an offline caller is not bricked by best-effort-loadable drift.
+Model files may live in a subdirectory (e.g. ``onnx/model.onnx``), so matching is done on the
+repo-relative path, not the bare filename.
+"""
+
+import json
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from fastembed.common import model_management
+from fastembed.common.model_management import ModelManagement
+
+REVISION = "0123456789abcdef0123456789abcdef01234567"
+
+
+def _seed_cache(cache: Path, files: dict[str, bytes], metadata: dict[str, Any]) -> Path:
+    snapshot = cache / "models--qdrant--fake-onnx"
+    for name, blob in files.items():
+        path = snapshot / name
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(blob)
+    snapshot.mkdir(parents=True, exist_ok=True)
+    (snapshot / ModelManagement.METADATA_FILE).write_text(json.dumps(metadata))
+    return snapshot
+
+
+def test_offline_probe_raises_when_model_file_is_corrupt(tmp_path: Path) -> None:
+    cache = tmp_path / "cache"
+    _seed_cache(
+        cache,
+        files={"model.onnx": b"x" * 10},
+        metadata={"model.onnx": {"size": 999_999, "blob_id": "deadbeef"}},
+    )
+
+    with pytest.raises(ValueError, match="corrupt"):
+        ModelManagement.download_files_from_huggingface(
+            "qdrant/fake-onnx",
+            cache_dir=str(cache),
+            extra_patterns=["model.onnx"],
+            local_files_only=True,
+        )
+
+
+def test_offline_probe_raises_when_subdir_model_file_is_corrupt(tmp_path: Path) -> None:
+    cache = tmp_path / "cache"
+    key = f"snapshots/{REVISION}/onnx/model.onnx"
+    _seed_cache(
+        cache,
+        files={key: b"x" * 10},
+        metadata={key: {"size": 999_999, "blob_id": "deadbeef"}},
+    )
+
+    with pytest.raises(ValueError, match="corrupt"):
+        ModelManagement.download_files_from_huggingface(
+            "qdrant/fake-onnx",
+            cache_dir=str(cache),
+            extra_patterns=["onnx/model.onnx"],
+            local_files_only=True,
+        )
+
+
+def test_offline_probe_tolerates_auxiliary_file_drift(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    cache = tmp_path / "cache"
+    snapshot = _seed_cache(
+        cache,
+        files={"model.onnx": b"x" * 10, "config.json": b"y" * 20},
+        metadata={
+            "model.onnx": {"size": 10, "blob_id": "modelblob"},
+            "config.json": {"size": 999, "blob_id": "configblob"},
+        },
+    )
+
+    # The model file matches metadata; only config.json drifted. The probe must NOT raise:
+    # it falls through to snapshot_download (mocked here to return the cached path).
+    monkeypatch.setattr(model_management, "snapshot_download", lambda **kwargs: str(snapshot))
+
+    result = ModelManagement.download_files_from_huggingface(
+        "qdrant/fake-onnx",
+        cache_dir=str(cache),
+        extra_patterns=["model.onnx"],
+        local_files_only=True,
+    )
+    assert result == str(snapshot)