Skip to content

ColBERT Retriever

rankify.retrievers.colbert_retriever

Run

Bases: object

Source code in rankify/utils/retrievers/colbert/colbert/infra/run.py 
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
class Run(object):
    _instance = None

    os.environ["TOKENIZERS_PARALLELISM"] = "true"  # NOTE: If a deadlock arises, switch to false!!

    def __new__(cls):
        """
        Singleton Pattern. See https://python-patterns.guide/gang-of-four/singleton/
        """
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.stack = []

            # TODO: Save a timestamp here! And re-use it! But allow the user to override it on calling Run().context a second time.
            run_config = RunConfig()
            run_config.assign_defaults()

            cls._instance.__append(run_config)

        # TODO: atexit.register(all_done)

        return cls._instance

    @property
    def config(self):
        return self.stack[-1]

    def __getattr__(self, name):
        if hasattr(self.config, name):
            return getattr(self.config, name)

        super().__getattr__(name)

    def __append(self, runconfig: RunConfig):
        # runconfig.disallow_writes(readonly=True)
        self.stack.append(runconfig)

    def __pop(self):
        self.stack.pop()

    @contextmanager
    def context(self, runconfig: RunConfig, inherit_config=True):
        if inherit_config:
            runconfig = RunConfig.from_existing(self.config, runconfig)

        self.__append(runconfig)

        try:
            yield
        finally:
            self.__pop()

    def open(self, path, mode='r'):
        path = os.path.join(self.path_, path)

        if not os.path.exists(self.path_):
            create_directory(self.path_)

        if ('w' in mode or 'a' in mode) and not self.overwrite:
            assert not os.path.exists(path), (self.overwrite, path)

            # create directory if it doesn't exist
            os.makedirs(os.path.dirname(path), exist_ok=True)

        return open(path, mode=mode)

    def print(self, *args):
        print_message("[" + str(self.rank) + "]", "\t\t", *args)

    def print_main(self, *args):
        if self.rank == 0:
            self.print(*args)

__new__()

Singleton Pattern. See https://python-patterns.guide/gang-of-four/singleton/

Source code in rankify/utils/retrievers/colbert/colbert/infra/run.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def __new__(cls):
    """
    Singleton Pattern. See https://python-patterns.guide/gang-of-four/singleton/
    """
    if cls._instance is None:
        cls._instance = super().__new__(cls)
        cls._instance.stack = []

        # TODO: Save a timestamp here! And re-use it! But allow the user to override it on calling Run().context a second time.
        run_config = RunConfig()
        run_config.assign_defaults()

        cls._instance.__append(run_config)

    # TODO: atexit.register(all_done)

    return cls._instance

RunConfig dataclass

Bases: BaseConfig, RunSettings

Source code in rankify/utils/retrievers/colbert/colbert/infra/config/config.py 
7
8
9
@dataclass
class RunConfig(BaseConfig, RunSettings):
    pass

ColBERTConfig dataclass

Bases: RunSettings, ResourceSettings, DocSettings, QuerySettings, TrainingSettings, IndexingSettings, SearchSettings, BaseConfig, TokenizerSettings

Source code in rankify/utils/retrievers/colbert/colbert/infra/config/config.py 
12
13
14
15
@dataclass
class ColBERTConfig(RunSettings, ResourceSettings, DocSettings, QuerySettings, TrainingSettings,
                    IndexingSettings, SearchSettings, BaseConfig, TokenizerSettings):
    pass

Searcher

Source code in rankify/utils/retrievers/colbert/colbert/searcher.py 
 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
class Searcher:
    def __init__(self, index, checkpoint=None, collection=None, config=None, index_root=None, verbose:int = 3):
        self.verbose = verbose
        if self.verbose > 1:
            print_memory_stats()

        initial_config = ColBERTConfig.from_existing(config, Run().config)

        default_index_root = initial_config.index_root_
        index_root = index_root if index_root else default_index_root
        self.index = os.path.join(index_root, index)
        self.index_config = ColBERTConfig.load_from_index(self.index)

        self.checkpoint = checkpoint or self.index_config.checkpoint
        self.checkpoint_config = ColBERTConfig.load_from_checkpoint(self.checkpoint)
        self.config = ColBERTConfig.from_existing(self.checkpoint_config, self.index_config, initial_config)

        #print(collection)
        #print(self.config)
        self.collection = Collection.cast(collection or self.config.collection)
        self.configure(checkpoint=self.checkpoint, collection=self.collection)

        self.checkpoint = Checkpoint(self.checkpoint, colbert_config=self.config, verbose=self.verbose)
        use_gpu = self.config.total_visible_gpus > 0
        if use_gpu:
            self.checkpoint = self.checkpoint.cuda()
        load_index_with_mmap = self.config.load_index_with_mmap
        if load_index_with_mmap and use_gpu:
            raise ValueError(f"Memory-mapped index can only be used with CPU!")
        self.ranker = IndexScorer(self.index, use_gpu, load_index_with_mmap)

        print_memory_stats()

    def configure(self, **kw_args):
        self.config.configure(**kw_args)

    def encode(self, text: TextQueries, full_length_search=False):
        queries = text if type(text) is list else [text]
        bsize = 128 if len(queries) > 128 else None

        self.checkpoint.query_tokenizer.query_maxlen = self.config.query_maxlen
        Q = self.checkpoint.queryFromText(queries, bsize=bsize, to_cpu=True, full_length_search=full_length_search)

        return Q

    def search(self, text: str, k=10, filter_fn=None, full_length_search=False, pids=None):
        Q = self.encode(text, full_length_search=full_length_search)
        return self.dense_search(Q, k, filter_fn=filter_fn, pids=pids)

    def search_all(self, queries: TextQueries, k=10, filter_fn=None, full_length_search=False, qid_to_pids=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)

    def _search_all_Q(self, queries, Q, k, filter_fn=None, qid_to_pids=None):
        qids = list(queries.keys())

        if qid_to_pids is None:
            qid_to_pids = {qid: None for qid in qids}

        all_scored_pids = [
            list(
                zip(
                    *self.dense_search(
                        Q[query_idx:query_idx+1],
                        k, filter_fn=filter_fn,
                        pids=qid_to_pids[qid]
                    )
                )
            )
            for query_idx, qid in tqdm(enumerate(qids))
        ]

        data = {qid: val for qid, val in zip(queries.keys(), all_scored_pids)}

        provenance = Provenance()
        provenance.source = 'Searcher::search_all'
        provenance.queries = queries.provenance()
        provenance.config = self.config.export()
        provenance.k = k

        return Ranking(data=data, provenance=provenance)

    def dense_search(self, Q: torch.Tensor, k=10, filter_fn=None, pids=None):
        if k <= 10:
            if self.config.ncells is None:
                self.configure(ncells=1)
            if self.config.centroid_score_threshold is None:
                self.configure(centroid_score_threshold=0.5)
            if self.config.ndocs is None:
                self.configure(ndocs=256)
        elif k <= 100:
            if self.config.ncells is None:
                self.configure(ncells=2)
            if self.config.centroid_score_threshold is None:
                self.configure(centroid_score_threshold=0.45)
            if self.config.ndocs is None:
                self.configure(ndocs=1024)
        else:
            if self.config.ncells is None:
                self.configure(ncells=4)
            if self.config.centroid_score_threshold is None:
                self.configure(centroid_score_threshold=0.4)
            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)

        return pids[:k], list(range(1, k+1)), scores[:k]

