Skip to content

LiT5 Reranker

rankify.models.lit5_reranker

BaseRanking

Bases: ABC

An abstract base class for implementing different ranking models.

This class defines the interface for all ranking models, ensuring that all subclasses implement the required methods.

Attributes:

Name Type Description
method str

The name of the ranking method.
model_name str

The name of the model being used for ranking.
api_key str

An optional API key for accessing remote models or services.
Source code in rankify/models/base.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
58
59
class BaseRanking(ABC):
    """
    An abstract base class for implementing different ranking models.

    This class defines the interface for all ranking models, ensuring that all subclasses implement the required methods.

    Attributes:
        method (str): The name of the ranking method.
        model_name (str): The name of the model being used for ranking.
        api_key (str, optional): An optional API key for accessing remote models or services.
    """

    @abstractmethod
    def __init__(self, method: str= None, model_name: str= None, api_key: str= None, **kwargs) ->None:
        """
        Initializes the base ranking model.

        Args:
            method (str, optional): The name of the ranking method. Defaults to None.
            model_name (str, optional): The name of the model being used for ranking. Defaults to None.
            api_key (str, optional): An optional API key for accessing remote models or services. Defaults to None.

        Example:
            ```python
            class MyRanking(BaseRanking):
                def __init__(self, method, model_name):
                    super().__init__(method, model_name)
            ```
        """
        pass

    @abstractmethod
    def rank(self, documents: list[Document] ):
        """
        Abstract method to rank a list of documents.

        Args:
            documents (list[Document]): A list of Document instances that need to be ranked.

        Raises:
            NotImplementedError: This method must be implemented by subclasses.

        Example:
            ```python
            class MyRanking(BaseRanking):
                def __init__(self, method, model_name):
                    super().__init__(method, model_name)

                def rank(self, documents):
                    # Ranking implementation here
                    pass
            ```
        """
        pass

__init__(method=None, model_name=None, api_key=None, **kwargs) abstractmethod

Initializes the base ranking model.

Parameters:

Name Type Description Default
method str

The name of the ranking method. Defaults to None.

 None
model_name str

The name of the model being used for ranking. Defaults to None.

 None
api_key str

An optional API key for accessing remote models or services. Defaults to None.

 None
Example 
class MyRanking(BaseRanking):
    def __init__(self, method, model_name):
        super().__init__(method, model_name)
Source code in rankify/models/base.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@abstractmethod
def __init__(self, method: str= None, model_name: str= None, api_key: str= None, **kwargs) ->None:
    """
    Initializes the base ranking model.

    Args:
        method (str, optional): The name of the ranking method. Defaults to None.
        model_name (str, optional): The name of the model being used for ranking. Defaults to None.
        api_key (str, optional): An optional API key for accessing remote models or services. Defaults to None.

    Example:
        ```python
        class MyRanking(BaseRanking):
            def __init__(self, method, model_name):
                super().__init__(method, model_name)
        ```
    """
    pass

rank(documents) abstractmethod

Abstract method to rank a list of documents.

Parameters:

Name Type Description Default
documents list[Document]

A list of Document instances that need to be ranked.

 required

Raises:

Type Description
NotImplementedError

This method must be implemented by subclasses.
Example 
class MyRanking(BaseRanking):
    def __init__(self, method, model_name):
        super().__init__(method, model_name)

    def rank(self, documents):
        # Ranking implementation here
        pass
Source code in rankify/models/base.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@abstractmethod
def rank(self, documents: list[Document] ):
    """
    Abstract method to rank a list of documents.

    Args:
        documents (list[Document]): A list of Document instances that need to be ranked.

    Raises:
        NotImplementedError: This method must be implemented by subclasses.

    Example:
        ```python
        class MyRanking(BaseRanking):
            def __init__(self, method, model_name):
                super().__init__(method, model_name)

            def rank(self, documents):
                # Ranking implementation here
                pass
        ```
    """
    pass

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}"

PromptMode

Bases: Enum

Source code in rankify/utils/models/rank_llm/rerank/rankllm.py 
15
16
17
18
19
20
21
22
23
24
class PromptMode(Enum):
    UNSPECIFIED = "unspecified"
    RANK_GPT = "rank_GPT"
    RANK_GPT_APEER = "rank_GPT_APEER"
    LRL = "LRL"
    MONOT5 = "monot5"
    LiT5 = "LiT5"

    def __str__(self):
        return self.value

Reranker

