RankFiD

rankify.models.rank_fid

Request dataclass

Source code in rankify/utils/models/rank_llm/data.py

@dataclass
class Request:
    query: Query
    candidates: List[Candidate] = field(default_factory=list)

Result dataclass

Source code in rankify/utils/models/rank_llm/data.py

@dataclass
class Result:
    query: Query
    candidates: list[Candidate] = field(default_factory=list)
    ranking_exec_summary: list[RankingExecInfo] = (field(default_factory=list),)

ListwiseRankLLM

Bases: RankLLM, ABC

All children of ListwiseRankLLM must implement these functions

• rerank_batched
• run_llm_batched
• run_llm
• create_prompt_batched
• create_prompt
• get_num_tokens
• cost_per_1k_token
• num_output_tokens

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

class ListwiseRankLLM(RankLLM, ABC):
    """
    All children of ListwiseRankLLM must implement these functions:
        - rerank_batched
        - run_llm_batched
        - run_llm
        - create_prompt_batched
        - create_prompt
        - get_num_tokens
        - cost_per_1k_token
        - num_output_tokens
    """

    def __init__(
        self,
        model: str,
        context_size: int,
        prompt_mode: PromptMode,
        num_few_shot_examples: int,
        window_size: int,
    ) -> None:
        super().__init__(model, context_size, prompt_mode)
        self._num_few_shot_examples = num_few_shot_examples
        self._window_size = window_size

    def get_output_filename(
        self,
        top_k_candidates: int,
        dataset_name: str,
        shuffle_candidates: bool,
        **kwargs: Any,
    ) -> str:
        _modelname = self._model.split("/")[-1]
        if _modelname.startswith("checkpoint"):
            _modelname = self._model.split("/")[-2] + "_" + _modelname
        name = (
            f"{_modelname}_{self._context_size}_{top_k_candidates}_{self._prompt_mode}"
        )
        if dataset_name:
            name = f"{name}_{dataset_name}"
        if self._num_few_shot_examples > 0:
            name += f"_{self._num_few_shot_examples}_shot"
        return (
            f"{name}_shuffled_{datetime.isoformat(datetime.now())}"
            if shuffle_candidates
            else f"{name}_{datetime.isoformat(datetime.now())}"
        )

    def max_tokens(self) -> int:
        """
        Returns the maximum number of tokens for a given model

        Returns:
            int: The maximum token count.
        """
        return self._context_size

    def permutation_pipeline_batched(
        self,
        results: List[Result],
        rank_start: int,
        rank_end: int,
        logging: bool = False,
        populate_exec_summary: bool = False,
        use_logits: bool = False,
        use_alpha: bool = False
    ) -> List[Result]:
        """
        Runs the permutation pipeline on a batch of result objects within the passed in rank range.

        Args:
            results (List[Result]): The list of result objects to process.
            rank_start (int): The start index for ranking.
            rank_end (int): The end index for ranking.
            logging (bool, optional): Flag to enable logging of operations. Defaults to False.

        Returns:
            List[Result]: The list of processed result objects after applying permutation.
        """
        prompts = []
        logger.info("Loading prompts.")
        prompts = self.create_prompt_batched(
            results, rank_start, rank_end, batch_size=32, use_alpha=use_alpha
        )
        if logging:
            for prompt in prompts:
                logger.debug(f"Prompt: {prompt[0]}\n")
        logger.info("Prompts loaded.")
        batched_results = self.run_llm_batched(
            [prompt for prompt, _ in prompts], current_window_size=rank_end - rank_start, use_logits=use_logits, use_alpha=use_alpha
        )

        for index, (result, (prompt, in_token_count)) in enumerate(
            zip(results, prompts)
        ):
            permutation, out_token_count = batched_results[index]
            if logging:
                logger.debug(f"output: {permutation}")
            if populate_exec_summary:
                if result.ranking_exec_summary is None:
                    result.ranking_exec_summary = []
                ranking_exec_info = RankingExecInfo(
                    prompt, permutation, in_token_count, out_token_count
                )
                result.ranking_exec_summary.append(ranking_exec_info)
            result = self.receive_permutation(result, permutation, rank_start, rank_end)

        return results

    def permutation_pipeline(
        self,
        result: Result,
        rank_start: int,
        rank_end: int,
        logging: bool = False,
        populate_exec_summary: bool = True,
        use_logits: bool = False,
        use_alpha: bool = False
    ) -> Result:
        """
        Runs the permutation pipeline on the passed in result set within the passed in rank range.

        Args:
            result (Result): The result object to process.
            rank_start (int): The start index for ranking.
            rank_end (int): The end index for ranking.
            logging (bool, optional): Flag to enable logging of operations. Defaults to False.

        Returns:
            Result: The processed result object after applying permutation.
        """
        prompt, in_token_count = self.create_prompt(result, rank_start, rank_end, use_alpha)
        if logging:
            logger.info(f"Prompt: {prompt}\n")
        permutation, out_token_count = self.run_llm(
            prompt, current_window_size=rank_end - rank_start, use_logits=use_logits, use_alpha=use_alpha
        )
        if logging:
            print(f"Output: {permutation}")
        if populate_exec_summary:
            ranking_exec_info = RankingExecInfo(
                prompt, permutation, in_token_count, out_token_count
            )
            result.ranking_exec_summary.append(ranking_exec_info)
        result = self.receive_permutation(result, permutation, rank_start, rank_end)
        return result

    def shuffle_and_rescore(
        rerank_results: List[Result], rank_start: int, rank_end: int
    ):
        """
        Shuffles candidates between rank_start and rank_end, and rescales scores based on new rank.

        Args:
            rerank_results (List[Result]): List of Result objects to process.
            rank_start (int): Start index for ranking.
            rank_end (int): End index for ranking.
        """
        for rerank_result in rerank_results:
            # Shuffle rerank_result hits between rank_start and rank_end
            rerank_result.candidates[rank_start:rank_end] = random.sample(
                rerank_result.candidates[rank_start:rank_end],
                len(rerank_result.candidates[rank_start:rank_end]),
            )
            # Rescore all candidates with 1/rank
            for i, cand in enumerate(rerank_result.candidates):
                cand["score"] = 1.0 / (i + 1)
                cand["rank"] = i + 1

    def sliding_windows_batched(
        self,
        requests: List[Request],
        rank_start: int,
        rank_end: int,
        window_size: int,
        step: int,
        shuffle_candidates: bool = False,
        logging: bool = False,
        populate_exec_summary: bool = False,
        use_logits: bool = False,
        use_alpha: bool = False
    ) -> List[Result]:
        """
        Applies the sliding window algorithm to the reranking process for a batch of result objects.
        Args:
            requests (List[Request]): The list of request objects to process.
            rank_start (int): The start index for ranking.
            rank_end (int): The end index for ranking.
            window_size (int): The size of each sliding window.
            step (int): The step size for moving the window.
            shuffle_candidates (bool, optional): Flag to shuffle candidates before processing. Defaults to False.
            logging (bool, optional): Flag to enable logging of operations. Defaults to False.
        Returns:
            List[Result]: The list of result objects after applying the sliding window technique.
        """
        rerank_results = [
            Result(
                query=copy.deepcopy(request.query),
                candidates=copy.deepcopy(request.candidates),
                ranking_exec_summary=[],
            )
            for request in requests
        ]
        if shuffle_candidates:
            self.shuffle_and_rescore(rerank_results, rank_start, rank_end)
        end_pos = rank_end
        start_pos = rank_end - window_size

        # end_pos > rank_start ensures that the list is non-empty while allowing last window to be smaller than window_size
        # start_pos + step != rank_start prevents processing of redundant windows (e.g. 0-20, followed by 0-10)
        while end_pos > rank_start and start_pos + step != rank_start:
            if logging:
                logger.info(f"start_pos: {start_pos}, end_pos: {end_pos}")
            start_pos = max(start_pos, rank_start)
            rerank_results = self.permutation_pipeline_batched(
                rerank_results, start_pos, end_pos, logging, populate_exec_summary, use_logits, use_alpha
            )
            end_pos = end_pos - step
            start_pos = start_pos - step
        return rerank_results

    def sliding_windows(
        self,
        request: Request,
        rank_start: int,
        rank_end: int,
        window_size: int,
        step: int,
        shuffle_candidates: bool = False,
        logging: bool = False,
        populate_exec_summary: bool = True,
        use_logits: bool = False,
        use_alpha: bool = False
    ) -> Result:
        """
        Applies the sliding window algorithm to the reranking process.

        Args:
            request (Request): The request object to process.
            rank_start (int): The start index for ranking.
            rank_end (int): The end index for ranking.
            window_size (int): The size of each sliding window.
            step (int): The step size for moving the window.
            shuffle_candidates (bool, optional): Flag to shuffle candidates before processing. Defaults to False.
            logging (bool, optional): Flag to enable logging of operations. Defaults to False.

        Returns:
            Result: The result object after applying the sliding window technique.
        """
        rerank_result = Result(
            query=copy.deepcopy(request.query),
            candidates=copy.deepcopy(request.candidates),
            ranking_exec_summary=[],
        )
        if shuffle_candidates:
            self.shuffle_and_rescore([rerank_result], rank_start, rank_end)
        end_pos = rank_end
        start_pos = rank_end - window_size
        # end_pos > rank_start ensures that the list is non-empty while allowing last window to be smaller than window_size
        # start_pos + step != rank_start prevents processing of redundant windows (e.g. 0-20, followed by 0-10)
        while end_pos > rank_start and start_pos + step != rank_start:
            start_pos = max(start_pos, rank_start)
            rerank_result = self.permutation_pipeline(
                rerank_result,
                start_pos,
                end_pos,
                logging,
                populate_exec_summary=populate_exec_summary,
                use_logits=use_logits,
                use_alpha=use_alpha
            )
            end_pos = end_pos - step
            start_pos = start_pos - step
        return rerank_result

    def get_ranking_cost_upperbound(
        self, num_q: int, rank_start: int, rank_end: int, window_size: int, step: int
    ) -> Tuple[float, int]:
        """
        Calculates the upper bound of the ranking cost for a given set of parameters.

        Args:
            num_q (int): The number of queries.
            rank_start (int): The start index for ranking.
            rank_end (int): The end index for ranking.
            window_size (int): The size of each sliding window.
            step (int): The step size for moving the window.

        Returns:
            Tuple[float, int]: A tuple object containing the cost and the total number of tokens used (input tokens + output tokens).
        """
        # For every prompt generated for every query assume the max context size is used.
        num_promt = (rank_end - rank_start - window_size) / step + 1
        input_token_count = (
            num_q * num_promt * (self._context_size - self.num_output_tokens())
        )
        output_token_count = num_q * num_promt * self.num_output_tokens()
        cost = (
            input_token_count * self.cost_per_1k_token(input_token=True)
            + output_token_count * self.cost_per_1k_token(input_token=False)
        ) / 1000.0
        return (cost, input_token_count + output_token_count)

    def get_ranking_cost(
        self,
        retrieved_results: List[Request],
        rank_start: int,
        rank_end: int,
        window_size: int,
        step: int,
    ) -> Tuple[float, int]:
        """
        Calculates the ranking cost based on actual token counts from generated prompts.

        Args:
            retrieved_results (List[Request]): A list of retrieved results for processing.
            rank_start (int): The start index for ranking.
            rank_end (int): The end index for ranking.
            window_size (int): The size of each sliding window.
            step (int): The step size for moving the window.

        Returns:
            Tuple[float, int]: A tuple object containing the calculated cost and the total number of tokens used (input tokens + output tokens).
        """
        input_token_count = 0
        output_token_count = 0
        # Go through the retrieval result using the sliding window and count the number of tokens for generated prompts.
        # This is an estimated cost analysis since the actual prompts' length will depend on the ranking.
        for result in retrieved_results:
            end_pos = rank_end
            start_pos = rank_end - window_size
            while start_pos >= rank_start:
                start_pos = max(start_pos, rank_start)
                prompt, _ = self.create_prompt(result, start_pos, end_pos)
                input_token_count += self.get_num_tokens(prompt)
                end_pos = end_pos - step
                start_pos = start_pos - step
                output_token_count += self.num_output_tokens()
        cost = (
            input_token_count * self.cost_per_1k_token(input_token=True)
            + output_token_count * self.cost_per_1k_token(input_token=False)
        ) / 1000.0
        return (cost, input_token_count + output_token_count)

    def _clean_response(self, response: str, use_alpha: bool = False) -> str:
        new_response = ""
        if use_alpha:
            for c in response:
                if not c.isalpha():
                    new_response += " "
                else:
                    new_response += str(ord(c) - ALPH_START_IDX)
            new_response = new_response.strip()
        else:
            for c in response:
                if not c.isdigit():
                    new_response += " "
                else:
                    new_response += c
            new_response = new_response.strip()

        return new_response

    def _remove_duplicate(self, response: List[int]) -> List[int]:
        new_response = []
        for c in response:
            if c not in new_response:
                new_response.append(c)
        return new_response

    def receive_permutation(
        self, result: Result, permutation: str, rank_start: int, rank_end: int
    ) -> Result:
        """
        Processes and applies a permutation to the ranking results.

        This function takes a permutation string, representing the new order of items,
        and applies it to a subset of the ranking results. It adjusts the ranks and scores in the
        'result' object based on this permutation.

        Args:
            result (Result): The result object containing the initial ranking results.
            permutation (str): A string representing the new order of items.
                            Each item in the string should correspond to a rank in the results.
            rank_start (int): The starting index of the range in the results to which the permutation is applied.
            rank_end (int): The ending index of the range in the results to which the permutation is applied.

        Returns:
            Result: The updated result object with the new ranking order applied.

        Note:
            This function assumes that the permutation string is a sequence of integers separated by spaces.
            Each integer in the permutation string corresponds to a 1-based index in the ranking results.
            The function first normalizes these to 0-based indices, removes duplicates, and then reorders
            the items in the specified range of the 'result.candidates' list according to the permutation.
            Items not mentioned in the permutation string remain in their original sequence but are moved after
            the permuted items.
        """

        # Parse and normalize the permutation indices
        response = self._clean_response(permutation)
        response = [int(x) - 1 for x in response.split()]
        response = self._remove_duplicate(response)

        # Extract the relevant candidates and create a mapping for new order
        cut_range = copy.deepcopy(result.candidates[rank_start:rank_end])
        original_rank = [tt for tt in range(len(cut_range))]
        response = [ss for ss in response if ss in original_rank]
        response = response + [tt for tt in original_rank if tt not in response]

        # Update candidates in the new order
        for j, x in enumerate(response):
            #print(cut_range[x])
            result.candidates[j + rank_start] = copy.deepcopy(cut_range[x])
            if result.candidates[j + rank_start]["score"] is not None:
                #print(cut_range[j])
                result.candidates[j + rank_start]["score"] = cut_range[j]["score"]

        return result

    def _replace_number(self, s: str) -> str:
        return re.sub(r"\[(\d+)\]", r"(\1)", s)

    def convert_doc_to_prompt_content(
        self, doc: Dict[str, Any], max_length: int
    ) -> str:
        if "text" in doc:
            content = doc["text"]
        elif "segment" in doc:
            content = doc["segment"]
        elif "contents" in doc:
            content = doc["contents"]
        elif "content" in doc:
            content = doc["content"]
        elif "body" in doc:
            content = doc["body"]
        else:
            content = doc["passage"]
        if "title" in doc and doc["title"]:
            content = "Title: " + doc["title"] + " " + "Content: " + content
        content = content.strip()
        content = fix_text(content)
        # For Japanese should cut by character: content = content[:int(max_length)]
        content = " ".join(content.split()[: int(max_length)])
        return self._replace_number(content)

