From fc5fc485dea76441d05b28ffcb64940df4a5920a Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Sat, 31 Oct 2026 07:07:02 -0700 Subject: [PATCH] Performance: load a cached Hugging Face embedding model without contacting the Hub Every chat question and AI suggestion request builds the embedding model, and sentence-transformers checks the Hub for a newer revision each time. With the model already cached that check cost about 2 seconds per request (0.2 seconds without it), and it fails outright when the Hub is unreachable. Load from the local cache first and only download when the model is missing. A full rebuild still checks the Hub, so a new model release is picked up together with re-embedding every document instead of silently mixing revisions. --- docs/administration.md | 2 ++ src/paperless_ai/embedding.py | 25 ++++++++++++--- src/paperless_ai/indexing.py | 2 +- src/paperless_ai/tests/test_ai_indexing.py | 6 ++++ src/paperless_ai/tests/test_embedding.py | 36 ++++++++++++++++++++++ 5 files changed, 65 insertions(+), 6 deletions(-) diff --git a/docs/administration.md b/docs/administration.md index 155716660e..80ca53eaaf 100644 --- a/docs/administration.md +++ b/docs/administration.md @@ -565,6 +565,8 @@ Specify `rebuild` to build the index from scratch from all documents in the data this the first time you enable the feature, or after changing the embedding backend or model. +When using the `huggingface` embedding backend, a model that is already downloaded is loaded from the local cache without contacting the Hugging Face Hub. Only `rebuild` checks the Hub for a newer revision of the model, so a model update is picked up together with re-embedding every document. If the model is not cached yet, it is downloaded on first use. + Specify `update` to incrementally index new and changed documents. This is what the scheduled task runs. diff --git a/src/paperless_ai/embedding.py b/src/paperless_ai/embedding.py index f4466686a1..8ab449c3c2 100644 --- a/src/paperless_ai/embedding.py +++ b/src/paperless_ai/embedding.py @@ -1,4 +1,5 @@ import re +from contextlib import suppress from typing import TYPE_CHECKING from django.conf import settings @@ -20,7 +21,17 @@ 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 embedding model for the configured backend. + + For Hugging Face, a cached model loads offline by default, saving seconds per + call. Pass ``refresh=True`` to check the Hub for a newer model revision. + """ match config.llm_embedding_backend: case LLMEmbeddingBackend.OPENAI_LIKE: from llama_index.embeddings.openai_like import OpenAILikeEmbedding @@ -52,11 +63,15 @@ 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 not refresh: + with suppress(OSError): + return HuggingFaceEmbedding(**hf_kwargs, local_files_only=True) + 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_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 7b2c199bb6..8bc13bd6ec 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -212,6 +212,10 @@ def test_update_llm_index( ai_config.assert_called_once() build_document_node.assert_called_once_with(real_document, chunk_size=512) + indexing.get_embedding_model.assert_called_once_with( + mock_config, + refresh=True, + ) @pytest.mark.django_db @@ -336,6 +340,8 @@ def test_update_llm_index_partial_update( indexing.update_llm_index(rebuild=False) + assert indexing.get_embedding_model.call_args.kwargs == {"refresh": False} + with indexing.get_vector_store() as store: assert store.table_exists(), ( "Expected the vector store table to exist after incremental update" diff --git a/src/paperless_ai/tests/test_embedding.py b/src/paperless_ai/tests/test_embedding.py index 3d2c04cb75..a408dfe1da 100644 --- a/src/paperless_ai/tests/test_embedding.py +++ b/src/paperless_ai/tests/test_embedding.py @@ -1,5 +1,6 @@ from unittest.mock import ANY from unittest.mock import MagicMock +from unittest.mock import call from unittest.mock import patch import pytest @@ -167,10 +168,45 @@ 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 +def test_get_embedding_model_huggingface_downloads_when_not_cached(mock_ai_config): + mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.HUGGINGFACE + mock_ai_config.return_value.llm_embedding_model = "some/model" + downloaded = MagicMock() + + with patch( + "llama_index.embeddings.huggingface.HuggingFaceEmbedding", + side_effect=[OSError("not cached"), downloaded], + ) as MockHuggingFaceEmbedding: + model = get_embedding_model(mock_ai_config.return_value) + + cache_folder = str(settings.DATA_DIR / "hf_cache") + assert MockHuggingFaceEmbedding.call_args_list == [ + call(model_name="some/model", cache_folder=cache_folder, local_files_only=True), + call(model_name="some/model", cache_folder=cache_folder), + ] + assert model is downloaded + + +def test_get_embedding_model_huggingface_refresh_checks_the_hub(mock_ai_config): + mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.HUGGINGFACE + mock_ai_config.return_value.llm_embedding_model = "some/model" + + with patch( + "llama_index.embeddings.huggingface.HuggingFaceEmbedding", + ) as MockHuggingFaceEmbedding: + get_embedding_model(mock_ai_config.return_value, refresh=True) + + MockHuggingFaceEmbedding.assert_called_once_with( + model_name="some/model", + cache_folder=str(settings.DATA_DIR / "hf_cache"), + ) + + def test_get_embedding_model_ollama(mock_ai_config): mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA mock_ai_config.return_value.llm_embedding_model = "embeddinggemma"