Source code in rankify/utils/models/rank_llm/rerank/reranker.py 
 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
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
class Reranker:
    def __init__(self, agent: Optional[RankLLM]) -> None:
        self._agent = agent

    def rerank_batch(
        self,
        requests: List[Request],
        rank_start: int = 0,
        rank_end: int = 100,
        shuffle_candidates: bool = False,
        logging: bool = False,
        **kwargs: Any,
    ) -> List[Result]:
        """
        Reranks a list of requests using the RankLLM agent.

        This function applies a sliding window algorithm to rerank the results.
        Each window of results is processed by the RankLLM agent to obtain a new ranking.

        Args:
            requests (List[Request]): The list of requests. Each request has a query and a candidates list.
            rank_start (int, optional): The starting rank for processing. Defaults to 0.
            rank_end (int, optional): The end rank for processing. Defaults to 100.
            window_size (int, optional): The size of each sliding window. Defaults to 20.
            step (int, optional): The step size for moving the window. Defaults to 10.
            shuffle_candidates (bool, optional): Whether to shuffle candidates before reranking. Defaults to False.
            logging (bool, optional): Enables logging of the reranking process. Defaults to False.
            vllm_batched (bool, optional): Whether to use VLLM batched processing. Defaults to False.
            sglang_batched (bool, optional): Whether to use SGLang batched processing. Defaults to False.
            tensorrt_batched (bool, optional): Whether to use TensorRT-LLM batched processing. Defaults to False.
            populate_exec_summary (bool, optional): Whether to populate the exec summary. Defaults to False.
            batched (bool, optional): Whether to use batched processing. Defaults to False.

        Returns:
            List[Result]: A list containing the reranked candidates.
        """
        return self._agent.rerank_batch(
            requests, rank_start, rank_end, shuffle_candidates, logging, **kwargs
        )

    def rerank(
        self,
        request: Request,
        rank_start: int = 0,
        rank_end: int = 100,
        shuffle_candidates: bool = False,
        logging: bool = False,
        **kwargs: Any,
    ) -> Result:
        """
        Reranks a request using the RankLLM agent.

        This function applies a sliding window algorithm to rerank the results.
        Each window of results is processed by the RankLLM agent to obtain a new ranking.

        Args:
            request (Request): The reranking request which has a query and a candidates list.
            rank_start (int, optional): The starting rank for processing. Defaults to 0.
            rank_end (int, optional): The end rank for processing. Defaults to 100.
            window_size (int, optional): The size of each sliding window. Defaults to 20.
            step (int, optional): The step size for moving the window. Defaults to 10.
            shuffle_candidates (bool, optional): Whether to shuffle candidates before reranking. Defaults to False.
            logging (bool, optional): Enables logging of the reranking process. Defaults to False.

        Returns:
            Result: the rerank result which contains the reranked candidates.
        """
        results = self.rerank_batch(
            requests=[request],
            rank_start=rank_start,
            rank_end=rank_end,
            shuffle_candidates=shuffle_candidates,
            logging=logging,
            **kwargs,
        )
        return results[0]

    def write_rerank_results(
        self,
        retrieval_method_name: str,
        results: List[Result],
        shuffle_candidates: bool = False,
        top_k_candidates: int = 100,
        dataset_name: str = None,
        rerank_results_dirname: str = "rerank_results",
        ranking_execution_summary_dirname: str = "ranking_execution_summary",
        vllm_batched: bool = False,
        sglang_batched: bool = False,
        tensorrt_batched: bool = False,
        **kwargs,
    ) -> str:
        """
        Writes the reranked results to files in specified formats.

        This function saves the reranked results in both TREC Eval format and JSON format.
        A summary of the ranking execution is saved as well.

        Args:
            retrieval_method_name (str): The name of the retrieval method.
            results (List[Result]): The reranked results to be written.
            shuffle_candidates (bool, optional): Indicates if the candidates were shuffled. Defaults to False.
            top_k_candidates (int, optional): The number of top candidates considered. Defaults to 100.
            pass_ct (int, optional): Pass count, if applicable. Defaults to None.
            window_size (int, optional): The window size used in reranking. Defaults to None.
            dataset_name (str, optional): The name of the dataset used. Defaults to None.
            vllm_batched (bool, optional): Indicates if vLLM inference backend used. Defaults to False.
            sglang_batched (bool, optional): Indicates if SGLang inference backend used. Defaults to False.

        Returns:
            str: The file name of the saved reranked results in TREC Eval format.

        Note:
            The function creates directories and files as needed. The file names are constructed based on the
            provided parameters and the current timestamp to ensure uniqueness so there are no collisions.
        """
        pass_ct: Optional[int] = kwargs.get("pass_ct", None)
        window_size: Optional[int] = kwargs.get("window_size", None)

        name = self._agent.get_output_filename(
            top_k_candidates, dataset_name, shuffle_candidates, **kwargs
        )

        if window_size is not None:
            name += f"_window_{window_size}"
        if pass_ct is not None:
            name += f"_pass_{pass_ct}"

        # Add vllm or sglang to rerank result file name if they are used
        if vllm_batched:
            name += "_vllm"
        if sglang_batched:
            name += "_sglang"
        if tensorrt_batched:
            name += "_tensorrt"

        # write rerank results
        writer = DataWriter(results)
        Path(f"{rerank_results_dirname}/{retrieval_method_name}/").mkdir(
            parents=True, exist_ok=True
        )
        result_file_name = (
            f"{rerank_results_dirname}/{retrieval_method_name}/{name}.txt"
        )
        writer.write_in_trec_eval_format(result_file_name)
        writer.write_in_jsonl_format(
            f"{rerank_results_dirname}/{retrieval_method_name}/{name}.jsonl"
        )
        # Write ranking execution summary
        Path(f"{ranking_execution_summary_dirname}/{retrieval_method_name}/").mkdir(
            parents=True, exist_ok=True
        )
        writer.write_ranking_exec_summary(
            f"{ranking_execution_summary_dirname}/{retrieval_method_name}/{name}.json"
        )
        return result_file_name

    def get_agent(self) -> RankLLM:
        return self._agent

    def create_agent(
        model_path: str,
        default_agent: RankLLM,
        interactive: bool,
        **kwargs: Any,
    ) -> RankLLM:
        """Construct rerank agent

        Keyword arguments:
        argument -- description
        model_path -- name of model
        default_agent -- used for interactive mode to pass in a pre-instantiated agent to use
        interactive -- whether to run retrieve_and_rerank in interactive mode, used by the API

        Return: rerank agent -- Option<RankLLM>
        """
        use_azure_openai: bool = kwargs.get("use_azure_openai", False)
        vllm_batched: bool = kwargs.get("vllm_batched", False)

        if interactive and default_agent is not None:
            # Default rerank agent
            agent = default_agent
        elif "gpt" in model_path or use_azure_openai:
            # GPT based reranking models

            keys_and_defaults = [
                ("context_size", 4096),
                ("prompt_mode", PromptMode.RANK_GPT),
                ("num_few_shot_examples", 0),
                ("window_size", 20),
            ]
            [
                context_size,
                prompt_mode,
                num_few_shot_examples,
                window_size,
            ] = extract_kwargs(keys_and_defaults, **kwargs)

            openai_keys = get_openai_api_key()
            agent = SafeOpenai(
                model=model_path,
                context_size=context_size,
                prompt_mode=prompt_mode,
                window_size=window_size,
                num_few_shot_examples=num_few_shot_examples,
                keys=openai_keys,
                **(get_azure_openai_args() if use_azure_openai else {}),
            )
        elif "vicuna" in model_path or "zephyr" in model_path:
            # RankVicuna or RankZephyr model suite
            print(f"Loading {model_path} ...")

            model_full_paths = {
                "rank_zephyr": "castorini/rank_zephyr_7b_v1_full",
                "rank_vicuna": "castorini/rank_vicuna_7b_v1",
            }

            keys_and_defaults = [
                ("context_size", 4096),
                ("prompt_mode", PromptMode.RANK_GPT),
                ("num_few_shot_examples", 0),
                ("device", "cuda"),
                ("num_gpus", 1),
                ("variable_passages", False),
                ("window_size", 20),
                ("system_message", None),
                ("vllm_batched", False),
                ("sglang_batched", False),
                ("tensorrt_batched", False),
                ("use_logits", False),
                ("use_alpha", False),
            ]
            [
                context_size,
                prompt_mode,
                num_few_shot_examples,
                device,
                num_gpus,
                variable_passages,
                window_size,
                system_message,
                vllm_batched,
                sglang_batched,
                tensorrt_batched,
                use_logits,
                use_alpha,
            ] = extract_kwargs(keys_and_defaults, **kwargs)

            agent = RankListwiseOSLLM(
                model=(
                    model_full_paths[model_path]
                    if model_path in model_full_paths
                    else model_path
                ),
                name=model_path,
                context_size=context_size,
                prompt_mode=prompt_mode,
                num_few_shot_examples=num_few_shot_examples,
                device=device,
                num_gpus=num_gpus,
                variable_passages=variable_passages,
                window_size=window_size,
                system_message=system_message,
                vllm_batched=vllm_batched,
                sglang_batched=sglang_batched,
                tensorrt_batched=tensorrt_batched,
                use_logits=use_logits,
                use_alpha=use_alpha,
            )

            print(f"Completed loading {model_path}")
        elif "monot5" in model_path:
            # using monot5
            print(f"Loading {model_path} ...")

            model_full_paths = {"monot5": "castorini/monot5-3b-msmarco-10k"}

            keys_and_defaults = [
                ("prompt_mode", PromptMode.MONOT5),
                ("context_size", 512),
                ("device", "cuda"),
                ("batch_size", 64),
            ]
            [prompt_mode, context_size, device, batch_size] = extract_kwargs(
                keys_and_defaults, **kwargs
            )

            agent = MonoT5(
                model=(
                    model_full_paths[model_path]
                    if model_path in model_full_paths
                    else model_path
                ),
                prompt_mode=prompt_mode,
                context_size=context_size,
                device=device,
                batch_size=batch_size,
            )

        elif "lit5-distill" in model_path.lower():
            keys_and_defaults = [
                ("context_size", 150),
                ("prompt_mode", PromptMode.LiT5),
                ("num_few_shot_examples", 0),
                ("window_size", 20),
                ("precision", "bfloat16"),
                ("device", "cuda"),
                # reuse this parameter, but its not for "vllm", but only for "batched"
                ("vllm_batched", False),
            ]
            (
                context_size,
                prompt_mode,
                num_few_shot_examples,
                window_size,
                precision,
                device,
                vllm_batched,
            ) = extract_kwargs(keys_and_defaults, **kwargs)

            agent = RankFiDDistill(
                model=model_path,
                context_size=context_size,
                prompt_mode=prompt_mode,
                num_few_shot_examples=num_few_shot_examples,
                window_size=window_size,
                precision=precision,
                device=device,
                batched=vllm_batched,
            )
            print(f"Completed loading {model_path}")
        elif "lit5-score" in model_path.lower():
            keys_and_defaults = [
                ("context_size", 150),
                ("prompt_mode", PromptMode.LiT5),
                ("num_few_shot_examples", 0),
                ("window_size", 100),
                ("precision", "bfloat16"),
                ("device", "cuda"),
                # reuse this parameter, but its not for "vllm", but only for "batched"
                ("vllm_batched", False),
            ]
            (
                context_size,
                prompt_mode,
                num_few_shot_examples,
                window_size,
                precision,
                device,
                vllm_batched,
            ) = extract_kwargs(keys_and_defaults, **kwargs)

            agent = RankFiDScore(
                model=model_path,
                context_size=context_size,
                prompt_mode=prompt_mode,
                num_few_shot_examples=num_few_shot_examples,
                window_size=window_size,
                precision=precision,
                device=device,
                batched=vllm_batched,
            )
            print(f"Completed loading {model_path}")
        elif vllm_batched:
            # supports loading models from huggingface
            print(f"Loading {model_path} ...")
            keys_and_defaults = [
                ("context_size", 4096),
                ("prompt_mode", PromptMode.RANK_GPT),
                ("num_few_shot_examples", 0),
                ("device", "cuda"),
                ("num_gpus", 1),
                ("variable_passages", False),
                ("window_size", 20),
                ("system_message", None),
                ("vllm_batched", True),
                ("use_logits", False),
                ("use_alpha", False),
            ]
            [
                context_size,
                prompt_mode,
                num_few_shot_examples,
                device,
                num_gpus,
                variable_passages,
                window_size,
                system_message,
                vllm_batched,
                use_logits,
                use_alpha,
            ] = extract_kwargs(keys_and_defaults, **kwargs)

            agent = RankListwiseOSLLM(
                model=(model_path),
                name=model_path,
                context_size=context_size,
                prompt_mode=prompt_mode,
                num_few_shot_examples=num_few_shot_examples,
                device=device,
                num_gpus=num_gpus,
                variable_passages=variable_passages,
                window_size=window_size,
                system_message=system_message,
                use_logits=use_logits,
                use_alpha=use_alpha,
                vllm_batched=vllm_batched,
            )

            print(f"Completed loading {model_path}")
        elif model_path in ["unspecified", "rank_random", "rank_identity"]:
            # NULL reranker
            agent = None
        else:
            raise ValueError(f"Unsupported model: {model_path}")

        if agent is None and model_path not in [
            "unspecified",
            "rank_random",
            "rank_identity",
        ]:
            raise ValueError(f"Unsupported model: {model_path}")
        return agent