max_tokens()

Returns the maximum number of tokens for a given model

Returns:

Name	Type	Description
int	int	The maximum token count.

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def max_tokens(self) -> int:
    """
    Returns the maximum number of tokens for a given model

    Returns:
        int: The maximum token count.
    """
    return self._context_size

permutation_pipeline_batched(results, rank_start, rank_end, logging=False, populate_exec_summary=False, use_logits=False, use_alpha=False)

Runs the permutation pipeline on a batch of result objects within the passed in rank range.

Parameters:

Name	Type	Description	Default
results	List[Result]	The list of result objects to process.	required
rank_start	int	The start index for ranking.	required
rank_end	int	The end index for ranking.	required
logging	bool	Flag to enable logging of operations. Defaults to False.	False

Returns:

Type	Description
List[Result]	List[Result]: The list of processed result objects after applying permutation.

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def permutation_pipeline_batched(
    self,
    results: List[Result],
    rank_start: int,
    rank_end: int,
    logging: bool = False,
    populate_exec_summary: bool = False,
    use_logits: bool = False,
    use_alpha: bool = False
) -> List[Result]:
    """
    Runs the permutation pipeline on a batch of result objects within the passed in rank range.

    Args:
        results (List[Result]): The list of result objects to process.
        rank_start (int): The start index for ranking.
        rank_end (int): The end index for ranking.
        logging (bool, optional): Flag to enable logging of operations. Defaults to False.

    Returns:
        List[Result]: The list of processed result objects after applying permutation.
    """
    prompts = []
    logger.info("Loading prompts.")
    prompts = self.create_prompt_batched(
        results, rank_start, rank_end, batch_size=32, use_alpha=use_alpha
    )
    if logging:
        for prompt in prompts:
            logger.debug(f"Prompt: {prompt[0]}\n")
    logger.info("Prompts loaded.")
    batched_results = self.run_llm_batched(
        [prompt for prompt, _ in prompts], current_window_size=rank_end - rank_start, use_logits=use_logits, use_alpha=use_alpha
    )

    for index, (result, (prompt, in_token_count)) in enumerate(
        zip(results, prompts)
    ):
        permutation, out_token_count = batched_results[index]
        if logging:
            logger.debug(f"output: {permutation}")
        if populate_exec_summary:
            if result.ranking_exec_summary is None:
                result.ranking_exec_summary = []
            ranking_exec_info = RankingExecInfo(
                prompt, permutation, in_token_count, out_token_count
            )
            result.ranking_exec_summary.append(ranking_exec_info)
        result = self.receive_permutation(result, permutation, rank_start, rank_end)

    return results

permutation_pipeline(result, rank_start, rank_end, logging=False, populate_exec_summary=True, use_logits=False, use_alpha=False)

Runs the permutation pipeline on the passed in result set within the passed in rank range.

Parameters:

Name	Type	Description	Default
result	Result	The result object to process.	required
rank_start	int	The start index for ranking.	required
rank_end	int	The end index for ranking.	required
logging	bool	Flag to enable logging of operations. Defaults to False.	False

Returns:

Name	Type	Description
Result	Result	The processed result object after applying permutation.

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def permutation_pipeline(
    self,
    result: Result,
    rank_start: int,
    rank_end: int,
    logging: bool = False,
    populate_exec_summary: bool = True,
    use_logits: bool = False,
    use_alpha: bool = False
) -> Result:
    """
    Runs the permutation pipeline on the passed in result set within the passed in rank range.

    Args:
        result (Result): The result object to process.
        rank_start (int): The start index for ranking.
        rank_end (int): The end index for ranking.
        logging (bool, optional): Flag to enable logging of operations. Defaults to False.

    Returns:
        Result: The processed result object after applying permutation.
    """
    prompt, in_token_count = self.create_prompt(result, rank_start, rank_end, use_alpha)
    if logging:
        logger.info(f"Prompt: {prompt}\n")
    permutation, out_token_count = self.run_llm(
        prompt, current_window_size=rank_end - rank_start, use_logits=use_logits, use_alpha=use_alpha
    )
    if logging:
        print(f"Output: {permutation}")
    if populate_exec_summary:
        ranking_exec_info = RankingExecInfo(
            prompt, permutation, in_token_count, out_token_count
        )
        result.ranking_exec_summary.append(ranking_exec_info)
    result = self.receive_permutation(result, permutation, rank_start, rank_end)
    return result

shuffle_and_rescore(rerank_results, rank_start, rank_end)

Shuffles candidates between rank_start and rank_end, and rescales scores based on new rank.

Parameters:

Name	Type	Description	Default
rerank_results	List[Result]	List of Result objects to process.	required
rank_start	int	Start index for ranking.	required
rank_end	int	End index for ranking.	required

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def shuffle_and_rescore(
    rerank_results: List[Result], rank_start: int, rank_end: int
):
    """
    Shuffles candidates between rank_start and rank_end, and rescales scores based on new rank.

    Args:
        rerank_results (List[Result]): List of Result objects to process.
        rank_start (int): Start index for ranking.
        rank_end (int): End index for ranking.
    """
    for rerank_result in rerank_results:
        # Shuffle rerank_result hits between rank_start and rank_end
        rerank_result.candidates[rank_start:rank_end] = random.sample(
            rerank_result.candidates[rank_start:rank_end],
            len(rerank_result.candidates[rank_start:rank_end]),
        )
        # Rescore all candidates with 1/rank
        for i, cand in enumerate(rerank_result.candidates):
            cand["score"] = 1.0 / (i + 1)
            cand["rank"] = i + 1

sliding_windows_batched(requests, rank_start, rank_end, window_size, step, shuffle_candidates=False, logging=False, populate_exec_summary=False, use_logits=False, use_alpha=False)

Applies the sliding window algorithm to the reranking process for a batch of result objects. Args: requests (List[Request]): The list of request objects to process. rank_start (int): The start index for ranking. rank_end (int): The end index for ranking. window_size (int): The size of each sliding window. step (int): The step size for moving the window. shuffle_candidates (bool, optional): Flag to shuffle candidates before processing. Defaults to False. logging (bool, optional): Flag to enable logging of operations. Defaults to False. Returns: List[Result]: The list of result objects after applying the sliding window technique.

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def sliding_windows_batched(
    self,
    requests: List[Request],
    rank_start: int,
    rank_end: int,
    window_size: int,
    step: int,
    shuffle_candidates: bool = False,
    logging: bool = False,
    populate_exec_summary: bool = False,
    use_logits: bool = False,
    use_alpha: bool = False
) -> List[Result]:
    """
    Applies the sliding window algorithm to the reranking process for a batch of result objects.
    Args:
        requests (List[Request]): The list of request objects to process.
        rank_start (int): The start index for ranking.
        rank_end (int): The end index for ranking.
        window_size (int): The size of each sliding window.
        step (int): The step size for moving the window.
        shuffle_candidates (bool, optional): Flag to shuffle candidates before processing. Defaults to False.
        logging (bool, optional): Flag to enable logging of operations. Defaults to False.
    Returns:
        List[Result]: The list of result objects after applying the sliding window technique.
    """
    rerank_results = [
        Result(
            query=copy.deepcopy(request.query),
            candidates=copy.deepcopy(request.candidates),
            ranking_exec_summary=[],
        )
        for request in requests
    ]
    if shuffle_candidates:
        self.shuffle_and_rescore(rerank_results, rank_start, rank_end)
    end_pos = rank_end
    start_pos = rank_end - window_size

    # end_pos > rank_start ensures that the list is non-empty while allowing last window to be smaller than window_size
    # start_pos + step != rank_start prevents processing of redundant windows (e.g. 0-20, followed by 0-10)
    while end_pos > rank_start and start_pos + step != rank_start:
        if logging:
            logger.info(f"start_pos: {start_pos}, end_pos: {end_pos}")
        start_pos = max(start_pos, rank_start)
        rerank_results = self.permutation_pipeline_batched(
            rerank_results, start_pos, end_pos, logging, populate_exec_summary, use_logits, use_alpha
        )
        end_pos = end_pos - step
        start_pos = start_pos - step
    return rerank_results

