File size: 3,629 Bytes
54461b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """Local embeddings via FastEmbed (ONNX, CPU, zero cost).
Nothing leaves the process. That is the compliance answer for a UAE government or legal
client — the corpus and every query stay inside the deployment boundary — and it is why
inference cost per query is dominated by one Claude call rather than by an embedding API.
`bge-small-en-v1.5` is asymmetric: queries are embedded with an instruction prefix and
documents without one. Getting that backwards costs real recall, so the two directions
are separate functions here rather than one `embed(text)` that callers must remember to
parameterise.
"""
from __future__ import annotations
import functools
import logging
import threading
from collections.abc import Iterable, Sequence
from fastembed import TextEmbedding
from app.core.settings import Settings, get_settings
logger = logging.getLogger(__name__)
_LOCK = threading.Lock()
_MODEL: TextEmbedding | None = None
def get_embedding_model(settings: Settings | None = None) -> TextEmbedding:
"""Process-wide embedding model.
Loaded once and shared: the ONNX session is thread-safe for inference, and holding a
single instance keeps a cold container from paying the load cost per request.
"""
global _MODEL # noqa: PLW0603 - deliberate process-wide singleton
if _MODEL is not None:
return _MODEL
cfg = settings or get_settings()
with _LOCK:
if _MODEL is None:
cfg.models_cache_dir.mkdir(parents=True, exist_ok=True)
logger.info("loading embedding model %s", cfg.embedding_model)
_MODEL = TextEmbedding(
model_name=cfg.embedding_model,
cache_dir=str(cfg.models_cache_dir),
)
return _MODEL
def embed_documents(texts: Sequence[str], settings: Settings | None = None) -> list[list[float]]:
"""Embed corpus passages. No instruction prefix — this is the document side."""
if not texts:
return []
model = get_embedding_model(settings)
return [vector.tolist() for vector in model.embed(list(texts))]
def embed_query(text: str, settings: Settings | None = None) -> list[float]:
"""Embed a search query, with the instruction prefix bge was trained to expect."""
cfg = settings or get_settings()
return _embed_query_cached(f"{cfg.embedding_query_prefix}{text}")
@functools.lru_cache(maxsize=512)
def _embed_query_cached(prefixed: str) -> list[float]:
"""Cache identical queries.
Repeated questions are the norm in a demo, and the eval harness embeds the same
query once per configuration. The key is the fully prefixed string, so changing the
prefix can never return a vector built with the old one.
``TextEmbedding.embed`` is used deliberately in place of ``query_embed``: measured
against this build of FastEmbed, ``query_embed(q)`` is bit-identical to ``embed(q)``
(cosine 1.000000) — it does *not* add bge's instruction prefix, despite the name.
Relying on it to do so would leave every query embedded on the document side of an
asymmetric model, losing recall with nothing in the logs to show for it.
"""
model = get_embedding_model()
vectors: Iterable[object] = model.embed([prefixed])
for vector in vectors:
return list(vector.tolist()) # type: ignore[attr-defined]
raise RuntimeError("embedding model returned no vector")
def reset_embedding_model() -> None:
"""Drop the cached model and query cache. Used by tests."""
global _MODEL # noqa: PLW0603 - mirrors get_embedding_model
with _LOCK:
_MODEL = None
_embed_query_cached.cache_clear()
|