209 lines
7.5 KiB
Python
209 lines
7.5 KiB
Python
"""
|
|
Embedding providers for the retrieval layer.
|
|
|
|
Two things live here:
|
|
|
|
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.
|
|
|
|
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.
|
|
|
|
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:
|
|
|
|
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.
|
|
"""
|
|
|
|
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 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="hashing"`` -> deterministic local fallback.
|
|
3. 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).
|
|
|
|
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.
|
|
"""
|
|
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()
|
|
|
|
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 "
|
|
"HashingEmbedder (deterministic, offline, non-semantic). Set the env vars "
|
|
"or pass embedder=RemoteEmbedder() for real semantic embeddings."
|
|
)
|
|
return HashingEmbedder()
|