rerank_batch(requests, rank_start=0, rank_end=100, shuffle_candidates=False, logging=False, **kwargs)

Reranks a list of requests using the RankLLM agent.

This function applies a sliding window algorithm to rerank the results. Each window of results is processed by the RankLLM agent to obtain a new ranking.

Parameters:

Name Type Description Default
requests List[Request]

The list of requests. Each request has a query and a candidates list.

 required
rank_start int

The starting rank for processing. Defaults to 0.

 0
rank_end int

The end rank for processing. Defaults to 100.

 100
window_size int

The size of each sliding window. Defaults to 20.

 required
step int

The step size for moving the window. Defaults to 10.

 required
shuffle_candidates bool

Whether to shuffle candidates before reranking. Defaults to False.

 False
logging bool

Enables logging of the reranking process. Defaults to False.

 False
vllm_batched bool

Whether to use VLLM batched processing. Defaults to False.

 required
sglang_batched bool

Whether to use SGLang batched processing. Defaults to False.

 required
tensorrt_batched bool

Whether to use TensorRT-LLM batched processing. Defaults to False.

 required
populate_exec_summary bool

Whether to populate the exec summary. Defaults to False.

 required
batched bool

Whether to use batched processing. Defaults to False.

 required

Returns:

Type Description
List[Result]

