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:
stumpylog committed 2026-11-09 11:36:42 -08:00
1 parent 936b1b3bfa
commit fc5fc485de
5 files changed
+65 -6

No files matched your search

+2
View File
@@ -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.
+20 -5
View File
@@ -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
+1 -1
View File
@@ -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"
+36
View File
@@ -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"