ApproximateSearcher
- class lightning_ir.retrieve.base.searcher.ApproximateSearcher(index_dir: Path | str, search_config: SearchConfig, module: BiEncoderModule, use_gpu: bool = True)[source]
Bases:
Searcher- __init__(index_dir: Path | str, search_config: SearchConfig, module: BiEncoderModule, use_gpu: bool = True) None
Initialize the Searcher.
- Parameters:
index_dir (Path | str) – The directory containing the index files.
search_config (SearchConfig) – The configuration for the search.
module (BiEncoderModule) – The bi-encoder module to use for scoring.
use_gpu (bool) – Whether to use GPU for computations. Defaults to True.
- Raises:
ValueError – If the document lengths do not match the index.
Methods
search(output)Search for documents based on the output of the bi-encoder model.
- search(output: BiEncoderOutput) Tuple[PackedTensor, List[List[str]]][source]
Search for documents based on the output of the bi-encoder model.
- Parameters:
output (BiEncoderOutput) – The output from the bi-encoder model containing query and document embeddings.
- Returns:
The top-k scores and corresponding document IDs.
- Return type:
Tuple[PackedTensor, List[List[str]]]