embed.py

"""Embeddings — bge-m3 (BAAI/bge-m3), 1024-dim dense, L2-normalized.

Single source of truth for the embedding space, and the reason scores are
reproducible: a caller that embeds the same text with the same open model gets
the same vector, and therefore the same cosine score against these news
vectors. If this file drifts from what a caller runs, cosine similarity between
the two is meaningless. The contract:

  - model:         BAAI/bge-m3
  - output:        dense vector, 1024 dims
  - normalization: L2 (unit length), so cosine == dot product
  - instruction:   none. bge-m3 dense retrieval is instruction-free: queries and
                   documents are embedded identically. Do not add a query
                   prefix or prompt.

Two backends, selected by `EMBED_BACKEND`:
  - "local" (default): runs bge-m3 in-process via sentence-transformers. Heavy
    (torch plus ~2.2GB of weights), lazy-loaded on first use.
  - "http": POSTs texts to a bge-m3 inference endpoint (`EMBED_ENDPOINT`) for
    hosts that shouldn't bundle torch. Same weights, same space.

Output is re-normalized either way, so both backends agree on unit length.
"""

import math
import os
import threading

# Tag written to the `model` column, identifying the embedding space a stored
# vector belongs to. Kept distinct from the weights id so the tag stays stable
# even if the weights path changes.
EMBED_MODEL = "bge-m3"
HF_MODEL_ID = "BAAI/bge-m3"
EMBED_DIM = 1024

# Local encode batch size. Modest so CPU-only boxes don't blow memory; a GPU
# box can raise it via ST_BATCH.
_BATCH = int(os.environ.get("ST_BATCH", "32"))


def _l2_normalize(vec: list[float]) -> list[float]:
    n = math.sqrt(sum(x * x for x in vec))
    return [x / n for x in vec] if n else vec


# --- local backend (sentence-transformers) ---------------------------------

_model = None  # lazy singleton — loading bge-m3 is expensive, do it once.
# The singleton is shared by the indexing loop and the /embed endpoint
# (embed_server.py). sentence-transformers' encode is not guaranteed safe under
# concurrent calls, so serialize them to prevent an interleaved-encode
# corruption. Contention is negligible in practice.
_encode_lock = threading.Lock()


def _local_model():
    global _model
    if _model is None:
        # Imported lazily so `import embed` (and the whole worker) stays cheap
        # and torch is only required when the local backend is actually used.
        from sentence_transformers import SentenceTransformer

        _model = SentenceTransformer(HF_MODEL_ID)
    return _model


def _embed_local(texts: list[str]) -> list[list[float]]:
    model = _local_model()
    with _encode_lock:
        embs = model.encode(
            texts,
            batch_size=_BATCH,
            normalize_embeddings=True,
            convert_to_numpy=True,
        )
    return [[float(x) for x in row] for row in embs]


# --- http backend (external bge-m3 endpoint) -------------------------------


def _embed_http(texts: list[str]) -> list[list[float]]:
    import httpx

    endpoint = os.environ.get("EMBED_ENDPOINT")
    if not endpoint:
        raise RuntimeError("EMBED_BACKEND=http but EMBED_ENDPOINT is not set")
    headers = {"Content-Type": "application/json"}
    token = os.environ.get("EMBED_ENDPOINT_TOKEN")
    if token:
        headers["Authorization"] = f"Bearer {token}"

    out: list[list[float]] = []
    for i in range(0, len(texts), 64):
        batch = texts[i:i + 64]
        res = httpx.post(
            endpoint,
            headers=headers,
            json={"inputs": batch, "model": HF_MODEL_ID},
            timeout=120.0,
        )
        if res.status_code != 200:
            raise RuntimeError(
                f"embed endpoint failed: {res.status_code} {res.text}"
            )
        out.extend(res.json())
    return out


# --- public interface ------------------------------------------------------


def embed_texts(texts: list[str]) -> list[list[float]]:
    """Embed N texts, returning vectors 1:1 with `texts` (same length/order),
    each an L2-normalized 1024-dim bge-m3 dense vector. Empty input → []."""
    if not texts:
        return []

    backend = os.environ.get("EMBED_BACKEND", "local")
    raw = _embed_http(texts) if backend == "http" else _embed_local(texts)

    # Normalize uniformly regardless of backend so cosine is dot-product and
    # every producer (worker, client) agrees on unit-length vectors.
    return [_l2_normalize(v) for v in raw]


def embed_text(text: str) -> list[float]:
    return embed_texts([text])[0]