studyhub/ingest/embedder.py

from __future__ import annotations

import hashlib
from typing import Protocol, runtime_checkable

import numpy as np


@runtime_checkable
class Embedder(Protocol):
    dim: int
    def encode(self, texts: list[str]) -> np.ndarray:  # (n, dim), L2-normalized float32
        ...


class FakeEmbedder:
    """Deterministic, dependency-free embedder for tests."""

    def __init__(self, dim: int = 32):
        self.dim = dim

    def encode(self, texts: list[str]) -> np.ndarray:
        out = np.zeros((len(texts), self.dim), dtype="float32")
        for i, t in enumerate(texts):
            digest = hashlib.sha256(t.encode("utf-8")).digest()
            reps = (self.dim // len(digest)) + 1
            raw = np.frombuffer((digest * reps)[: self.dim], dtype=np.uint8).astype("float32")
            v = raw - raw.mean()
            norm = float(np.linalg.norm(v))
            out[i] = v / norm if norm > 0 else v
        return out


class BGEM3Embedder:
    """Real embedder (BAAI/bge-m3). Lazy import so tests need no model download.

    Used both for offline index-building and for query embedding on the Space
    (weights are downloaded from the Hub on first use). An int8-ONNX build to
    speed cold starts is a documented follow-up, not yet implemented.
    """

    def __init__(self, model_name: str = "BAAI/bge-m3"):
        from FlagEmbedding import BGEM3FlagModel
        import torch
        # fp16 only on GPU — fp16 matmul isn't implemented for CPU inference, which
        # is what CI runners and the free Space use; fp32 is correct (just more RAM).
        self._model = BGEM3FlagModel(model_name, use_fp16=torch.cuda.is_available())
        self.dim = 1024

    def encode(self, texts: list[str]) -> np.ndarray:
        dense = self._model.encode(texts, return_dense=True)["dense_vecs"]
        v = np.asarray(dense, dtype="float32")
        norms = np.linalg.norm(v, axis=1, keepdims=True)
        norms[norms == 0] = 1.0
        return v / norms