src/endopath/precomputed_retrieval.py

"""CPU-only serve-time visual retrieval (issue #21).

The deployed app does not load ColPali or torch. Instead it scores precomputed
page vectors against precomputed field-query vectors with the same
late-interaction (MaxSim / ColBERT) rule colpali_retrieval uses live, but in
plain numpy. Both the page images and the nine field queries are known ahead
of time, so their vectors are precomputed offline on GPU (issue #19) and this
module just scores them, which is cheap CPU matrix math.

The scoring reproduces colpali's `processor.score_multi_vector`: for a query's
token vectors Q [Tq, D] and a page's patch vectors P [Tp, D], the page score is
sum over query tokens of the max over page patches of their dot product
(sum_q max_p (q . p)). A slow test checks this numpy version against the real
score_multi_vector so the two cannot drift.

Storage format (what issue #19 writes, what this loads):
  - data/embeddings/query_embeddings.npz : one array per field name, each [Tq, D].
  - data/embeddings/{patient_filename}.npz : one array per page, keyed
    page_000, page_001, ... each [Tp, D] (per-page keys, so pages with
    different patch counts are handled).
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional

import numpy as np

EMBEDDINGS_DIR = Path("data") / "embeddings"
QUERY_EMBEDDINGS_NAME = "query_embeddings.npz"


def score_pages(query_emb: np.ndarray, page_embs: list[np.ndarray]) -> np.ndarray:
    """MaxSim score of a query against each page. query_emb is [Tq, D]; each
    page in page_embs is [Tp, D] (patch counts may differ per page). Returns
    one score per page. Computed in float32 so float16-stored vectors don't
    lose precision in the dot products."""
    query = query_emb.astype(np.float32, copy=False)
    scores = np.empty(len(page_embs), dtype=np.float32)
    for i, page in enumerate(page_embs):
        # [Tq, Tp] token-vs-patch similarities; max over patches per query
        # token, summed over query tokens.
        sim = query @ page.astype(np.float32, copy=False).T
        scores[i] = sim.max(axis=1).sum()
    return scores


def rank_pages_with_margin(
    query_emb: np.ndarray, page_embs: list[np.ndarray]
) -> tuple[int, float, float]:
    """(1-indexed best page, top score, margin), mirroring
    colpali_retrieval.rank_pages_with_margin. Margin is best minus second-best
    (the query-relative confidence signal #9 gates on); a single-page case has
    no second page, so its top score is returned as the margin."""
    if not page_embs:
        raise ValueError("no page embeddings provided")
    scores = score_pages(query_emb, page_embs)
    order = np.argsort(scores)[::-1]
    best = int(order[0])
    top_score = float(scores[best])
    margin = top_score - float(scores[order[1]]) if len(scores) > 1 else top_score
    return best + 1, top_score, margin


def rank_pages(query_emb: np.ndarray, page_embs: list[np.ndarray]) -> tuple[int, float]:
    page_number, top_score, _ = rank_pages_with_margin(query_emb, page_embs)
    return page_number, top_score


# --- model-output -> storage-format conversion -----------------------------
# Duck-typed (calls .detach()/.cpu() only if present) so this module stays
# torch-free: the precompute script (issue #19) passes real torch tensors, the
# tests pass numpy arrays, and both work.


def _to_numpy(array_like) -> np.ndarray:
    if hasattr(array_like, "detach"):
        array_like = array_like.detach().cpu()
    return np.asarray(array_like, dtype=np.float32)


def pages_from_model_output(page_output) -> list[np.ndarray]:
    """Convert embed_pages' [n_pages, n_patches, D] output to the per-page
    numpy arrays this module stores."""
    return [_to_numpy(page_output[i]) for i in range(len(page_output))]


def query_from_model_output(query_output) -> np.ndarray:
    """Convert embed_query's [1, Tq, D] output to a [Tq, D] numpy array."""
    return _to_numpy(query_output[0])


# --- storage format I/O ----------------------------------------------------


def _case_path(patient_filename: str, embeddings_dir: Optional[Path] = None) -> Path:
    return (embeddings_dir or EMBEDDINGS_DIR) / f"{patient_filename}.npz"


def save_case_page_embeddings(
    patient_filename: str, page_embs: list[np.ndarray], embeddings_dir: Optional[Path] = None
) -> Path:
    path = _case_path(patient_filename, embeddings_dir)
    path.parent.mkdir(parents=True, exist_ok=True)
    arrays = {f"page_{i:03d}": page for i, page in enumerate(page_embs)}
    np.savez(path, **arrays)
    return path


def load_case_page_embeddings(
    patient_filename: str, embeddings_dir: Optional[Path] = None
) -> list[np.ndarray]:
    path = _case_path(patient_filename, embeddings_dir)
    with np.load(path) as data:
        return [data[key] for key in sorted(data.files)]


def has_case_embeddings(patient_filename: str, embeddings_dir: Optional[Path] = None) -> bool:
    return _case_path(patient_filename, embeddings_dir).exists()


def save_query_embeddings(
    query_embs: dict[str, np.ndarray], embeddings_dir: Optional[Path] = None
) -> Path:
    base = embeddings_dir or EMBEDDINGS_DIR
    base.mkdir(parents=True, exist_ok=True)
    path = base / QUERY_EMBEDDINGS_NAME
    np.savez(path, **query_embs)
    return path


def load_query_embeddings(embeddings_dir: Optional[Path] = None) -> dict[str, np.ndarray]:
    path = (embeddings_dir or EMBEDDINGS_DIR) / QUERY_EMBEDDINGS_NAME
    with np.load(path) as data:
        return {key: data[key] for key in data.files}