ely-eye/bench/systems.py

"""Retrieval arms for the ablation. Every arm reads the same merged baseline database.

The arms are the product's own signals, switched on one at a time: sparse BM25, dense
cosine, the evidence graph, the product's weighted fusion, and fusion plus a cross-encoder
reranker. Rankings are lists of source ids: an atom votes for its source, and a source is
ranked by its best atom.
"""

from __future__ import annotations

import json
import re
import sqlite3
import time
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from bench.sets_common import ROOT

RUNS = ROOT / "runs" / "baseline"
INDEX_DIR = RUNS / "index"
MODELS = ROOT / "models"
DENSE_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
RERANK_MODEL = "BAAI/bge-reranker-base"


def local_model_dir(models: Path, repo_id: str) -> Path:
    """Where a fetched model lives: one flat directory per repo, named by its repo name only."""
    return models / repo_id.split("/")[-1]
DENSE_DIM = 384
FUSION_WEIGHTS = {"sparse": 0.55, "dense": 0.35, "graph": 0.10}
RERANK_CANDIDATES = 50
RERANK_TEXT_CHARS = 400
ARMS = ("bm25_only", "dense_only", "graph_only", "no_rerank", "fusion_rerank")


@dataclass(frozen=True)
class Atom:
    atom_id: str
    source_id: str
    text: str


def load_atoms(db_path: Path) -> list[Atom]:
    """Every atom in the merged database, in a stable order."""
    if not db_path.is_file():
        raise SystemExit(f"{db_path}: no merged baseline database, run python -m bench.ingest_baseline first")
    connection = sqlite3.connect(db_path)
    try:
        rows = connection.execute("SELECT atom_id, source_id, text FROM atoms ORDER BY atom_id").fetchall()
    except sqlite3.OperationalError as error:
        raise SystemExit(f"{db_path}: cannot read the atoms table ({error})") from error
    finally:
        connection.close()
    if not rows:
        raise SystemExit(f"{db_path}: the atoms table is empty")
    return [Atom(atom_id=str(atom_id), source_id=str(source_id), text=str(text or "")) for atom_id, source_id, text in rows]


def tokenize(text: str) -> list[str]:
    """The product's tokenizer: word characters plus CJK, lowercased."""
    return re.findall(r"[\w\u4e00-\u9fff]+", text.lower())


class SparseIndex:
    """BM25 over atom text, built once; the product rebuilds it on every query."""

    def __init__(self, atoms: Sequence[Atom]) -> None:
        from rank_bm25 import BM25Okapi

        corpus = [tokenize(f"{atom.text} {atom.source_id}") for atom in atoms]
        self._model = BM25Okapi(corpus)
        self._count = len(atoms)

    def scores(self, query: str) -> np.ndarray:
        tokens = tokenize(query)
        if not tokens:
            return np.zeros(self._count, dtype=np.float64)
        return np.asarray(self._model.get_scores(tokens), dtype=np.float64)


class DenseIndex:
    """Brute-force cosine over cached sentence embeddings, one vector per atom."""

    def __init__(self, atoms: Sequence[Atom], model_dir: Path, batch_size: int = 256, cache_name: str = "") -> None:
        self._cache_name = cache_name or model_dir.name
        self._vectors = self._load_or_build(atoms, model_dir, batch_size)

    def _vectors_path(self, model_dir: Path) -> Path:
        return INDEX_DIR / f"{self._cache_name}.npy"

    def _load_or_build(self, atoms: Sequence[Atom], model_dir: Path, batch_size: int) -> np.ndarray:
        path = self._vectors_path(model_dir)
        order_path = INDEX_DIR / f"{self._cache_name}.atoms.json"
        atom_ids = [atom.atom_id for atom in atoms]
        if path.is_file() and order_path.is_file():
            vectors = np.load(path)
            cached_ids = json.loads(order_path.read_text(encoding="utf-8"))
            if cached_ids == atom_ids and vectors.shape == (len(atoms), DENSE_DIM):
                return vectors
            if cached_ids != atom_ids:
                raise SystemExit(
                    f"{path}: cached vectors cover {len(cached_ids)} atoms but the database has {len(atoms)}; "
                    "delete bench/runs/baseline/index and rebuild"
                )
        INDEX_DIR.mkdir(parents=True, exist_ok=True)
        model = _sentence_model(model_dir)
        texts = [atom.text[:RERANK_TEXT_CHARS] if atom.text else atom.source_id for atom in atoms]
        vectors = model.encode(
            texts,
            batch_size=batch_size,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=len(texts) > 2000,
        ).astype(np.float32)
        if vectors.shape[1] != DENSE_DIM:
            raise SystemExit(f"{model_dir}: expected {DENSE_DIM} dimensions, got {vectors.shape[1]}")
        np.save(path, vectors)
        order_path.write_text(json.dumps(atom_ids), encoding="utf-8")
        return vectors

    def scores(self, query: str, model_dir: Path) -> np.ndarray:
        model = _sentence_model(model_dir)
        vector = model.encode([query], normalize_embeddings=True, convert_to_numpy=True).astype(np.float32)
        return np.asarray(self._vectors @ vector[0], dtype=np.float64)