BaseRetriever

Bases: ABC

Abstract base class for all retrieval methods in the rankify framework.

This class defines the common interface that all retrievers must implement, ensuring consistency across different retrieval methods (BM25, DPR, etc.).

Source code in rankify/retrievers/base_retriever.py 
 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
class BaseRetriever(ABC):
    """
    Abstract base class for all retrieval methods in the rankify framework.

    This class defines the common interface that all retrievers must implement,
    ensuring consistency across different retrieval methods (BM25, DPR, etc.).
    """

    def __init__(self, n_docs: int = 10, batch_size: int = 36, threads: int = 30, **kwargs):
        """
        Initialize the base retriever.

        Args:
            n_docs (int): Number of documents to retrieve per query
            batch_size (int): Number of queries to process in a batch
            threads (int): Number of parallel threads for retrieval
            **kwargs: Additional parameters specific to each retriever
        """
        self.n_docs = n_docs
        self.batch_size = batch_size
        self.threads = threads
        self.config = kwargs

    @abstractmethod
    def _initialize_searcher(self) -> Any:
        """Initialize the specific searcher for this retriever type."""
        pass

    @abstractmethod
    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve relevant contexts for the given documents.

        Args:
            documents (List[Document]): List of documents containing queries

        Returns:
            List[Document]: Documents updated with retrieved contexts
        """
        pass

    def _preprocess_query(self, query: str) -> str:
        """
        Preprocess query text. Can be overridden by specific retrievers.

        Args:
            query (str): Raw query text

        Returns:
            str: Preprocessed query text
        """
        return query.strip()

__init__(n_docs=10, batch_size=36, threads=30, **kwargs)

Initialize the base retriever.

Parameters:

Name Type Description Default
n_docs int

Number of documents to retrieve per query

 10
batch_size int

Number of queries to process in a batch

 36
threads int

Number of parallel threads for retrieval

 30
**kwargs

Additional parameters specific to each retriever

 {}
Source code in rankify/retrievers/base_retriever.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def __init__(self, n_docs: int = 10, batch_size: int = 36, threads: int = 30, **kwargs):
    """
    Initialize the base retriever.

    Args:
        n_docs (int): Number of documents to retrieve per query
        batch_size (int): Number of queries to process in a batch
        threads (int): Number of parallel threads for retrieval
        **kwargs: Additional parameters specific to each retriever
    """
    self.n_docs = n_docs
    self.batch_size = batch_size
    self.threads = threads
    self.config = kwargs

retrieve(documents) abstractmethod

Retrieve relevant contexts for the given documents.

Parameters:

Name Type Description Default
documents List[Document]

List of documents containing queries

 required

Returns:

Type Description
List[Document]

List[Document]: Documents updated with retrieved contexts
Source code in rankify/retrievers/base_retriever.py 
34
35
36
37
38
39
40
41
42
43
44
45
@abstractmethod
def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve relevant contexts for the given documents.

    Args:
        documents (List[Document]): List of documents containing queries

    Returns:
        List[Document]: Documents updated with retrieved contexts
    """
    pass

IndexManager

Manages downloading, caching, and loading of retrieval indexes.

Handles both prebuilt indexes (wiki, msmarco) and custom user indexes.

Source code in rankify/retrievers/index_manager.py 
 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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
