247 lines
9.1 KiB
Python
247 lines
9.1 KiB
Python
"""
|
|
Embedding providers for the retrieval layer.
|
|
|
|
Three providers, all behind the same ``Embedder`` interface:
|
|
|
|
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. ``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.
|
|
|
|
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.
|
|
|
|
``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
|
|
|
|
import hashlib
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
|
|
# The vector dimension the hashing fallback projects into. The remote provider
|
|
# returns its own native dimension; the index is built to whatever the embedder
|
|
# reports via ``dimension``.
|
|
HASHING_DIM = 256
|
|
|
|
|
|
class Embedder(ABC):
|
|
"""Abstract embedding provider."""
|
|
|
|
#: Dimensionality of the vectors this provider emits.
|
|
DIM: int
|
|
|
|
@abstractmethod
|
|
def embed(self, texts: Sequence[str]) -> np.ndarray:
|
|
"""Embed a batch of texts, returning an (N, DIM) float32 array."""
|
|
raise NotImplementedError
|
|
|
|
def embed_one(self, text: str) -> np.ndarray:
|
|
"""Embed a single text, returning a (DIM,) float32 array."""
|
|
return self.embed([text])[0]
|
|
|
|
|
|
class HashingEmbedder(Embedder):
|
|
"""
|
|
Deterministic, dependency-free embedding via the hashing trick.
|
|
|
|
We tokenize into character 3-grams, hash each one into a bucket, and write
|
|
+1/-1 (sign hashed from the gram) into a fixed-size vector, then L2-normalize.
|
|
|
|
This is *not* semantically meaningful (synonyms are not mapped together), but
|
|
it is stable, fast, and shares enough lexical overlap with the query for
|
|
near-neighbors to cluster around the right topics. Perfect for exercising
|
|
the ANN pipeline without an API key.
|
|
"""
|
|
|
|
DIM = HASHING_DIM
|
|
_NGRAM = 3
|
|
_SEED = 0xC0FFEE
|
|
|
|
def __init__(self, dim: int = HASHING_DIM, ngram: int = 3):
|
|
self.DIM = dim
|
|
self._ngram = ngram
|
|
|
|
@staticmethod
|
|
def _grams(text: str, ngram: int = 3) -> list[str]:
|
|
t = text.lower()
|
|
return [t[i:i + ngram] for i in range(max(0, len(t) - ngram + 1))]
|
|
|
|
def embed(self, texts: Sequence[str]) -> np.ndarray:
|
|
out = np.zeros((len(texts), self.DIM), dtype=np.float32)
|
|
ngram = self._ngram
|
|
for i, text in enumerate(texts):
|
|
for gram in self._grams(text or "", ngram):
|
|
h = int.from_bytes(
|
|
hashlib.sha1(gram.encode("utf-8")).digest()[:8],
|
|
"little",
|
|
)
|
|
bucket = h % self.DIM
|
|
sign = 1.0 if (h >> 63) & 1 else -1.0
|
|
out[i, bucket] += sign
|
|
norm = np.linalg.norm(out[i])
|
|
if norm > 0:
|
|
out[i] /= norm
|
|
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.
|
|
|
|
Configured entirely through the environment:
|
|
|
|
EMBEDDING_API_BASE base URL, e.g. https://api.openai.com/v1
|
|
EMBEDDING_API_KEY API key (required)
|
|
EMBEDDING_MODEL model name, default text-embedding-3-small
|
|
|
|
Any endpoint with the OpenAI embeddings shape works (OpenAI, Together,
|
|
Groq, a local vLLM/Ollama server, etc.).
|
|
"""
|
|
|
|
_DEFAULT_MODEL = "text-embedding-3-small"
|
|
_DEFAULT_DIM = 1536 # text-embedding-3-small's default output dim
|
|
|
|
def __init__(
|
|
self,
|
|
api_base: str | None = None,
|
|
api_key: str | None = None,
|
|
model: str | None = None,
|
|
dim: int | None = None,
|
|
):
|
|
import requests # imported lazily so the fallback needs no network stack
|
|
|
|
self._requests = requests
|
|
self._api_base = (api_base or os.environ.get("EMBEDDING_API_BASE", "")).rstrip("/")
|
|
self._api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "")
|
|
self._model = model or os.environ.get("EMBEDDING_MODEL", self._DEFAULT_MODEL)
|
|
self._dim = dim or int(
|
|
os.environ.get("EMBEDDING_DIM", self._DEFAULT_DIM)
|
|
)
|
|
if not self._api_base:
|
|
raise ValueError(
|
|
"EMBEDDING_API_BASE is not set; configure it (or use the local "
|
|
"HashingEmbedder fallback)."
|
|
)
|
|
if not self._api_key:
|
|
raise ValueError(
|
|
"EMBEDDING_API_KEY is not set; set it, or use the local "
|
|
"HashingEmbedder fallback."
|
|
)
|
|
|
|
@property
|
|
def DIM(self) -> int:
|
|
return self._dim
|
|
|
|
def embed(self, texts: Sequence[str]) -> np.ndarray:
|
|
resp = self._requests.post(
|
|
f"{self._api_base}/embeddings",
|
|
headers={"Authorization": f"Bearer {self._api_key}"},
|
|
json={"model": self._model, "input": list(texts), "dimensions": self._dim},
|
|
timeout=60,
|
|
)
|
|
resp.raise_for_status()
|
|
payload = resp.json()
|
|
items = payload["data"]
|
|
# Some backends don't preserve input order; sort by the returned index.
|
|
items = sorted(items, key=lambda d: d.get("index", 0))
|
|
return np.array([d["embedding"] for d in items], dtype=np.float32)
|
|
|
|
|
|
def get_embedder(embedder: str | Embedder | None = None) -> Embedder:
|
|
"""
|
|
Resolve an embedder from a name, an instance, or the environment.
|
|
|
|
Resolution order:
|
|
|
|
1. ``embedder="remote"`` or an ``Embedder`` instance -> use it directly.
|
|
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).
|
|
5. Otherwise -> the local semantic model (with a printed note).
|
|
|
|
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):
|
|
return embedder
|
|
if name == "hashing":
|
|
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")
|
|
if api_base and api_key:
|
|
print(f"[embed] using remote provider: {api_base} (model={os.environ.get('EMBEDDING_MODEL')})")
|
|
return RemoteEmbedder()
|
|
|
|
print(
|
|
"[embed] no EMBEDDING_API_BASE/KEY configured -> using local "
|
|
"sentence-transformers model (semantic). Set EMBEDDING_LOCAL_MODEL to "
|
|
"override, or embedder='hashing' for a dependency-free fallback."
|
|
)
|
|
return LocalEmbedder()
|