Add local semantic embedder (sentence-transformers) as default local path
This commit is contained in:
+71
-33
@@ -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,13 +212,14 @@ 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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user