Skip to content

First Reranker

rankify.models.first_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}"

RankListwiseOSLLM

Bases: RankLLM

Source code in rankify/utils/models/rank_listwise_os_llm.py 
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
class RankListwiseOSLLM(RankLLM):
    def __init__(
        self,
        model: str,
        context_size: int = 4096,
        prompt_mode: PromptMode = PromptMode.RANK_GPT,
        num_few_shot_examples: int = 0,
        device: str = "cuda",
        num_gpus: int = 1,
        variable_passages: bool = False,
        window_size: int = 20,
        system_message: str = None,
        batched: bool = False,
        gpu_memory_utilization: float = 0.9,  # Increased GPU memory usage
        max_model_len: int = 29168  # Limit max model length within cache limit
    ) -> None:
        super().__init__(model, context_size, prompt_mode, num_few_shot_examples)
        self._device = device
        if self._device == "cuda":
            assert torch.cuda.is_available(), "CUDA is not available on this device"
        if prompt_mode != PromptMode.RANK_GPT:
            raise ValueError(
                f"Unsupported prompt mode: {prompt_mode}. Only RANK_GPT is supported."
            )

        # Initialize the LLM model with custom memory and length configurations
        self._llm = LLM(
            model=model, 
            max_logprobs=30, 
            enforce_eager=False,
            gpu_memory_utilization=gpu_memory_utilization,  # Set memory utilization
            max_model_len=max_model_len  # Set max sequence length
        )

        self._tokenizer = self._llm.get_tokenizer()
        self.system_message_supported = "system" in self._tokenizer.chat_template
        self._batched = batched
        self._variable_passages = variable_passages
        self._window_size = window_size
        self._system_message = system_message
        self._output_token_estimate = None

        if num_few_shot_examples > 0:
            with open("data/output_v2_aug_filtered.jsonl", "r") as json_file:
                self._examples = list(json_file)[1:-1]

    def _evaluate_logits(self, logits: Dict[str, 'Logit'], use_alpha: bool, total: Tuple[int, int]) -> Tuple[str, Dict[int, float]]:
        if use_alpha:
            evaluations = {
                ord(logit.decoded_token): logit.logprob
                for logit in logits.values()
                if len(logit.decoded_token) == 1 and 
                   logit.decoded_token.isalpha() and 
                   ALPH_START_IDX + 1 <= ord(logit.decoded_token) <= ALPH_START_IDX + self._window_size
            }
            sorted_evaluations = sorted(evaluations.items(), key=lambda x: -x[1])
            result_string = ">".join([f"[{chr(x)}]" for x, y in sorted_evaluations])
        else:
            evaluations = {
                int(logit.decoded_token): logit.logprob
                for logit in logits.values()
                if logit.decoded_token.isnumeric() and
                   not unicodedata.name(logit.decoded_token).startswith(('SUPERSCRIPT', 'VULGAR FRACTION', 'SUBSCRIPT')) and
                   total[0] <= int(logit.decoded_token) <= total[1]
            }
            sorted_evaluations = sorted(evaluations.items(), key=lambda x: -x[1])
            result_string = ">".join([f"[{x}]" for x, y in sorted_evaluations])

        return result_string, evaluations

    def _get_logits_single_digit(self, output: RequestOutput, use_alpha: bool = False, effective_location: int = 1, total: Tuple[int, int] = (1, 9)):
        logits = output.outputs[0].logprobs[effective_location]
        return self._evaluate_logits(logits, use_alpha, total)

    def run_llm_batched(
        self,
        prompts: List[Union[str, List[Dict[str, str]]]],
        current_window_size: Optional[int] = None,
        use_logits: bool = False,
        use_alpha: bool = False,
    ) -> List[Tuple[str, int]]:
        if current_window_size is None:
            current_window_size = self._window_size

        temp = 0.0
        if use_logits:
            params = SamplingParams(
                min_tokens=2,
                max_tokens=2, 
                temperature=temp,
                logprobs=30,
            )
            outputs = self._llm.generate(prompts, sampling_params=params, use_tqdm=True)
            arr = [self._get_logits_single_digit(output, use_alpha=use_alpha) for output in outputs]
            return [(s, len(s)) for s, __ in arr]
        else:
            params = SamplingParams(
                temperature=temp,
                max_tokens=self.num_output_tokens(use_alpha, current_window_size),
                min_tokens=self.num_output_tokens(use_alpha, current_window_size),
            )
            outputs = self._llm.generate(prompts, sampling_params=params, use_tqdm=True)
            return [
                (output.outputs[0].text, len(output.outputs[0].token_ids))
                for output in outputs
            ]

    def run_llm(
        self, prompt: str, current_window_size: Optional[int] = None, use_logits: bool = False, use_alpha: bool = False
    ) -> Tuple[str, int]:
        if current_window_size is None:
            current_window_size = self._window_size

        temp = 0.0
        if use_logits:
            params = SamplingParams(min_tokens=1, max_tokens=1, temperature=temp, logprobs=30)
            output = self._llm.generate([prompt+"["], sampling_params=params, use_tqdm=False)[0]
            s, _ = self._get_logits_single_digit(output, effective_location=0, use_alpha=use_alpha)
            return s, len(s)
        else:
            max_new_tokens = self.num_output_tokens(use_alpha, current_window_size)
            params = SamplingParams(min_tokens=max_new_tokens, max_tokens=max_new_tokens, temperature=temp)
            output = self._llm.generate([prompt], sampling_params=params, use_tqdm=False)[0]
            return output.outputs[0].text, len(output.outputs[0].text)

    def num_output_tokens(self, use_alpha: bool, current_window_size: Optional[int] = None) -> int:
        if current_window_size is None:
            current_window_size = self._window_size

        if self._output_token_estimate and self._window_size == current_window_size:
            return self._output_token_estimate

        if use_alpha:
            token_str = " > ".join([f"[{i+1}]" for i in range(current_window_size)])
        else:
            token_str = " > ".join([f"[{chr(ALPH_START_IDX+i+1)}]" for i in range(current_window_size)])

        _output_token_estimate = len(self._tokenizer.encode(token_str)) - 1

        if self._window_size == current_window_size:
            self._output_token_estimate = _output_token_estimate

        return _output_token_estimate

    def _add_prefix_prompt(self, use_alpha, query: str, num: int) -> str:
        if use_alpha:
            return f"I will provide you with {num} passages, each indicated by a alphabetical identifier []. Rank the passages based on their relevance to the search query: {query}.\n"
        else:
            return f"I will provide you with {num} passages, each indicated by a numerical identifier []. Rank the passages based on their relevance to the search query: {query}.\n"

    def _add_post_prompt(self, use_alpha, query: str, num: int) -> str:
        if use_alpha:
            example_ordering = "[B] > [A]" if self._variable_passages else "[D] > [B]"
        else:
            example_ordering = "[2] > [1]" if self._variable_passages else "[4] > [2]"
        return f"Search Query: {query}.\nRank the {num} passages above based on their relevance to the search query. All the passages should be included and listed using identifiers, in descending order of relevance. The output format should be [] > [], e.g., {example_ordering}, Only respond with the ranking results, do not say any word or explain."

    def _add_few_shot_examples(self, conv):
        for _ in range(self._num_few_shot_examples):
            ex = random.choice(self._examples)
            obj = json.loads(ex)
            prompt = obj["conversations"][0]["value"]
            response = obj["conversations"][1]["value"]
            conv.append_message(conv.roles[0], prompt)
            conv.append_message(conv.roles[1], response)
        return conv

    def _add_few_shot_examples_messages(self, messages):
        for _ in range(self._num_few_shot_examples):
            ex = random.choice(self._examples)
            obj = json.loads(ex)
            prompt = obj["conversations"][0]["value"]
            response = obj["conversations"][1]["value"]
            messages.append({"role": "user", "content": prompt})
            messages.append({"role": "assistant", "content": response})
        return messages

    def create_prompt(self, result, use_alpha: bool, rank_start: int, rank_end: int) -> Tuple[str, int]:
        query = result.query
        num = len(result.hits[rank_start:rank_end])
        max_length = 300
        while True:
            messages = list()
            if self._system_message and self.system_message_supported:
                messages.append({"role": "system", "content": self._system_message})
            messages = self._add_few_shot_examples_messages(messages)
            prefix = self._add_prefix_prompt(use_alpha, query, num)
            rank = 0
            input_context = f"{prefix}\n"
            for hit in result.hits[rank_start:rank_end]:
                rank += 1
                content = hit["content"].replace("Title: Content: ", "").strip()
                content = " ".join(content.split()[:max_length])
                identifier = chr(ALPH_START_IDX + rank) if use_alpha else str(rank)
                input_context += f"[{identifier}] {self._replace_number(content, use_alpha)}\n"
            input_context += self._add_post_prompt(use_alpha, query, num)
            messages.append({"role": "user", "content": input_context})
            if self._system_message and not self.system_message_supported:
                messages[0]["content"] = self._system_message + "\n " + messages[0]["content"]
            prompt = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            prompt = fix_text(prompt)
            num_tokens = self.get_num_tokens(prompt)
            if num_tokens <= self.max_tokens() - self.num_output_tokens(rank_end - rank_start):
                break
            else:
                max_length -= max(
                    1,
                    (
                        num_tokens - self.max_tokens() + self.num_output_tokens(rank_end - rank_start)
                    ) // ((rank_end - rank_start) * 4),
                )
        return prompt, num_tokens

    def create_prompt_batched(
        self,
        results,
        use_alpha: bool,
        rank_start: int,
        rank_end: int,
        batch_size: int = 32,
    ) -> List[Tuple[str, int]]:
        def chunks(lst, n):
            """Yield successive n-sized chunks from lst."""
            for i in range(0, len(lst), n):
                yield lst[i : i + n]

        all_completed_prompts = []

        with ThreadPoolExecutor() as executor:
            for batch in chunks(results, batch_size):
                completed_prompts = list(
                    executor.map(
                        lambda result: self.create_prompt(result, use_alpha, rank_start, rank_end),
                        batch,
                    )
                )
                all_completed_prompts.extend(completed_prompts)
        return all_completed_prompts

    def get_num_tokens(self, prompt: str) -> int:
        return len(self._tokenizer.encode(prompt))

    def cost_per_1k_token(self, input_token: bool) -> float:
        return 0

