1"""Base searcher class and configuration for retrieval tasks."""
2
3from __future__ import annotations
4
5from abc import ABC, abstractmethod
6from pathlib import Path
7from typing import TYPE_CHECKING, Literal
8
9import torch
10
11from ...bi_encoder.bi_encoder_model import BiEncoderEmbedding, SingleVectorBiEncoderConfig
12from .packed_tensor import PackedTensor
13
14if TYPE_CHECKING:
15 from ...bi_encoder import BiEncoderModule, BiEncoderOutput
16
17
[docs]
18def cat_arange(arange_starts: torch.Tensor, arange_ends: torch.Tensor) -> torch.Tensor:
19 """Concatenates arange tensors into a single tensor.
20
21 Args:
22 arange_starts (torch.Tensor): The start indices of the ranges.
23 arange_ends (torch.Tensor): The end indices of the ranges.
24 Returns:
25 torch.Tensor: A tensor containing the concatenated ranges.
26 """
27 arange_lengths = arange_ends - arange_starts
28 offsets = torch.cumsum(arange_lengths, dim=0) - arange_lengths - arange_starts
29 return torch.arange(arange_lengths.sum()) - torch.repeat_interleave(offsets, arange_lengths)
30
31
[docs]
32class Searcher(ABC):
33 """Base class for searchers in the Lightning IR framework."""
34
[docs]
35 def __init__(
36 self, index_dir: Path | str, search_config: SearchConfig, module: BiEncoderModule, use_gpu: bool = True
37 ) -> None:
38 """Initialize the Searcher.
39
40 Args:
41 index_dir (Path | str): The directory containing the index files.
42 search_config (SearchConfig): The configuration for the search.
43 module (BiEncoderModule): The bi-encoder module to use for scoring.
44 use_gpu (bool): Whether to use GPU for computations. Defaults to True.
45 Raises:
46 ValueError: If the document lengths do not match the index.
47 """
48 super().__init__()
49 self.index_dir = Path(index_dir)
50 self.search_config = search_config
51 self.use_gpu = use_gpu
52 self.module = module
53 self.device = torch.device("cuda") if use_gpu and torch.cuda.is_available() else torch.device("cpu")
54
55 self.doc_ids = (self.index_dir / "doc_ids.txt").read_text().split()
56 self.doc_lengths = torch.load(self.index_dir / "doc_lengths.pt", weights_only=True)
57
58 self.to_gpu()
59
60 self.num_docs = len(self.doc_ids)
61 self.cumulative_doc_lengths = torch.cumsum(self.doc_lengths, dim=0)
62 self.num_embeddings = int(self.cumulative_doc_lengths[-1].item())
63
64 self.doc_is_single_vector = self.num_docs == self.num_embeddings
65 self.query_is_single_vector = isinstance(module.config, SingleVectorBiEncoderConfig)
66
67 if self.doc_lengths.shape[0] != self.num_docs or self.doc_lengths.sum() != self.num_embeddings:
68 raise ValueError("doc_lengths do not match index")
69
[docs]
70 def to_gpu(self) -> None:
71 """Move the searcher to the GPU if available."""
72 self.doc_lengths = self.doc_lengths.to(self.device)
73
74 def _filter_and_sort(
75 self, doc_scores: PackedTensor, doc_idcs: PackedTensor, k: int | None = None
76 ) -> tuple[PackedTensor, PackedTensor]:
77 """Filter and sort the document scores and indices.
78
79 Args:
80 doc_scores (PackedTensor): The document scores.
81 doc_idcs (PackedTensor): The document indices.
82 k (int | None): The number of top documents to return. If None, use the configured k from search_config.
83 Defaults to None.
84 Returns:
85 tuple[PackedTensor, PackedTensor]: The filtered and sorted document scores and indices.
86 """
87 k = k or self.search_config.k
88 per_query_doc_scores = torch.split(doc_scores, doc_scores.lengths)
89 per_query_doc_idcs = torch.split(doc_idcs, doc_idcs.lengths)
90 num_docs = []
91 new_doc_scores = []
92 new_doc_idcs = []
93 for _scores, _idcs in zip(per_query_doc_scores, per_query_doc_idcs):
94 _k = min(k, _scores.shape[0])
95 top_values, top_idcs = torch.topk(_scores, _k)
96 new_doc_scores.append(top_values)
97 new_doc_idcs.append(_idcs[top_idcs])
98 num_docs.append(_k)
99 return PackedTensor(torch.cat(new_doc_scores), lengths=num_docs), PackedTensor(
100 torch.cat(new_doc_idcs), lengths=num_docs
101 )
102
[docs]
103 @abstractmethod
104 def search(self, output: BiEncoderOutput) -> tuple[PackedTensor, list[list[str]]]:
105 """Search for documents based on the output of the bi-encoder model.
106
107 Args:
108 output (BiEncoderOutput): The output from the bi-encoder model containing query and document embeddings.
109 Returns:
110 tuple[PackedTensor, list[list[str]]]: The top-k scores and corresponding document IDs.
111 """
112 ...
113
114
[docs]
115class ExactSearcher(Searcher):
116 """Searcher that retrieves documents using exact matching of query embeddings."""
117
[docs]
118 def search(self, output: BiEncoderOutput) -> tuple[PackedTensor, list[list[str]]]:
119 """Search for documents based on the output of the bi-encoder model.
120
121 Args:
122 output (BiEncoderOutput): The output from the bi-encoder model containing query and document embeddings.
123 Returns:
124 tuple[PackedTensor, list[list[str]]]: The top-k scores and corresponding document IDs.
125 """
126 query_embeddings = output.query_embeddings
127 if query_embeddings is None:
128 raise ValueError("Expected query_embeddings in BiEncoderOutput")
129 query_embeddings = query_embeddings.to(self.device)
130
131 scores = self._score(query_embeddings)
132
133 # aggregate doc token scores
134 if not self.doc_is_single_vector:
135 scores = torch.scatter_reduce(
136 torch.zeros(scores.shape[0], self.num_docs, device=scores.device),
137 1,
138 self.doc_token_idcs[None].long().expand_as(scores),
139 scores,
140 "amax",
141 )
142
143 # aggregate query token scores
144 if not self.query_is_single_vector:
145 if query_embeddings.scoring_mask is None:
146 raise ValueError("Expected scoring_mask in multi-vector query_embeddings")
147 query_lengths = query_embeddings.scoring_mask.sum(-1)
148 query_token_idcs = torch.arange(query_lengths.shape[0]).to(query_lengths).repeat_interleave(query_lengths)
149 scores = torch.scatter_reduce(
150 torch.zeros(query_lengths.shape[0], self.num_docs, device=scores.device),
151 0,
152 query_token_idcs[:, None].expand_as(scores),
153 scores,
154 self.module.config.query_aggregation_function,
155 )
156 top_scores, top_idcs = torch.topk(scores, self.search_config.k)
157 doc_ids = [[self.doc_ids[idx] for idx in _doc_idcs] for _doc_idcs in top_idcs.tolist()]
158 return PackedTensor(top_scores.view(-1), lengths=[self.search_config.k] * len(doc_ids)), doc_ids
159
160 @property
161 def doc_token_idcs(self) -> torch.Tensor:
162 """Get the document token indices for scoring.
163
164 Returns:
165 torch.Tensor: The document token indices.
166 """
167 if not hasattr(self, "_doc_token_idcs"):
168 self._doc_token_idcs = (
169 torch.arange(self.doc_lengths.shape[0])
170 .to(device=self.doc_lengths.device)
171 .repeat_interleave(self.doc_lengths)
172 )
173 return self._doc_token_idcs
174
175 @abstractmethod
176 def _score(self, query_embeddings: BiEncoderEmbedding) -> torch.Tensor: ...
177
178
[docs]
179class ApproximateSearcher(Searcher):
[docs]
180 def search(self, output: BiEncoderOutput) -> tuple[PackedTensor, list[list[str]]]:
181 """Search for documents based on the output of the bi-encoder model.
182
183 Args:
184 output (BiEncoderOutput): The output from the bi-encoder model containing query and document embeddings.
185 Returns:
186 tuple[PackedTensor, list[list[str]]]: The top-k scores and corresponding document IDs.
187 """
188 query_embeddings = output.query_embeddings
189 if query_embeddings is None:
190 raise ValueError("Expected query_embeddings in BiEncoderOutput")
191 query_embeddings = query_embeddings.to(self.device)
192
193 candidate_scores, candidate_idcs = self._candidate_retrieval(query_embeddings)
194 scores, doc_idcs = self._aggregate_doc_scores(candidate_scores, candidate_idcs, query_embeddings)
195 scores = self._aggregate_query_scores(scores, query_embeddings)
196 scores, doc_idcs = self._filter_and_sort(scores, doc_idcs)
197 doc_ids = [
198 [self.doc_ids[doc_idx] for doc_idx in _doc_ids.tolist()] for _doc_ids in doc_idcs.split(doc_idcs.lengths)
199 ]
200
201 return scores, doc_ids
202
203 def _aggregate_doc_scores(
204 self, candidate_scores: PackedTensor, candidate_idcs: PackedTensor, query_embeddings: BiEncoderEmbedding
205 ) -> tuple[PackedTensor, PackedTensor]:
206 if self.doc_is_single_vector:
207 return candidate_scores, candidate_idcs
208
209 query_lengths = query_embeddings.scoring_mask.sum(-1)
210 num_query_vecs = query_lengths.sum()
211
212 # map vec_idcs to doc_idcs
213 candidate_doc_idcs = torch.searchsorted(
214 self.cumulative_doc_lengths,
215 candidate_idcs.to(self.cumulative_doc_lengths.device),
216 side="right",
217 )
218
219 # convert candidate_scores `num_query_vecs x candidate_k` to `num_query_doc_pairs x num_query_vecs`
220 # and aggregate the maximum doc_vector score per query_vector
221 max_query_length = query_lengths.max()
222 num_docs_per_query_candidate = torch.tensor(candidate_scores.lengths)
223
224 query_idcs = (
225 torch.arange(query_lengths.shape[0], device=query_lengths.device)
226 .repeat_interleave(query_lengths)
227 .repeat_interleave(num_docs_per_query_candidate)
228 )
229 query_vector_idcs = cat_arange(torch.zeros_like(query_lengths), query_lengths).repeat_interleave(
230 num_docs_per_query_candidate
231 )
232
233 stacked = torch.stack([query_idcs, candidate_doc_idcs])
234 unique_idcs, ranking_doc_idcs = stacked.unique(return_inverse=True, dim=1)
235 num_docs = unique_idcs[0].bincount()
236 doc_idcs = PackedTensor(unique_idcs[1], lengths=num_docs.tolist())
237 total_num_docs = num_docs.sum()
238
239 unpacked_scores = torch.full((total_num_docs * max_query_length,), float("nan"), device=query_lengths.device)
240 index = ranking_doc_idcs * max_query_length + query_vector_idcs
241 unpacked_scores = torch.scatter_reduce(
242 unpacked_scores, 0, index, candidate_scores, "max", include_self=False
243 ).view(total_num_docs, max_query_length)
244
245 # impute the missing values
246 if self.search_config.imputation_strategy == "gather":
247 # reconstruct the doc embeddings and re-compute the scores
248 imputation_values = torch.empty_like(unpacked_scores)
249 doc_embeddings = self._reconstruct_doc_embeddings(doc_idcs)
250 similarity = self.module.model.compute_similarity(query_embeddings, doc_embeddings, doc_idcs.lengths)
251 unpacked_scores = self.module.model._aggregate(
252 similarity, doc_embeddings.scoring_mask, "max", dim=-1
253 ).squeeze(-1)
254 elif self.search_config.imputation_strategy == "min":
255 per_query_vec_min = torch.scatter_reduce(
256 torch.empty(num_query_vecs),
257 0,
258 torch.arange(query_lengths.sum()).repeat_interleave(num_docs_per_query_candidate),
259 candidate_scores,
260 "min",
261 include_self=False,
262 )
263 imputation_values = torch.nn.utils.rnn.pad_sequence(
264 per_query_vec_min.split(query_lengths.tolist()), batch_first=True
265 ).repeat_interleave(num_docs, dim=0)
266 elif self.search_config.imputation_strategy == "zero":
267 imputation_values = torch.zeros_like(unpacked_scores)
268 else:
269 raise ValueError(f"Invalid imputation strategy: {self.search_config.imputation_strategy}")
270
271 is_nan = torch.isnan(unpacked_scores)
272 unpacked_scores[is_nan] = imputation_values[is_nan]
273
274 return PackedTensor(unpacked_scores, lengths=num_docs.tolist()), doc_idcs
275
276 def _aggregate_query_scores(self, scores: PackedTensor, query_embeddings: BiEncoderEmbedding) -> PackedTensor:
277 if self.query_is_single_vector:
278 return scores
279 query_scoring_mask = query_embeddings.scoring_mask.repeat_interleave(torch.tensor(scores.lengths), dim=0)
280 scores = PackedTensor(
281 self.module.model._aggregate(
282 scores, query_scoring_mask, self.module.config.query_aggregation_function, dim=1
283 ).squeeze(-1),
284 lengths=scores.lengths,
285 )
286 return scores
287
288 @abstractmethod
289 def _candidate_retrieval(self, query_embeddings: BiEncoderEmbedding) -> tuple[PackedTensor, PackedTensor]:
290 """Retrieves initial candidates using the query embeddings. Returns candidate scores and candidate vector
291 indices of shape `num_query_vecs x candidate_k` (packed). Candidate indices are None if all doc vectors are
292 scored.
293
294 Args:
295 query_embeddings (BiEncoderEmbedding): The query embeddings to use for candidate retrieval.
296 Returns:
297 tuple[PackedTensor, PackedTensor]: The candidate scores and candidate vector indices.
298 """
299 ...
300
301 @abstractmethod
302 def _gather_doc_embeddings(self, idcs: torch.Tensor) -> torch.Tensor:
303 """Gather document embeddings based on the provided indices.
304
305 Args:
306 idcs (torch.Tensor): The indices of the document embeddings to gather.
307 Returns:
308 torch.Tensor: The gathered document embeddings.
309 """
310 ...
311
312 def _reconstruct_doc_embeddings(self, doc_idcs: PackedTensor) -> BiEncoderEmbedding:
313 """Reconstruct document embeddings based on the provided document indices.
314
315 Args:
316 doc_idcs (PackedTensor): The packed tensor containing document indices.
317 Returns:
318 BiEncoderEmbedding: The reconstructed document embeddings.
319 """
320 # unique doc_idcs per query
321 unique_doc_idcs, inverse_idcs = torch.unique(doc_idcs, return_inverse=True)
322
323 # gather all vectors for unique doc_idcs
324 doc_lengths = self.doc_lengths[unique_doc_idcs]
325 start_doc_idcs = self.cumulative_doc_lengths[unique_doc_idcs - 1]
326 start_doc_idcs[unique_doc_idcs == 0] = 0
327 all_doc_idcs = cat_arange(start_doc_idcs, start_doc_idcs + doc_lengths)
328 all_doc_embeddings = self._gather_doc_embeddings(all_doc_idcs)
329 unique_embeddings = torch.nn.utils.rnn.pad_sequence(
330 list(torch.split(all_doc_embeddings, doc_lengths.tolist())),
331 batch_first=True,
332 ).to(inverse_idcs.device)
333 embeddings = unique_embeddings[inverse_idcs]
334
335 # mask out padding
336 doc_lengths = doc_lengths[inverse_idcs]
337 scoring_mask = torch.arange(embeddings.shape[1], device=embeddings.device) < doc_lengths[:, None]
338 doc_embeddings = BiEncoderEmbedding(embeddings=embeddings, scoring_mask=scoring_mask, encoding=None)
339 return doc_embeddings
340
341
[docs]
342class SearchConfig:
343 """Configuration class for searchers in the Lightning IR framework."""
344
345 search_class: type[Searcher]
346
347 SUPPORTED_MODELS: set[str]
348
[docs]
349 def __init__(self, k: int = 10) -> None:
350 """Initialize the SearchConfig.
351
352 Args:
353 k (int): The number of top documents to retrieve. Defaults to 10.
354 """
355 self.k = k
356
357
[docs]
358class ExactSearchConfig(SearchConfig):
359 """Configuration class for exact searchers in the Lightning IR framework."""
360
361 search_class = ExactSearcher
362
363
[docs]
364class ApproximateSearchConfig(SearchConfig):
365 """Configuration class for approximate searchers in the Lightning IR framework."""
366
367 search_class = ApproximateSearcher
368
[docs]
369 def __init__(
370 self, k: int = 10, candidate_k: int = 100, imputation_strategy: Literal["min", "gather", "zero"] = "gather"
371 ) -> None:
372 """Initialize the ApproximateSearchConfig.
373
374 Args:
375 k (int): The number of top documents to retrieve. Defaults to 10.
376 candidate_k (int): The number of candidate documents to consider for scoring. Defaults to 100.
377 imputation_strategy (Literal["min", "gather", "zero"]): Strategy for imputing missing scores. Defaults to
378 "gather".
379 """
380 super().__init__(k)
381 self.k = k
382 self.candidate_k = candidate_k
383 self.imputation_strategy = imputation_strategy
