Add arxiv-complete retrieval layer (semantic + keyword)
This commit is contained in:
@@ -0,0 +1,208 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user