List[Result]: A list containing the reranked candidates.
Source code in rankify/utils/models/rank_llm/rerank/reranker.py 
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
def rerank_batch(
    self,
    requests: List[Request],
    rank_start: int = 0,
    rank_end: int = 100,
    shuffle_candidates: bool = False,
    logging: bool = False,
    **kwargs: Any,
) -> List[Result]:
    """
    Reranks a list of requests using the RankLLM agent.

    This function applies a sliding window algorithm to rerank the results.
    Each window of results is processed by the RankLLM agent to obtain a new ranking.

    Args:
        requests (List[Request]): The list of requests. Each request has a query and a candidates list.
        rank_start (int, optional): The starting rank for processing. Defaults to 0.
        rank_end (int, optional): The end rank for processing. Defaults to 100.
        window_size (int, optional): The size of each sliding window. Defaults to 20.
        step (int, optional): The step size for moving the window. Defaults to 10.
        shuffle_candidates (bool, optional): Whether to shuffle candidates before reranking. Defaults to False.
        logging (bool, optional): Enables logging of the reranking process. Defaults to False.
        vllm_batched (bool, optional): Whether to use VLLM batched processing. Defaults to False.
        sglang_batched (bool, optional): Whether to use SGLang batched processing. Defaults to False.
        tensorrt_batched (bool, optional): Whether to use TensorRT-LLM batched processing. Defaults to False.
        populate_exec_summary (bool, optional): Whether to populate the exec summary. Defaults to False.
        batched (bool, optional): Whether to use batched processing. Defaults to False.

    Returns:
        List[Result]: A list containing the reranked candidates.
    """
    return self._agent.rerank_batch(
        requests, rank_start, rank_end, shuffle_candidates, logging, **kwargs
    )

rerank(request, rank_start=0, rank_end=100, shuffle_candidates=False, logging=False, **kwargs)

Reranks a request using the RankLLM agent.

This function applies a sliding window algorithm to rerank the results. Each window of results is processed by the RankLLM agent to obtain a new ranking.

Parameters:

Name Type Description Default
request Request

The reranking request which has a query and a candidates list.

 required
rank_start int

The starting rank for processing. Defaults to 0.

 0
rank_end int

The end rank for processing. Defaults to 100.

 100
window_size int

The size of each sliding window. Defaults to 20.

 required
step int

The step size for moving the window. Defaults to 10.

 required
shuffle_candidates bool

Whether to shuffle candidates before reranking. Defaults to False.

 False
logging bool

Enables logging of the reranking process. Defaults to False.

 False

Returns:

Name Type Description
Result Result

the rerank result which contains the reranked candidates.
Source code in rankify/utils/models/rank_llm/rerank/reranker.py 
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
def rerank(
    self,
    request: Request,
    rank_start: int = 0,
    rank_end: int = 100,
    shuffle_candidates: bool = False,
    logging: bool = False,
    **kwargs: Any,
) -> Result:
    """
    Reranks a request using the RankLLM agent.

    This function applies a sliding window algorithm to rerank the results.
    Each window of results is processed by the RankLLM agent to obtain a new ranking.

    Args:
        request (Request): The reranking request which has a query and a candidates list.
        rank_start (int, optional): The starting rank for processing. Defaults to 0.
        rank_end (int, optional): The end rank for processing. Defaults to 100.
        window_size (int, optional): The size of each sliding window. Defaults to 20.
        step (int, optional): The step size for moving the window. Defaults to 10.
        shuffle_candidates (bool, optional): Whether to shuffle candidates before reranking. Defaults to False.
        logging (bool, optional): Enables logging of the reranking process. Defaults to False.

    Returns:
        Result: the rerank result which contains the reranked candidates.
    """
    results = self.rerank_batch(
        requests=[request],
        rank_start=rank_start,
        rank_end=rank_end,
        shuffle_candidates=shuffle_candidates,
        logging=logging,
        **kwargs,
    )
    return results[0]

write_rerank_results(retrieval_method_name, results, shuffle_candidates=False, top_k_candidates=100, dataset_name=None, rerank_results_dirname='rerank_results', ranking_execution_summary_dirname='ranking_execution_summary', vllm_batched=False, sglang_batched=False, tensorrt_batched=False, **kwargs)

Writes the reranked results to files in specified formats.

This function saves the reranked results in both TREC Eval format and JSON format. A summary of the ranking execution is saved as well.

Parameters:

Name Type Description Default
retrieval_method_name str

The name of the retrieval method.

 required
results List[Result]

The reranked results to be written.

 required
shuffle_candidates bool

Indicates if the candidates were shuffled. Defaults to False.

 False
top_k_candidates int

The number of top candidates considered. Defaults to 100.

 100
pass_ct int

Pass count, if applicable. Defaults to None.

 required
window_size int

The window size used in reranking. Defaults to None.

 required
dataset_name str

The name of the dataset used. Defaults to None.

 None
vllm_batched bool