PromptMode

Bases: Enum

Source code in rankify/utils/models/rank_listwise_os_llm.py 
20
21
22
23
24
25
26
class PromptMode(Enum):
    UNSPECIFIED = "unspecified"
    RANK_GPT = "rank_GPT"
    LRL = "LRL"

    def __str__(self):
        return self.value

Result

Source code in rankify/utils/models/rank_listwise_os_llm.py 
705
706
707
708
709
710
711
712
713
714
715
716
717
class Result:
    def __init__(
        self,
        query: str,
        hits: List[Dict[str, Any]],
        ranking_exec_summary: List[RankingExecInfo] = None,
    ):
        self.query = query
        self.hits = hits
        self.ranking_exec_summary = ranking_exec_summary

    def __repr__(self):
        return str(self.__dict__)

FirstReranker

Source code in rankify/utils/models/rank_listwise_os_llm.py 
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
class FirstReranker:
    def __init__(self, agent: RankListwiseOSLLM):
        self.agent = agent

    def rerank(self, retrieved_result: Result, use_logits, use_alpha, rank_start, rank_end, window_size, step, logging=False, batched=False) -> Result:
        if batched:
            return self.agent.sliding_windows_batched(
                retrieved_results=[retrieved_result],
                use_logits=use_logits,
                use_alpha=use_alpha,
                rank_start=rank_start,
                rank_end=rank_end,
                window_size=window_size,
                step=step,
                logging=logging
            )[0]

        return self.agent.sliding_windows(
            retrieved_result=retrieved_result,
            use_logits=use_logits,
            use_alpha=use_alpha,
            rank_start=rank_start,
            rank_end=rank_end,
            window_size=window_size,
            step=step,
            logging=logging
        )