sliding_windows(request, rank_start, rank_end, window_size, step, shuffle_candidates=False, logging=False, populate_exec_summary=True, use_logits=False, use_alpha=False)

Applies the sliding window algorithm to the reranking process.

Parameters:

Name	Type	Description	Default
request	Request	The request object to process.	required
rank_start	int	The start index for ranking.	required
rank_end	int	The end index for ranking.	required
window_size	int	The size of each sliding window.	required
step	int	The step size for moving the window.	required
shuffle_candidates	bool	Flag to shuffle candidates before processing. Defaults to False.	False
logging	bool	Flag to enable logging of operations. Defaults to False.	False

Returns:

Name	Type	Description
Result	Result	The result object after applying the sliding window technique.

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def sliding_windows(
    self,
    request: Request,
    rank_start: int,
    rank_end: int,
    window_size: int,
    step: int,
    shuffle_candidates: bool = False,
    logging: bool = False,
    populate_exec_summary: bool = True,
    use_logits: bool = False,
    use_alpha: bool = False
) -> Result:
    """
    Applies the sliding window algorithm to the reranking process.

    Args:
        request (Request): The request object to process.
        rank_start (int): The start index for ranking.
        rank_end (int): The end index for ranking.
        window_size (int): The size of each sliding window.
        step (int): The step size for moving the window.
        shuffle_candidates (bool, optional): Flag to shuffle candidates before processing. Defaults to False.
        logging (bool, optional): Flag to enable logging of operations. Defaults to False.

    Returns:
        Result: The result object after applying the sliding window technique.
    """
    rerank_result = Result(
        query=copy.deepcopy(request.query),
        candidates=copy.deepcopy(request.candidates),
        ranking_exec_summary=[],
    )
    if shuffle_candidates:
        self.shuffle_and_rescore([rerank_result], rank_start, rank_end)
    end_pos = rank_end
    start_pos = rank_end - window_size
    # end_pos > rank_start ensures that the list is non-empty while allowing last window to be smaller than window_size
    # start_pos + step != rank_start prevents processing of redundant windows (e.g. 0-20, followed by 0-10)
    while end_pos > rank_start and start_pos + step != rank_start:
        start_pos = max(start_pos, rank_start)
        rerank_result = self.permutation_pipeline(
            rerank_result,
            start_pos,
            end_pos,
            logging,
            populate_exec_summary=populate_exec_summary,
            use_logits=use_logits,
            use_alpha=use_alpha
        )
        end_pos = end_pos - step
        start_pos = start_pos - step
    return rerank_result

get_ranking_cost_upperbound(num_q, rank_start, rank_end, window_size, step)

Calculates the upper bound of the ranking cost for a given set of parameters.

Parameters:

Name	Type	Description	Default
num_q	int	The number of queries.	required
rank_start	int	The start index for ranking.	required
rank_end	int	The end index for ranking.	required
window_size	int	The size of each sliding window.	required
step	int	The step size for moving the window.	required

Returns:

Type	Description
Tuple[float, int]	Tuple[float, int]: A tuple object containing the cost and the total number of tokens used (input tokens + output tokens).

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def get_ranking_cost_upperbound(
    self, num_q: int, rank_start: int, rank_end: int, window_size: int, step: int
) -> Tuple[float, int]:
    """
    Calculates the upper bound of the ranking cost for a given set of parameters.

    Args:
        num_q (int): The number of queries.
        rank_start (int): The start index for ranking.
        rank_end (int): The end index for ranking.
        window_size (int): The size of each sliding window.
        step (int): The step size for moving the window.

    Returns:
        Tuple[float, int]: A tuple object containing the cost and the total number of tokens used (input tokens + output tokens).
    """
    # For every prompt generated for every query assume the max context size is used.
    num_promt = (rank_end - rank_start - window_size) / step + 1
    input_token_count = (
        num_q * num_promt * (self._context_size - self.num_output_tokens())
    )
    output_token_count = num_q * num_promt * self.num_output_tokens()
    cost = (
        input_token_count * self.cost_per_1k_token(input_token=True)
        + output_token_count * self.cost_per_1k_token(input_token=False)
    ) / 1000.0
    return (cost, input_token_count + output_token_count)

get_ranking_cost(retrieved_results, rank_start, rank_end, window_size, step)

Calculates the ranking cost based on actual token counts from generated prompts.

Parameters:

Name	Type	Description	Default
retrieved_results	List[Request]	A list of retrieved results for processing.	required
rank_start	int	The start index for ranking.	required
rank_end	int	The end index for ranking.	required
window_size	int	The size of each sliding window.	required
step	int	The step size for moving the window.	required

Returns:

Type	Description
Tuple[float, int]	Tuple[float, int]: A tuple object containing the calculated cost and the total number of tokens used (input tokens + output tokens).

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def get_ranking_cost(
    self,
    retrieved_results: List[Request],
    rank_start: int,
    rank_end: int,
    window_size: int,
    step: int,
) -> Tuple[float, int]:
    """
    Calculates the ranking cost based on actual token counts from generated prompts.

    Args:
        retrieved_results (List[Request]): A list of retrieved results for processing.
        rank_start (int): The start index for ranking.
        rank_end (int): The end index for ranking.
        window_size (int): The size of each sliding window.
        step (int): The step size for moving the window.

    Returns:
        Tuple[float, int]: A tuple object containing the calculated cost and the total number of tokens used (input tokens + output tokens).
    """
    input_token_count = 0
    output_token_count = 0
    # Go through the retrieval result using the sliding window and count the number of tokens for generated prompts.
    # This is an estimated cost analysis since the actual prompts' length will depend on the ranking.
    for result in retrieved_results:
        end_pos = rank_end
        start_pos = rank_end - window_size
        while start_pos >= rank_start:
            start_pos = max(start_pos, rank_start)
            prompt, _ = self.create_prompt(result, start_pos, end_pos)
            input_token_count += self.get_num_tokens(prompt)
            end_pos = end_pos - step
            start_pos = start_pos - step
            output_token_count += self.num_output_tokens()
    cost = (
        input_token_count * self.cost_per_1k_token(input_token=True)
        + output_token_count * self.cost_per_1k_token(input_token=False)
    ) / 1000.0
    return (cost, input_token_count + output_token_count)

receive_permutation(result, permutation, rank_start, rank_end)

Processes and applies a permutation to the ranking results.

This function takes a permutation string, representing the new order of items, and applies it to a subset of the ranking results. It adjusts the ranks and scores in the 'result' object based on this permutation.

Parameters:

Name	Type	Description	Default
result	Result	The result object containing the initial ranking results.	required
permutation	str	A string representing the new order of items. Each item in the string should correspond to a rank in the results.	required
rank_start	int	The starting index of the range in the results to which the permutation is applied.	required
rank_end	int	The ending index of the range in the results to which the permutation is applied.	required

Returns:

Name	Type	Description
Result	Result	The updated result object with the new ranking order applied.

Note

This function assumes that the permutation string is a sequence of integers separated by spaces. Each integer in the permutation string corresponds to a 1-based index in the ranking results. The function first normalizes these to 0-based indices, removes duplicates, and then reorders the items in the specified range of the 'result.candidates' list according to the permutation. Items not mentioned in the permutation string remain in their original sequence but are moved after the permuted items.

Source code in rankify/utils/models/rank_llm/rerank/listwise/listwise_rankllm.py

def receive_permutation(
    self, result: Result, permutation: str, rank_start: int, rank_end: int
) -> Result:
    """
    Processes and applies a permutation to the ranking results.

    This function takes a permutation string, representing the new order of items,
    and applies it to a subset of the ranking results. It adjusts the ranks and scores in the
    'result' object based on this permutation.

    Args:
        result (Result): The result object containing the initial ranking results.
        permutation (str): A string representing the new order of items.
                        Each item in the string should correspond to a rank in the results.
        rank_start (int): The starting index of the range in the results to which the permutation is applied.
        rank_end (int): The ending index of the range in the results to which the permutation is applied.

    Returns:
        Result: The updated result object with the new ranking order applied.

    Note:
        This function assumes that the permutation string is a sequence of integers separated by spaces.
        Each integer in the permutation string corresponds to a 1-based index in the ranking results.
        The function first normalizes these to 0-based indices, removes duplicates, and then reorders
        the items in the specified range of the 'result.candidates' list according to the permutation.
        Items not mentioned in the permutation string remain in their original sequence but are moved after
        the permuted items.
    """

    # Parse and normalize the permutation indices
    response = self._clean_response(permutation)
    response = [int(x) - 1 for x in response.split()]
    response = self._remove_duplicate(response)

    # Extract the relevant candidates and create a mapping for new order
    cut_range = copy.deepcopy(result.candidates[rank_start:rank_end])
    original_rank = [tt for tt in range(len(cut_range))]
    response = [ss for ss in response if ss in original_rank]
    response = response + [tt for tt in original_rank if tt not in response]

    # Update candidates in the new order
    for j, x in enumerate(response):
        #print(cut_range[x])
        result.candidates[j + rank_start] = copy.deepcopy(cut_range[x])
        if result.candidates[j + rank_start]["score"] is not None:
            #print(cut_range[j])
            result.candidates[j + rank_start]["score"] = cut_range[j]["score"]

    return result

FiD

Bases: T5ForConditionalGeneration

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