Indicates if vLLM inference backend used. Defaults to False.

 False
sglang_batched bool

Indicates if SGLang inference backend used. Defaults to False.

 False

Returns:

Name Type Description
str str

The file name of the saved reranked results in TREC Eval format.
Note

The function creates directories and files as needed. The file names are constructed based on the provided parameters and the current timestamp to ensure uniqueness so there are no collisions.

Source code in rankify/utils/models/rank_llm/rerank/reranker.py 
 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
def write_rerank_results(
    self,
    retrieval_method_name: str,
    results: List[Result],
    shuffle_candidates: bool = False,
    top_k_candidates: int = 100,
    dataset_name: str = None,
    rerank_results_dirname: str = "rerank_results",
    ranking_execution_summary_dirname: str = "ranking_execution_summary",
    vllm_batched: bool = False,
    sglang_batched: bool = False,
    tensorrt_batched: bool = False,
    **kwargs,
) -> str:
    """
    Writes the reranked results to files in specified formats.

    This function saves the reranked results in both TREC Eval format and JSON format.
    A summary of the ranking execution is saved as well.

    Args:
        retrieval_method_name (str): The name of the retrieval method.
        results (List[Result]): The reranked results to be written.
        shuffle_candidates (bool, optional): Indicates if the candidates were shuffled. Defaults to False.
        top_k_candidates (int, optional): The number of top candidates considered. Defaults to 100.
        pass_ct (int, optional): Pass count, if applicable. Defaults to None.
        window_size (int, optional): The window size used in reranking. Defaults to None.
        dataset_name (str, optional): The name of the dataset used. Defaults to None.
        vllm_batched (bool, optional): Indicates if vLLM inference backend used. Defaults to False.
        sglang_batched (bool, optional): Indicates if SGLang inference backend used. Defaults to False.

    Returns:
        str: The file name of the saved reranked results in TREC Eval format.

    Note:
        The function creates directories and files as needed. The file names are constructed based on the
        provided parameters and the current timestamp to ensure uniqueness so there are no collisions.
    """
    pass_ct: Optional[int] = kwargs.get("pass_ct", None)
    window_size: Optional[int] = kwargs.get("window_size", None)

    name = self._agent.get_output_filename(
        top_k_candidates, dataset_name, shuffle_candidates, **kwargs
    )

    if window_size is not None:
        name += f"_window_{window_size}"
    if pass_ct is not None:
        name += f"_pass_{pass_ct}"

    # Add vllm or sglang to rerank result file name if they are used
    if vllm_batched:
        name += "_vllm"
    if sglang_batched:
        name += "_sglang"
    if tensorrt_batched:
        name += "_tensorrt"

    # write rerank results
    writer = DataWriter(results)
    Path(f"{rerank_results_dirname}/{retrieval_method_name}/").mkdir(
        parents=True, exist_ok=True
    )
    result_file_name = (
        f"{rerank_results_dirname}/{retrieval_method_name}/{name}.txt"
    )
    writer.write_in_trec_eval_format(result_file_name)
    writer.write_in_jsonl_format(
        f"{rerank_results_dirname}/{retrieval_method_name}/{name}.jsonl"
    )
    # Write ranking execution summary
    Path(f"{ranking_execution_summary_dirname}/{retrieval_method_name}/").mkdir(
        parents=True, exist_ok=True
    )
    writer.write_ranking_exec_summary(
        f"{ranking_execution_summary_dirname}/{retrieval_method_name}/{name}.json"
    )
    return result_file_name

create_agent(model_path, default_agent, interactive, **kwargs)

Construct rerank agent

Keyword arguments: argument -- description model_path -- name of model default_agent -- used for interactive mode to pass in a pre-instantiated agent to use interactive -- whether to run retrieve_and_rerank in interactive mode, used by the API

Return: rerank agent -- Option

