Download an embedding model whose cached snapshot lacks its graph

The chunker caches only a model's tokenizer.json in the embeddings cache.
FastEmbed counts any cached snapshot as the model, so it never downloaded
the ONNX graph and every load failed with NO_SUCHFILE. The loader now
checks the files FastEmbed needs and fetches them first when they are
missing; offline it leaves them for FastEmbed to report.

docsgpt/vectorstore/embeddings_local.py

@@ -20,6 +20,7 @@ from the repository itself rather than assumed.
 
 import json
 import logging
+import os
 import threading
 from dataclasses import replace
 from typing import Any, List, Optional
@@ -39,6 +40,16 @@ _register_lock = threading.Lock()
 _FALLBACK_ONNX_FILE = "onnx/model.onnx"
 _FALLBACK_POOLING = "mean"
 
+# The files FastEmbed fetches with every model besides its graph (see
+# ``ModelManagement.download_files_from_huggingface``).
+_FASTEMBED_SUPPORT_FILES = (
+    "config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "preprocessor_config.json",
+)
+
 # Sentence-transformers records how a model turns token vectors into one
 # vector, and whether it normalises the result, as files in the repository.
 # Reading them is the difference between running a model and running something
@@ -244,6 +255,47 @@ def _spec_for(model_name: str) -> EmbeddingModel:
     )
 
 
+def _complete_model_cache(repo: str, cache_dir: Optional[str]) -> None:
+    """Download a model's graph when its cached snapshot lacks it.
+
+    FastEmbed first loads a model from the cache and only downloads when that
+    fails, and a snapshot of the repository counts as cached whatever it
+    holds. The chunker caches only ``tokenizer.json`` there (see
+    ``docsgpt/parser/tokenization.py``), so FastEmbed then finds no graph and
+    fails to load. Fetching the files it needs first makes either order work.
+    Offline, nothing is fetched and FastEmbed reports what is missing.
+
+    Args:
+        repo: The model's FastEmbed name.
+        cache_dir: ``EMBEDDINGS_CACHE_DIR``, or None for FastEmbed's default.
+    """
+    from fastembed import TextEmbedding
+    from fastembed.common.utils import define_cache_dir
+    from huggingface_hub import hf_hub_download, snapshot_download
+
+    lowered = repo.lower()
+    description = next(
+        (d for d in TextEmbedding._list_supported_models() if str(getattr(d, "model", "")).lower() == lowered),
+        None,
+    )
+    source = getattr(getattr(description, "sources", None), "hf", None)
+    model_file = getattr(description, "model_file", None)
+    if not source or not isinstance(model_file, str):
+        return  # Not fetched from Hugging Face: FastEmbed handles it.
+    needed = [model_file, *(getattr(description, "additional_files", None) or [])]
+    cache = str(define_cache_dir(cache_dir))
+    try:
+        for filename in needed:
+            hf_hub_download(source, filename, cache_dir=cache, local_files_only=True)
+        return
+    except Exception:  # noqa: BLE001 -- any miss means the snapshot is incomplete
+        pass
+    if os.environ.get("HF_HUB_OFFLINE", "").strip().upper() in {"1", "TRUE", "YES", "ON"}:
+        return
+    logger.warning("The cached %s has no %s; downloading the model files.", source, ", ".join(needed))
+    snapshot_download(repo_id=source, allow_patterns=[*_FASTEMBED_SUPPORT_FILES, *needed], cache_dir=cache)
+
+
 def _pad_to_longest_in_batch(model: Any) -> None:
     """Undo a fixed padding width baked into a model's ``tokenizer.json``."""
     # FastEmbed enables padding only when the tokenizer declares none, so a
@@ -295,6 +347,7 @@ class EmbeddingsWrapper:
             cache_dir = settings.EMBEDDINGS_CACHE_DIR
             if cache_dir:
                 init_kwargs["cache_dir"] = cache_dir
+            _complete_model_cache(self.spec.repo, cache_dir or None)
             self.model = TextEmbedding(**init_kwargs)
         except Exception as exc:
             raise RuntimeError(

tests/vectorstore/test_embeddings_local.py

@@ -413,3 +413,63 @@ class TestPoolingOverrides:
         with patch.object(embeddings_local.settings, "EMBEDDINGS_POOLING", "banana"):
             spec = embeddings_local._spec_for("some-org/plain-onnx-export")
         assert spec.pooling == embeddings_local._FALLBACK_POOLING
+
+
+class TestIncompleteModelCache:
+    """The chunker caches only a model's ``tokenizer.json`` in the same
+    directory. FastEmbed treats any cached snapshot as the model and then
+    fails to open its ONNX graph, so the loader completes the snapshot first."""
+
+    @staticmethod
+    def _description():
+        from types import SimpleNamespace
+
+        return SimpleNamespace(
+            model="sentence-transformers/all-mpnet-base-v2",
+            model_file="onnx/model.onnx",
+            additional_files=[],
+            sources=SimpleNamespace(hf="sentence-transformers/all-mpnet-base-v2"),
+        )
+
+    def _complete(self, monkeypatch, tmp_path, cached: set, offline: str = ""):
+        from docsgpt.vectorstore import embeddings_local
+
+        downloads = []
+
+        def hf_hub_download(repo_id, filename, cache_dir=None, local_files_only=False):
+            if filename not in cached:
+                raise FileNotFoundError(filename)
+            return f"{cache_dir}/{filename}"
+
+        def snapshot_download(**kwargs):
+            downloads.append(kwargs)
+            return "/snapshot"
+
+        monkeypatch.setenv("HF_HUB_OFFLINE", offline)
+        with patch("fastembed.TextEmbedding._list_supported_models", return_value=[self._description()]), \
+                patch("huggingface_hub.hf_hub_download", side_effect=hf_hub_download), \
+                patch("huggingface_hub.snapshot_download", side_effect=snapshot_download):
+            embeddings_local._complete_model_cache("sentence-transformers/all-mpnet-base-v2", str(tmp_path))
+        return downloads
+
+    def test_a_snapshot_with_only_the_tokenizer_gets_its_model(self, monkeypatch, tmp_path):
+        downloads = self._complete(monkeypatch, tmp_path, cached={"tokenizer.json"})
+        assert len(downloads) == 1
+        assert downloads[0]["repo_id"] == "sentence-transformers/all-mpnet-base-v2"
+        assert downloads[0]["cache_dir"] == str(tmp_path)
+        assert "onnx/model.onnx" in downloads[0]["allow_patterns"]
+        assert "tokenizer_config.json" in downloads[0]["allow_patterns"]
+
+    def test_a_complete_snapshot_downloads_nothing(self, monkeypatch, tmp_path):
+        assert self._complete(monkeypatch, tmp_path, cached={"tokenizer.json", "onnx/model.onnx"}) == []
+
+    def test_offline_never_downloads(self, monkeypatch, tmp_path):
+        assert self._complete(monkeypatch, tmp_path, cached={"tokenizer.json"}, offline="1") == []
+
+    def test_loading_completes_the_cache_first(self, fake_fastembed, monkeypatch):
+        from docsgpt.vectorstore import embeddings_local
+
+        calls = []
+        monkeypatch.setattr(embeddings_local, "_complete_model_cache", lambda repo, cache: calls.append(repo))
+        embeddings_local.EmbeddingsWrapper("huggingface_sentence-transformers/all-mpnet-base-v2")
+        assert calls == ["sentence-transformers/all-mpnet-base-v2"]