class FiD(T5ForConditionalGeneration):
    _keys_to_ignore_on_load_missing = [
        r"encoder\.embed_tokens\.weight",
        r"decoder\.embed_tokens\.weight",
        r"lm_head\.weight",
    ]
    _keys_to_ignore_on_load_unexpected = [
        r"decoder\.block\.0\.layer\.1\.EncDecAttention\.relative_attention_bias\.weight",
    ]

    ANSWER_EOS_TOKEN = 1

    def __init__(self, config):
        super().__init__(config)
        self.model_dim = config.d_model

        self.shared = nn.Embedding(config.vocab_size, config.d_model)

        encoder_config = copy.deepcopy(config)
        encoder_config.is_decoder = False
        encoder_config.use_cache = False
        encoder_config.is_encoder_decoder = False

        self.encoder = FiDStack(encoder_config, self.shared)

        decoder_config = copy.deepcopy(config)
        decoder_config.is_decoder = True
        decoder_config.is_encoder_decoder = False
        decoder_config.use_cache = True
        decoder_config.num_layers = config.num_decoder_layers
        self.decoder = FiDStack(decoder_config, self.shared)

        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

        # Model parallel
        self.model_parallel = False
        self.device_map = None

    def reset_n_passages(self, n_passages: int):
        self.encoder.reset_n_passages(n_passages)
        self.decoder.reset_n_passages(n_passages)

    def set_checkpoint(self, use_checkpoint):
        """
        Enable or disable checkpointing in the encoder.
        See https://pytorch.org/docs/stable/checkpoint.html
        """
        for mod in self.encoder.encoder.block:
            mod.use_checkpoint = use_checkpoint

    def reset_score_storage(self):
        """
        Reset score storage, only used when cross-attention scores are saved
        to train a retriever.
        """
        for mod in self.decoder.block:
            mod.layer[1].EncDecAttention.normalized_score_storage = None

    @torch.no_grad()
    def get_crossattention_scores(
        self, n_passages, mask, ids, mask_query=None, output_sequence_lengths=[]
    ):
        """
        Cross-attention scores are aggregated to obtain a single scalar per
        passage. This scalar can be seen as a similarity score between the
        question and the input passage. It is obtained by averaging the
        cross-attention scores obtained on the first decoded token over heads,
        layers, and tokens of the input passage.

        More details in Distilling Knowledge from Reader to Retriever:
        https://arxiv.org/abs/2012.04584.
        """
        norms = []
        for mod in self.decoder.block:
            norms.append(mod.layer[1].EncDecAttention.normalized_score_storage)
        norms = torch.stack(norms)

        output = {}
        self.aggregate_value(
            norms,
            mask,
            n_passages,
            ids,
            mask_query,
            output,
            prefix="norms",
            output_sequence_lengths=output_sequence_lengths,
        )
        return output

    def aggregate_value(
        self,
        scores,
        mask,
        n_passages,
        ids,
        mask_query=None,
        output={},
        prefix="",
        output_sequence_lengths=[],
    ):
        n_layers, bsz, n_tokens, total_tokens = scores.size()

        ids = ids.view(bsz, n_passages, -1)
        scores = scores.view(n_layers, bsz, n_tokens, n_passages, -1)
        mask = mask.view(bsz, n_passages, -1)
        scores = scores.masked_fill(~mask[None, :, None], 0.0)

        scores = scores.sum(dim=[0])

        scores_woquery = None
        # Compute scores based on scores without query
        if not mask_query is None:
            output[f"{prefix}woquery"] = self.get_woquery_score(
                scores,
                mask_query,
                mask,
                n_layers,
                output_sequence_lengths=output_sequence_lengths,
            )

        return output

    def get_woquery_score(
        self, scores, mask_query, mask, n_layers, output_sequence_lengths
    ):
        if scores.size(-1) > mask_query.size(-1):
            zero_padding = torch.zeros(
                [mask_query.size(0), scores.size(-1) - mask_query.size(-1)],
                device=mask_query.device,
                dtype=torch.bool,
            )
            mask_query = torch.cat([mask_query, zero_padding], dim=-1)
        mask_query = mask * (~mask_query[:, None])
        scores_woquery = scores.masked_fill(~mask_query[:, None], 0.0)

        ntokens_woquery = 256 * n_layers

        # zero out scores after EOS token. This is needed when batching results in sequences with different lengths.
        for i in range(len(scores_woquery)):
            scores_woquery[i, output_sequence_lengths[i] :, :, :] = 0

        scores_woquery = scores_woquery.sum(dim=[1, 3])
        return scores_woquery / ntokens_woquery

    def overwrite_forward_crossattention(self):
        """
        Replace cross-attention forward function, only used to save
        cross-attention scores.
        """
        for mod in self.decoder.block:
            xattn = mod.layer[1].EncDecAttention
            xattn.forward = types.MethodType(cross_attention_forward, xattn)

    def create_crossattention_storage(self):
        for mod in self.decoder.block:
            xattn = mod.layer[1].EncDecAttention
            xattn.normalized_score_storage = None

set_checkpoint(use_checkpoint)

Enable or disable checkpointing in the encoder. See https://pytorch.org/docs/stable/checkpoint.html

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

def set_checkpoint(self, use_checkpoint):
    """
    Enable or disable checkpointing in the encoder.
    See https://pytorch.org/docs/stable/checkpoint.html
    """
    for mod in self.encoder.encoder.block:
        mod.use_checkpoint = use_checkpoint

reset_score_storage()

Reset score storage, only used when cross-attention scores are saved to train a retriever.

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

def reset_score_storage(self):
    """
    Reset score storage, only used when cross-attention scores are saved
    to train a retriever.
    """
    for mod in self.decoder.block:
        mod.layer[1].EncDecAttention.normalized_score_storage = None

get_crossattention_scores(n_passages, mask, ids, mask_query=None, output_sequence_lengths=[])

Cross-attention scores are aggregated to obtain a single scalar per passage. This scalar can be seen as a similarity score between the question and the input passage. It is obtained by averaging the cross-attention scores obtained on the first decoded token over heads, layers, and tokens of the input passage.

More details in Distilling Knowledge from Reader to Retriever: https://arxiv.org/abs/2012.04584.

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

@torch.no_grad()
def get_crossattention_scores(
    self, n_passages, mask, ids, mask_query=None, output_sequence_lengths=[]
):
    """
    Cross-attention scores are aggregated to obtain a single scalar per
    passage. This scalar can be seen as a similarity score between the
    question and the input passage. It is obtained by averaging the
    cross-attention scores obtained on the first decoded token over heads,
    layers, and tokens of the input passage.

    More details in Distilling Knowledge from Reader to Retriever:
    https://arxiv.org/abs/2012.04584.
    """
    norms = []
    for mod in self.decoder.block:
        norms.append(mod.layer[1].EncDecAttention.normalized_score_storage)
    norms = torch.stack(norms)

    output = {}
    self.aggregate_value(
        norms,
        mask,
        n_passages,
        ids,
        mask_query,
        output,
        prefix="norms",
        output_sequence_lengths=output_sequence_lengths,
    )
    return output

overwrite_forward_crossattention()

Replace cross-attention forward function, only used to save cross-attention scores.

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

def overwrite_forward_crossattention(self):
    """
    Replace cross-attention forward function, only used to save
    cross-attention scores.
    """
    for mod in self.decoder.block:
        xattn = mod.layer[1].EncDecAttention
        xattn.forward = types.MethodType(cross_attention_forward, xattn)

FiDCrossAttentionScore

Bases: T5ForConditionalGeneration

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

class FiDCrossAttentionScore(T5ConditionalGenerationCrossAttentionScore):
    _keys_to_ignore_on_load_missing = [
        r"encoder\.embed_tokens\.weight",
        r"decoder\.embed_tokens\.weight",
        r"lm_head\.weight",
    ]
    _keys_to_ignore_on_load_unexpected = [
        r"decoder\.block\.0\.layer\.1\.EncDecAttention\.relative_attention_bias\.weight",
    ]

    ANSWER_EOS_TOKEN = 1

    def __init__(self, config):
        super().__init__(config)
        self.model_dim = config.d_model

        self.shared = nn.Embedding(config.vocab_size, config.d_model)

        encoder_config = copy.deepcopy(config)
        encoder_config.is_decoder = False
        encoder_config.use_cache = False
        encoder_config.is_encoder_decoder = False

        self.encoder = FiDStackCrossAttentionScore(encoder_config, self.shared)

        decoder_config = copy.deepcopy(config)
        decoder_config.is_decoder = True
        decoder_config.is_encoder_decoder = False
        decoder_config.use_cache = True
        decoder_config.num_layers = config.num_decoder_layers
        self.decoder = FiDStackCrossAttentionScore(decoder_config, self.shared)

        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

        # Model parallel
        self.model_parallel = False
        self.device_map = None

    def set_checkpoint(self, use_checkpoint):
        """
        Enable or disable checkpointing in the encoder.
        See https://pytorch.org/docs/stable/checkpoint.html
        """
        for mod in self.encoder.encoder.block:
            mod.use_checkpoint = use_checkpoint

    def reset_score_storage(self):
        """
        Reset score storage, only used when cross-attention scores are saved
        to train a retriever.
        """
        for mod in self.decoder.block:
            mod.layer[1].EncDecAttention.normalized_score_storage = None

    @torch.no_grad()
    def get_crossattention_scores(
        self, n_passages, mask, ids, mask_query=None, output_sequence_lengths=[]
    ):
        """
        Cross-attention scores are aggregated to obtain a single scalar per
        passage. This scalar can be seen as a similarity score between the
        question and the input passage. It is obtained by averaging the
        cross-attention scores obtained on the first decoded token over heads,
        layers, and tokens of the input passage.

        More details in Distilling Knowledge from Reader to Retriever:
        https://arxiv.org/abs/2012.04584.
        """
        norms = []
        for mod in self.decoder.block:
            norms.append(mod.layer[1].EncDecAttention.normalized_score_storage)
        norms = torch.stack(norms)

        output = {}
        self.aggregate_value(
            norms,
            mask,
            n_passages,
            ids,
            mask_query,
            output,
            prefix="norms",
            output_sequence_lengths=output_sequence_lengths,
        )
        return output

    def aggregate_value(
        self,
        scores,
        mask,
        n_passages,
        ids,
        mask_query=None,
        output={},
        prefix="",
        output_sequence_lengths=[],
    ):
        n_layers, bsz, n_tokens, total_tokens = scores.size()

        ids = ids.view(bsz, n_passages, -1)
        scores = scores.view(n_layers, bsz, n_tokens, n_passages, -1)
        mask = mask.view(bsz, n_passages, -1)
        scores = scores.masked_fill(~mask[None, :, None], 0.0)

        scores = scores.sum(dim=[0])

        scores_woquery = None
        # Compute scores based on scores without query
        if not mask_query is None:
            output[f"{prefix}woquery"] = self.get_woquery_score(
                scores,
                mask_query,
                mask,
                n_layers,
                output_sequence_lengths=output_sequence_lengths,
            )

        return output

    def get_woquery_score(
        self, scores, mask_query, mask, n_layers, output_sequence_lengths
    ):
        if scores.size(-1) > mask_query.size(-1):
            zero_padding = torch.zeros(
                [mask_query.size(0), scores.size(-1) - mask_query.size(-1)],
                device=mask_query.device,
                dtype=torch.bool,
            )
            mask_query = torch.cat([mask_query, zero_padding], dim=-1)
        mask_query = mask * (~mask_query[:, None])
        scores_woquery = scores.masked_fill(~mask_query[:, None], 0.0)

        ntokens_woquery = 256 * n_layers

        # zero out scores after EOS token. This is needed when batching results in sequences with different lengths.
        for i in range(len(scores_woquery)):
            scores_woquery[i, output_sequence_lengths[i] :, :, :] = 0

        scores_woquery = scores_woquery.sum(dim=[1, 3])
        return scores_woquery / ntokens_woquery

    def overwrite_forward_crossattention(self):
        """
        Replace cross-attention forward function, only used to save
        cross-attention scores.
        """
        for mod in self.decoder.block:
            xattn = mod.layer[1].EncDecAttention
            xattn.forward = types.MethodType(cross_attention_forward, xattn)

    def create_crossattention_storage(self):
        for mod in self.decoder.block:
            xattn = mod.layer[1].EncDecAttention
            xattn.normalized_score_storage = None

set_checkpoint(use_checkpoint)

Enable or disable checkpointing in the encoder. See https://pytorch.org/docs/stable/checkpoint.html

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

def set_checkpoint(self, use_checkpoint):
    """
    Enable or disable checkpointing in the encoder.
    See https://pytorch.org/docs/stable/checkpoint.html
    """
    for mod in self.encoder.encoder.block:
        mod.use_checkpoint = use_checkpoint

reset_score_storage()

Reset score storage, only used when cross-attention scores are saved to train a retriever.

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

def reset_score_storage(self):
    """
    Reset score storage, only used when cross-attention scores are saved
    to train a retriever.
    """
    for mod in self.decoder.block:
        mod.layer[1].EncDecAttention.normalized_score_storage = None