Source code in rankify/utils/models/rank_llm/rerank/reranker.py 
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
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
def create_agent(
    model_path: str,
    default_agent: RankLLM,
    interactive: bool,
    **kwargs: Any,
) -> RankLLM:
    """Construct rerank agent

    Keyword arguments:
    argument -- description
    model_path -- name of model
    default_agent -- used for interactive mode to pass in a pre-instantiated agent to use
    interactive -- whether to run retrieve_and_rerank in interactive mode, used by the API

    Return: rerank agent -- Option<RankLLM>
    """
    use_azure_openai: bool = kwargs.get("use_azure_openai", False)
    vllm_batched: bool = kwargs.get("vllm_batched", False)

    if interactive and default_agent is not None:
        # Default rerank agent
        agent = default_agent
    elif "gpt" in model_path or use_azure_openai:
        # GPT based reranking models

        keys_and_defaults = [
            ("context_size", 4096),
            ("prompt_mode", PromptMode.RANK_GPT),
            ("num_few_shot_examples", 0),
            ("window_size", 20),
        ]
        [
            context_size,
            prompt_mode,
            num_few_shot_examples,
            window_size,
        ] = extract_kwargs(keys_and_defaults, **kwargs)

        openai_keys = get_openai_api_key()
        agent = SafeOpenai(
            model=model_path,
            context_size=context_size,
            prompt_mode=prompt_mode,
            window_size=window_size,
            num_few_shot_examples=num_few_shot_examples,
            keys=openai_keys,
            **(get_azure_openai_args() if use_azure_openai else {}),
        )
    elif "vicuna" in model_path or "zephyr" in model_path:
        # RankVicuna or RankZephyr model suite
        print(f"Loading {model_path} ...")

        model_full_paths = {
            "rank_zephyr": "castorini/rank_zephyr_7b_v1_full",
            "rank_vicuna": "castorini/rank_vicuna_7b_v1",
        }

        keys_and_defaults = [
            ("context_size", 4096),
            ("prompt_mode", PromptMode.RANK_GPT),
            ("num_few_shot_examples", 0),
            ("device", "cuda"),
            ("num_gpus", 1),
            ("variable_passages", False),
            ("window_size", 20),
            ("system_message", None),
            ("vllm_batched", False),
            ("sglang_batched", False),
            ("tensorrt_batched", False),
            ("use_logits", False),
            ("use_alpha", False),
        ]
        [
            context_size,
            prompt_mode,
            num_few_shot_examples,
            device,
            num_gpus,
            variable_passages,
            window_size,
            system_message,
            vllm_batched,
            sglang_batched,
            tensorrt_batched,
            use_logits,
            use_alpha,
        ] = extract_kwargs(keys_and_defaults, **kwargs)

        agent = RankListwiseOSLLM(
            model=(
                model_full_paths[model_path]
                if model_path in model_full_paths
                else model_path
            ),
            name=model_path,
            context_size=context_size,
            prompt_mode=prompt_mode,
            num_few_shot_examples=num_few_shot_examples,
            device=device,
            num_gpus=num_gpus,
            variable_passages=variable_passages,
            window_size=window_size,
            system_message=system_message,
            vllm_batched=vllm_batched,
            sglang_batched=sglang_batched,
            tensorrt_batched=tensorrt_batched,
            use_logits=use_logits,
            use_alpha=use_alpha,
        )

        print(f"Completed loading {model_path}")
    elif "monot5" in model_path:
        # using monot5
        print(f"Loading {model_path} ...")

        model_full_paths = {"monot5": "castorini/monot5-3b-msmarco-10k"}

        keys_and_defaults = [
            ("prompt_mode", PromptMode.MONOT5),
            ("context_size", 512),
            ("device", "cuda"),
            ("batch_size", 64),
        ]
        [prompt_mode, context_size, device, batch_size] = extract_kwargs(
            keys_and_defaults, **kwargs
        )

        agent = MonoT5(
            model=(
                model_full_paths[model_path]
                if model_path in model_full_paths
                else model_path
            ),
            prompt_mode=prompt_mode,
            context_size=context_size,
            device=device,
            batch_size=batch_size,
        )

    elif "lit5-distill" in model_path.lower():
        keys_and_defaults = [
            ("context_size", 150),
            ("prompt_mode", PromptMode.LiT5),
            ("num_few_shot_examples", 0),
            ("window_size", 20),
            ("precision", "bfloat16"),
            ("device", "cuda"),
            # reuse this parameter, but its not for "vllm", but only for "batched"
            ("vllm_batched", False),
        ]
        (
            context_size,
            prompt_mode,
            num_few_shot_examples,
            window_size,
            precision,
            device,
            vllm_batched,
        ) = extract_kwargs(keys_and_defaults, **kwargs)

        agent = RankFiDDistill(
            model=model_path,
            context_size=context_size,
            prompt_mode=prompt_mode,
            num_few_shot_examples=num_few_shot_examples,
            window_size=window_size,
            precision=precision,
            device=device,
            batched=vllm_batched,
        )
        print(f"Completed loading {model_path}")
    elif "lit5-score" in model_path.lower():
        keys_and_defaults = [
            ("context_size", 150),
            ("prompt_mode", PromptMode.LiT5),
            ("num_few_shot_examples", 0),
            ("window_size", 100),
            ("precision", "bfloat16"),
            ("device", "cuda"),
            # reuse this parameter, but its not for "vllm", but only for "batched"
            ("vllm_batched", False),
        ]
        (
            context_size,
            prompt_mode,
            num_few_shot_examples,
            window_size,
            precision,
            device,
            vllm_batched,
        ) = extract_kwargs(keys_and_defaults, **kwargs)

        agent = RankFiDScore(
            model=model_path,
            context_size=context_size,
            prompt_mode=prompt_mode,
            num_few_shot_examples=num_few_shot_examples,
            window_size=window_size,
            precision=precision,
            device=device,
            batched=vllm_batched,
        )
        print(f"Completed loading {model_path}")
    elif vllm_batched:
        # supports loading models from huggingface
        print(f"Loading {model_path} ...")
        keys_and_defaults = [
            ("context_size", 4096),
            ("prompt_mode", PromptMode.RANK_GPT),
            ("num_few_shot_examples", 0),
            ("device", "cuda"),
            ("num_gpus", 1),
            ("variable_passages", False),
            ("window_size", 20),
            ("system_message", None),
            ("vllm_batched", True),
            ("use_logits", False),
            ("use_alpha", False),
        ]
        [
            context_size,
            prompt_mode,
            num_few_shot_examples,
            device,
            num_gpus,
            variable_passages,
            window_size,
            system_message,
            vllm_batched,
            use_logits,
            use_alpha,
        ] = extract_kwargs(keys_and_defaults, **kwargs)

        agent = RankListwiseOSLLM(
            model=(model_path),
            name=model_path,
            context_size=context_size,
            prompt_mode=prompt_mode,
            num_few_shot_examples=num_few_shot_examples,
            device=device,
            num_gpus=num_gpus,
            variable_passages=variable_passages,
            window_size=window_size,
            system_message=system_message,
            use_logits=use_logits,
            use_alpha=use_alpha,
            vllm_batched=vllm_batched,
        )

        print(f"Completed loading {model_path}")
    elif model_path in ["unspecified", "rank_random", "rank_identity"]:
        # NULL reranker
        agent = None
    else:
        raise ValueError(f"Unsupported model: {model_path}")

    if agent is None and model_path not in [
        "unspecified",
        "rank_random",
        "rank_identity",
    ]:
        raise ValueError(f"Unsupported model: {model_path}")
    return agent

Request dataclass

Source code in rankify/utils/models/rank_llm/data.py 
21
22
23
24
@dataclass
class Request:
    query: Query
    candidates: List[Candidate] = field(default_factory=list)

LiT5DistillReranker

Bases: BaseRanking

Implements LiT5-Distill, a listwise reranker based on T5 encoder-decoder models.

LiT5-Distill is designed for zero-shot ranking, leveraging sequence-to-sequence architectures to improve retrieval performance with efficient inference.

References
  • Tamber et al. (2023): Scaling Down, LiTting Up: Efficient Zero-Shot Listwise Reranking with Seq2seq Encoder-Decoder Models. Paper

Attributes:

Name Type Description
context_size int

The maximum number of tokens used in ranking (default: 300).
window_size int