class IndexManager:
    """
    Manages downloading, caching, and loading of retrieval indexes.

    Handles both prebuilt indexes (wiki, msmarco) and custom user indexes.
    """

    def __init__(self, cache_dir: str = None):
        self.cache_dir = cache_dir or os.environ.get("RERANKING_CACHE_DIR", "./cache")
        self.index_configs = self._load_index_configs()

    def _load_index_configs(self) -> Dict:
        """Load index configuration mappings."""
        return {
            "bm25": {
                "wiki": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bm25_wiki.zip",
                    "prebuilt": "wikipedia-dpr-100w"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bm25_index_msmarco.zip?download=true",
                    "prebuilt": None
                }
            },
            "dpr-multi": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.dpr-multi",
                    "encoder": "facebook/dpr-question_encoder-multiset-base"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_passage_dpr_multi.zip",
                    "encoder": "facebook/dpr-question_encoder-multiset-base"
                }
            },
            "dpr-single": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.dpr-single-nq",
                    "encoder": "facebook/dpr-question_encoder-single-nq-base"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_passage_dpr_single.zip",
                    "encoder": "facebook/dpr-question_encoder-single-nq-base"
                }
            },
            "ance-multi": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.ance-multi",
                    "encoder": "castorini/ance-dpr-question-multi"
                },
                "msmarco": {
                    "prebuilt": "msmarco-v1-passage.ance", 
                    "encoder": "castorini/ance-msmarco-passage"
                }
            },

            "bpr-single": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.bpr-single-nq",
                    "encoder": "castorini/bpr-nq-question-encoder"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_passage_bpr.zip",
                    "encoder": "castorini/bpr-nq-question-encoder"
                }
            },
            "bge": {
                "wiki": {
                    "urls": [
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bgb_index.tar.gz.part1?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bgb_index.tar.gz.part2?download=true",
                    ],
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true",
                    "encoder": "BAAI/bge-large-en-v1.5"
                },
                "msmarco": {
                    "urls": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_embeddings_bgb.zip?download=true",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true",
                    "encoder": "BAAI/bge-large-en-v1.5"
                }
            },
            "colbert": {
                "wiki": {
                    "urls": [
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.001?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.002?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.003?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.004?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.005?download=true"
                    ],
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true",
                    "model": "colbert-ir/colbertv2.0"
                },
                "msmarco": {
                    "urls": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_embeddings_colbert.zip?download=true",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true",
                    "model": "colbert-ir/colbertv2.0"
                }
            },
            "contriever": {
                "wiki": {
                    "url": "https://dl.fbaipublicfiles.com/contriever/embeddings/contriever-msmarco/wikipedia_embeddings.tar",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true",
                    "model": "facebook/contriever-msmarco",
                    "vector_size": 768
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_embeddings_contriever.zip?download=true",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true",
                    "model": "facebook/contriever-msmarco",
                    "vector_size": 768
                }
            },
            "online": {
                "web": {
                    "search_provider": "web",
                    "chunk_size": 500,
                    "chunk_overlap": 50,
                    "requires_api_key": True
                }
            },
            "hyde": {
                "wiki": {
                    "base_retriever": "contriever",
                    "base_model": "facebook/contriever-msmarco",
                    "base_index_type": "wiki",
                    "task": "web search",
                    "llm_model": "gpt-3.5-turbo-0125",
                    "num_generated_docs": 1,
                    "max_token_generated_docs": 512,
                    "temperature": 0.7,
                    "requires_api_key": True
                },
                "msmarco": {
                    "base_retriever": "contriever", 
                    "base_model": "facebook/contriever-msmarco",
                    "base_index_type": "msmarco",
                    "task": "web search",
                    "llm_model": "gpt-3.5-turbo-0125",
                    "num_generated_docs": 1,
                    "max_token_generated_docs": 512,
                    "temperature": 0.7,
                    "requires_api_key": True
                }
            }
        }

    def get_index_path(self, method: str, index_type: str, custom_path: str = None) -> str:
        """
        Get the path to the index for a given method and index type.

        Args:
            method (str): Retrieval method (e.g., 'bm25', 'dpr-multi', 'ance')
            index_type (str): Index type ('wiki', 'msmarco', or 'custom')
            custom_path (str): Path to custom index if index_type is 'custom'

        Returns:
            str: Path to the index
        """
        if custom_path:
            return custom_path

        if method not in self.index_configs:
            raise ValueError(f"Unsupported method: {method}")

        if index_type not in self.index_configs[method]:
            raise ValueError(f"Unsupported index type '{index_type}' for method '{method}'")

        config = self.index_configs[method][index_type]

        # If it's a prebuilt index, return the identifier
        if "prebuilt" in config and config["prebuilt"]:
            return config["prebuilt"]

        # Otherwise, download and return local path
        return self._ensure_index_downloaded(method, index_type)

    def _ensure_index_downloaded(self, method: str, index_type: str) -> str:
        """Download and extract index if not already available."""
        config = self.index_configs[method][index_type]
        url = config["url"]

        # Create local directory path
        index_name = f"{method}_{index_type}"
        local_dir = os.path.join(self.cache_dir, "index", index_name)

        if not os.path.exists(local_dir):
            print(f"Downloading {method} index for {index_type}...")
            self._download_and_extract(url, local_dir)

        return local_dir

    def _download_and_extract(self, url: str, destination: str):
        """Download and extract a ZIP file."""
        os.makedirs(destination, exist_ok=True)

        zip_name = os.path.basename(url).split("?")[0]
        zip_path = os.path.join(self.cache_dir, "temp", zip_name)

        os.makedirs(os.path.dirname(zip_path), exist_ok=True)

        if not os.path.exists(zip_path):
            response = requests.get(url, stream=True)
            response.raise_for_status()

            with open(zip_path, "wb") as f:
                for chunk in tqdm(response.iter_content(chunk_size=1024), 
                                desc=f"Downloading {zip_name}"):
                    f.write(chunk)

        print(f"Extracting {zip_name}...")
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(destination)

        # Clean up ZIP file
        os.remove(zip_path)
        print("Extraction complete.")

    def load_id_mapping(self, index_path: str) -> Optional[Dict]:
        """Load ID mapping if available."""
        mapping_path = os.path.join(index_path, "id_mapping.json")
        if os.path.exists(mapping_path):
            with open(mapping_path, "r", encoding="utf-8") as f:
                mapping = json.load(f)
                return {v: k for k, v in mapping.items()}  # Reverse mapping
        return None

    # def load_corpus(self, index_path: str) -> Dict:
    #     """Load corpus data from index path."""
    #     corpus_file = os.path.join(index_path, "corpus.jsonl")
    #     if not os.path.exists(corpus_file):
    #         return {}

    #     corpus = {}
    #     with open(corpus_file, "r", encoding="utf-8") as f:
    #         for line in f:
    #             doc = json.loads(line.strip())
    #             doc_id = doc.get("docid") or doc.get("id")
    #             contents = doc.get("contents", "")
    #             title = doc.get("title", contents[:100] if contents else "No Title")
    #             corpus[str(doc_id)] = {"contents": contents, "title": title}

    #     return corpus
    def load_corpus(self, index_path: str) -> Dict:
        """Load corpus data from index path (supports multiple formats)."""
        # Try different corpus file formats in order of preference
        corpus_files = [
            "corpus_metadata.json",  # Your custom format
            "corpus.jsonl",          # Standard JSONL format
            "corpus.json"            # Alternative JSON format
        ]

        for corpus_filename in corpus_files:
            corpus_file = os.path.join(index_path, corpus_filename)
            if os.path.exists(corpus_file):
                print(f"📚 Loading corpus from {corpus_filename}")

                if corpus_filename.endswith('.jsonl'):
                    # JSONL format (one JSON object per line)
                    return self._load_corpus_jsonl(corpus_file)
                else:
                    # JSON format (single JSON object)
                    return self._load_corpus_json(corpus_file)

        print(f"❌ No corpus file found in {index_path}")
        return {}

    def _load_corpus_jsonl(self, corpus_file: str) -> Dict:
        """Load corpus from JSONL format (one JSON object per line)."""
        corpus = {}
        with open(corpus_file, "r", encoding="utf-8") as f:
            for line in f:
                doc = json.loads(line.strip())
                doc_id = doc.get("docid") or doc.get("id")
                contents = doc.get("contents", "")
                title = doc.get("title", contents[:100] if contents else "No Title")
                corpus[str(doc_id)] = {"contents": contents, "title": title}
        return corpus

    def _load_corpus_json(self, corpus_file: str) -> Dict:
        """Load corpus from JSON format (single JSON object)."""
        corpus = {}
        with open(corpus_file, "r", encoding="utf-8") as f:
            data = json.load(f)

        # Handle different JSON structures
        if isinstance(data, dict):
            # If it's a dict, iterate through it
            for doc_id, doc_data in data.items():
                if isinstance(doc_data, dict):
                    contents = doc_data.get("contents") or doc_data.get("text") or ""
                    title = doc_data.get("title", contents[:100] if contents else "No Title")
                else:
                    # If doc_data is a string, use it as contents
                    contents = str(doc_data)
                    title = contents[:100] if contents else "No Title"
                corpus[str(doc_id)] = {"contents": contents, "title": title}

        elif isinstance(data, list):
            # If it's a list, iterate through documents
            for doc in data:
                doc_id = doc.get("docid") or doc.get("id")
                contents = doc.get("contents", "")
                title = doc.get("title", contents[:100] if contents else "No Title")
                corpus[str(doc_id)] = {"contents": contents, "title": title}

        return corpus
    def download_and_extract_index(self, url: str) -> str:
        """Download and extract index from URL, return local path."""
        # This is used by DenseRetriever pattern
        import tempfile
        import zipfile

        # Create temporary file
        with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as tmp_file:
            tmp_path = tmp_file.name

        try:
            # Download
            print(f"Downloading from {url}...")
            response = requests.get(url, stream=True)
            response.raise_for_status()

            with open(tmp_path, 'wb') as f:
                for chunk in tqdm(response.iter_content(chunk_size=1024)):
                    f.write(chunk)

            # Extract to cache directory
            extract_dir = os.path.join(self.cache_dir, "downloaded_index")
            os.makedirs(extract_dir, exist_ok=True)

            with zipfile.ZipFile(tmp_path, 'r') as zip_ref:
                zip_ref.extractall(extract_dir)

            return extract_dir

        finally:
            # Clean up temporary file
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

