Add local semantic embedder (sentence-transformers) as default local path

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
+1. ``RemoteEmbedder`` -- OpenAI-compatible remote endpoint.  This is the
-   OpenAI-compatible remote provider.  This is the **production** path: it hits
+   **production** path: it hits an HTTP endpoint and returns real dense vectors.
-   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
+2. ``LocalEmbedder`` -- real dense embeddings from a small local
-   no model weights, no GPU, no API key and no internet.  It projects a bag of
+   sentence-transformers model.  This is the **default local** path when no API
-   character n-grams into a fixed-dimension vector via the hashing trick.  This
+   key is present: genuine semantic vectors (synonyms/paraphrases map together),
-   lets the whole ANN pipeline (FAISS build -> ANN search -> ranking) run
+   no API cost, modest RAM.  ``paraphrase-MiniLM-L3-v2`` by default (~90 MB).
-   end-to-end offline, which is exactly what we want for demos and CI.
+   Slower than hashing (~1-2 h for the full 3.15M corpus) but far better quality.
 
-WHY NOT A LOCAL EMBEDDING MODEL?
+3. ``HashingEmbedder`` -- deterministic, dependency-free fallback.  No model
---------------------------------
+   weights, no GPU, no API key, no internet.  Projects a bag of character n-grams
-This box is a Strix Halo AP with ~104 GiB of unified RAM.  It has *tried* to
+   into a fixed-dimension vector via the hashing trick.  Not semantic (synonyms
-load a large local embedding/LLM model before and OOM-killed.  So we do not
+   are not mapped together), but stable and fast -- handy for CI or a quick
-import torch or try to run sentence-transformers here.  If you have the RAM and
+   offline smoke test.
-want real embeddings without paying for an API, drop in a local model:
 
-    class LocalEmbedder(Embedder):
+``get_embedder()`` resolves one of the above: an explicit instance/name wins,
-        DIM = 768
+then the API key if set, otherwise the local semantic model.
-        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.
 """
 
 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.
+      2. ``embedder="local"`` -> local sentence-transformers model (semantic).
-      3. If ``EMBEDDING_API_BASE`` and ``EMBEDDING_API_KEY`` are both set ->
+      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
+      This means the CLI works out of the box with no configuration: it uses real
-    degrades to the reproducible local embedding when no API key is present.
+      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 "
+        "sentence-transformers model (semantic). Set EMBEDDING_LOCAL_MODEL to "
-        "or pass embedder=RemoteEmbedder() for real semantic embeddings."
+        "override, or embedder='hashing' for a dependency-free fallback."
     )
-    return HashingEmbedder()
+    return LocalEmbedder()