The window size for processing candidate passages (default: 20).
_reranker Reranker

The LiT5-Distill reranking agent.
Example 
from rankify.dataset.dataset import Document, Question, Context
from rankify.models.reranking import Reranking

question = Question("What are the effects of climate change?")
contexts = [
    Context(text="Rising temperatures are causing ice caps to melt.", id=0),
    Context(text="Many species face extinction due to habitat loss.", id=1),
    Context(text="Ocean acidification is increasing.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize LiT5-Distill reranker
model = Reranking(method='lit5distill', model_name='castorini/LiT5-Distill-base')
model.rank([document])

# Print reordered contexts
print("Reordered Contexts:")
for context in document.reorder_contexts:
    print(context.text)
Source code in rankify/models/lit5_reranker.py 
 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
class LiT5DistillReranker(BaseRanking):
    """
    Implements **LiT5-Distill**, a **listwise reranker** based on **T5 encoder-decoder models**.


    **LiT5-Distill** is designed for **zero-shot ranking**, leveraging **sequence-to-sequence architectures** 
    to improve **retrieval performance** with **efficient inference**.

    References:
        - **Tamber et al. (2023)**: *Scaling Down, LiTting Up: Efficient Zero-Shot Listwise Reranking with Seq2seq Encoder-Decoder Models*.
          [Paper](https://arxiv.org/abs/2312.16098)

    Attributes:
        context_size (int): The **maximum number of tokens** used in ranking (default: `300`).
        window_size (int): The **window size** for processing candidate passages (default: `20`).
        _reranker (Reranker): The **LiT5-Distill** reranking agent.

    Example:
        ```python
        from rankify.dataset.dataset import Document, Question, Context
        from rankify.models.reranking import Reranking

        question = Question("What are the effects of climate change?")
        contexts = [
            Context(text="Rising temperatures are causing ice caps to melt.", id=0),
            Context(text="Many species face extinction due to habitat loss.", id=1),
            Context(text="Ocean acidification is increasing.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize LiT5-Distill reranker
        model = Reranking(method='lit5distill', model_name='castorini/LiT5-Distill-base')
        model.rank([document])

        # Print reordered contexts
        print("Reordered Contexts:")
        for context in document.reorder_contexts:
            print(context.text)
        ```
    """
    def __init__(self,method: str= None, model_name: str = "castorini/LiT5-Distill-base", api_key: str= None, **kwargs):
        """
        Initializes the **LiT5-Distill reranker**.

        Args:
            method (str, optional): The **reranking method name**.
            model_name (str, optional): The name of the **pre-trained LiT5-Distill model** (default: `"castorini/LiT5-Distill-base"`).
            api_key (str, optional): API key for **authentication** (if needed).
        """
        self.context_size = 300
        self.window_size = 20


        self._reranker = Reranker(
            Reranker.create_agent(model_name, default_agent=None, interactive=False)
        )

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks each document's **contexts** using the **LiT5-Distill model**.

        Args:
            documents (List[Document]): A list of **Document** instances containing contexts to rerank.

        Returns:
            List[Document]: The reranked list of **Document** instances with updated `reorder_contexts`.
        """
        for document in tqdm(documents, desc="Reranking Documents"):

            document = self._rerank_document(document)
        return documents

    def _rerank_document(self, document: Document):
        """
        Reranks a **single document** using the **LiT5-Distill model**.

        Args:
            document (Document): A **Document** instance containing query and contexts.

        Returns:
            Document: The document with **reranked** contexts in `reorder_contexts`.
        """
        # Prepare request data structure for reranking
        request = Request(
            query={"text": document.question.question},  # Extract `text` from `Question`
            candidates=[
                {"docid": ctx.id, "doc": {"text": ctx.text} , "score": ctx.score } for ctx in document.contexts
            ]
        )
        rank_start= 0
        rank_end = 100
        window_size = 20
        step= 10
        shuffle_candidates = False
        logging = False,

        # Rerank using the agent
        reranked_result = self._reranker.rerank(
            request=request,
            rank_start=rank_start,
            rank_end=rank_end,
            window_size=window_size,
            step=step,
            shuffle_candidates=shuffle_candidates,
            #logging=logging,
        )
            # Create a mapping from docid to the original context

        contexts = copy.deepcopy(document.contexts)
        docid_to_context = {str(ctx.id): ctx for ctx in contexts}
        #print("Hello")
        #print(docid_to_context)
        #print(reranked_result.candidates)
        # Reorder contexts based on reranked_result
        reorder_contexts = []
        for candidate in reranked_result.candidates:
            #print(candidate)
            d = docid_to_context[str(candidate["docid"])]
            #print(d)
            d.score = candidate["score"]
            reorder_contexts.append(d)
        document.reorder_contexts = reorder_contexts

        #print(document.reorder_contexts)
        return document

__init__(method=None, model_name='castorini/LiT5-Distill-base', api_key=None, **kwargs)

Initializes the LiT5-Distill reranker.

Parameters:

Name Type Description Default
method str

The reranking method name.

 None
model_name str

The name of the pre-trained LiT5-Distill model (default: "castorini/LiT5-Distill-base").

 'castorini/LiT5-Distill-base'
api_key str

API key for authentication (if needed).

 None
Source code in rankify/models/lit5_reranker.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(self,method: str= None, model_name: str = "castorini/LiT5-Distill-base", api_key: str= None, **kwargs):
    """
    Initializes the **LiT5-Distill reranker**.

    Args:
        method (str, optional): The **reranking method name**.
        model_name (str, optional): The name of the **pre-trained LiT5-Distill model** (default: `"castorini/LiT5-Distill-base"`).
        api_key (str, optional): API key for **authentication** (if needed).
    """
    self.context_size = 300
    self.window_size = 20


    self._reranker = Reranker(
        Reranker.create_agent(model_name, default_agent=None, interactive=False)
    )

rank(documents)

Reranks each document's contexts using the LiT5-Distill model.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document instances containing contexts to rerank.

 required

Returns:

Type Description
List[Document]

List[Document]: The reranked list of Document instances with updated reorder_contexts.
Source code in rankify/models/lit5_reranker.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks each document's **contexts** using the **LiT5-Distill model**.

    Args:
        documents (List[Document]): A list of **Document** instances containing contexts to rerank.

    Returns:
        List[Document]: The reranked list of **Document** instances with updated `reorder_contexts`.
    """
    for document in tqdm(documents, desc="Reranking Documents"):

        document = self._rerank_document(document)
    return documents

LiT5ScoreReranker

Bases: BaseRanking

Implements LiT5-Score, a listwise reranker that generates direct ranking scores for passages.

LiT5-Score assigns numerical scores to each passage relative to a query, allowing for zero-shot ranking without fine-tuning.

References
  • Tamber et al. (2023): Scaling Down, LiTting Up: Efficient Zero-Shot Listwise Reranking with Seq2seq Encoder-Decoder Models. Paper

Attributes:

Name Type Description
context_size int

The maximum number of tokens used in ranking (default: 300).
window_size int

The window size for processing candidate passages (default: 20).
_reranker Reranker

The LiT5-Score reranking agent.
Source code in rankify/models/lit5_reranker.py 
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
class LiT5ScoreReranker(BaseRanking):
    """
    Implements **LiT5-Score**, a **listwise reranker** that generates **direct ranking scores** for passages.

    **LiT5-Score** assigns **numerical scores** to **each passage** relative to a **query**, 
    allowing for **zero-shot ranking** without fine-tuning.

    References:
        - **Tamber et al. (2023)**: *Scaling Down, LiTting Up: Efficient Zero-Shot Listwise Reranking with Seq2seq Encoder-Decoder Models*.
          [Paper](https://arxiv.org/abs/2312.16098)

    Attributes:
        context_size (int): The **maximum number of tokens** used in ranking (default: `300`).
        window_size (int): The **window size** for processing candidate passages (default: `20`).
        _reranker (Reranker): The **LiT5-Score** reranking agent.
    """
    def __init__(self, method: str = None, model_name: str = "castorini/LiT5-Score-base", api_key: str = None, **kwargs):
        """
        Initializes the **LiT5-Score reranker**.

        Args:
            method (str, optional): The **reranking method name**.
            model_name (str, optional): The name of the **pre-trained LiT5-Score model** (default: `"castorini/LiT5-Score-base"`).
            api_key (str, optional): API key for **authentication** (if needed).
        """
        self.context_size = 300
        self.window_size = 20

        self._reranker = Reranker(
            Reranker.create_agent(model_name, default_agent=None, interactive=False)
        )

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks each document's **contexts** using the **LiT5-Score model**.

        Args:
            documents (List[Document]): A list of **Document** instances containing contexts to rerank.

        Returns:
            List[Document]: The reranked list of **Document** instances with updated `reorder_contexts`.
        """
        for document in tqdm(documents, desc="Reranking Documents"):
            document = self._rerank_document(document)
        return documents

    def _rerank_document(self, document: Document):
        """
        Reranks a **single document** using the **LiT5-Score model**.

        Args:
            document (Document): A **Document** instance containing query and contexts.

        Returns:
            Document: The document with **reranked** contexts in `reorder_contexts`.
        """
        # Prepare request data structure for reranking
        request = Request(
            query={"text": document.question.question},  # Extract `text` from `Question`
            candidates=[
                {"docid": ctx.id, "doc": {"text": ctx.text}, "score": ctx.score} for ctx in document.contexts
            ]
        )

        # Reranking parameters
        rank_start = 0
        rank_end = 100
        window_size = 20
        step = 10
        shuffle_candidates = False
        logging = False

        # Rerank using the agent
        reranked_result = self._reranker.rerank(
            request=request,
            rank_start=rank_start,
            rank_end=rank_end,
            window_size=window_size,
            step=step,
            shuffle_candidates=shuffle_candidates,
            logging=logging,
        )

        # Create a mapping from docid to the original context
        contexts = copy.deepcopy(document.contexts)
        docid_to_context = {str(ctx.id): ctx for ctx in contexts}
        # Reorder contexts based on reranked_result
        reorder_contexts = []
        for candidate in reranked_result.candidates:
            d = docid_to_context[str(candidate["docid"])]
            d.score = candidate["score"]
            reorder_contexts.append(d)
        document.reorder_contexts = reorder_contexts
        return document

__init__(method=None, model_name='castorini/LiT5-Score-base', api_key=None, **kwargs)

Initializes the LiT5-Score reranker.

Parameters:

Name Type Description Default
method str

The reranking method name.

 None
model_name str

The name of the pre-trained LiT5-Score model (default: "castorini/LiT5-Score-base").

 'castorini/LiT5-Score-base'
api_key str

API key for authentication (if needed).

 None
Source code in rankify/models/lit5_reranker.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def __init__(self, method: str = None, model_name: str = "castorini/LiT5-Score-base", api_key: str = None, **kwargs):
    """
    Initializes the **LiT5-Score reranker**.

    Args:
        method (str, optional): The **reranking method name**.
        model_name (str, optional): The name of the **pre-trained LiT5-Score model** (default: `"castorini/LiT5-Score-base"`).
        api_key (str, optional): API key for **authentication** (if needed).
    """
    self.context_size = 300
    self.window_size = 20

    self._reranker = Reranker(
        Reranker.create_agent(model_name, default_agent=None, interactive=False)
    )

rank(documents)

Reranks each document's contexts using the LiT5-Score model.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document instances containing contexts to rerank.

 required

Returns:

Type Description
List[Document]

List[Document]: The reranked list of Document instances with updated reorder_contexts.
Source code in rankify/models/lit5_reranker.py 
172
173
174
175
176
177
178
179
180
181
182
183
184
def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks each document's **contexts** using the **LiT5-Score model**.

    Args:
        documents (List[Document]): A list of **Document** instances containing contexts to rerank.

    Returns:
        List[Document]: The reranked list of **Document** instances with updated `reorder_contexts`.
    """
    for document in tqdm(documents, desc="Reranking Documents"):
        document = self._rerank_document(document)
    return documents