FirstModelReranker

Bases: BaseRanking

Implements FIRST: Faster Improved Listwise Reranking with Single Token Decoding.

FIRST is a listwise reranking model that employs a window-based decoding approach to efficiently rank passages using single-token predictions instead of full-text generation.

This scalable and efficient reranking method improves ranking speed and accuracy while maintaining high retrieval effectiveness.

Attributes:

Name Type Description
method str

The reranking method name.
model_name str

The name of the model used for reranking.
api_key str

API key for accessing remote models (if applicable).
context_size int

Maximum input length for the reranking model (default: 4096 tokens).
top_k int

Number of top-ranked passages retained after reranking (default: 20).
window_size int

Size of the sliding window used in listwise ranking (default: 9).
step_size int

Step size for moving the ranking window (default: 9).
use_logits bool

Whether to use logits-based scoring (default: False).
use_alpha bool

Whether to apply adaptive alpha scaling in ranking (default: False).
batched bool

Whether to use batched ranking for efficiency (default: False).
device str

Computing device ("cuda" if available, otherwise "cpu").
agent RankListwiseOSLLM

The ranking model instance used for passage ranking.
References
  • Gangi Reddy et al. FIRST: Faster Improved Listwise Reranking with Single Token Decoding
    Paper
See Also
  • Reranking: Main interface for reranking models, including FirstModelReranker.

Examples:

Basic usage with the Reranking class: 

from rankify.dataset.dataset import Document, Question, Answer, Context
from rankify.models.reranking import Reranking

