diff --git a/docsgpt/vectorstore/embeddings_local.py b/docsgpt/vectorstore/embeddings_local.py index 5a236e3b..88573cd5 100644 --- a/docsgpt/vectorstore/embeddings_local.py +++ b/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( diff --git a/tests/vectorstore/test_embeddings_local.py b/tests/vectorstore/test_embeddings_local.py index 5d61a72b..b3090a1c 100644 --- a/tests/vectorstore/test_embeddings_local.py +++ b/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"]