get_index_path(method, index_type, custom_path=None)

Get the path to the index for a given method and index type.

Parameters:

Name Type Description Default
method str

Retrieval method (e.g., 'bm25', 'dpr-multi', 'ance')

 required
index_type str

Index type ('wiki', 'msmarco', or 'custom')

 required
custom_path str

Path to custom index if index_type is 'custom'

 None

Returns:

Name Type Description
str str

Path to the index
Source code in rankify/retrievers/index_manager.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def get_index_path(self, method: str, index_type: str, custom_path: str = None) -> str:
    """
    Get the path to the index for a given method and index type.

    Args:
        method (str): Retrieval method (e.g., 'bm25', 'dpr-multi', 'ance')
        index_type (str): Index type ('wiki', 'msmarco', or 'custom')
        custom_path (str): Path to custom index if index_type is 'custom'

    Returns:
        str: Path to the index
    """
    if custom_path:
        return custom_path

    if method not in self.index_configs:
        raise ValueError(f"Unsupported method: {method}")

    if index_type not in self.index_configs[method]:
        raise ValueError(f"Unsupported index type '{index_type}' for method '{method}'")

    config = self.index_configs[method][index_type]

    # If it's a prebuilt index, return the identifier
    if "prebuilt" in config and config["prebuilt"]:
        return config["prebuilt"]

    # Otherwise, download and return local path
    return self._ensure_index_downloaded(method, index_type)

load_id_mapping(index_path)

Load ID mapping if available.

Source code in rankify/retrievers/index_manager.py 
227
228
229
230
231
232
233
234
def load_id_mapping(self, index_path: str) -> Optional[Dict]:
    """Load ID mapping if available."""
    mapping_path = os.path.join(index_path, "id_mapping.json")
    if os.path.exists(mapping_path):
        with open(mapping_path, "r", encoding="utf-8") as f:
            mapping = json.load(f)
            return {v: k for k, v in mapping.items()}  # Reverse mapping
    return None

load_corpus(index_path)

Load corpus data from index path (supports multiple formats).

Source code in rankify/retrievers/index_manager.py 
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def load_corpus(self, index_path: str) -> Dict:
    """Load corpus data from index path (supports multiple formats)."""
    # Try different corpus file formats in order of preference
    corpus_files = [
        "corpus_metadata.json",  # Your custom format
        "corpus.jsonl",          # Standard JSONL format
        "corpus.json"            # Alternative JSON format
    ]

    for corpus_filename in corpus_files:
        corpus_file = os.path.join(index_path, corpus_filename)
        if os.path.exists(corpus_file):
            print(f"📚 Loading corpus from {corpus_filename}")

            if corpus_filename.endswith('.jsonl'):
                # JSONL format (one JSON object per line)
                return self._load_corpus_jsonl(corpus_file)
            else:
                # JSON format (single JSON object)
                return self._load_corpus_json(corpus_file)

    print(f"❌ No corpus file found in {index_path}")
    return {}

download_and_extract_index(url)

Download and extract index from URL, return local path.