get_crossattention_scores(n_passages, mask, ids, mask_query=None, output_sequence_lengths=[])

Cross-attention scores are aggregated to obtain a single scalar per passage. This scalar can be seen as a similarity score between the question and the input passage. It is obtained by averaging the cross-attention scores obtained on the first decoded token over heads, layers, and tokens of the input passage.

More details in Distilling Knowledge from Reader to Retriever: https://arxiv.org/abs/2012.04584.

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

@torch.no_grad()
def get_crossattention_scores(
    self, n_passages, mask, ids, mask_query=None, output_sequence_lengths=[]
):
    """
    Cross-attention scores are aggregated to obtain a single scalar per
    passage. This scalar can be seen as a similarity score between the
    question and the input passage. It is obtained by averaging the
    cross-attention scores obtained on the first decoded token over heads,
    layers, and tokens of the input passage.

    More details in Distilling Knowledge from Reader to Retriever:
    https://arxiv.org/abs/2012.04584.
    """
    norms = []
    for mod in self.decoder.block:
        norms.append(mod.layer[1].EncDecAttention.normalized_score_storage)
    norms = torch.stack(norms)

    output = {}
    self.aggregate_value(
        norms,
        mask,
        n_passages,
        ids,
        mask_query,
        output,
        prefix="norms",
        output_sequence_lengths=output_sequence_lengths,
    )
    return output

overwrite_forward_crossattention()

Replace cross-attention forward function, only used to save cross-attention scores.

Source code in rankify/utils/models/rank_llm/rerank/listwise/lit5/model.py

def overwrite_forward_crossattention(self):
    """
    Replace cross-attention forward function, only used to save
    cross-attention scores.
    """
    for mod in self.decoder.block:
        xattn = mod.layer[1].EncDecAttention
        xattn.forward = types.MethodType(cross_attention_forward, xattn)

PromptMode

Bases: Enum

Source code in rankify/utils/models/rank_llm/rerank/rankllm.py

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

RankFiDDistill

Bases: ListwiseRankLLM

Implements RankFiDDistill, a listwise ranking approach leveraging Fusion-in-Decoder (FiD) for effective retriever-reader knowledge distillation.

RankFiDDistill utilizes Fusion-in-Decoder (FiD) for reranking retrieved passages using multi-document cross-attention. The model is optimized for ranking efficiency and distilling knowledge from reader-based models.

References

• Izacard, G. & Grave, E. (2020): Distilling Knowledge from Reader to Retriever for Question Answering. Paper

Attributes:

Name	Type	Description
model	str	The name or path of the pre-trained RankFiDDistill model.
context_size	int	The maximum number of passages used for ranking.
prompt_mode	PromptMode	Defines the prompt template for FiD.
num_few_shot_examples	int	Number of few-shot examples for ranking.
window_size	int	The window size for ranking multiple documents at a time.
step_size	int	The step size for sliding window ranking.
precision	str	Precision mode ("float32", "bfloat16", "float16").
device	str	The device to use ("cuda" or "cpu").
batched	bool	Whether to enable batch processing.

Example

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