def _sentence_model(model_dir: Path) -> Any:
    """sentence-transformers ships no stubs, so the model is typed Any rather than a fake object."""
    from sentence_transformers import SentenceTransformer

    if not model_dir.is_dir():
        raise SystemExit(f"{model_dir}: model not downloaded, run python -m bench.fetch_models first")
    return SentenceTransformer(str(model_dir), local_files_only=True)


def top_sources(scores: np.ndarray, atoms: Sequence[Atom], top_k: int) -> list[str]:
    """Rank sources by their best-scoring atom; ties break on source id."""
    if len(scores) != len(atoms):
        raise SystemExit(f"scores has {len(scores)} entries for {len(atoms)} atoms")
    best: dict[str, float] = {}
    for atom, score in zip(atoms, scores, strict=True):
        current = best.get(atom.source_id)
        if current is None or score > current:
            best[atom.source_id] = float(score)
    ranked = sorted(best.items(), key=lambda item: (-item[1], item[0]))
    return [source_id for source_id, _ in ranked[:top_k]]


def fuse(sparse: np.ndarray, dense: np.ndarray, graph: np.ndarray) -> np.ndarray:
    """The product's weighted mix, with each signal scaled to [0, 1]."""
    return (
        FUSION_WEIGHTS["sparse"] * _unit(sparse)
        + FUSION_WEIGHTS["dense"] * _unit(dense)
        + FUSION_WEIGHTS["graph"] * _unit(graph)
    )


def _unit(values: np.ndarray) -> np.ndarray:
    if values.size == 0:
        return values
    low = float(values.min())
    high = float(values.max())
    if high - low <= 1e-12:
        return np.zeros_like(values)
    return (values - low) / (high - low)


class Reranker:
    """Cross-encoder over the fusion candidate list."""

    def __init__(self, model_dir: Path) -> None:
        if not model_dir.is_dir():
            raise SystemExit(f"{model_dir}: model not downloaded, run python -m bench.fetch_models first")
        from sentence_transformers import CrossEncoder

        self._model = CrossEncoder(str(model_dir), local_files_only=True)

    def scores(self, query: str, texts: Sequence[str]) -> np.ndarray:
        if not texts:
            return np.zeros(0, dtype=np.float64)
        pairs = [(query, text[:RERANK_TEXT_CHARS] or " ") for text in texts]
        return np.asarray(self._model.predict(pairs, batch_size=16), dtype=np.float64)


def arm_needs(arm: str, models: Path) -> bool:
    """Whether an arm's model files are present, so a run can skip it instead of crashing."""
    if arm not in ARMS:
        raise SystemExit(f"{arm}: unknown arm, expected one of {ARMS}")
    required: list[Path] = []
    if arm in ("dense_only", "no_rerank", "fusion_rerank"):
        required.append(local_model_dir(models, DENSE_MODEL))
    if arm == "fusion_rerank":
        required.append(local_model_dir(models, RERANK_MODEL))
    return all(directory.is_dir() for directory in required)