Source code in rankify/retrievers/index_manager.py 
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
def download_and_extract_index(self, url: str) -> str:
    """Download and extract index from URL, return local path."""
    # This is used by DenseRetriever pattern
    import tempfile
    import zipfile

    # Create temporary file
    with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as tmp_file:
        tmp_path = tmp_file.name

    try:
        # Download
        print(f"Downloading from {url}...")
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(tmp_path, 'wb') as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024)):
                f.write(chunk)

        # Extract to cache directory
        extract_dir = os.path.join(self.cache_dir, "downloaded_index")
        os.makedirs(extract_dir, exist_ok=True)

        with zipfile.ZipFile(tmp_path, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)

        return extract_dir

    finally:
        # Clean up temporary file
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

Document

Represents a document consisting of a question, answers, and contexts.

Attributes:

Name Type Description
question Question

The question associated with the document.
answers Answer

The answers to the question.
contexts list[Context]

A list of related contexts.
reorder_contexts list[Context] or None

A reordered list of contexts based on relevance.
Source code in rankify/dataset/dataset.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
class Document:
    """
    Represents a document consisting of a question, answers, and contexts.

    Attributes:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context]): A list of related contexts.
        reorder_contexts (list[Context] or None): A reordered list of contexts based on relevance.
    """
    def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
        """
        Initializes a Document instance.

        Args:
            question (Question): The question associated with the document.
            answers (Answer): The answers to the question.
            contexts (list[Context], optional): A list of contexts related to the question.

        Example:
            ```python
            q = Question("What is the capital of France?")
            a = Answer(["Paris"])
            c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
            d = Document(question=q, answers=a, contexts=[c1, c2])
            print(d)
            ```
        """
        self.question: Question = question
        self.answers: Answer = answers
        self.contexts: List[Context] = contexts
        self.reorder_contexts: List[Context] = None
        self.id = str(id) 

    @classmethod
    def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
        """
        Creates a Document instance from a dictionary.

        Args:
            data (dict): A dictionary containing the question, answers, and contexts.
            n_docs (int, optional): The number of contexts to include. Defaults to 100.

        Returns:
            Document: A new Document instance.

        Example:
            ```python
            data = {
                "question": "What is the capital of France?",
                "answers": ["Paris"],
                "ctxs": [
                    {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                    {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
                ]
            }
            d = Document.from_dict(data)
            print(d.question)
            ```
        """
        question = Question(data["question"])
        if "answers" in data:
            answers = Answer(data["answers"])
        else:
            answers =Answer('')

        if "query_id" in data:
            id = data["query_id"]
        else:
            id = None
        contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
        return cls(question, answers, contexts, id=id)

    def to_dict(self) -> Dict[str, Optional[object]]:
        """
        Converts the document into a dictionary representation.

        Returns:
            dict: A dictionary containing the question, answers, and contexts.
        """
        return {
            "question": self.question.question,
            "answers": self.answers.answers,
            "contexts": [ctx.to_dict() for ctx in self.contexts]
        }
    def to_dict_reoreder(self) -> Dict[str,Optional[object]]:
        return {
            "question" : self.question.question,
            "answers" : self.answers.answers,
            "contexts" : [ctx.to_dict() for ctx in self.reorder_contexts]
        }
    def __str__(self) -> str:
        """
        Returns a string representation of the Document instance.

        Returns:
            str: The formatted document information.

        Example:
            ```python
            d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
            print(d)
            ```
        """
        contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
        reorder_contexts_str= ''
        if self.reorder_contexts is not None:
            reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
        return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

__init__(question, answers, contexts=None, id=None)

Initializes a Document instance.

Parameters:

Name Type Description Default
question Question

The question associated with the document.

 required
answers Answer

The answers to the question.

 required
contexts list[Context]

A list of contexts related to the question.

 None
Example 
q = Question("What is the capital of France?")
a = Answer(["Paris"])
c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
d = Document(question=q, answers=a, contexts=[c1, c2])
print(d)
Source code in rankify/dataset/dataset.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
    """
    Initializes a Document instance.

    Args:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context], optional): A list of contexts related to the question.

    Example:
        ```python
        q = Question("What is the capital of France?")
        a = Answer(["Paris"])
        c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
        d = Document(question=q, answers=a, contexts=[c1, c2])
        print(d)
        ```
    """
    self.question: Question = question
    self.answers: Answer = answers
    self.contexts: List[Context] = contexts
    self.reorder_contexts: List[Context] = None
    self.id = str(id)

from_dict(data, n_docs=100) classmethod

Creates a Document instance from a dictionary.

Parameters:

Name Type Description Default
data dict

A dictionary containing the question, answers, and contexts.

 required
n_docs int

The number of contexts to include. Defaults to 100.

 100

Returns:

Name Type Description
Document Document