question = Question("Who invented the first light bulb?")
answers = Answer(["Thomas Edison is credited with inventing the first practical light bulb."])
contexts = [
    Context(text="Nikola Tesla contributed to AC electricity.", id=0),
    Context(text="Thomas Edison patented the first practical light bulb.", id=1),
    Context(text="Light bulbs use tungsten filaments.", id=2),
    Context(text="The Wright brothers invented the airplane.", id=3),
]
document = Document(question=question, answers=answers, contexts=contexts)

# Initialize Reranking with FirstModelReranker
model = Reranking(method='first_ranker', model_name='base')
model.rank([document])

# Print reordered contexts
print("Reordered Contexts:")
for context in document.reorder_contexts:
    print(context.text)

Notes
  • FIRST uses a windowed approach to efficiently process large query-document pairs.
  • Integrated into the Reranking class, meaning users should use Reranking instead of FirstModelReranker directly.
  • Uses single-token decoding to achieve fast and effective ranking.
Source code in rankify/models/first_reranker.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
class FirstModelReranker(BaseRanking):
    """
    Implements **FIRST: Faster Improved Listwise Reranking with Single Token Decoding**.



    **FIRST** is a listwise reranking model that employs a **window-based decoding approach** 
    to efficiently rank passages using **single-token predictions** instead of full-text generation.

    This **scalable and efficient reranking method** improves ranking speed and accuracy while 
    maintaining high retrieval effectiveness.

    Attributes:
        method (str, optional): The reranking method name.
        model_name (str): The name of the model used for reranking.
        api_key (str, optional): API key for accessing remote models (if applicable).
        context_size (int): Maximum input length for the reranking model (default: `4096` tokens).
        top_k (int): Number of **top-ranked passages** retained after reranking (default: `20`).
        window_size (int): Size of the **sliding window** used in listwise ranking (default: `9`).
        step_size (int): Step size for moving the **ranking window** (default: `9`).
        use_logits (bool): Whether to use **logits-based scoring** (default: `False`).
        use_alpha (bool): Whether to apply **adaptive alpha scaling** in ranking (default: `False`).
        batched (bool): Whether to use **batched ranking** for efficiency (default: `False`).
        device (str): Computing device (`"cuda"` if available, otherwise `"cpu"`).
        agent (RankListwiseOSLLM): The ranking model instance used for passage ranking.

    References:
        - **Gangi Reddy et al.** *FIRST: Faster Improved Listwise Reranking with Single Token Decoding*  
          [Paper](https://arxiv.org/abs/2406.15657)

    See Also:
        - `Reranking`: Main interface for reranking models, including `FirstModelReranker`.

    Examples:
        **Basic usage with the `Reranking` class:**
        ```python
        from rankify.dataset.dataset import Document, Question, Answer, Context
        from rankify.models.reranking import Reranking

        question = Question("Who invented the first light bulb?")
        answers = Answer(["Thomas Edison is credited with inventing the first practical light bulb."])
        contexts = [
            Context(text="Nikola Tesla contributed to AC electricity.", id=0),
            Context(text="Thomas Edison patented the first practical light bulb.", id=1),
            Context(text="Light bulbs use tungsten filaments.", id=2),
            Context(text="The Wright brothers invented the airplane.", id=3),
        ]
        document = Document(question=question, answers=answers, contexts=contexts)

        # Initialize Reranking with FirstModelReranker
        model = Reranking(method='first_ranker', model_name='base')
        model.rank([document])

        # Print reordered contexts
        print("Reordered Contexts:")
        for context in document.reorder_contexts:
            print(context.text)
        ```

    Notes:
        - **FIRST** uses a **windowed approach** to efficiently process large query-document pairs.
        - Integrated into the `Reranking` class, meaning users should use `Reranking` instead of `FirstModelReranker` directly.
        - Uses **single-token decoding** to achieve **fast and effective ranking**.
    """

    def __init__(self, method: str = None, model_name: str = None, api_key: str = None, **kwargs):
        """
        Initializes the FIRST model reranker.

        Args:
            method (str, optional): The reranking method name.
            model_name (str): The name of the reranking model.
            api_key (str, optional): API key for remote access (if applicable).
            **kwargs: Additional parameters for model configuration.

        Example:
            ```python
            model = FirstModelReranker(method='first_ranker', model_name='base')
            ```
        """
        self.method = method
        self.model_name = model_name
        self.api_key = api_key
        self.context_size = kwargs.get("context_size", 4096)
        self.top_k = kwargs.get("top_k", 20)
        self.window_size = kwargs.get("window_size", 4)
        self.step_size = kwargs.get("step_size", 2)
        self.use_logits = kwargs.get("use_logits", False)
        self.use_alpha = kwargs.get("use_alpha", False)
        self.batched = kwargs.get("batched", False)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.agent = self._initialize_agent()

    def _initialize_agent(self):
        """
        Initializes the **RankListwiseOSLLM** agent for reranking with the specified model and parameters.

        Returns:
            RankListwiseOSLLM: A listwise reranking model instance.

        Example:
            ```python
            agent = model._initialize_agent()
            ```
        """
        return RankListwiseOSLLM(
            model=self.model_name,
            context_size=self.context_size,
            prompt_mode=PromptMode.RANK_GPT,
            num_few_shot_examples=0,
            device=self.device,
            num_gpus=1,
            variable_passages=True,
            window_size=self.window_size,
            system_message="You are RankLLM, an intelligent assistant that can rank passages based on their relevancy to the query",
            batched=self.batched,
            max_model_len=8192
        )

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks a list of documents using the **FIRST reranking model**.

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

        Returns:
            List[Document]: Documents with updated `reorder_contexts` after reranking.

        Raises:
            ValueError: If no contexts are provided for reranking.

        Example:
            ```python
            reranked_docs = model.rank(documents)
            ```
        """
        for document in tqdm(documents, desc="Reranking Documents"):
            self._rerank_document(document)
        return documents

    def _rerank_document(self, document: Document):
        """
        Applies the **FIRST reranking model** to reorder the document contexts.

        Args:
            document (Document): A `Document` instance to be reranked.

        Returns:
            None: Updates `document.reorder_contexts` in place.

        Example:
            ```python
            model._rerank_document(document)
            ```
        """
        result = Result(
            query=document.question.question,
            hits=[{"docid": ctx.id, "content": ctx.text , "rank":  0, 'score':0} for ctx in document.contexts]
        )
        # Perform reranking using `FirstReranker`
        reranker = FirstReranker(agent=self.agent)
        reranked_result = reranker.rerank(
            retrieved_result=result,
            use_logits=self.use_logits,
            use_alpha=self.use_alpha,
            rank_start=0,
            rank_end=self.top_k,
            window_size=self.window_size,
            step=self.step_size,
            logging=False,
            batched=self.batched
        )
        #print(reranked_result)
        # Update `Document.reorder_contexts` based on reranked `hits`
        context_dict = {ctx.id: ctx for ctx in document.contexts}

        document.reorder_contexts = [
            context_dict[hit["docid"]] for hit in reranked_result.hits if hit["docid"] in context_dict
        ]

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

Initializes the FIRST model reranker.

Parameters:

Name Type Description Default
method str

The reranking method name.

 None
model_name str

The name of the reranking model.

 None
api_key str

API key for remote access (if applicable).

 None
**kwargs

Additional parameters for model configuration.

 {}
Example 
model = FirstModelReranker(method='first_ranker', model_name='base')
Source code in rankify/models/first_reranker.py 
 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
def __init__(self, method: str = None, model_name: str = None, api_key: str = None, **kwargs):
    """
    Initializes the FIRST model reranker.

    Args:
        method (str, optional): The reranking method name.
        model_name (str): The name of the reranking model.
        api_key (str, optional): API key for remote access (if applicable).
        **kwargs: Additional parameters for model configuration.

    Example:
        ```python
        model = FirstModelReranker(method='first_ranker', model_name='base')
        ```
    """
    self.method = method
    self.model_name = model_name
    self.api_key = api_key
    self.context_size = kwargs.get("context_size", 4096)
    self.top_k = kwargs.get("top_k", 20)
    self.window_size = kwargs.get("window_size", 4)
    self.step_size = kwargs.get("step_size", 2)
    self.use_logits = kwargs.get("use_logits", False)
    self.use_alpha = kwargs.get("use_alpha", False)
    self.batched = kwargs.get("batched", False)
    self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.agent = self._initialize_agent()

rank(documents)

Reranks a list of documents using the FIRST reranking model.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document instances to rerank.

 required

Returns:

Type Description
List[Document]

List[Document]: Documents with updated reorder_contexts after reranking.

Raises:

Type Description
ValueError

If no contexts are provided for reranking.
Example 
reranked_docs = model.rank(documents)
Source code in rankify/models/first_reranker.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks a list of documents using the **FIRST reranking model**.

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

    Returns:
        List[Document]: Documents with updated `reorder_contexts` after reranking.

    Raises:
        ValueError: If no contexts are provided for reranking.

    Example:
        ```python
        reranked_docs = model.rank(documents)
        ```
    """
    for document in tqdm(documents, desc="Reranking Documents"):
        self._rerank_document(document)
    return documents