class ArmRunner:
    """One arm, ranked queries in batches, with wall-clock timings recorded."""

    def __init__(self, arm: str, atoms: Sequence[Atom], models: Path) -> None:
        if arm not in ARMS:
            raise SystemExit(f"{arm}: unknown arm, expected one of {ARMS}")
        self.arm = arm
        self.atoms = atoms
        self.models = models
        self.timings: list[float] = []
        self._sparse: SparseIndex | None = None
        self._dense: DenseIndex | None = None
        self._reranker: Reranker | None = None
        if arm in ("bm25_only", "no_rerank", "fusion_rerank"):
            self._sparse = SparseIndex(atoms)
        if arm in ("dense_only", "no_rerank", "fusion_rerank"):
            self._dense = DenseIndex(atoms, local_model_dir(models, DENSE_MODEL))
        if arm == "fusion_rerank":
            self._reranker = Reranker(local_model_dir(models, RERANK_MODEL))

    def rank(self, query: str, top_k: int = 10) -> list[str]:
        started = time.perf_counter()
        ranked = self._rank(query, top_k)
        self.timings.append(time.perf_counter() - started)
        return ranked

    def _rank(self, query: str, top_k: int) -> list[str]:
        if self.arm == "graph_only":
            # The product never populates atom.relations, so the graph signal is a constant.
            graph = np.zeros(len(self.atoms), dtype=np.float64)
            return top_sources(graph, self.atoms, top_k)
        sparse = self._sparse.scores(query) if self._sparse is not None else np.zeros(len(self.atoms))
        if self.arm == "bm25_only":
            return top_sources(sparse, self.atoms, top_k)
        dense_dir = local_model_dir(self.models, DENSE_MODEL)
        dense = (
            self._dense.scores(query, dense_dir) if self._dense is not None else np.zeros(len(self.atoms))
        )
        if self.arm == "dense_only":
            return top_sources(dense, self.atoms, top_k)
        graph = np.zeros(len(self.atoms), dtype=np.float64)
        fused = fuse(sparse, dense, graph)
        if self.arm == "no_rerank":
            return top_sources(fused, self.atoms, top_k)
        candidates = top_sources(fused, self.atoms, RERANK_CANDIDATES)
        by_source = _best_atom_text(self.atoms, candidates)
        if self._reranker is None:
            raise SystemExit(f"{self.arm}: the reranker model is missing, run python -m bench.fetch_models")
        rerank_scores = self._reranker.scores(query, [text for _, text in by_source])
        order = sorted(zip(candidates, rerank_scores, strict=True), key=lambda item: (-item[1], item[0]))
        return [source_id for source_id, _ in order[:top_k]]

    def percentile(self, fraction: float) -> float:
        if not self.timings:
            raise SystemExit(f"{self.arm}: no timings recorded")
        if not 0.0 <= fraction <= 1.0:
            raise SystemExit(f"percentile must be within [0, 1], got {fraction}")
        return float(np.percentile(np.asarray(self.timings), fraction * 100))


def _best_atom_text(atoms: Sequence[Atom], source_ids: Sequence[str]) -> list[tuple[str, str]]:
    longest: dict[str, str] = {}
    for atom in atoms:
        if atom.source_id in longest:
            if len(atom.text) > len(longest[atom.source_id]):
                longest[atom.source_id] = atom.text
        else:
            longest[atom.source_id] = atom.text
    return [(source_id, longest.get(source_id, "")) for source_id in source_ids]


def live_product_rankings(db_home: Path, queries: Sequence[str], top_k: int = 10) -> tuple[list[list[str]], list[float]]:
    """Rankings from the unmodified RetrievalService, with per-query wall time."""
    from ely_eye.config import Settings
    from ely_eye.db import Database
    from ely_eye.retrieval import RetrievalService

    settings = Settings(home=db_home)
    service = RetrievalService(settings=settings, database=Database(settings))
    rankings: list[list[str]] = []
    timings: list[float] = []
    for query in queries:
        started = time.perf_counter()
        hits = service.search(query, top_k=top_k)
        timings.append(time.perf_counter() - started)
        # The product ranks atoms; collapse to sources by first (best) hit so the arms stay comparable.
        sources: list[str] = []
        for hit in hits:
            source_id = hit.atom.source_id
            if source_id not in sources:
                sources.append(source_id)
        rankings.append(sources[:top_k])
    return rankings, timings