From be22cc2a9992a5c8b22284148bb983ee74aded36 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sun, 20 Sep 2026 10:06:34 -0700 Subject: [PATCH] Add local semantic embedder (sentence-transformers) as default local path --- retrieval/embed.py | 106 ++++++++++++++++++++++++++++++--------------- 1 file changed, 72 insertions(+), 34 deletions(-) diff --git a/retrieval/embed.py b/retrieval/embed.py index 2404849..25ab3d6 100644 --- a/retrieval/embed.py +++ b/retrieval/embed.py @@ -1,35 +1,27 @@ """ Embedding providers for the retrieval layer. -Two things live here: +Three providers, all behind the same ``Embedder`` interface: -1. A *pluggable* embedding interface (``Embedder``) plus a concrete - OpenAI-compatible remote provider. This is the **production** path: it hits - an HTTP endpoint and returns real dense vectors. +1. ``RemoteEmbedder`` -- OpenAI-compatible remote endpoint. This is the + **production** path: it hits an HTTP endpoint and returns real dense vectors. + Configured via ``EMBEDDING_API_BASE`` / ``EMBEDDING_API_KEY`` / + ``EMBEDDING_MODEL``. -2. A **deterministic local hashing fallback** (``HashingEmbedder``) that needs - no model weights, no GPU, no API key and no internet. It projects a bag of - character n-grams into a fixed-dimension vector via the hashing trick. This - lets the whole ANN pipeline (FAISS build -> ANN search -> ranking) run - end-to-end offline, which is exactly what we want for demos and CI. +2. ``LocalEmbedder`` -- real dense embeddings from a small local + sentence-transformers model. This is the **default local** path when no API + key is present: genuine semantic vectors (synonyms/paraphrases map together), + no API cost, modest RAM. ``paraphrase-MiniLM-L3-v2`` by default (~90 MB). + Slower than hashing (~1-2 h for the full 3.15M corpus) but far better quality. -WHY NOT A LOCAL EMBEDDING MODEL? --------------------------------- -This box is a Strix Halo AP with ~104 GiB of unified RAM. It has *tried* to -load a large local embedding/LLM model before and OOM-killed. So we do not -import torch or try to run sentence-transformers here. If you have the RAM and -want real embeddings without paying for an API, drop in a local model: +3. ``HashingEmbedder`` -- deterministic, dependency-free fallback. No model + weights, no GPU, no API key, no internet. Projects a bag of character n-grams + into a fixed-dimension vector via the hashing trick. Not semantic (synonyms + are not mapped together), but stable and fast -- handy for CI or a quick + offline smoke test. - class LocalEmbedder(Embedder): - DIM = 768 - def __init__(self): - from sentence_transformers import SentenceTransformer - self._model = SentenceTransformer("all-MiniLM-L6-v2") - def embed(self, texts): - return self._model.encode(texts, normalize_embeddings=True) - -and pass it to the index/query builders via the ``embedder=`` argument. The -rest of the pipeline is agnostic to where the vectors come from. +``get_embedder()`` resolves one of the above: an explicit instance/name wins, +then the API key if set, otherwise the local semantic model. """ from __future__ import annotations @@ -107,6 +99,49 @@ class HashingEmbedder(Embedder): return out +class LocalEmbedder(Embedder): + """ + Real dense embeddings from a local sentence-transformers model. + + Uses a small CPU model (``paraphrase-MiniLM-L3-v2`` by default, ~90 MB) so it + runs on a laptop without an API key and without the Strix Halo OOM history we + had loading big models. This is the *production local* path: unlike + ``HashingEmbedder`` it produces genuinely semantic vectors (synonyms and + paraphrases map close together). + + Trade-off is speed: embedding the full 3.15M-row corpus takes ~1-2 hours on + this box. For a quick demo, index a single ``--category`` subset instead. + + Any model name accepted by ``SentenceTransformer`` works; set + ``EMBEDDING_LOCAL_MODEL`` to override (e.g. ``all-MiniLM-L6-v2`` for better + quality at ~3x the time). + """ + + DIM = 384 # paraphrase-MiniLM-L3-v2's dimensionality + _DEFAULT_MODEL = "paraphrase-MiniLM-L3-v2" + + def __init__(self, model: str | None = None): + from sentence_transformers import SentenceTransformer + + self._model_name = model or os.environ.get( + "EMBEDDING_LOCAL_MODEL", self._DEFAULT_MODEL + ) + self._model = SentenceTransformer(self._model_name) + # Cache the real dim so the index/registry match the model. + DIM = getattr(self._model, "get_embedding_dimension", None) or getattr( + self._model, "get_sentence_embedding_dimension", None + ) + self.DIM = int(DIM()) + self._normalize = True + + def embed(self, texts: Sequence[str]) -> np.ndarray: + return self._model.encode( + list(texts), + normalize_embeddings=self._normalize, + show_progress_bar=False, + ) + + class RemoteEmbedder(Embedder): """ OpenAI-compatible remote embedding provider. @@ -177,14 +212,15 @@ def get_embedder(embedder: str | Embedder | None = None) -> Embedder: Resolution order: 1. ``embedder="remote"`` or an ``Embedder`` instance -> use it directly. - 2. ``embedder="hashing"`` -> deterministic local fallback. - 3. If ``EMBEDDING_API_BASE`` and ``EMBEDDING_API_KEY`` are both set -> + 2. ``embedder="local"`` -> local sentence-transformers model (semantic). + 3. ``embedder="hashing"`` -> deterministic local fallback (no deps). + 4. If ``EMBEDDING_API_BASE`` and ``EMBEDDING_API_KEY`` are both set -> remote provider (the production path). - 4. Otherwise -> the local hashing fallback (with a printed note). + 5. Otherwise -> the local semantic model (with a printed note). - This means the CLI works out of the box with no configuration: it silently - degrades to the reproducible local embedding when no API key is present. - """ + This means the CLI works out of the box with no configuration: it uses real + semantic embeddings from a small local model when no API key is present. + """ name = (embedder or "").lower() if isinstance(embedder, str) else None if isinstance(embedder, Embedder): @@ -193,6 +229,8 @@ def get_embedder(embedder: str | Embedder | None = None) -> Embedder: return HashingEmbedder() if name == "remote": return RemoteEmbedder() + if name == "local": + return LocalEmbedder() api_base = os.environ.get("EMBEDDING_API_BASE") api_key = os.environ.get("EMBEDDING_API_KEY") @@ -202,7 +240,7 @@ def get_embedder(embedder: str | Embedder | None = None) -> Embedder: print( "[embed] no EMBEDDING_API_BASE/KEY configured -> using local " - "HashingEmbedder (deterministic, offline, non-semantic). Set the env vars " - "or pass embedder=RemoteEmbedder() for real semantic embeddings." + "sentence-transformers model (semantic). Set EMBEDDING_LOCAL_MODEL to " + "override, or embedder='hashing' for a dependency-free fallback." ) - return HashingEmbedder() + return LocalEmbedder()