A new Document instance.
Example 
data = {
    "question": "What is the capital of France?",
    "answers": ["Paris"],
    "ctxs": [
        {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
        {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
    ]
}
d = Document.from_dict(data)
print(d.question)
Source code in rankify/dataset/dataset.py 
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
@classmethod
def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
    """
    Creates a Document instance from a dictionary.

    Args:
        data (dict): A dictionary containing the question, answers, and contexts.
        n_docs (int, optional): The number of contexts to include. Defaults to 100.

    Returns:
        Document: A new Document instance.

    Example:
        ```python
        data = {
            "question": "What is the capital of France?",
            "answers": ["Paris"],
            "ctxs": [
                {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
            ]
        }
        d = Document.from_dict(data)
        print(d.question)
        ```
    """
    question = Question(data["question"])
    if "answers" in data:
        answers = Answer(data["answers"])
    else:
        answers =Answer('')

    if "query_id" in data:
        id = data["query_id"]
    else:
        id = None
    contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
    return cls(question, answers, contexts, id=id)

to_dict()

Converts the document into a dictionary representation.

Returns:

Name Type Description
dict Dict[str, Optional[object]]

A dictionary containing the question, answers, and contexts.
Source code in rankify/dataset/dataset.py 
269
270
271
272
273
274
275
276
277
278
279
280
def to_dict(self) -> Dict[str, Optional[object]]:
    """
    Converts the document into a dictionary representation.

    Returns:
        dict: A dictionary containing the question, answers, and contexts.
    """
    return {
        "question": self.question.question,
        "answers": self.answers.answers,
        "contexts": [ctx.to_dict() for ctx in self.contexts]
    }

__str__()

Returns a string representation of the Document instance.

Returns:

Name Type Description
str str

The formatted document information.
Example 
d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
print(d)
Source code in rankify/dataset/dataset.py 
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def __str__(self) -> str:
    """
    Returns a string representation of the Document instance.

    Returns:
        str: The formatted document information.

    Example:
        ```python
        d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
        print(d)
        ```
    """
    contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
    reorder_contexts_str= ''
    if self.reorder_contexts is not None:
        reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
    return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

Context

Represents a context with metadata such as score and title.

Attributes:

Name Type Description
score float

The relevance score of the context.
has_answer bool

Whether the context contains an answer.
id int

The identifier of the context.
title str

The title of the context.
text str

The text of the context.
Source code in rankify/dataset/dataset.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
class Context:
    """
    Represents a context with metadata such as score and title.

    Attributes:
        score (float, optional): The relevance score of the context.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.
    """
    def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
        """
        Initializes a Context instance.

        Args:
            score (float, optional): The relevance score.
            has_answer (bool, optional): Whether the context contains an answer.
            id (int, optional): The identifier of the context.
            title (str, optional): The title of the context.
            text (str, optional): The text of the context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c)
            ```
        """
        self.score: Optional[float] = score
        self.has_answer: Optional[bool] = has_answer
        self.id: Optional[str] = id
        self.title: Optional[str] = title
        self.text: Optional[str] = text

    def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

        """
        Converts the Context instance to a dictionary.

        Args:
            save_text (bool): Whether to include text in the output dictionary.

        Returns:
            dict: The context data.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c.to_dict())
            ```
        """
        context_dict = {
            "score": float(self.score) if self.score is not None else None,
            "has_answer": self.has_answer,
            "id": self.id,
            }

        # Include 'text' only if save_text is True
        if save_text:
            context_dict["text"] = self.text
            context_dict["title"] =  self.title

        return context_dict
    def __str__(self) -> str:
        """
        Returns a string representation of the Context instance.

        Returns:
            str: The formatted context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(str(c))
            ```
        """
        return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

__init__(score=None, has_answer=None, id=None, title=None, text=None)

Initializes a Context instance.

Parameters:

Name Type Description Default
score float

The relevance score.

 None
has_answer bool

Whether the context contains an answer.

 None
id int

The identifier of the context.

 None
title str

The title of the context.

 None
text str

The text of the context.

 None
Example 
c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c)
Source code in rankify/dataset/dataset.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
    """
    Initializes a Context instance.

    Args:
        score (float, optional): The relevance score.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c)
        ```
    """
    self.score: Optional[float] = score
    self.has_answer: Optional[bool] = has_answer
    self.id: Optional[str] = id
    self.title: Optional[str] = title
    self.text: Optional[str] = text

to_dict(save_text=False)

Converts the Context instance to a dictionary.

Parameters:

Name Type Description Default
save_text bool

Whether to include text in the output dictionary.

 False

Returns:

Name Type Description
dict Dict[str, Optional[object]]

The context data.
Example 
c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c.to_dict())
Source code in rankify/dataset/dataset.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

    """
    Converts the Context instance to a dictionary.

    Args:
        save_text (bool): Whether to include text in the output dictionary.

    Returns:
        dict: The context data.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c.to_dict())
        ```
    """
    context_dict = {
        "score": float(self.score) if self.score is not None else None,
        "has_answer": self.has_answer,
        "id": self.id,
        }

    # Include 'text' only if save_text is True
    if save_text:
        context_dict["text"] = self.text
        context_dict["title"] =  self.title

    return context_dict

__str__()

Returns a string representation of the Context instance.

Returns:

Name Type Description
str str

The formatted context.
Example 
c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(str(c))
Source code in rankify/dataset/dataset.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def __str__(self) -> str:
    """
    Returns a string representation of the Context instance.

    Returns:
        str: The formatted context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(str(c))
        ```
    """
    return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

ColBERTRetriever

Bases: BaseRetriever

ColBERT retriever with backward compatibility for prebuilt indices.

Supports two modes: 1. Prebuilt indices (wiki, msmarco) - Original format with passages.tsv 2. Custom indices - New format with collection.tsv + ID mappings

Source code in rankify/retrievers/colbert_retriever.py 
 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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