# Define a query and contexts
question = Question("What are the effects of climate change?")
contexts = [
    Context(text="Climate change leads to rising sea levels.", id=0),
    Context(text="Artificial intelligence is transforming industries.", id=1),
    Context(text="Global temperatures are increasing due to CO2 emissions.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize RankFiDDistill Reranker
model = Reranking(method='lit5dist', model_name='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/rank_fid.py

class RankFiDDistill(ListwiseRankLLM):
    """
    Implements **RankFiDDistill**, a **listwise ranking approach** leveraging 
    **Fusion-in-Decoder (FiD)** for effective **retriever-reader knowledge distillation**.



    RankFiDDistill utilizes **Fusion-in-Decoder (FiD)** for **reranking retrieved passages** 
    using **multi-document cross-attention**. The model is optimized for **ranking efficiency** 
    and **distilling knowledge** from reader-based models.

    References:
        - **Izacard, G. & Grave, E. (2020)**: *Distilling Knowledge from Reader to Retriever for Question Answering*.
          [Paper](https://arxiv.org/abs/2012.04584)

    Attributes:
        model (str): The **name or path** of the pre-trained **RankFiDDistill** model.
        context_size (int): The **maximum number of passages** used for ranking.
        prompt_mode (PromptMode): Defines the **prompt template** for FiD.
        num_few_shot_examples (int): Number of **few-shot examples** for ranking.
        window_size (int): The **window size** for ranking multiple documents at a time.
        step_size (int): The **step size** for sliding window ranking.
        precision (str): Precision mode (`"float32"`, `"bfloat16"`, `"float16"`).
        device (str): The device to use (`"cuda"` or `"cpu"`).
        batched (bool): Whether to **enable batch processing**.

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

        # Define a query and contexts
        question = Question("What are the effects of climate change?")
        contexts = [
            Context(text="Climate change leads to rising sea levels.", id=0),
            Context(text="Artificial intelligence is transforming industries.", id=1),
            Context(text="Global temperatures are increasing due to CO2 emissions.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize RankFiDDistill Reranker
        model = Reranking(method='lit5dist', model_name='LiT5-Distill-base')
        model.rank([document])

        # Print reordered contexts
        print("Reordered Contexts:")
        for context in document.reorder_contexts:
            print(context.text)
        ```
    """
    def _post_init(self):
        self._to_precision(self._precision)

    def _tokenize(self, s: str):
        return self._tokenizer(s)

    def _to_precision(self, precision: str) -> None:
        """
        We don't support python12 for now, after python 12, the code should be changed into
        """
        if precision == "float32":
            self._llm = self._llm.float()
        elif precision == "bfloat16":
            self._llm = self._llm.bfloat16()
        elif precision == "float16":
            self._llm = self._llm.float16()

    def __init__(
        self,
        model: str,
        context_size: int = 150,
        prompt_mode: PromptMode = PromptMode.LiT5,  # Placeholder for actual mode
        num_few_shot_examples: int = 0,
        window_size: int = 20,
        step_size: int = 10,
        precision: str = "bfloat16",
        device: str = "cuda",
        batched: bool = False,
    ) -> None:
        """
        Initializes RankFiDDistill for reranking.

        Args:
            model (str): Path or name of the **RankFiDDistill** model.
            context_size (int, optional): Number of passages used for ranking.
            prompt_mode (PromptMode, optional): Defines the **FiD prompt mode**.
            num_few_shot_examples (int, optional): Number of **few-shot examples** used for ranking.
            window_size (int, optional): Defines the **window size** for ranking.
            step_size (int, optional): Defines the **step size** for sliding window ranking.
            precision (str, optional): Precision format (`"float32"`, `"bfloat16"`, `"float16"`).
            device (str, optional): The device for computation (`"cuda"` or `"cpu"`).
            batched (bool, optional): Whether to use **batch processing**.
        """
        super().__init__(
            model=model,
            context_size=context_size,
            prompt_mode=prompt_mode,
            num_few_shot_examples=num_few_shot_examples,
            window_size=window_size,
        )
        self._precision = precision
        self._tokenizer = T5Tokenizer.from_pretrained(model)
        self._llm = FiD.from_pretrained(model).to(device).eval()

        self._device = device

        self._window_size = window_size
        self._stride = step_size

        self._batched = batched

        self._answer_maxlength = len(
            " > ".join(map(lambda x: f"[{x}]", range(1, window_size + 1)))
        )

        self._output_token_estimate = None

        self._post_init()

    def _run_llm_by_length_unified( self, batch_prompts: List[List[str]]) -> List[Tuple[str, int]]:
        if len(batch_prompts) == 0:
            return []

        self._llm.eval()

        batch_size = len(batch_prompts)
        n_passages = len(batch_prompts[0])

        # single batch, unsqueeze
        inputs = {
            k: v.reshape(batch_size, -1).to(self._device)
            for k, v in self._tokenizer(
                [prompt for prompts in batch_prompts for prompt in prompts],
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=self.max_tokens(),
            ).items()
        }

        with torch.no_grad():
            self._llm.reset_n_passages(n_passages=n_passages)
            outputs = self._llm.generate(
                **inputs,
                max_length=self._answer_maxlength,
                do_sample=False,
            )

        decoded_outputs = [
            self._tokenizer.decode(outputs[i], skip_special_tokens=True)
            for i in range(outputs.shape[0])
        ]

        # all token size should be equal
        return [
            (decoded_output, outputs.shape[1]) for decoded_output in decoded_outputs
        ]

    def rerank_batch(
        self,
        requests: List[Request],
        rank_start: int = 0,
        rank_end: int = 100,
        shuffle_candidates: bool = False,
        logging: bool = False,
        use_logits: bool = False,
        use_alpha: bool = False,
        **kwargs
    ) -> List[Result]:
        """
        Reranks documents in batch using RankFiDDistill.

        Args:
            requests (List[Request]): List of **query and candidate passages** for ranking.
            rank_start (int, optional): The **starting rank index**.
            rank_end (int, optional): The **ending rank index**.
            shuffle_candidates (bool, optional): Whether to **shuffle candidates** before ranking.
            logging (bool, optional): Enable logging for debugging.

        Returns:
            List[Result]: **Ranked list of documents**.
        """
        top_k_retrieve: int = kwargs.get("top_k_retrieve", 100)

        window_size: int = kwargs.get("window_size", self._window_size)
        window_size = min(window_size, top_k_retrieve)
        step: int = kwargs.get("step_size", self._stride)

        populate_exec_summary: bool = kwargs.get("populate_exec_summary", False)

        batch_size = kwargs.get("batch_size", 1)
        #print("|||||||||||||||||||||||||||||||")
        if self._batched:
            # reranking using vllm
            if len(set([len(req.candidates) for req in requests])) != 1:
                raise ValueError(
                    "Batched requests must have the same number of candidates"
                )

            result = []

            #with tqdm(range(0, len(requests))) as bar:
            for i in range(0, len(requests), batch_size):
                batch = requests[i : min(i + batch_size, len(requests))]
                batch_result = self.sliding_windows_batched(
                    batch,
                    rank_start=max(rank_start, 0),
                    rank_end=min(
                        rank_end, len(requests[0].candidates)
                    ),  # TODO: Fails arbitrary hit sizes
                    window_size=window_size,
                    step=step,
                    shuffle_candidates=shuffle_candidates,
                    logging=logging,
                    populate_exec_summary=populate_exec_summary,
                )
                result.extend(batch_result)
                #bar.update(len(batch))

            return result
        else:
            # Normal operation mode
            results = []
            for request in requests:
                result = self.sliding_windows(
                    request,
                    rank_start=max(rank_start, 0),
                    rank_end=min(rank_end, len(request.candidates)),
                    window_size=window_size,
                    step=step,
                    shuffle_candidates=shuffle_candidates,
                    logging=logging,
                    populate_exec_summary=populate_exec_summary,
                )
                results.append(result)
            return results

    def run_llm_batched(
        self, prompts: List[List[Dict[str, str]]], **kwargs
    ) -> List[Tuple[str, int]]:

        if len(prompts) == 0:
            return []

        # unfortunately, we are not allowed to use VLLM on T5. However, we could unify the prompts by passage size
        #   (which is commonly the same) then rerank stuff having same passage sizes

        prompt_infos = [list(map(lambda x: x["text"], prompt)) for prompt in prompts]

        return self._run_llm_by_length_unified(prompt_infos)

    def create_prompt_batched(
        self, results: List[Result], rank_start: int, rank_end: int, batch_size: int
    ) -> List[Tuple[List[Dict[str, str]], int]]:
        return [self.create_prompt(result, rank_start, rank_end) for result in results]

    def run_llm(self, prompts: List[Dict[str, str]], **kwargs) -> Tuple[str, int]:
        """
        Runs RankFiDDistill to generate ranking predictions.

        Args:
            prompts (List[Dict[str, str]]): **List of query-context pairs** formatted for FiD ranking.

        Returns:
            Tuple[str, int]: **Ranked list of passages**.
        """

        return self._run_llm_by_length_unified(
            [list(map(lambda x: x["text"], prompts))]
        )[0]

    def create_prompt(
        self, result: Result, rank_start: int, rank_end: int, use_alpha: bool= False
    ) -> Tuple[List[Dict[str, str]], int]:
        """
        Create a prompt based on the result and given ranking range.
        """

        # For now, we concat the prompt, because it seems LiT5 is also concatting the stuff
        prompts = [
            {
                "text": self._gen_passage(
                    result.query["text"],
                    i + 1 - rank_start,
                    self.convert_doc_to_prompt_content(
                        result.candidates[i]["doc"], self.max_tokens()
                    ),
                )
            }
            for i in range(rank_start, rank_end)
        ]

        return prompts, sum(self.get_num_tokens(prompt["text"]) for prompt in prompts)

    def get_num_tokens(self, prompt: Union[str, List[Dict[str, str]]]) -> int:
        """
        Abstract method to calculate the number of tokens contained in the given prompt.
        """
        if isinstance(prompt, str):
            return len(self._tokenizer.encode(prompt))
        elif isinstance(prompt, list):
            return sum(len(self._tokenizer.encode(item["text"])) for item in prompt)
        else:
            raise ValueError(
                "Prompt must be a string or a list of dictionaries with a 'text' key."
            )

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

    def num_output_tokens(self, current_window_size: Optional[int] = None) -> int:
        if current_window_size is None:
            current_window_size = self._window_size
        if (
            self._output_token_estimate is not None
            and self._window_size == current_window_size
        ):
            return self._output_token_estimate
        else:
            output_token_estimate = (
                len(
                    self._tokenizer.encode(
                        " > ".join([f"[{i + 1}]" for i in range(current_window_size)])
                    )
                )
                - 1
            )
            if (
                self._output_token_estimate is None
                and self._window_size == current_window_size
            ):
                self._output_token_estimate = output_token_estimate

            return output_token_estimate

    @staticmethod
    def _gen_passage(query: str, index: int, passage: str) -> str:
        """
        Formats passages for the RankFiDDistill prompt.

        Parameters
        ----------
        query : str
            The search query.
        index : int
            The passage index in the ranking.
        passage : str
            The passage text.

        Returns
        -------
        str
            The formatted passage.
        """
        return f"Search Query: {query} Passage: [{index}] {passage} Relevance Ranking: "

__init__(model, context_size=150, prompt_mode=PromptMode.LiT5, num_few_shot_examples=0, window_size=20, step_size=10, precision='bfloat16', device='cuda', batched=False)

Initializes RankFiDDistill for reranking.

Parameters:

Name	Type	Description	Default
model	str	Path or name of the RankFiDDistill model.	required
context_size	int	Number of passages used for ranking.	150
prompt_mode	PromptMode	Defines the FiD prompt mode.	LiT5
num_few_shot_examples	int	Number of few-shot examples used for ranking.	0
window_size	int	Defines the window size for ranking.	20
step_size	int	Defines the step size for sliding window ranking.	10
precision	str	Precision format ("float32", "bfloat16", "float16").	'bfloat16'
device	str	The device for computation ("cuda" or "cpu").	'cuda'
batched	bool	Whether to use batch processing.	False

Source code in rankify/models/rank_fid.py

def __init__(
    self,
    model: str,
    context_size: int = 150,
    prompt_mode: PromptMode = PromptMode.LiT5,  # Placeholder for actual mode
    num_few_shot_examples: int = 0,
    window_size: int = 20,
    step_size: int = 10,
    precision: str = "bfloat16",
    device: str = "cuda",
    batched: bool = False,
) -> None:
    """
    Initializes RankFiDDistill for reranking.

    Args:
        model (str): Path or name of the **RankFiDDistill** model.
        context_size (int, optional): Number of passages used for ranking.
        prompt_mode (PromptMode, optional): Defines the **FiD prompt mode**.
        num_few_shot_examples (int, optional): Number of **few-shot examples** used for ranking.
        window_size (int, optional): Defines the **window size** for ranking.
        step_size (int, optional): Defines the **step size** for sliding window ranking.
        precision (str, optional): Precision format (`"float32"`, `"bfloat16"`, `"float16"`).
        device (str, optional): The device for computation (`"cuda"` or `"cpu"`).
        batched (bool, optional): Whether to use **batch processing**.
    """
    super().__init__(
        model=model,
        context_size=context_size,
        prompt_mode=prompt_mode,
        num_few_shot_examples=num_few_shot_examples,
        window_size=window_size,
    )
    self._precision = precision
    self._tokenizer = T5Tokenizer.from_pretrained(model)
    self._llm = FiD.from_pretrained(model).to(device).eval()

    self._device = device

    self._window_size = window_size
    self._stride = step_size

    self._batched = batched

    self._answer_maxlength = len(
        " > ".join(map(lambda x: f"[{x}]", range(1, window_size + 1)))
    )

    self._output_token_estimate = None

    self._post_init()

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

Reranks documents in batch using RankFiDDistill.

Parameters:

Name	Type	Description	Default
requests	List[Request]	List of query and candidate passages for ranking.	required
rank_start	int	The starting rank index.	0
rank_end	int	The ending rank index.	100
shuffle_candidates	bool	Whether to shuffle candidates before ranking.	False
logging	bool	Enable logging for debugging.	False

Returns:

Type	Description
List[Result]	List[Result]: Ranked list of documents.

Source code in rankify/models/rank_fid.py

def rerank_batch(
    self,
    requests: List[Request],
    rank_start: int = 0,
    rank_end: int = 100,
    shuffle_candidates: bool = False,
    logging: bool = False,
    use_logits: bool = False,
    use_alpha: bool = False,
    **kwargs
) -> List[Result]:
    """
    Reranks documents in batch using RankFiDDistill.

    Args:
        requests (List[Request]): List of **query and candidate passages** for ranking.
        rank_start (int, optional): The **starting rank index**.
        rank_end (int, optional): The **ending rank index**.
        shuffle_candidates (bool, optional): Whether to **shuffle candidates** before ranking.
        logging (bool, optional): Enable logging for debugging.

    Returns:
        List[Result]: **Ranked list of documents**.
    """
    top_k_retrieve: int = kwargs.get("top_k_retrieve", 100)

    window_size: int = kwargs.get("window_size", self._window_size)
    window_size = min(window_size, top_k_retrieve)
    step: int = kwargs.get("step_size", self._stride)

    populate_exec_summary: bool = kwargs.get("populate_exec_summary", False)

    batch_size = kwargs.get("batch_size", 1)
    #print("|||||||||||||||||||||||||||||||")
    if self._batched:
        # reranking using vllm
        if len(set([len(req.candidates) for req in requests])) != 1:
            raise ValueError(
                "Batched requests must have the same number of candidates"
            )

        result = []

        #with tqdm(range(0, len(requests))) as bar:
        for i in range(0, len(requests), batch_size):
            batch = requests[i : min(i + batch_size, len(requests))]
            batch_result = self.sliding_windows_batched(
                batch,
                rank_start=max(rank_start, 0),
                rank_end=min(
                    rank_end, len(requests[0].candidates)
                ),  # TODO: Fails arbitrary hit sizes
                window_size=window_size,
                step=step,
                shuffle_candidates=shuffle_candidates,
                logging=logging,
                populate_exec_summary=populate_exec_summary,
            )
            result.extend(batch_result)
            #bar.update(len(batch))

        return result
    else:
        # Normal operation mode
        results = []
        for request in requests:
            result = self.sliding_windows(
                request,
                rank_start=max(rank_start, 0),
                rank_end=min(rank_end, len(request.candidates)),
                window_size=window_size,
                step=step,
                shuffle_candidates=shuffle_candidates,
                logging=logging,
                populate_exec_summary=populate_exec_summary,
            )
            results.append(result)
        return results

run_llm(prompts, **kwargs)

Runs RankFiDDistill to generate ranking predictions.

Parameters:

Name	Type	Description	Default
prompts	List[Dict[str, str]]	List of query-context pairs formatted for FiD ranking.	required

Returns:

Type	Description
Tuple[str, int]	Tuple[str, int]: Ranked list of passages.

Source code in rankify/models/rank_fid.py

def run_llm(self, prompts: List[Dict[str, str]], **kwargs) -> Tuple[str, int]:
    """
    Runs RankFiDDistill to generate ranking predictions.

    Args:
        prompts (List[Dict[str, str]]): **List of query-context pairs** formatted for FiD ranking.

    Returns:
        Tuple[str, int]: **Ranked list of passages**.
    """

    return self._run_llm_by_length_unified(
        [list(map(lambda x: x["text"], prompts))]
    )[0]

create_prompt(result, rank_start, rank_end, use_alpha=False)

Create a prompt based on the result and given ranking range.

Source code in rankify/models/rank_fid.py

def create_prompt(
    self, result: Result, rank_start: int, rank_end: int, use_alpha: bool= False
) -> Tuple[List[Dict[str, str]], int]:
    """
    Create a prompt based on the result and given ranking range.
    """

    # For now, we concat the prompt, because it seems LiT5 is also concatting the stuff
    prompts = [
        {
            "text": self._gen_passage(
                result.query["text"],
                i + 1 - rank_start,
                self.convert_doc_to_prompt_content(
                    result.candidates[i]["doc"], self.max_tokens()
                ),
            )
        }
        for i in range(rank_start, rank_end)
    ]

    return prompts, sum(self.get_num_tokens(prompt["text"]) for prompt in prompts)

get_num_tokens(prompt)

Abstract method to calculate the number of tokens contained in the given prompt.

Source code in rankify/models/rank_fid.py

def get_num_tokens(self, prompt: Union[str, List[Dict[str, str]]]) -> int:
    """
    Abstract method to calculate the number of tokens contained in the given prompt.
    """
    if isinstance(prompt, str):
        return len(self._tokenizer.encode(prompt))
    elif isinstance(prompt, list):
        return sum(len(self._tokenizer.encode(item["text"])) for item in prompt)
    else:
        raise ValueError(
            "Prompt must be a string or a list of dictionaries with a 'text' key."
        )

RankFiDScore

Bases: ListwiseRankLLM

Implements RankFiDScore [18]_, a listwise ranking approach leveraging Fusion-in-Decoder (FiD) with cross-attention scoring for accurate ranking.

.. _[18]: https://arxiv.org/abs/2012.04584

RankFiDScore utilizes Fusion-in-Decoder (FiD) models optimized for zero-shot listwise ranking by leveraging cross-attention weights for precise passage relevance estimation.

References

• Izacard, G. & Grave, E. (2020): Distilling Knowledge from Reader to Retriever for Question Answering. Paper

Attributes:

Name	Type	Description
model	str	The name or path of the pre-trained RankFiDScore model.
context_size	int	The maximum number of passages used for ranking.
prompt_mode	PromptMode	Defines the prompt template for FiD.
num_few_shot_examples	int	Number of few-shot examples for ranking.
window_size	int	The window size for ranking multiple documents at a time.
step_size	int	The step size for sliding window ranking.
precision	str	Precision mode ("float32", "bfloat16", "float16").
device	str	The device to use ("cuda" or "cpu").

Examples:

Basic Usage:

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

# Define a query and contexts
question = Question("What are the effects of climate change?")
contexts = [
    Context(text="Climate change leads to rising sea levels.", id=0),
    Context(text="Artificial intelligence is transforming industries.", id=1),
    Context(text="Global temperatures are increasing due to CO2 emissions.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize RankFiDScore Reranker
model = Reranking(method='lit5score', model_name='LiT5-Score-base')
model.rank([document])

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

Source code in rankify/models/rank_fid.py

class RankFiDScore(ListwiseRankLLM):
    """
    Implements **RankFiDScore** `[18]_`, a **listwise ranking approach** leveraging 
    **Fusion-in-Decoder (FiD)** with **cross-attention scoring** for accurate ranking.

    .. _[18]: https://arxiv.org/abs/2012.04584

    RankFiDScore utilizes **Fusion-in-Decoder (FiD) models** optimized for **zero-shot listwise ranking** 
    by leveraging **cross-attention weights** for precise **passage relevance estimation**.

    References:
        - **Izacard, G. & Grave, E. (2020)**: *Distilling Knowledge from Reader to Retriever for Question Answering*.
          [Paper](https://arxiv.org/abs/2012.04584)

    Attributes:
        model (str): The **name or path** of the pre-trained **RankFiDScore** model.
        context_size (int): The **maximum number of passages** used for ranking.
        prompt_mode (PromptMode): Defines the **prompt template** for FiD.
        num_few_shot_examples (int): Number of **few-shot examples** for ranking.
        window_size (int): The **window size** for ranking multiple documents at a time.
        step_size (int): The **step size** for sliding window ranking.
        precision (str): Precision mode (`"float32"`, `"bfloat16"`, `"float16"`).
        device (str): The device to use (`"cuda"` or `"cpu"`).

    Examples:
        **Basic Usage:**
        ```python
        from rankify.dataset.dataset import Document, Question, Context
        from rankify.models.reranking import Reranking

        # Define a query and contexts
        question = Question("What are the effects of climate change?")
        contexts = [
            Context(text="Climate change leads to rising sea levels.", id=0),
            Context(text="Artificial intelligence is transforming industries.", id=1),
            Context(text="Global temperatures are increasing due to CO2 emissions.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize RankFiDScore Reranker
        model = Reranking(method='lit5score', model_name='LiT5-Score-base')
        model.rank([document])

        # Print reordered contexts
        print("Reordered Contexts:")
        for context in document.reorder_contexts:
            print(context.text)
        ```
    """
    def _post_init(self):
        # set the overwrite forward cross attention
        self._llm.overwrite_forward_crossattention()
        self._to_precision(self._precision)

    def _tokenize(self, s: str):
        return self._tokenizer(s)

    def _to_precision(self, precision: str) -> None:
        """
        We don't support python12 for now, after python 12, the code should be changed into
        """
        if precision == "float32":
            self._llm = self._llm.float()
        elif precision == "bfloat16":
            self._llm = self._llm.bfloat16()
        elif precision == "float16":
            self._llm = self._llm.float16()

    def __init__(
        self,
        model: str,
        context_size: int = 150,
        prompt_mode: PromptMode = PromptMode.LiT5,  # Placeholder for actual mode
        num_few_shot_examples: int = 0,
        window_size: int = 20,
        step_size: int = 10,
        precision: str = "bfloat16",
        device: str = "cuda",
        batched: bool = False,
    ) -> None:
        """
        Initializes RankFiDScore for reranking.

        Args:
            model (str): Path or name of the **RankFiDScore** model.
            context_size (int, optional): Number of passages to use for ranking.
            prompt_mode (PromptMode, optional): Defines the **FiD prompt mode**.
            num_few_shot_examples (int, optional): Number of **few-shot examples** used for ranking.
            window_size (int, optional): Defines the **window size** for ranking.
            step_size (int, optional): Defines the **step size** for sliding window ranking.
            precision (str, optional): Precision format (`"float32"`, `"bfloat16"`, `"float16"`).
            device (str, optional): The device for computation (`"cuda"` or `"cpu"`).
            batched (bool, optional): Whether to use **batch processing**.
        """

        super().__init__(
            model=model,
            context_size=context_size,
            prompt_mode=prompt_mode,
            num_few_shot_examples=num_few_shot_examples,
            window_size=window_size,
        )
        self._precision = precision
        self._tokenizer = T5Tokenizer.from_pretrained(model)
        self._llm = FiDCrossAttentionScore.from_pretrained(model).to(device).eval()

        self._device = device
        self._window_size = window_size
        self._stride = step_size

        self._batched = batched

        self._output_token_estimate = None

        self._post_init()

    def _run_llm_by_length_unified(
        self, batch_prompts: List[List[Tuple[str, str]]]
    ) -> List[Tuple[str, int]]:
        if len(batch_prompts) == 0:
            return []

        # get arbitrary query (they should be the same)
        queries = [prompts[0][0] for prompts in batch_prompts]
        batch_size = len(batch_prompts)
        n_passages = len(batch_prompts[0])

        inputs = {
            k: v.reshape(batch_size, -1).to(self._device)
            for k, v in self._tokenizer(
                [prompt for prompts in batch_prompts for (_, prompt) in prompts],
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=self.max_tokens(),
            ).items()
        }

        passage_ids = inputs["input_ids"]
        passage_mask = inputs["attention_mask"]

        with torch.no_grad():
            self._llm.reset_score_storage()

            outputs = self._llm.generate(
                **inputs, max_length=20, do_sample=False, n_passages=n_passages
            )

        output_sequence_lengths = []

        for output in outputs:
            output_length = 0
            for j in range(output.shape[0]):
                if output[j] == FiDCrossAttentionScore.ANSWER_EOS_TOKEN:
                    output_length = j
                    break
            else:
                output_length = outputs.shape[1]
            output_sequence_lengths.append(output_length)

        query_mask_reader = self._tokenizer(
            queries,
            max_length=self.max_tokens(),
            padding="longest",
            truncation=True,
            return_tensors="pt",
            add_special_tokens=False,
        )["attention_mask"].bool()

        with torch.no_grad():
            crossattention_scores = self._llm.get_crossattention_scores(
                n_passages,
                ids=passage_ids.to(self._device),
                mask=passage_mask.bool().to(self._device),
                mask_query=query_mask_reader.to(self._device),
                output_sequence_lengths=output_sequence_lengths,
            )
            # only supports normswoquery for now
            crossattention_score: torch.Tensor = crossattention_scores["normswoquery"]
            sorted, idxes = torch.sort(crossattention_score, dim=-1, descending=True)
            idxes = idxes.detach().cpu()

        return [
            (
                " > ".join([f"[{x + 1}]" for x in idxes[i].tolist()]),
                output_sequence_lengths[i] + crossattention_score.shape[1],
            )
            for i in range(idxes.shape[0])
        ]

    def rerank_batch(
        self,
        requests: List[Request],
        rank_start: int = 0,
        rank_end: int = 100,
        shuffle_candidates: bool = False,
        logging: bool = False,
        use_logits: bool = False,
        use_alpha: bool = False,
        **kwargs
    ) -> List[Result]:
        """
        Reranks documents in batch using RankFiDScore.

        Args:
            requests (List[Request]): List of requests containing queries and candidate passages.
            rank_start (int, optional): The starting rank index.
            rank_end (int, optional): The ending rank index.
            shuffle_candidates (bool, optional): Whether to shuffle candidate passages before ranking.
            logging (bool, optional): Enable logging for debugging.

        Returns:
            List[Result]: The reranked documents.
        """
        top_k_retrieve: int = kwargs.get("top_k_retrieve", 100)

        window_size: int = kwargs.get("window_size", self._window_size)
        window_size = min(window_size, top_k_retrieve)
        step: int = kwargs.get("step_size", self._stride)

        populate_exec_summary: bool = kwargs.get("populate_exec_summary", False)

        batch_size = kwargs.get("batch_size", 1)

        if self._batched:
            # reranking using vllm
            if len(set([len(req.candidates) for req in requests])) != 1:
                raise ValueError(
                    "Batched requests must have the same number of candidates"
                )

            result = []

            with tqdm(range(0, len(requests))) as bar:
                for i in range(0, len(requests), batch_size):
                    batch = requests[i : min(i + batch_size, len(requests))]
                    batch_result = self.sliding_windows_batched(
                        batch,
                        rank_start=max(rank_start, 0),
                        rank_end=min(
                            rank_end, len(requests[0].candidates)
                        ),  # TODO: Fails arbitrary hit sizes
                        window_size=window_size,
                        step=step,
                        shuffle_candidates=shuffle_candidates,
                        logging=logging,
                        populate_exec_summary=populate_exec_summary,
                    )
                    result.extend(batch_result)
                    bar.update(len(batch))

            return result
        else:
            # Normal operation mode
            results = []
            for request in requests: #tqdm(
                result = self.sliding_windows(
                    request,
                    rank_start=max(rank_start, 0),
                    rank_end=min(rank_end, len(request.candidates)),
                    window_size=window_size,
                    step=step,
                    shuffle_candidates=shuffle_candidates,
                    logging=logging,
                    populate_exec_summary=populate_exec_summary,
                )
                results.append(result)
            return results

    def run_llm_batched(
        self, prompts: List[List[Dict[str, str]]], **kwargs
    ) -> List[Tuple[str, int]]:
        if len(prompts) == 0:
            return []

        # unfortunately, we are not allowed to use VLLM on T5. However, we could unify the prompts by passage size
        #   (which is commonly the same) then rerank stuff having same passage sizes

        processed_prompts = [
            [(x["query"], x["text"]) for x in prmpt] for prmpt in prompts
        ]

        return self._run_llm_by_length_unified(processed_prompts)

    def create_prompt_batched(
        self, results: List[Result], rank_start: int, rank_end: int, batch_size: int
    ) -> List[Tuple[List[Dict[str, str]], int]]:
        return [self.create_prompt(result, rank_start, rank_end) for result in results]

    def run_llm(self, prompts: List[Dict[str, str]], **kwargs) -> Tuple[str, int]:
        """
        Runs RankFiDScore to generate ranking predictions.

        Args:
            prompts (List[Dict[str, str]]): **List of query-context pairs** formatted for FiD ranking.

        Returns:
            Tuple[str, int]: **Ranked list of passages**.
        """
        return self._run_llm_by_length_unified(
            [[(x["query"], x["text"]) for x in prompts]]
        )[0]

    def create_prompt(
        self, result: Result, rank_start: int, rank_end: int, use_alpha: bool= False
    ) -> Tuple[List[Dict[str, str]], int]:
        """
        Creates a **prompt** based on the result and the specified **ranking range**.

        Args:
            result (Result): The result object containing **query and candidate passages**.
            rank_start (int): The **starting rank index**.
            rank_end (int): The **ending rank index**.
            use_alpha (bool, optional): Whether to **apply alpha weighting**.

        Returns:
            Tuple[List[Dict[str, str]], int]: A **list of formatted prompts** and their **token count**.
        """
        query = result.query["text"]
        results = []

        sum_token = 0

        for i in range(rank_start, rank_end):
            results.append(
                {
                    "query": f"question: {query}",
                    "text": self._gen_passage(
                        query,
                        self.convert_doc_to_prompt_content(
                            result.candidates[i]["doc"], self.max_tokens()
                        ),
                    ),
                }
            )
            sum_token += len(self._tokenizer.encode(results[-1]["text"]))

        return results, sum_token

    def get_num_tokens(self, prompt: str) -> int:
        """
        Computes the number of tokens in a given **prompt string**.

        Args:
            prompt (str): The input prompt text.

        Returns:
            int: The number of tokens in the **prompt**.
        """
        return len(self._tokenizer.encode(prompt))

    def cost_per_1k_token(self, input_token: bool) -> float:
        """
        Returns the estimated **cost per 1,000 tokens**.

        Args:
            input_token (bool): Whether to compute for **input tokens**.

        Returns:
            float: The cost per **1,000 tokens**.
        """
        return 0.0

    def num_output_tokens(self, current_window_size: Optional[int] = None) -> int:
        """
        Computes the **number of output tokens** for the current **window size**.

        Args:
            current_window_size (Optional[int], optional): The **size of the current ranking window**.

        Returns:
            int: The estimated **output token count**.
        """
        if current_window_size is None:
            current_window_size = self._window_size
        if (
            self._output_token_estimate is not None
            and self._window_size == current_window_size
        ):
            return self._output_token_estimate
        else:
            output_token_estimate = (
                len(
                    self._tokenizer.encode(
                        " > ".join([f"[{i + 1}]" for i in range(current_window_size)])
                    )
                )
                - 1
            )
            if (
                self._output_token_estimate is None
                and self._window_size == current_window_size
            ):
                self._output_token_estimate = output_token_estimate

            return output_token_estimate

    @staticmethod
    def _gen_passage(query: str, passage: str) -> str:
        """
        Formats passages for the RankFiDScore prompt.

        Args:
            query (str): The search query.
            passage (str): The passage text.

        Returns:
            str: The formatted passage in **RankFiDScore format**.
        """
        return f"question: {query} context: {passage}"

__init__(model, context_size=150, prompt_mode=PromptMode.LiT5, num_few_shot_examples=0, window_size=20, step_size=10, precision='bfloat16', device='cuda', batched=False)

Initializes RankFiDScore for reranking.

Parameters:

Name	Type	Description	Default
model	str	Path or name of the RankFiDScore model.	required
context_size	int	Number of passages to use for ranking.	150
prompt_mode	PromptMode	Defines the FiD prompt mode.	LiT5
num_few_shot_examples	int	Number of few-shot examples used for ranking.	0
window_size	int	Defines the window size for ranking.	20
step_size	int	Defines the step size for sliding window ranking.	10
precision	str	Precision format ("float32", "bfloat16", "float16").	'bfloat16'
device	str	The device for computation ("cuda" or "cpu").	'cuda'
batched	bool	Whether to use batch processing.	False

Source code in rankify/models/rank_fid.py

def __init__(
    self,
    model: str,
    context_size: int = 150,
    prompt_mode: PromptMode = PromptMode.LiT5,  # Placeholder for actual mode
    num_few_shot_examples: int = 0,
    window_size: int = 20,
    step_size: int = 10,
    precision: str = "bfloat16",
    device: str = "cuda",
    batched: bool = False,
) -> None:
    """
    Initializes RankFiDScore for reranking.

    Args:
        model (str): Path or name of the **RankFiDScore** model.
        context_size (int, optional): Number of passages to use for ranking.
        prompt_mode (PromptMode, optional): Defines the **FiD prompt mode**.
        num_few_shot_examples (int, optional): Number of **few-shot examples** used for ranking.
        window_size (int, optional): Defines the **window size** for ranking.
        step_size (int, optional): Defines the **step size** for sliding window ranking.
        precision (str, optional): Precision format (`"float32"`, `"bfloat16"`, `"float16"`).
        device (str, optional): The device for computation (`"cuda"` or `"cpu"`).
        batched (bool, optional): Whether to use **batch processing**.
    """

    super().__init__(
        model=model,
        context_size=context_size,
        prompt_mode=prompt_mode,
        num_few_shot_examples=num_few_shot_examples,
        window_size=window_size,
    )
    self._precision = precision
    self._tokenizer = T5Tokenizer.from_pretrained(model)
    self._llm = FiDCrossAttentionScore.from_pretrained(model).to(device).eval()

    self._device = device
    self._window_size = window_size
    self._stride = step_size

    self._batched = batched

    self._output_token_estimate = None

    self._post_init()

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

Reranks documents in batch using RankFiDScore.

Parameters:

Name	Type	Description	Default
requests	List[Request]	List of requests containing queries and candidate passages.	required
rank_start	int	The starting rank index.	0
rank_end	int	The ending rank index.	100
shuffle_candidates	bool	Whether to shuffle candidate passages before ranking.	False
logging	bool	Enable logging for debugging.	False

Returns:

Type	Description
List[Result]	List[Result]: The reranked documents.

Source code in rankify/models/rank_fid.py

def rerank_batch(
    self,
    requests: List[Request],
    rank_start: int = 0,
    rank_end: int = 100,
    shuffle_candidates: bool = False,
    logging: bool = False,
    use_logits: bool = False,
    use_alpha: bool = False,
    **kwargs
) -> List[Result]:
    """
    Reranks documents in batch using RankFiDScore.

    Args:
        requests (List[Request]): List of requests containing queries and candidate passages.
        rank_start (int, optional): The starting rank index.
        rank_end (int, optional): The ending rank index.
        shuffle_candidates (bool, optional): Whether to shuffle candidate passages before ranking.
        logging (bool, optional): Enable logging for debugging.

    Returns:
        List[Result]: The reranked documents.
    """
    top_k_retrieve: int = kwargs.get("top_k_retrieve", 100)

    window_size: int = kwargs.get("window_size", self._window_size)
    window_size = min(window_size, top_k_retrieve)
    step: int = kwargs.get("step_size", self._stride)

    populate_exec_summary: bool = kwargs.get("populate_exec_summary", False)

    batch_size = kwargs.get("batch_size", 1)

    if self._batched:
        # reranking using vllm
        if len(set([len(req.candidates) for req in requests])) != 1:
            raise ValueError(
                "Batched requests must have the same number of candidates"
            )

        result = []

        with tqdm(range(0, len(requests))) as bar:
            for i in range(0, len(requests), batch_size):
                batch = requests[i : min(i + batch_size, len(requests))]
                batch_result = self.sliding_windows_batched(
                    batch,
                    rank_start=max(rank_start, 0),
                    rank_end=min(
                        rank_end, len(requests[0].candidates)
                    ),  # TODO: Fails arbitrary hit sizes
                    window_size=window_size,
                    step=step,
                    shuffle_candidates=shuffle_candidates,
                    logging=logging,
                    populate_exec_summary=populate_exec_summary,
                )
                result.extend(batch_result)
                bar.update(len(batch))

        return result
    else:
        # Normal operation mode
        results = []
        for request in requests: #tqdm(
            result = self.sliding_windows(
                request,
                rank_start=max(rank_start, 0),
                rank_end=min(rank_end, len(request.candidates)),
                window_size=window_size,
                step=step,
                shuffle_candidates=shuffle_candidates,
                logging=logging,
                populate_exec_summary=populate_exec_summary,
            )
            results.append(result)
        return results

run_llm(prompts, **kwargs)

Runs RankFiDScore to generate ranking predictions.

Parameters:

Name	Type	Description	Default
prompts	List[Dict[str, str]]	List of query-context pairs formatted for FiD ranking.	required

Returns:

Type	Description
Tuple[str, int]	Tuple[str, int]: Ranked list of passages.

Source code in rankify/models/rank_fid.py

def run_llm(self, prompts: List[Dict[str, str]], **kwargs) -> Tuple[str, int]:
    """
    Runs RankFiDScore to generate ranking predictions.

    Args:
        prompts (List[Dict[str, str]]): **List of query-context pairs** formatted for FiD ranking.

    Returns:
        Tuple[str, int]: **Ranked list of passages**.
    """
    return self._run_llm_by_length_unified(
        [[(x["query"], x["text"]) for x in prompts]]
    )[0]

create_prompt(result, rank_start, rank_end, use_alpha=False)

Creates a prompt based on the result and the specified ranking range.

Parameters:

Name	Type	Description	Default
result	Result	The result object containing query and candidate passages.	required
rank_start	int	The starting rank index.	required
rank_end	int	The ending rank index.	required
use_alpha	bool	Whether to apply alpha weighting.	False

Returns:

Type	Description
Tuple[List[Dict[str, str]], int]	Tuple[List[Dict[str, str]], int]: A list of formatted prompts and their token count.

Source code in rankify/models/rank_fid.py

def create_prompt(
    self, result: Result, rank_start: int, rank_end: int, use_alpha: bool= False
) -> Tuple[List[Dict[str, str]], int]:
    """
    Creates a **prompt** based on the result and the specified **ranking range**.

    Args:
        result (Result): The result object containing **query and candidate passages**.
        rank_start (int): The **starting rank index**.
        rank_end (int): The **ending rank index**.
        use_alpha (bool, optional): Whether to **apply alpha weighting**.

    Returns:
        Tuple[List[Dict[str, str]], int]: A **list of formatted prompts** and their **token count**.
    """
    query = result.query["text"]
    results = []

    sum_token = 0

    for i in range(rank_start, rank_end):
        results.append(
            {
                "query": f"question: {query}",
                "text": self._gen_passage(
                    query,
                    self.convert_doc_to_prompt_content(
                        result.candidates[i]["doc"], self.max_tokens()
                    ),
                ),
            }
        )
        sum_token += len(self._tokenizer.encode(results[-1]["text"]))

    return results, sum_token

get_num_tokens(prompt)

Computes the number of tokens in a given prompt string.

Parameters:

Name	Type	Description	Default
prompt	str	The input prompt text.	required

Returns:

Name	Type	Description
int	int	The number of tokens in the prompt.

Source code in rankify/models/rank_fid.py

def get_num_tokens(self, prompt: str) -> int:
    """
    Computes the number of tokens in a given **prompt string**.

    Args:
        prompt (str): The input prompt text.

    Returns:
        int: The number of tokens in the **prompt**.
    """
    return len(self._tokenizer.encode(prompt))

cost_per_1k_token(input_token)

Returns the estimated cost per 1,000 tokens.

Parameters:

Name	Type	Description	Default
input_token	bool	Whether to compute for input tokens.	required

Returns:

Name	Type	Description
float	float	The cost per 1,000 tokens.

Source code in rankify/models/rank_fid.py

def cost_per_1k_token(self, input_token: bool) -> float:
    """
    Returns the estimated **cost per 1,000 tokens**.

    Args:
        input_token (bool): Whether to compute for **input tokens**.

    Returns:
        float: The cost per **1,000 tokens**.
    """
    return 0.0

num_output_tokens(current_window_size=None)

Computes the number of output tokens for the current window size.

Parameters:

Name	Type	Description	Default
current_window_size	Optional[int]	The size of the current ranking window.	None

Returns:

Name	Type	Description
int	int	The estimated output token count.

Source code in rankify/models/rank_fid.py

def num_output_tokens(self, current_window_size: Optional[int] = None) -> int:
    """
    Computes the **number of output tokens** for the current **window size**.

    Args:
        current_window_size (Optional[int], optional): The **size of the current ranking window**.

    Returns:
        int: The estimated **output token count**.
    """
    if current_window_size is None:
        current_window_size = self._window_size
    if (
        self._output_token_estimate is not None
        and self._window_size == current_window_size
    ):
        return self._output_token_estimate
    else:
        output_token_estimate = (
            len(
                self._tokenizer.encode(
                    " > ".join([f"[{i + 1}]" for i in range(current_window_size)])
                )
            )
            - 1
        )
        if (
            self._output_token_estimate is None
            and self._window_size == current_window_size
        ):
            self._output_token_estimate = output_token_estimate

        return output_token_estimate