mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-11-10 00:18:41 +00:00
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.
This commit is contained in:
1 parent
936b1b3bfa
commit
fc5fc485de
5 files changed
+65
-6
No files matched your search
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in new issue
Block a user