diff --git a/src/paperless_ai/embedding.py b/src/paperless_ai/embedding.py index f4466686a1..f32c4d02ef 100644 --- a/src/paperless_ai/embedding.py +++ b/src/paperless_ai/embedding.py @@ -20,7 +20,18 @@ OCR_LEADER_REGEX = re.compile(r"[._\-\u00b7]{4,}") HORIZONTAL_WHITESPACE_REGEX = re.compile(r"[ \t\u00a0]+") -def get_embedding_model(config: AIConfig) -> "BaseEmbedding": +def get_embedding_model( + config: AIConfig, + *, + refresh: bool = False, +) -> "BaseEmbedding": + """Build the configured embedding model. + + For the Hugging Face backend, a model already in the local cache is loaded + without contacting the Hub, which saves seconds per call. ``refresh=True`` + checks the Hub for a newer revision first; index rebuilds use it so a new + model release is picked up together with re-embedding every document. + """ match config.llm_embedding_backend: case LLMEmbeddingBackend.OPENAI_LIKE: from llama_index.embeddings.openai_like import OpenAILikeEmbedding @@ -52,11 +63,18 @@ def get_embedding_model(config: AIConfig) -> "BaseEmbedding": case LLMEmbeddingBackend.HUGGINGFACE: from llama_index.embeddings.huggingface import HuggingFaceEmbedding - return HuggingFaceEmbedding( - model_name=config.llm_embedding_model + hf_kwargs = { + "model_name": config.llm_embedding_model or "sentence-transformers/all-MiniLM-L6-v2", - cache_folder=str(settings.DATA_DIR / "hf_cache"), - ) + "cache_folder": str(settings.DATA_DIR / "hf_cache"), + } + if refresh: + return HuggingFaceEmbedding(**hf_kwargs) + try: + return HuggingFaceEmbedding(**hf_kwargs, local_files_only=True) + except OSError: + # Not cached yet, download it + return HuggingFaceEmbedding(**hf_kwargs) case LLMEmbeddingBackend.OLLAMA: from llama_index.embeddings.ollama import OllamaEmbedding from ollama import AsyncClient diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 91a2e7e304..9e2377f34e 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -470,7 +470,7 @@ def update_llm_index( logger.warning("No documents found to index.") chunk_size = config.llm_embedding_chunk_size - embed_model = get_embedding_model(config) + embed_model = get_embedding_model(config, refresh=rebuild) with write_store(embed_model_name=model_name) as store: if rebuild or not store.table_exists(): diff --git a/src/paperless_ai/tests/test_embedding.py b/src/paperless_ai/tests/test_embedding.py index 3d2c04cb75..cffbc8b34c 100644 --- a/src/paperless_ai/tests/test_embedding.py +++ b/src/paperless_ai/tests/test_embedding.py @@ -167,6 +167,7 @@ def test_get_embedding_model_huggingface(mock_ai_config): MockHuggingFaceEmbedding.assert_called_once_with( model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=str(settings.DATA_DIR / "hf_cache"), + local_files_only=True, ) assert model == MockHuggingFaceEmbedding.return_value diff --git a/src/paperless_ai/tests/test_embedding_hf_local_files.py b/src/paperless_ai/tests/test_embedding_hf_local_files.py new file mode 100644 index 0000000000..84d3ad6fb1 --- /dev/null +++ b/src/paperless_ai/tests/test_embedding_hf_local_files.py @@ -0,0 +1,106 @@ +from pathlib import Path +from unittest.mock import MagicMock +from unittest.mock import call +from unittest.mock import patch + +import pytest +import pytest_mock +from django.conf import settings + +from paperless.models import LLMEmbeddingBackend +from paperless_ai import indexing +from paperless_ai.embedding import get_embedding_model +from paperless_ai.tests.conftest import FakeEmbedding +from paperless_testing.factories import DocumentFactory + +MODEL = "sentence-transformers/all-MiniLM-L6-v2" +CACHE = str(settings.DATA_DIR / "hf_cache") + + +@pytest.fixture +def hf_config() -> MagicMock: + config = MagicMock() + config.llm_embedding_backend = LLMEmbeddingBackend.HUGGINGFACE + config.llm_embedding_model = MODEL + return config + + +@pytest.fixture +def hf_embedding(mocker: pytest_mock.MockerFixture) -> MagicMock: + return mocker.patch("llama_index.embeddings.huggingface.HuggingFaceEmbedding") + + +class TestHuggingFaceLocalFirst: + def test_uses_cached_model_without_hub( + self, + hf_config: MagicMock, + hf_embedding: MagicMock, + ) -> None: + model = get_embedding_model(hf_config) + + hf_embedding.assert_called_once_with( + model_name=MODEL, + cache_folder=CACHE, + local_files_only=True, + ) + assert model is hf_embedding.return_value + + def test_downloads_when_not_cached( + self, + hf_config: MagicMock, + hf_embedding: MagicMock, + ) -> None: + downloaded = MagicMock() + hf_embedding.side_effect = [OSError("not cached"), downloaded] + + model = get_embedding_model(hf_config) + + assert hf_embedding.call_args_list == [ + call(model_name=MODEL, cache_folder=CACHE, local_files_only=True), + call(model_name=MODEL, cache_folder=CACHE), + ] + assert model is downloaded + + def test_refresh_checks_the_hub( + self, + hf_config: MagicMock, + hf_embedding: MagicMock, + ) -> None: + get_embedding_model(hf_config, refresh=True) + + hf_embedding.assert_called_once_with(model_name=MODEL, cache_folder=CACHE) + + +@pytest.mark.django_db +class TestIndexUpdateRefresh: + @pytest.fixture(autouse=True) + def _index_inputs(self, temp_llm_index_dir: Path) -> None: + DocumentFactory.create() + + def _update(self, mock_embed_model: FakeEmbedding, *, rebuild: bool) -> MagicMock: + mock_config = MagicMock() + mock_config.llm_embedding_chunk_size = 512 + with ( + patch("paperless_ai.indexing.AIConfig", return_value=mock_config), + patch("paperless_ai.indexing.build_document_node", return_value=[]), + patch( + "paperless_ai.indexing.get_embedding_model", + return_value=mock_embed_model, + ) as get_model, + ): + indexing.update_llm_index(rebuild=rebuild) + get_model.assert_called_once() + return get_model + + def test_rebuild_refreshes(self, mock_embed_model: FakeEmbedding) -> None: + get_model = self._update(mock_embed_model, rebuild=True) + + assert get_model.call_args.kwargs == {"refresh": True} + + def test_incremental_update_stays_local( + self, + mock_embed_model: FakeEmbedding, + ) -> None: + get_model = self._update(mock_embed_model, rebuild=False) + + assert get_model.call_args.kwargs == {"refresh": False}