File size: 5,523 Bytes
eea689d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """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}
|