From 5da4da3cfe4a383afce63c34fbad297cbb468191 Mon Sep 17 00:00:00 2001 From: Vishal Bakshi Date: Sun, 27 Jul 2025 18:07:27 -0700 Subject: [PATCH 1/4] Create ROADMAP.md --- ROADMAP.md | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 ROADMAP.md diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 00000000..c6066192 --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,36 @@ +# stanford-futuredata/ColBERT Roadmap + +## Guiding Philosophy +stanford-futuredata/ColBERT should remain a stable, canonical reference implementation of late interaction, especially for newcomers to late interaction and the GPU-poor*. Folks looking for the bleeding-edge of late interaction should check out the fantastic [lightonai/PyLate library](https://github.com/lightonai/pylate). + +\*stanford-futuredata/ColBERT's coupled model-index design enables [batched encoding](https://github.com/stanford-futuredata/ColBERT/blob/8627585ad290c21720eaa54e325e7c8c301d15f6/colbert/indexing/collection_indexer.py#L376) with immediate compression, maintaining [sub-5GB memory usage](https://vishalbakshi.github.io/blog/posts/2025-02-14-RAGatouille-ColBERT-Memory-Profiling/#profiling-results) even for multi-million document collections. + +## Immediate Goal: Dependency Stabilization (~3 months) + +- Upgrade PyTorch to 2.x +- Upgrade transformers (remove deprecated AdamW) +- Replace faiss with [fastkmeans](https://github.com/AnswerDotAI/fastkmeans) +- Test Python 3.9-3.12 compatibility +- Resolve crypt.h/ninja errors. +- Merge distributed training fix in [#258](https://github.com/stanford-futuredata/ColBERT/pull/258/files#diff-12632f8041dc63139b026f92118749d36110bc0fbbbd6180206b3109fc694c7f) (potentially related: [#132](https://github.com/stanford-futuredata/ColBERT/issues/132) and [#233](https://github.com/stanford-futuredata/ColBERT/issues/233)) +- Replace git-python with GitPython in PyPI (already changed in repo in commit [736f88b](https://github.com/stanford-futuredata/ColBERT/commit/736f88b981078a2c8687c8ee33c0f390e01284cd)) + +## Medium-Term Goals: Documentation and Bug Fixes (~6 months) + +- Update documentation + - Address documentation updates in issues/PRs ([#316](https://github.com/stanford-futuredata/ColBERT/pull/316), [#153](https://github.com/stanford-futuredata/ColBERT/issues/153), [#167](https://github.com/stanford-futuredata/ColBERT/issues/167), etc.). + - Create an llms.txt and llms_ctx.txt for the repo. +- Investigate issues: + - Bug ([#159](https://github.com/stanford-futuredata/ColBERT/issues/159), [#317](https://github.com/stanford-futuredata/ColBERT/issues/317), [#360](https://github.com/stanford-futuredata/ColBERT/issues/360), etc.) + - Training ([#262](https://github.com/stanford-futuredata/ColBERT/issues/262), [#265](https://github.com/stanford-futuredata/ColBERT/issues/265), [#291](https://github.com/stanford-futuredata/ColBERT/issues/291), etc.). + - IndexUpdater ([#180](https://github.com/stanford-futuredata/ColBERT/issues/180), [#261](https://github.com/stanford-futuredata/ColBERT/issues/261), [#276](https://github.com/stanford-futuredata/ColBERT/issues/276), etc.). + - Multi-GPU ([#158](https://github.com/stanford-futuredata/ColBERT/issues/158), [#265](https://github.com/stanford-futuredata/ColBERT/issues/265), [#318](https://github.com/stanford-futuredata/ColBERT/issues/318), etc.). + - Ready-to-close issues after review/repro ([#139](https://github.com/stanford-futuredata/ColBERT/issues/139), [#179](https://github.com/stanford-futuredata/ColBERT/issues/179), [#335](https://github.com/stanford-futuredata/ColBERT/issues/335), etc.). + - etc. + +## Long-Term Goals: Feature Requests (~3 months) + +- Resuming training from checkpoint ([#307](https://github.com/stanford-futuredata/ColBERT/issues/307)). +- Allow string pids ([#326](https://github.com/stanford-futuredata/ColBERT/pull/326)). +- Explore [batch size handling options](https://github.com/stanford-futuredata/ColBERT/blob/8627585ad290c21720eaa54e325e7c8c301d15f6/colbert/search/index_storage.py#L121) to resolve OOM during search. +- etc. From 81218db160f61af8ae30611aa2fe62e38ddf8977 Mon Sep 17 00:00:00 2001 From: Vishal Bakshi Date: Sun, 10 Aug 2025 19:29:58 -0700 Subject: [PATCH 2/4] Update version to 0.2.22 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 9f9162cc..cd7aa30d 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ setuptools.setup( name="colbert-ai", - version="0.2.20", + version="0.2.22", author="Omar Khattab", author_email="okhattab@stanford.edu", description="Efficient and Effective Passage Search via Contextualized Late Interaction over BERT", From ec1a9d2d730a466bc6c56cc98330d69106597352 Mon Sep 17 00:00:00 2001 From: robinnarsinghranabhat Date: Wed, 10 Sep 2025 21:58:36 -0500 Subject: [PATCH 3/4] bugfix : Incorrect sample division in multi-gpu ddp training --- colbert/data/examples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colbert/data/examples.py b/colbert/data/examples.py index 543f6f86..074659ac 100644 --- a/colbert/data/examples.py +++ b/colbert/data/examples.py @@ -41,7 +41,7 @@ def tolist(self, rank=None, nranks=None): if rank or nranks: assert rank in range(nranks), (rank, nranks) - return [self.data[idx] for idx in range(0, len(self.data), nranks)] # if line_idx % nranks == rank + return [self.data[idx + rank] for idx in range(0, len(self.data), nranks) if idx + rank < len(self.data)] # if line_idx % nranks == rank return list(self.data) From 546f0882e5da4e6b6153e03e0f54cf871c8cf7a4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 17 May 2026 06:42:23 +0000 Subject: [PATCH 4/4] Add in-line filtered late-interaction search Introduces an inverted-list-based filter that is traversed in parallel with IVF cluster lists during candidate generation, replacing the existing post-filter pattern with a true in-line one. The downstream sort/unique step then operates on the already-filtered packed PIDs rather than on the unfiltered union, which is the main savings. New modules: - colbert/search/filter_index.py FilterIndex (term -> sorted PID list, packed in a StridedTensor), FilterExpression (Term / And / Or), MaterializedFilter (BitmapFilter / SortedListFilter), each exposing filter_packed(packed_pids, lengths) and a per-cell merge-join reference path. - colbert/search/filtered_search.py filtered_ivf_lookup (vectorised) and filtered_ivf_lookup_per_cell (reference per-block layout for a future CUDA kernel). Hooks (additive, backward-compatible): CandidateGeneration.generate_candidates, IndexScorer.retrieve/rank, and Searcher.search/search_all/dense_search/_search_all_Q all accept a new optional materialized_filter kwarg. The existing filter_fn callback path is unchanged. GPU plan: the filter_packed primitive is the natural lowering target for a fused CUDA kernel that reads from the IVF tensor and the filter bitmap (or sorted list) in one pass; for the full-VRAM mode the BitmapFilter path is already branch-free and coalesced. Tests: colbert/tests/inline_filter_test.py exercises FilterIndex construction, expression materialisation, both filter representations, the vectorised and per-cell lookup paths, and randomised equivalence between in-line filtered candidate gen and the legacy post-filter path. 27 tests, all passing on CPU. https://claude.ai/code/session_01M8ykB5AJXoBNknM8dsH8tm --- colbert/search/candidate_generation.py | 13 +- colbert/search/filter_index.py | 441 +++++++++++++++++++++ colbert/search/filtered_search.py | 165 ++++++++ colbert/search/index_storage.py | 27 +- colbert/searcher.py | 24 +- colbert/tests/inline_filter_test.py | 513 +++++++++++++++++++++++++ 6 files changed, 1170 insertions(+), 13 deletions(-) create mode 100644 colbert/search/filter_index.py create mode 100644 colbert/search/filtered_search.py create mode 100644 colbert/tests/inline_filter_test.py diff --git a/colbert/search/candidate_generation.py b/colbert/search/candidate_generation.py index 96a1840f..bce931a9 100644 --- a/colbert/search/candidate_generation.py +++ b/colbert/search/candidate_generation.py @@ -42,7 +42,18 @@ def generate_candidate_scores(self, Q, eids): E = E.cuda() return (Q.unsqueeze(0) @ E.unsqueeze(2)).squeeze(-1).T - def generate_candidates(self, config, Q): + def generate_candidates(self, config, Q, materialized_filter=None): + """Generate candidate PIDs for query ``Q``. + + When ``materialized_filter`` is not ``None`` the filter's allowed-PID + set is applied *during* the IVF cell-list traversal — the downstream + sort/unique runs on the already-filtered packed PIDs rather than on + the unfiltered union. See ``colbert.search.filtered_search``. + """ + if materialized_filter is not None: + from colbert.search.filtered_search import generate_filtered_candidates + return generate_filtered_candidates(self, config, Q, materialized_filter) + ncells = config.ncells assert isinstance(self.ivf, StridedTensor) diff --git a/colbert/search/filter_index.py b/colbert/search/filter_index.py new file mode 100644 index 00000000..e36b0fd5 --- /dev/null +++ b/colbert/search/filter_index.py @@ -0,0 +1,441 @@ +""" +Inverted-list filter index for in-line filtered late-interaction search. + +This module provides: + + - ``FilterIndex`` : maps ``(field, value) -> sorted PID list`` via a + ``StridedTensor`` over int32 PIDs. + - ``FilterExpression`` : logical query (``Term`` / ``And`` / ``Or``). + - ``MaterializedFilter``: query-time representation, either ``BitmapFilter`` + or ``SortedListFilter``. Both expose + ``filter_packed(packed_pids, lengths)``, which is + the primitive invoked while walking the IVF cell + lists during candidate generation. + +The materialised filter is computed once per query and reused across every +centroid cell touched by that query, so the per-query setup cost amortises over +``ncells * query_maxlen`` per-cell intersections. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union + +import torch + +from colbert.search.strided_tensor import StridedTensor + + +PID_DTYPE = torch.int32 + + +def _to_pid_tensor(values: Iterable[int], device: Optional[torch.device] = None) -> torch.Tensor: + if isinstance(values, torch.Tensor): + t = values.to(dtype=PID_DTYPE) + else: + t = torch.tensor(list(values), dtype=PID_DTYPE) + if device is not None: + t = t.to(device) + return t + + +class FilterIndex: + """Inverted lists keyed by ``(field, value)``. + + A passage may be associated with multiple values per field (multi-valued + fields are supported by passing ``Sequence[Any]`` as the per-passage value). + Posting lists are stored sorted-ascending and packed into a single + ``StridedTensor`` for cheap term lookup. + """ + + def __init__( + self, + vocab: Dict[Tuple[str, Any], int], + posting_lists: StridedTensor, + num_passages: int, + ): + self._vocab = vocab + self._posting_lists = posting_lists + self._num_passages = int(num_passages) + + @property + def num_passages(self) -> int: + return self._num_passages + + @property + def num_terms(self) -> int: + return len(self._vocab) + + @property + def fields(self) -> List[str]: + return sorted({f for f, _ in self._vocab.keys()}) + + def has_term(self, field: str, value: Any) -> bool: + return (field, value) in self._vocab + + def term_id(self, field: str, value: Any) -> Optional[int]: + return self._vocab.get((field, value)) + + def term_pids(self, field: str, value: Any) -> torch.Tensor: + """Sorted-ascending int32 tensor of PIDs that match ``field == value``. + + Returns an empty tensor if the term is unknown. + """ + tid = self._vocab.get((field, value)) + if tid is None: + return torch.empty(0, dtype=PID_DTYPE) + pids, _ = self._posting_lists.lookup(torch.tensor([tid], dtype=torch.long)) + return pids.to(PID_DTYPE) + + @classmethod + def from_field_data( + cls, + field_data: Dict[str, Sequence[Any]], + num_passages: int, + use_gpu: bool = False, + ) -> "FilterIndex": + """Build an index from a dict of ``field_name -> per-passage value(s)``. + + Each entry may be a flat sequence (single-valued field) or a sequence + of sequences (multi-valued: ``per_pid_values[pid] -> Sequence[value]``). + Missing values are skipped — if a passage has no entry for a field, it + is simply absent from that field's posting lists. + """ + vocab: Dict[Tuple[str, Any], int] = {} + term_to_pids: Dict[int, List[int]] = {} + + for field, per_pid in field_data.items(): + if len(per_pid) > num_passages: + raise ValueError( + f"field {field!r} has {len(per_pid)} entries but num_passages={num_passages}" + ) + for pid, raw in enumerate(per_pid): + if raw is None: + continue + values: Sequence[Any] + if isinstance(raw, (list, tuple, set, frozenset)): + values = list(raw) + else: + values = [raw] + for v in values: + key = (field, v) + tid = vocab.get(key) + if tid is None: + tid = len(vocab) + vocab[key] = tid + term_to_pids[tid] = [] + term_to_pids[tid].append(pid) + + if len(vocab) == 0: + packed = torch.empty(0, dtype=PID_DTYPE) + lengths = torch.empty(0, dtype=torch.long) + else: + packed_parts: List[torch.Tensor] = [] + lengths_list: List[int] = [] + for tid in range(len(vocab)): + pids_for_term = sorted(set(term_to_pids[tid])) + packed_parts.append(_to_pid_tensor(pids_for_term)) + lengths_list.append(len(pids_for_term)) + packed = torch.cat(packed_parts) if packed_parts else torch.empty(0, dtype=PID_DTYPE) + lengths = torch.tensor(lengths_list, dtype=torch.long) + + posting_lists = StridedTensor(packed, lengths, use_gpu=use_gpu) + return cls(vocab=vocab, posting_lists=posting_lists, num_passages=num_passages) + + +# --------------------------------------------------------------------------- +# Query expressions +# --------------------------------------------------------------------------- + + +class FilterExpression(ABC): + """Logical filter expression. Compose with ``Term``, ``And``, ``Or``.""" + + @abstractmethod + def _collect_sorted_pids(self, fi: FilterIndex) -> torch.Tensor: + """Internal: return the (sorted-ascending) allowed-PID tensor.""" + + def materialize( + self, + filter_index: FilterIndex, + mode: str = "auto", + device: Optional[torch.device] = None, + bitmap_density_threshold: float = 0.02, + ) -> "MaterializedFilter": + """Compile the expression into a runtime ``MaterializedFilter``. + + ``mode``: + - ``'sorted'``: always return a ``SortedListFilter``. + - ``'bitmap'``: always return a ``BitmapFilter``. + - ``'auto'`` : pick by density. Bitmap memory is ``O(N_passages)``, + sorted memory is ``O(N_allowed)``. Bitmap also gives + branch-free O(1) membership, which is much friendlier + for GPU. So we pick ``BitmapFilter`` whenever the + allowed set is at least ``bitmap_density_threshold`` + of the collection, OR whenever target device is GPU. + """ + allowed = self._collect_sorted_pids(filter_index) + if device is not None: + allowed = allowed.to(device) + N = filter_index.num_passages + chosen = mode + if chosen == "auto": + on_gpu = (device is not None) and (torch.device(device).type == "cuda") + density = allowed.numel() / max(N, 1) + chosen = "bitmap" if (on_gpu or density >= bitmap_density_threshold) else "sorted" + if chosen == "bitmap": + return BitmapFilter.from_sorted_pids(allowed, N) + if chosen == "sorted": + return SortedListFilter(allowed, N) + raise ValueError(f"unknown materialization mode: {mode!r}") + + +class Term(FilterExpression): + def __init__(self, field: str, value: Any): + self.field = field + self.value = value + + def _collect_sorted_pids(self, fi: FilterIndex) -> torch.Tensor: + return fi.term_pids(self.field, self.value) + + def __repr__(self) -> str: + return f"Term({self.field!r}, {self.value!r})" + + +class Or(FilterExpression): + def __init__(self, *children: FilterExpression): + if not children: + raise ValueError("Or requires at least one child") + self.children = list(children) + + def _collect_sorted_pids(self, fi: FilterIndex) -> torch.Tensor: + parts = [c._collect_sorted_pids(fi) for c in self.children] + if not parts: + return torch.empty(0, dtype=PID_DTYPE) + cat = torch.cat(parts) + # torch.unique returns sorted-ascending unique values + return torch.unique(cat).to(PID_DTYPE) + + def __repr__(self) -> str: + return "Or(" + ", ".join(repr(c) for c in self.children) + ")" + + +class And(FilterExpression): + def __init__(self, *children: FilterExpression): + if not children: + raise ValueError("And requires at least one child") + self.children = list(children) + + def _collect_sorted_pids(self, fi: FilterIndex) -> torch.Tensor: + parts = [c._collect_sorted_pids(fi) for c in self.children] + # Intersect: sort by length and progressively reduce + parts.sort(key=lambda t: t.numel()) + result = parts[0] + for p in parts[1:]: + if result.numel() == 0: + break + # both sorted-ascending; use searchsorted for vectorized membership + idx = torch.searchsorted(p, result) + idx = idx.clamp(max=p.numel() - 1) if p.numel() > 0 else idx + if p.numel() == 0: + result = result.new_empty(0) + break + mask = p[idx] == result + result = result[mask] + return result.to(PID_DTYPE) + + def __repr__(self) -> str: + return "And(" + ", ".join(repr(c) for c in self.children) + ")" + + +# --------------------------------------------------------------------------- +# Materialised filters +# --------------------------------------------------------------------------- + + +def _segment_count(mask_long: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: + """Return per-segment sum of ``mask_long`` given per-segment ``lengths``. + + Equivalent to ``[mask_long[off:off+L].sum() for off, L in zip(offsets, lengths)]`` + but vectorised via ``scatter_add_``. + """ + if lengths.numel() == 0: + return torch.zeros(0, dtype=lengths.dtype, device=lengths.device) + seg_ids = torch.repeat_interleave( + torch.arange(lengths.numel(), device=lengths.device), lengths.to(lengths.device) + ) + out = torch.zeros(lengths.numel(), dtype=mask_long.dtype, device=mask_long.device) + out.scatter_add_(0, seg_ids.to(out.device), mask_long) + return out.to(lengths.dtype) + + +class MaterializedFilter(ABC): + """Runtime filter — answers ``is_pid_allowed`` plus a batched ``filter_packed``.""" + + def __init__(self, num_passages: int): + self._num_passages = int(num_passages) + + @property + def num_passages(self) -> int: + return self._num_passages + + @property + @abstractmethod + def num_allowed(self) -> int: ... + + @property + @abstractmethod + def device(self) -> torch.device: ... + + @abstractmethod + def contains(self, pids: torch.Tensor) -> torch.Tensor: + """Boolean membership: returns ``mask`` such that ``mask[i] == True`` iff + ``pids[i]`` passes the filter. Shape preserved.""" + + def filter_pids(self, pids: torch.Tensor) -> torch.Tensor: + """Filter a 1-D PID tensor. Useful as a drop-in for the legacy + ``filter_fn`` callback used in :class:`IndexScorer.rank`.""" + if pids.numel() == 0: + return pids + mask = self.contains(pids) + return pids[mask] + + def filter_packed( + self, packed_pids: torch.Tensor, lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Vectorised in-line filtering pass. + + Args: + packed_pids: 1-D int tensor [total_len]. Concatenation of per-cell + (sorted-ascending) PID lists. + lengths : 1-D long tensor [num_cells]. Length of each per-cell list. + + Returns: + ``(filtered_packed, filtered_lengths)`` with the same packed/lengths + semantics — each per-cell slice is now restricted to PIDs that pass + the filter, preserving order. + + The "in-line" property holds because we never materialise the deduped + union of the unfiltered cell lists; the dedup downstream operates on + the already-filtered packed tensor. + """ + assert packed_pids.dim() == 1 + assert lengths.dim() == 1 + if packed_pids.numel() == 0: + return packed_pids, lengths + mask = self.contains(packed_pids) + filtered_packed = packed_pids[mask] + filtered_lengths = _segment_count(mask.long(), lengths) + return filtered_packed, filtered_lengths + + def filter_per_cell_reference( + self, packed_pids: torch.Tensor, lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-cell merge-join reference path. + + Walks each cell list element-by-element and emits only filter-passing + PIDs. This is the canonical "traversed while doing hit collection" + idiom — slow in Python, but it's the textbook algorithm and the + reference layout for a future per-block CUDA kernel (one cell per + thread block, warp-cooperative merge inside). + + Output is identical to :meth:`filter_packed`. + """ + offsets = torch.cumsum(lengths, dim=0) - lengths + out_parts: List[torch.Tensor] = [] + out_lengths: List[int] = [] + for i in range(lengths.numel()): + off = int(offsets[i].item()) + L = int(lengths[i].item()) + if L == 0: + out_lengths.append(0) + continue + cell = packed_pids[off : off + L] + m = self.contains(cell) + kept = cell[m] + out_parts.append(kept) + out_lengths.append(int(kept.numel())) + if out_parts: + filtered_packed = torch.cat(out_parts) + else: + filtered_packed = packed_pids.new_empty(0) + filtered_lengths = torch.tensor(out_lengths, dtype=lengths.dtype, device=lengths.device) + return filtered_packed, filtered_lengths + + +class BitmapFilter(MaterializedFilter): + """Boolean mask of length ``num_passages``. + + Membership is O(1), branch-free, and trivially coalesced on GPU. Memory is + ``num_passages`` bits (here implemented as a ``torch.bool`` tensor so it's + one byte per passage; a packed-bit version is a straightforward swap if + needed). Always preferred when the index lives on GPU. + """ + + def __init__(self, bitmap: torch.Tensor, num_passages: int): + super().__init__(num_passages) + assert bitmap.dtype == torch.bool + assert bitmap.numel() == num_passages + self._bitmap = bitmap + + @classmethod + def from_sorted_pids(cls, sorted_pids: torch.Tensor, num_passages: int) -> "BitmapFilter": + bitmap = torch.zeros(num_passages, dtype=torch.bool, device=sorted_pids.device) + if sorted_pids.numel() > 0: + bitmap[sorted_pids.long()] = True + return cls(bitmap, num_passages) + + @property + def num_allowed(self) -> int: + return int(self._bitmap.sum().item()) + + @property + def device(self) -> torch.device: + return self._bitmap.device + + @property + def bitmap(self) -> torch.Tensor: + return self._bitmap + + def contains(self, pids: torch.Tensor) -> torch.Tensor: + return self._bitmap[pids.long()] + + +class SortedListFilter(MaterializedFilter): + """Sorted-ascending allowed-PID tensor; membership via ``searchsorted``. + + Memory is ``O(N_allowed)`` — preferred when the filter is selective enough + that the bitmap is wasteful. The internal sorted list is also the natural + input to a per-cell merge-join CUDA kernel. + """ + + def __init__(self, sorted_pids: torch.Tensor, num_passages: int): + super().__init__(num_passages) + if sorted_pids.numel() > 0: + assert sorted_pids.dim() == 1 + self._sorted_pids = sorted_pids.contiguous() + + @property + def num_allowed(self) -> int: + return int(self._sorted_pids.numel()) + + @property + def device(self) -> torch.device: + return self._sorted_pids.device + + @property + def sorted_pids(self) -> torch.Tensor: + return self._sorted_pids + + def contains(self, pids: torch.Tensor) -> torch.Tensor: + if self._sorted_pids.numel() == 0: + return torch.zeros(pids.numel(), dtype=torch.bool, device=pids.device) + # searchsorted requires matching dtypes + haystack = self._sorted_pids + if haystack.dtype != pids.dtype: + haystack = haystack.to(pids.dtype) + idx = torch.searchsorted(haystack, pids) + idx_clamped = idx.clamp(max=haystack.numel() - 1) + in_bounds = idx < haystack.numel() + return in_bounds & (haystack[idx_clamped] == pids) diff --git a/colbert/search/filtered_search.py b/colbert/search/filtered_search.py new file mode 100644 index 00000000..0db7925f --- /dev/null +++ b/colbert/search/filtered_search.py @@ -0,0 +1,165 @@ +""" +Filtered candidate generation: walk the IVF cell lists with a filter applied +*during* the traversal, never materialising the deduped union of the +unfiltered cell lists. + +Two equivalent code paths are provided: + + - :func:`filtered_ivf_lookup` : vectorised, batched over all cells. + Single pass over packed PIDs. + Easy to lower to a fused CUDA + kernel. + - :func:`filtered_ivf_lookup_per_cell` : per-cell merge-join reference + implementation. Slow in Python + but exists as the canonical + algorithm and as the layout + template for a per-block kernel. + +Both consume an ``ivf`` :class:`StridedTensor` (centroid -> sorted PID list) +and a :class:`MaterializedFilter`. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from colbert.search.filter_index import MaterializedFilter +from colbert.search.strided_tensor import StridedTensor + + +def _gather_packed_cells( + ivf: StridedTensor, cell_ids: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Slice the underlying packed IVF tensor for the given cells. + + Returns ``(packed_pids, lengths)`` with one segment per ``cell_ids[i]``, + in input order. The returned ``packed_pids`` matches what + :meth:`StridedTensor.lookup` would produce in ``'packed'`` mode but without + going through the C++ segmented_lookup extension — that lets the filtered + path stay pure-torch and load-extension-free for unit tests. + """ + cell_ids = cell_ids.long().cpu() + lengths = ivf.lengths[cell_ids] + offsets = ivf.offsets[cell_ids] + if cell_ids.numel() == 0: + return ivf.tensor.new_empty(0), lengths + starts = offsets.tolist() + lens = lengths.tolist() + parts = [ivf.tensor[s : s + L] for s, L in zip(starts, lens) if L > 0] + if parts: + packed = torch.cat(parts) + else: + packed = ivf.tensor.new_empty(0) + return packed, lengths + + +def filtered_ivf_lookup( + ivf: StridedTensor, + cell_ids: torch.Tensor, + materialized_filter: MaterializedFilter, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Vectorised filtered IVF lookup. + + Walks the (packed) per-cell PID lists once, applies the filter pointwise, + and returns the compacted output along with per-cell filtered lengths. + + Args: + ivf : :class:`StridedTensor` indexed by centroid id. + cell_ids : 1-D int tensor of centroid ids to retrieve. + materialized_filter : a :class:`MaterializedFilter`. + + Returns: + ``(filtered_packed_pids, filtered_lengths)``. + """ + packed_pids, lengths = _gather_packed_cells(ivf, cell_ids) + target_device = materialized_filter.device + if packed_pids.device != target_device: + packed_pids = packed_pids.to(target_device) + if lengths.device != target_device: + lengths = lengths.to(target_device) + return materialized_filter.filter_packed(packed_pids, lengths) + + +def filtered_ivf_lookup_per_cell( + ivf: StridedTensor, + cell_ids: torch.Tensor, + materialized_filter: MaterializedFilter, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-cell merge-join reference implementation. + + Functionally identical to :func:`filtered_ivf_lookup` but loops over + individual cells. Use only as a correctness/teaching reference; for the + real CPU path use the vectorised form. For the GPU full-VRAM path, a + future CUDA kernel can mirror this loop with one block per cell. + """ + packed_pids, lengths = _gather_packed_cells(ivf, cell_ids) + target_device = materialized_filter.device + if packed_pids.device != target_device: + packed_pids = packed_pids.to(target_device) + if lengths.device != target_device: + lengths = lengths.to(target_device) + return materialized_filter.filter_per_cell_reference(packed_pids, lengths) + + +def generate_filtered_candidate_pids( + candidate_generator, + Q: torch.Tensor, + ncells: int, + materialized_filter: MaterializedFilter, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Equivalent of :meth:`CandidateGeneration.generate_candidate_pids` but + with in-line filtering applied during the IVF traversal. + + Returns the (still-duplicated, still-unsorted-across-cells) + filtered packed PIDs plus the centroid score matrix. The downstream + dedup/sort step in :meth:`CandidateGeneration.generate_candidates` will + operate on this filtered output, so it never sees the unfiltered union. + """ + cells, scores = candidate_generator.get_cells(Q, ncells) + pids, _filtered_lengths = filtered_ivf_lookup( + candidate_generator.ivf, cells, materialized_filter + ) + if candidate_generator.use_gpu: + pids = pids.cuda() + return pids, scores + + +def generate_filtered_candidates( + candidate_generator, + config, + Q: torch.Tensor, + materialized_filter: MaterializedFilter, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Drop-in replacement for :meth:`CandidateGeneration.generate_candidates`. + + Same shape of return values, but the filter is applied while collecting + hits from the cluster (centroid) lists rather than as a post-step. + """ + ncells = config.ncells + + assert isinstance(candidate_generator.ivf, StridedTensor) + + Q = Q.squeeze(0) + if candidate_generator.use_gpu: + Q = Q.cuda().half() + assert Q.dim() == 2 + + pids, centroid_scores = generate_filtered_candidate_pids( + candidate_generator, Q, ncells, materialized_filter + ) + + if pids.numel() == 0: + empty = pids + counts = empty.new_empty(0, dtype=torch.long) + if candidate_generator.use_gpu: + empty, counts = empty.cuda(), counts.cuda() + return empty, centroid_scores + + sorter = pids.sort() + pids = sorter.values + pids, pids_counts = torch.unique_consecutive(pids, return_counts=True) + if candidate_generator.use_gpu: + pids, pids_counts = pids.cuda(), pids_counts.cuda() + return pids, centroid_scores diff --git a/colbert/search/index_storage.py b/colbert/search/index_storage.py index e8deb209..b4a336cf 100644 --- a/colbert/search/index_storage.py +++ b/colbert/search/index_storage.py @@ -74,9 +74,11 @@ def set_embeddings_strided(self): def lookup_pids(self, passage_ids, out_device='cuda', return_mask=False): return self.embeddings_strided.lookup_pids(passage_ids, out_device) - def retrieve(self, config, Q): + def retrieve(self, config, Q, materialized_filter=None): Q = Q[:, :config.query_maxlen] # NOTE: Candidate generation uses only the query tokens - pids, centroid_scores = self.generate_candidates(config, Q) + pids, centroid_scores = self.generate_candidates( + config, Q, materialized_filter=materialized_filter + ) return pids, centroid_scores @@ -84,10 +86,24 @@ def embedding_ids_to_pids(self, embedding_ids): all_pids = torch.unique(self.emb2pid[embedding_ids.long()].cuda(), sorted=False) return all_pids - def rank(self, config, Q, filter_fn=None, pids=None): + def rank(self, config, Q, filter_fn=None, pids=None, materialized_filter=None): + """Score the top docs for query ``Q``. + + Args: + materialized_filter: optional :class:`MaterializedFilter`. When set, + the filter is applied *in-line* during candidate generation — see + ``colbert.search.filtered_search``. Mutually preferred over + ``filter_fn`` for selective filters because the candidate set + never includes filter-rejected PIDs. + filter_fn: legacy post-candidate filter callback. Still applied + after candidate generation, including the in-line-filtered case + (for stacking custom logic on top). + """ with torch.inference_mode(): if pids is None: - pids, centroid_scores = self.retrieve(config, Q) + pids, centroid_scores = self.retrieve( + config, Q, materialized_filter=materialized_filter + ) else: pids = torch.tensor(pids, dtype=torch.int32, device=Q.device) centroid_scores = None @@ -101,6 +117,9 @@ def rank(self, config, Q, filter_fn=None, pids=None): if len(pids) == 0: return [], [] + if len(pids) == 0: + return [], [] + scores, pids = self.score_pids(config, Q, pids, centroid_scores) scores_sorter = scores.sort(descending=True) diff --git a/colbert/searcher.py b/colbert/searcher.py index 8bc07c50..65ca58a7 100644 --- a/colbert/searcher.py +++ b/colbert/searcher.py @@ -62,19 +62,24 @@ def encode(self, text: TextQueries, full_length_search=False): return Q - def search(self, text: str, k=10, filter_fn=None, full_length_search=False, pids=None): + def search(self, text: str, k=10, filter_fn=None, full_length_search=False, pids=None, + materialized_filter=None): Q = self.encode(text, full_length_search=full_length_search) - return self.dense_search(Q, k, filter_fn=filter_fn, pids=pids) + return self.dense_search(Q, k, filter_fn=filter_fn, pids=pids, + materialized_filter=materialized_filter) - def search_all(self, queries: TextQueries, k=10, filter_fn=None, full_length_search=False, qid_to_pids=None): + def search_all(self, queries: TextQueries, k=10, filter_fn=None, full_length_search=False, + qid_to_pids=None, materialized_filter=None): queries = Queries.cast(queries) queries_ = list(queries.values()) Q = self.encode(queries_, full_length_search=full_length_search) - return self._search_all_Q(queries, Q, k, filter_fn=filter_fn, qid_to_pids=qid_to_pids) + return self._search_all_Q(queries, Q, k, filter_fn=filter_fn, qid_to_pids=qid_to_pids, + materialized_filter=materialized_filter) - def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None): + def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None, + materialized_filter=None): qids = list(queries.keys()) if qid_to_pids is None: @@ -86,7 +91,8 @@ def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None): *self.dense_search( Q[query_idx:query_idx+1], k, filter_fn=filter_fn, - pids=qid_to_pids[qid] + pids=qid_to_pids[qid], + materialized_filter=materialized_filter, ) ) ) @@ -103,7 +109,8 @@ def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None): return Ranking(data=data, provenance=provenance) - def dense_search(self, Q: torch.Tensor, k=10, filter_fn=None, pids=None): + def dense_search(self, Q: torch.Tensor, k=10, filter_fn=None, pids=None, + materialized_filter=None): if k <= 10: if self.config.ncells is None: self.configure(ncells=1) @@ -126,6 +133,7 @@ def dense_search(self, Q: torch.Tensor, k=10, filter_fn=None, pids=None): if self.config.ndocs is None: self.configure(ndocs=max(k * 4, 4096)) - pids, scores = self.ranker.rank(self.config, Q, filter_fn=filter_fn, pids=pids) + pids, scores = self.ranker.rank(self.config, Q, filter_fn=filter_fn, pids=pids, + materialized_filter=materialized_filter) return pids[:k], list(range(1, k+1)), scores[:k] diff --git a/colbert/tests/inline_filter_test.py b/colbert/tests/inline_filter_test.py new file mode 100644 index 00000000..a6c825b0 --- /dev/null +++ b/colbert/tests/inline_filter_test.py @@ -0,0 +1,513 @@ +"""Unit tests for in-line filtered late-interaction candidate generation. + +These tests deliberately avoid the full ``Indexer``/``Searcher`` stack so they +can run on CPU without a model or a real index. They exercise: + + - ``FilterIndex.from_field_data`` (inverted-list construction) + - ``Term`` / ``And`` / ``Or`` expression materialisation + - ``BitmapFilter`` and ``SortedListFilter`` filter primitives + - ``filtered_ivf_lookup`` (vectorised path) + - ``filtered_ivf_lookup_per_cell`` (reference per-cell merge-join path) + - End-to-end equivalence of in-line vs post-filter candidate generation, + against a brute-force Python reference. +""" + +from __future__ import annotations + +import random +from typing import Dict, List, Set, Tuple + +import pytest +import torch + +from colbert.search.filter_index import ( + And, + BitmapFilter, + FilterIndex, + MaterializedFilter, + Or, + SortedListFilter, + Term, +) +from colbert.search.filtered_search import ( + filtered_ivf_lookup, + filtered_ivf_lookup_per_cell, +) +from colbert.search.strided_tensor import StridedTensor + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _build_ivf(per_centroid_pids: List[List[int]]) -> StridedTensor: + """Build a (sorted-ascending) IVF StridedTensor from a list of cell lists.""" + sorted_lists = [sorted(set(c)) for c in per_centroid_pids] + flat: List[int] = [] + lengths: List[int] = [] + for c in sorted_lists: + flat.extend(c) + lengths.append(len(c)) + packed = torch.tensor(flat, dtype=torch.int32) + lens = torch.tensor(lengths, dtype=torch.long) + return StridedTensor(packed, lens, use_gpu=False) + + +def _brute_force_filter( + per_centroid_pids: List[List[int]], + cell_ids: List[int], + allowed_pids: Set[int], +) -> Tuple[List[int], List[int]]: + """Reference: per-cell intersect with ``allowed_pids``.""" + packed: List[int] = [] + lengths: List[int] = [] + for c in cell_ids: + cell = sorted(set(per_centroid_pids[c])) + kept = [p for p in cell if p in allowed_pids] + packed.extend(kept) + lengths.append(len(kept)) + return packed, lengths + + +def _dedup_sort(pids: List[int]) -> List[int]: + return sorted(set(pids)) + + +# --------------------------------------------------------------------------- +# FilterIndex construction +# --------------------------------------------------------------------------- + + +def test_filter_index_basic_single_value_field(): + """Single-valued field maps each pid -> one term.""" + fi = FilterIndex.from_field_data( + field_data={"lang": ["en", "fr", "en", "de", "fr"]}, + num_passages=5, + ) + assert fi.num_passages == 5 + assert fi.num_terms == 3 + assert fi.has_term("lang", "en") + assert not fi.has_term("lang", "es") + assert _dedup_sort(fi.term_pids("lang", "en").tolist()) == [0, 2] + assert _dedup_sort(fi.term_pids("lang", "fr").tolist()) == [1, 4] + assert _dedup_sort(fi.term_pids("lang", "de").tolist()) == [3] + + +def test_filter_index_multi_value_field(): + """Multi-valued field: a passage can be in multiple posting lists.""" + fi = FilterIndex.from_field_data( + field_data={"tag": [["a", "b"], ["b"], ["a", "c"], []]}, + num_passages=4, + ) + assert fi.num_terms == 3 + assert _dedup_sort(fi.term_pids("tag", "a").tolist()) == [0, 2] + assert _dedup_sort(fi.term_pids("tag", "b").tolist()) == [0, 1] + assert _dedup_sort(fi.term_pids("tag", "c").tolist()) == [2] + + +def test_filter_index_missing_values_skipped(): + fi = FilterIndex.from_field_data( + field_data={"region": ["us", None, "eu", None, "us"]}, + num_passages=5, + ) + assert _dedup_sort(fi.term_pids("region", "us").tolist()) == [0, 4] + assert _dedup_sort(fi.term_pids("region", "eu").tolist()) == [2] + + +def test_filter_index_unknown_term_returns_empty(): + fi = FilterIndex.from_field_data( + field_data={"lang": ["en"]}, + num_passages=1, + ) + assert fi.term_pids("lang", "zz").numel() == 0 + assert fi.term_pids("missing_field", "anything").numel() == 0 + + +# --------------------------------------------------------------------------- +# Expression materialisation +# --------------------------------------------------------------------------- + + +@pytest.fixture +def small_index(): + # 10 passages, two fields + return FilterIndex.from_field_data( + field_data={ + "lang": ["en", "fr", "en", "de", "fr", "en", "en", "de", "fr", "en"], + "year": [2020, 2021, 2021, 2020, 2022, 2022, 2020, 2021, 2022, 2022], + }, + num_passages=10, + ) + + +def test_term_materialize_matches_index_lookup(small_index): + mat = Term("lang", "en").materialize(small_index, mode="sorted") + expected = sorted({0, 2, 5, 6, 9}) + assert _dedup_sort(mat.sorted_pids.tolist()) == expected + assert mat.num_allowed == len(expected) + + +def test_or_materialize_is_union(small_index): + expr = Or(Term("lang", "en"), Term("year", 2022)) + mat = expr.materialize(small_index, mode="sorted") + expected = sorted({0, 2, 5, 6, 9} | {4, 5, 8, 9}) + assert _dedup_sort(mat.sorted_pids.tolist()) == expected + + +def test_and_materialize_is_intersection(small_index): + expr = And(Term("lang", "en"), Term("year", 2022)) + mat = expr.materialize(small_index, mode="sorted") + expected = sorted({0, 2, 5, 6, 9} & {4, 5, 8, 9}) + assert _dedup_sort(mat.sorted_pids.tolist()) == expected + + +def test_and_with_no_matches(small_index): + expr = And(Term("lang", "en"), Term("lang", "fr")) + mat = expr.materialize(small_index, mode="sorted") + assert mat.num_allowed == 0 + + +def test_complex_expression(small_index): + # (lang == en AND year == 2020) OR lang == de + expr = Or( + And(Term("lang", "en"), Term("year", 2020)), + Term("lang", "de"), + ) + mat = expr.materialize(small_index, mode="sorted") + en_2020 = {0, 6} # lang_en = {0,2,5,6,9}, year_2020 = {0,3,6} + de = {3, 7} + expected = sorted(en_2020 | de) + assert _dedup_sort(mat.sorted_pids.tolist()) == expected + + +# --------------------------------------------------------------------------- +# Auto-mode selection +# --------------------------------------------------------------------------- + + +def test_auto_picks_sorted_for_sparse_filter(): + # 1000-passage index; filter selects 5 passages → 0.5% density → sorted + fi = FilterIndex.from_field_data( + field_data={"f": ["x"] * 5 + ["y"] * 995}, + num_passages=1000, + ) + mat = Term("f", "x").materialize(fi, mode="auto") + assert isinstance(mat, SortedListFilter) + + +def test_auto_picks_bitmap_for_dense_filter(): + # 1000-passage index; filter selects 500 → 50% density → bitmap + fi = FilterIndex.from_field_data( + field_data={"f": ["a"] * 500 + ["b"] * 500}, + num_passages=1000, + ) + mat = Term("f", "a").materialize(fi, mode="auto") + assert isinstance(mat, BitmapFilter) + + +# --------------------------------------------------------------------------- +# MaterializedFilter primitives +# --------------------------------------------------------------------------- + + +@pytest.fixture +def made_filters(): + """Same allowed-PID set, materialised both ways.""" + allowed = sorted({2, 3, 5, 7, 11, 13}) + N = 20 + sorted_t = torch.tensor(allowed, dtype=torch.int32) + return { + "allowed": allowed, + "N": N, + "sorted": SortedListFilter(sorted_t, N), + "bitmap": BitmapFilter.from_sorted_pids(sorted_t, N), + } + + +def test_contains_matches(made_filters): + pids = torch.tensor([0, 2, 3, 4, 5, 11, 19, 13], dtype=torch.int32) + expected = torch.tensor([False, True, True, False, True, True, False, True]) + for name in ("sorted", "bitmap"): + mat = made_filters[name] + got = mat.contains(pids) + assert torch.equal(got, expected), name + + +def test_filter_packed_matches_per_cell(made_filters): + # 4 cells, each sorted ascending + cells = [[0, 1, 2, 3, 5], [4, 6, 7, 11], [9, 13, 15], [12, 14, 16, 17, 19]] + packed = torch.tensor(sum(cells, []), dtype=torch.int32) + lengths = torch.tensor([len(c) for c in cells], dtype=torch.long) + for name in ("sorted", "bitmap"): + mat = made_filters[name] + v_pack, v_lens = mat.filter_packed(packed, lengths) + r_pack, r_lens = mat.filter_per_cell_reference(packed, lengths) + assert torch.equal(v_pack, r_pack), name + assert torch.equal(v_lens, r_lens), name + # also compare against pure-python expected + allowed = set(made_filters["allowed"]) + expected_per_cell = [[p for p in c if p in allowed] for c in cells] + expected_packed = sum(expected_per_cell, []) + expected_lens = [len(c) for c in expected_per_cell] + assert v_pack.tolist() == expected_packed, name + assert v_lens.tolist() == expected_lens, name + + +def test_filter_packed_handles_empty_inputs(made_filters): + mat = made_filters["bitmap"] + packed = torch.empty(0, dtype=torch.int32) + lengths = torch.empty(0, dtype=torch.long) + p, L = mat.filter_packed(packed, lengths) + assert p.numel() == 0 + assert L.numel() == 0 + + +def test_filter_packed_all_empty_cells(made_filters): + """Every cell has length 0 — output must preserve the 0-length cells.""" + mat = made_filters["bitmap"] + packed = torch.empty(0, dtype=torch.int32) + lengths = torch.tensor([0, 0, 0], dtype=torch.long) + p, L = mat.filter_packed(packed, lengths) + assert p.numel() == 0 + assert torch.equal(L, torch.tensor([0, 0, 0], dtype=torch.long)) + + +def test_filter_packed_no_match_keeps_cell_count(made_filters): + """All PIDs rejected — output keeps the same number of (zero-length) cells.""" + mat = made_filters["bitmap"] + cells = [[0, 1], [4, 6], [8, 9]] # none of these are in allowed = {2,3,5,7,11,13} + packed = torch.tensor(sum(cells, []), dtype=torch.int32) + lengths = torch.tensor([len(c) for c in cells], dtype=torch.long) + p, L = mat.filter_packed(packed, lengths) + assert p.numel() == 0 + assert L.tolist() == [0, 0, 0] + + +# --------------------------------------------------------------------------- +# filtered_ivf_lookup end-to-end +# --------------------------------------------------------------------------- + + +def test_filtered_ivf_lookup_matches_brute_force(small_index): + per_centroid = [ + [0, 2, 5, 9], # centroid 0 + [1, 4, 8], # centroid 1 + [3, 6, 7], # centroid 2 + [0, 3, 5, 7], # centroid 3 + [], # centroid 4 (empty) + [1, 2, 8, 9], # centroid 5 + ] + ivf = _build_ivf(per_centroid) + + # filter: lang == en (pids {0, 2, 5, 6, 9}) + expr = Term("lang", "en") + cells_to_probe = torch.tensor([0, 2, 3, 4, 5], dtype=torch.long) + + for mode in ("sorted", "bitmap"): + mat = expr.materialize(small_index, mode=mode) + v_pack, v_lens = filtered_ivf_lookup(ivf, cells_to_probe, mat) + r_pack, r_lens = filtered_ivf_lookup_per_cell(ivf, cells_to_probe, mat) + + exp_pack, exp_lens = _brute_force_filter( + per_centroid, cells_to_probe.tolist(), {0, 2, 5, 6, 9} + ) + + assert v_pack.tolist() == exp_pack, mode + assert v_lens.tolist() == exp_lens, mode + assert r_pack.tolist() == exp_pack, mode + assert r_lens.tolist() == exp_lens, mode + + +def test_filtered_ivf_lookup_no_cells(): + per_centroid = [[0, 1], [2, 3]] + ivf = _build_ivf(per_centroid) + fi = FilterIndex.from_field_data({"f": ["x", "x", "y", "y"]}, num_passages=4) + mat = Term("f", "x").materialize(fi, mode="sorted") + p, L = filtered_ivf_lookup(ivf, torch.tensor([], dtype=torch.long), mat) + assert p.numel() == 0 + assert L.numel() == 0 + + +def test_filtered_ivf_lookup_filter_matches_nothing(small_index): + per_centroid = [[0, 1, 2], [3, 4, 5]] + ivf = _build_ivf(per_centroid) + # filter for an unknown value -> empty allowed set + mat = Term("lang", "zz").materialize(small_index, mode="sorted") + p, L = filtered_ivf_lookup(ivf, torch.tensor([0, 1], dtype=torch.long), mat) + assert p.numel() == 0 + assert L.tolist() == [0, 0] + + +# --------------------------------------------------------------------------- +# In-line == post-filter equivalence (the core correctness claim) +# --------------------------------------------------------------------------- + + +def _simulate_postfilter_candidate_gen( + ivf: StridedTensor, cells: torch.Tensor, allowed: Set[int] +) -> List[int]: + """Reference: ColBERT's current path -- union all cells, sort+unique, then filter.""" + parts: List[int] = [] + for c in cells.tolist(): + L = int(ivf.lengths[c].item()) + off = int(ivf.offsets[c].item()) + parts.extend(ivf.tensor[off : off + L].tolist()) + deduped = sorted(set(parts)) + return [p for p in deduped if p in allowed] + + +def _simulate_inline_candidate_gen( + ivf: StridedTensor, cells: torch.Tensor, mat: MaterializedFilter +) -> List[int]: + """New path: filter during cell traversal, then sort+unique.""" + packed, _lens = filtered_ivf_lookup(ivf, cells, mat) + if packed.numel() == 0: + return [] + sorted_packed, _ = packed.sort() + deduped, _ = torch.unique_consecutive(sorted_packed, return_counts=True) + return deduped.tolist() + + +@pytest.mark.parametrize("seed", [0, 1, 2, 3, 4, 5]) +def test_inline_equals_postfilter_randomized(seed): + """Randomised correctness: for arbitrary IVF + filter + cell selection, + in-line filtered candidate gen must yield the same deduped candidate PID + set as the existing post-filter path.""" + rng = random.Random(seed) + N = 200 + num_centroids = 40 + avg_list_len = 20 + # Random IVF: each centroid gets ~avg_list_len random pids (sorted unique) + per_centroid = [] + for _ in range(num_centroids): + k = rng.randint(0, avg_list_len * 2) + c = sorted({rng.randrange(N) for _ in range(k)}) + per_centroid.append(c) + ivf = _build_ivf(per_centroid) + + # Random filter: each passage tagged with one of {a,b,c}; we filter for {a} OR {b} + field = [rng.choice(["a", "b", "c"]) for _ in range(N)] + fi = FilterIndex.from_field_data({"tag": field}, num_passages=N) + + expr = Or(Term("tag", "a"), Term("tag", "b")) + allowed_set = {p for p, v in enumerate(field) if v in ("a", "b")} + + # Sample a random subset of centroids to probe + num_probe = rng.randint(1, num_centroids) + cells = torch.tensor( + sorted(rng.sample(range(num_centroids), num_probe)), dtype=torch.long + ) + + expected = _simulate_postfilter_candidate_gen(ivf, cells, allowed_set) + + for mode in ("sorted", "bitmap"): + mat = expr.materialize(fi, mode=mode) + got = _simulate_inline_candidate_gen(ivf, cells, mat) + assert got == expected, f"seed={seed} mode={mode}" + + +# --------------------------------------------------------------------------- +# Wiring through CandidateGeneration +# --------------------------------------------------------------------------- + + +class _StubCodec: + """Minimal stand-in for ResidualCodec; ``CandidateGeneration.get_cells`` + only needs ``codec.centroids``.""" + def __init__(self, centroids: torch.Tensor): + self.centroids = centroids + + +class _StubConfig: + def __init__(self, ncells: int): + self.ncells = ncells + + +class _StubCandidateGen: + """Minimal harness exposing exactly the attributes + :meth:`CandidateGeneration.generate_candidates` needs.""" + use_gpu = False + + def __init__(self, ivf: StridedTensor, centroids: torch.Tensor): + self.ivf = ivf + self.codec = _StubCodec(centroids) + + +def test_generate_candidates_inline_vs_postfilter(): + """End-to-end: ``CandidateGeneration.generate_candidates`` with and without + ``materialized_filter`` must produce the same deduped candidate PID set + (modulo filtering).""" + from colbert.search.candidate_generation import CandidateGeneration + + rng = random.Random(123) + N = 80 + num_centroids = 12 + per_centroid = [ + sorted({rng.randrange(N) for _ in range(rng.randint(2, 15))}) + for _ in range(num_centroids) + ] + ivf = _build_ivf(per_centroid) + + # Random centroid vectors (low-dim for speed) -- exact values don't matter + # for correctness; we just need ``get_cells`` to pick *some* cells. + # Q has shape (1, num_query_tokens, dim) like the real ColBERT path; after + # the .squeeze(0) inside generate_candidates it becomes (num_query_tokens, dim). + dim = 4 + num_query_tokens = 6 + centroids = torch.randn(num_centroids, dim) + Q = torch.randn(1, num_query_tokens, dim) + + cg = _StubCandidateGen(ivf, centroids) + # bind real methods from CandidateGeneration onto the stub + cg.get_cells = CandidateGeneration.get_cells.__get__(cg, _StubCandidateGen) + cg.generate_candidate_pids = CandidateGeneration.generate_candidate_pids.__get__( + cg, _StubCandidateGen + ) + cg.generate_candidates = CandidateGeneration.generate_candidates.__get__( + cg, _StubCandidateGen + ) + + cfg = _StubConfig(ncells=4) + + # Path 1: unfiltered baseline + unfiltered_pids, _ = cg.generate_candidates(cfg, Q) + unfiltered_set = set(unfiltered_pids.tolist()) + + # Path 2: post-filter callback (legacy mechanism) + field = [rng.choice(["x", "y", "z"]) for _ in range(N)] + fi = FilterIndex.from_field_data({"f": field}, num_passages=N) + allowed = {p for p, v in enumerate(field) if v == "x"} + postfiltered = sorted(unfiltered_set & allowed) + + # Path 3: in-line filtered candidate gen + for mode in ("sorted", "bitmap"): + mat = Term("f", "x").materialize(fi, mode=mode) + inline_pids, _ = cg.generate_candidates(cfg, Q, materialized_filter=mat) + inline_set = sorted(set(inline_pids.tolist())) + assert inline_set == postfiltered, mode + + +def test_generate_candidates_inline_empty_filter_returns_no_pids(): + from colbert.search.candidate_generation import CandidateGeneration + + per_centroid = [[0, 1, 2], [3, 4, 5]] + ivf = _build_ivf(per_centroid) + centroids = torch.randn(2, 4) + Q = torch.randn(1, 3, 4) + + cg = _StubCandidateGen(ivf, centroids) + cg.get_cells = CandidateGeneration.get_cells.__get__(cg, _StubCandidateGen) + cg.generate_candidate_pids = CandidateGeneration.generate_candidate_pids.__get__( + cg, _StubCandidateGen + ) + cg.generate_candidates = CandidateGeneration.generate_candidates.__get__( + cg, _StubCandidateGen + ) + cfg = _StubConfig(ncells=2) + + fi = FilterIndex.from_field_data({"f": ["x"] * 6}, num_passages=6) + # Filter for a value that doesn't exist -> empty allowed set + mat = Term("f", "nope").materialize(fi, mode="sorted") + pids, _ = cg.generate_candidates(cfg, Q, materialized_filter=mat) + assert pids.numel() == 0