Module search 

Query-time search over a built Index.

The flow mirrors PLAID’s reference implementation:

1. For every query token, find the n_probe top dot-product coarse centroids (matching S_c,q = C · Q^T). This is paper Stage 1.
2. Optional centroid pruning: drop probed centroids whose best score across all query tokens sits below centroid_score_threshold. This is used both here (to shrink the reachable centroid set before postings are gathered) and below, at token level, when scoring D̃.
3. Union the doc-level IVF postings for the surviving centroids to get a set of candidate documents.
4. Optional centroid interaction. When n_candidate_docs is set, the cascade runs in two refinement passes:
   • Stage 2 (paper §4.3): approximate MaxSim with the pruned mask applied at token level; keep top n_candidate_docs (paper’s ndocs).
   • Stage 3 (paper §4.2): approximate MaxSim without pruning on the Stage 2 survivors; keep top n_candidate_docs / 4 (paper’s ndocs/4 empirical heuristic), clamped to at least top_k.
5. Decode the survivors’ stored residual codes into approximate embeddings and compute exact MaxSim against the query.
6. Sort by exact MaxSim score and return the top-top_k.

The pruning and interaction stages share a precomputed [n_query_tokens, n_centroids] query-centroid score matrix that step 1 would conceptually compute anyway, and together they drastically shrink the set of candidates reaching the expensive decode in step 5. Callers that want the legacy behaviour (every probed candidate decodes, no pruning) can pass n_candidate_docs = None and centroid_score_threshold = None.

Structs

SearchParams

Tunable knobs for a single search call.

SearchResult

One entry in a ranked search result list.

Functions

search

Run a PLAID-style search over index and return the top-top_k documents ranked by ColBERT MaxSim against query_tokens.

top_n_centroids

Indices of the n centroids with the highest dot-product against point, most relevant first.