class ColBERTRetriever(BaseRetriever):
    """
    ColBERT retriever with backward compatibility for prebuilt indices.

    Supports two modes:
    1. Prebuilt indices (wiki, msmarco) - Original format with passages.tsv
    2. Custom indices - New format with collection.tsv + ID mappings
    """

    def __init__(self, model: str = "colbert-ir/colbertv2.0", index_type: str = "wiki", 
                 index_folder: str = None, **kwargs):
        super().__init__(**kwargs)
        self.model = model
        self.index_type = index_type
        self.index_folder = index_folder
        self.tokenizer = SimpleTokenizer()

        # Initialize index manager
        self.index_manager = IndexManager()

        # Initialize flag (will be properly set in _setup_paths)
        self.is_custom_index = False

        # Setup paths and determine index type
        self._setup_paths()

        # Load ID mappings (only for custom indices)
        self.id_mappings = self._load_id_mappings()

        # Initialize searcher
        self.searcher = self._initialize_searcher()

        # Load passages
        self.passages = self._load_passages()

    def _setup_paths(self):
        """Setup file paths based on whether it's a custom or prebuilt index."""
        if self.index_folder:
            # Check if it's actually a custom index by looking for custom files
            collection_file = os.path.join(self.index_folder, "collection.tsv")
            id_mapping_file = os.path.join(self.index_folder, "id_mapping.json")

            if os.path.exists(collection_file) and os.path.exists(id_mapping_file):
                # True custom index
                self.is_custom_index = True
                self.index_path = self.index_folder
                self.passages_file = os.path.join(self.index_folder, "passages.tsv")
                self.collection_file = collection_file
                self.id_mapping_file = id_mapping_file
                print(f"Using custom index: {self.index_folder}")
            else:
                # Prebuilt index with index_folder specified
                self.is_custom_index = False
                self.index_path = self.index_folder
                passages_file = os.path.join(self.index_folder, "passages.tsv")
                if os.path.exists(passages_file):
                    self.passages_file = passages_file
                else:
                    # Try to find passages file in parent directories
                    self._ensure_index_and_passages_downloaded()
                    self.passages_file = self._get_passages_file_path()
                self.collection_file = self.passages_file  # Same file for prebuilt indices
                self.id_mapping_file = None
                print(f"Using prebuilt index at: {self.index_folder}")
        else:
            # Standard prebuilt index
            self.is_custom_index = False
            self.index_path = os.path.join(self.index_manager.cache_dir, "index", "colbert")
            self._ensure_index_and_passages_downloaded()
            self.passages_file = self._get_passages_file_path()
            self.collection_file = self.passages_file  # Same file for prebuilt indices
            self.id_mapping_file = None
            print(f"Using prebuilt index: {self.index_type}")

    def _load_id_mappings(self):
        """Load ID mappings for custom indices only."""
        if not self.is_custom_index or not self.id_mapping_file:
            return None

        if not os.path.exists(self.id_mapping_file):
            print(f"Warning: ID mapping file not found: {self.id_mapping_file}")
            return None

        try:
            with open(self.id_mapping_file, 'r', encoding='utf-8') as f:
                mappings = json.load(f)

            # Convert string keys back to integers for sequential_to_original
            sequential_to_original = {int(k): v for k, v in mappings.get("sequential_to_original", {}).items()}
            original_to_sequential = mappings.get("original_to_sequential", {})

            print(f"✅ Loaded ID mappings: {len(original_to_sequential)} entries")
            return {
                "sequential_to_original": sequential_to_original,
                "original_to_sequential": original_to_sequential
            }

        except Exception as e:
            print(f"Warning: Error loading ID mappings: {e}")
            return None

    def _initialize_searcher(self):
        """Initialize ColBERT searcher with backward compatibility."""
        with Run().context(RunConfig(nranks=1, experiment="colbert")):

            if self.is_custom_index:
                # Custom index - use our format
                index_root = self.index_folder
                collection_path = self.collection_file

                # Verify collection file exists
                if not os.path.exists(collection_path):
                    raise FileNotFoundError(f"ColBERT collection file not found: {collection_path}")

                print(f"Initializing custom ColBERT searcher...")
                print(f"  Index root: {index_root}")
                print(f"  Collection: {collection_path}")

            else:
                # Prebuilt index - use original format
                if self.index_folder:
                    index_root = self.index_folder
                else:
                    index_root = os.path.join(self.index_path, self.index_type)
                collection_path = self.passages_file

                # Verify passages file exists (this is what prebuilt indices use)
                if not os.path.exists(collection_path):
                    raise FileNotFoundError(f"Passages file not found: {collection_path}")

                print(f"Initializing prebuilt ColBERT searcher...")
                print(f"  Index root: {index_root}")
                print(f"  Collection: {collection_path}")

            config = ColBERTConfig(
                root=index_root,
                collection=collection_path
            )

            return Searcher(index=index_root, config=config)

    def _ensure_index_and_passages_downloaded(self):
        """Download and extract ColBERT index and passages if needed (prebuilt indices only)."""
        if self.is_custom_index or self.index_folder:
            return  # Skip for custom indices or when index_folder is provided

        os.makedirs(self.index_path, exist_ok=True)

        if self.index_type not in self.index_manager.index_configs.get("colbert", {}):
            raise ValueError(f"Unsupported ColBERT index type: {self.index_type}")

        config = self.index_manager.index_configs["colbert"][self.index_type]

        # Check if index exists
        index_dir = os.path.join(self.index_path, self.index_type)
        if not os.path.exists(index_dir):
            urls = config.get("urls")
            if isinstance(urls, list):
                # Multi-part download
                for url in urls:
                    filename = self._extract_filename_from_url(url)
                    local_path = os.path.join(self.index_path, filename)
                    if not os.path.exists(local_path):
                        self._download_file(url, local_path)
                self._extract_multi_part_zip()
            else:
                # Single file download
                zip_path = os.path.join(self.index_path, "index.zip")
                if not os.path.exists(zip_path):
                    self._download_file(urls, zip_path)
                self._extract_zip_files()

        # Download passages if needed
        passages_url = config.get("passages_url")
        if passages_url:
            passages_filename = self._extract_filename_from_url(passages_url)
            passages_path = os.path.join(self.index_manager.cache_dir, passages_filename)
            if not os.path.exists(passages_path):
                self._download_file(passages_url, passages_path)

    def _get_passages_file_path(self):
        """Get path to passages file for prebuilt indices."""
        if self.index_folder and not self.is_custom_index:
            # Direct path provided for prebuilt index
            passages_path = os.path.join(self.index_folder, "passages.tsv")
            if os.path.exists(passages_path):
                return passages_path

        if self.index_type not in self.index_manager.index_configs.get("colbert", {}):
            raise ValueError(f"Unsupported ColBERT index type: {self.index_type}")

        config = self.index_manager.index_configs["colbert"][self.index_type]
        passages_url = config.get("passages_url")
        if passages_url:
            passages_filename = self._extract_filename_from_url(passages_url)
            return os.path.join(self.index_manager.cache_dir, passages_filename)
        return None

    def _extract_filename_from_url(self, url):
        """Extract filename from URL, ignoring query parameters."""
        parsed_url = urlparse(url)
        return os.path.basename(parsed_url.path).split('?')[0]

    def _download_file(self, url: str, save_path: str):
        """Download file from URL with progress bar."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(save_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024), 
                            desc=f"Downloading {os.path.basename(save_path)}"):
                f.write(chunk)

    def _extract_zip_files(self):
        """Extract ZIP files in the index folder."""
        zip_files = [f for f in os.listdir(self.index_path) if f.endswith(".zip")]

        for zip_file in zip_files:
            zip_path = os.path.join(self.index_path, zip_file)

            print(f"Extracting {zip_file}...")
            with zipfile.ZipFile(zip_path, "r") as zip_ref:
                zip_ref.extractall(self.index_path)

            os.remove(zip_path)

    def _extract_multi_part_zip(self):
        """Extract multi-part ZIP files by combining and decompressing them."""
        # Find all parts (handle both .001, .002 format and other naming patterns)
        all_files = os.listdir(self.index_path)
        zip_parts = []

        # Look for wiki.zip.001, wiki.zip.002, etc.
        for f in all_files:
            if f.startswith("wiki.zip.") and f.split('.')[-1].isdigit():
                zip_parts.append(f)

        if not zip_parts:
            print("No multi-part ZIP files found")
            return

        # Sort by part number
        zip_parts.sort(key=lambda x: int(x.split('.')[-1]))

        combined_zip_path = os.path.join(self.index_path, "combined.zip")

        print(f"Combining {len(zip_parts)} parts into {combined_zip_path}...")

        # Combine all parts
        with open(combined_zip_path, "wb") as combined:
            for part in zip_parts:
                part_path = os.path.join(self.index_path, part)
                print(f"Adding {part} to combined archive...")
                with open(part_path, "rb") as part_file:
                    # Use shutil.copyfileobj for better memory efficiency
                    import shutil
                    shutil.copyfileobj(part_file, combined)

        # Extract the combined ZIP
        print(f"Extracting combined ZIP file: {combined_zip_path}...")
        try:
            with zipfile.ZipFile(combined_zip_path, "r") as zip_ref:
                zip_ref.extractall(self.index_path)
            print("Extraction completed successfully")
        except zipfile.BadZipFile as e:
            print(f"Error: Combined file is not a valid ZIP: {e}")
            raise
        finally:
            # Clean up: remove combined file and parts
            if os.path.exists(combined_zip_path):
                os.remove(combined_zip_path)
            for part in zip_parts:
                part_path = os.path.join(self.index_path, part)
                if os.path.exists(part_path):
                    os.remove(part_path)

    def _load_passages(self):
        """Load passages with format detection."""
        if not self.passages_file or not os.path.exists(self.passages_file):
            print("Warning: Passages file not found, using empty passages dict")
            return {}

        passages = {}
        print(f"Loading passages from {self.passages_file}...")

        # Detect if the first line has integer or string IDs
        with open(self.passages_file, "r", encoding="utf-8") as f:
            first_line = f.readline().strip()
            f.seek(0)  # Reset to beginning

            # Check if header exists
            has_header = first_line.startswith('id\t') or first_line.lower().startswith('id\t')

            if has_header:
                next(f)  # Skip header

            # Process lines
            for line_num, line in enumerate(f, 1):
                try:
                    line = line.strip()
                    if not line:
                        continue

                    parts = line.split("\t")
                    if len(parts) >= 2:
                        pid = parts[0]  # Keep as string (works for both int and string IDs)
                        text = parts[1] if len(parts) > 1 else ""
                        title = parts[2] if len(parts) > 2 else ""

                        passages[pid] = {"text": text, "title": title}
                    else:
                        print(f"Warning: Malformed line {line_num} in passages file")

                except Exception as e:
                    print(f"Error processing line {line_num}: {e}")
                    continue

        print(f"✅ Loaded {len(passages)} passages")
        return passages

    def _convert_result_id(self, result_id: int) -> str:
        """Convert ColBERT result ID to the appropriate format."""
        if self.is_custom_index and self.id_mappings:
            # Custom index: convert sequential ID back to original string ID
            return self.id_mappings["sequential_to_original"].get(result_id, str(result_id))
        else:
            # Prebuilt index: IDs are already in correct format
            return str(result_id)

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using ColBERT."""
        print(f"Retrieving {len(documents)} documents with ColBERT...")
        print(f"Mode: {'Custom index' if self.is_custom_index else f'Prebuilt index ({self.index_type})'}")

        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            query = document.question.question

            try:
                # ColBERT search returns (document_ids, ranks, scores)
                results = self.searcher.search(query, k=self.n_docs)
                document_ids, ranks, scores = results

                contexts = []
                for result_id, rank, score in zip(document_ids, ranks, scores):
                    try:
                        context = self._create_context_from_result(result_id, score, document)
                        if context:
                            contexts.append(context)
                    except Exception as e:
                        print(f"Error processing result {result_id}: {e}")

                document.contexts = contexts

            except Exception as e:
                print(f"Error searching for document {i}: {e}")
                document.contexts = []

        return documents

    def _create_context_from_result(self, result_id: int, score: float, document: Document) -> Context:
        """Create Context object from ColBERT search result."""
        # Convert result ID to appropriate format
        passage_id = self._convert_result_id(result_id)

        # Find passage by ID
        if passage_id in self.passages:
            passage = self.passages[passage_id]
            return Context(
                id=passage_id,
                title=passage["title"],
                text=passage["text"],
                score=float(score),
                has_answer=has_answers(passage["text"], document.answers.answers, self.tokenizer)
            )
        else:
            print(f"Warning: Passage {passage_id} (result_id: {result_id}) not found in passages dict")
            return Context(
                id=passage_id,
                title="Passage not found",
                text=f"Passage {passage_id} not found in corpus",
                score=float(score),
                has_answer=False
            )

retrieve(documents)

Retrieve relevant contexts using ColBERT.

Source code in rankify/retrievers/colbert_retriever.py 
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using ColBERT."""
    print(f"Retrieving {len(documents)} documents with ColBERT...")
    print(f"Mode: {'Custom index' if self.is_custom_index else f'Prebuilt index ({self.index_type})'}")

    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        query = document.question.question

        try:
            # ColBERT search returns (document_ids, ranks, scores)
            results = self.searcher.search(query, k=self.n_docs)
            document_ids, ranks, scores = results

            contexts = []
            for result_id, rank, score in zip(document_ids, ranks, scores):
                try:
                    context = self._create_context_from_result(result_id, score, document)
                    if context:
                        contexts.append(context)
                except Exception as e:
                    print(f"Error processing result {result_id}: {e}")

            document.contexts = contexts

        except Exception as e:
            print(f"Error searching for document {i}: {e}")
            document.contexts = []

    return documents