Vicuna Reranker

rankify.models.vicuna_reranker

BaseRanking

Bases: ABC

An abstract base class for implementing different ranking models.

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

Attributes:

Name	Type	Description
method	str	The name of the ranking method.
model_name	str	The name of the model being used for ranking.
api_key	str	An optional API key for accessing remote models or services.

Source code in rankify/models/base.py

class BaseRanking(ABC):
    """
    An abstract base class for implementing different ranking models.

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

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

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

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

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

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

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

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

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

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

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

Initializes the base ranking model.

Parameters:

Name	Type	Description	Default
method	str	The name of the ranking method. Defaults to None.	None
model_name	str	The name of the model being used for ranking. Defaults to None.	None
api_key	str	An optional API key for accessing remote models or services. Defaults to None.	None

Example

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

Source code in rankify/models/base.py

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

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

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

rank(documents) abstractmethod

Abstract method to rank a list of documents.

Parameters:

Name	Type	Description	Default
documents	list[Document]	A list of Document instances that need to be ranked.	required

Raises:

Type	Description
NotImplementedError	This method must be implemented by subclasses.

Example

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

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

Source code in rankify/models/base.py

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

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

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

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

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

Document

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

Attributes:

Name	Type	Description
question	Question	The question associated with the document.
answers	Answer	The answers to the question.
contexts	list[Context]	A list of related contexts.
reorder_contexts	list[Context] or None	A reordered list of contexts based on relevance.

Source code in rankify/dataset/dataset.py

class Document:
    """
    Represents a document consisting of a question, answers, and contexts.

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

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

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

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

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

        Returns:
            Document: A new Document instance.

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

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

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

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

        Returns:
            str: The formatted document information.

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

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

Initializes a Document instance.

Parameters:

Name	Type	Description	Default
question	Question	The question associated with the document.	required
answers	Answer	The answers to the question.	required
contexts	list[Context]	A list of contexts related to the question.	None

Example

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

Source code in rankify/dataset/dataset.py

def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
    """
    Initializes a Document instance.

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

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

from_dict(data, n_docs=100) classmethod

Creates a Document instance from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict	A dictionary containing the question, answers, and contexts.	required
n_docs	int	The number of contexts to include. Defaults to 100.	100

Returns:

Name	Type	Description
Document	Document	A new Document instance.

Example

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

Source code in rankify/dataset/dataset.py

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

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

    Returns:
        Document: A new Document instance.

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

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

to_dict()

Converts the document into a dictionary representation.

Returns:

Name	Type	Description
dict	Dict[str, Optional[object]]	A dictionary containing the question, answers, and contexts.

Source code in rankify/dataset/dataset.py

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

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

__str__()

Returns a string representation of the Document instance.

Returns:

Name	Type	Description
str	str	The formatted document information.

Example

d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
print(d)

Source code in rankify/dataset/dataset.py

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

    Returns:
        str: The formatted document information.

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

Request dataclass

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

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

RankListwiseOSLLM

Bases: ListwiseRankLLM

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

class RankListwiseOSLLM(ListwiseRankLLM):
    def __init__(
        self,
        model: str,
        name: str = "",
        context_size: int = 4096,
        prompt_mode: PromptMode = PromptMode.RANK_GPT,
        num_few_shot_examples: int = 0,
        device: str = "cuda",
        num_gpus: int = 1,
        variable_passages: bool = False,
        window_size: int = 20,
        system_message: str = None,
        vllm_batched: bool = False,
        sglang_batched: bool = False,
    ) -> None:
        """
         Creates instance of the RankListwiseOSLLM class, an extension of RankLLM designed for performing listwise ranking of passages using a specified language model. Advanced configurations are supported such as GPU acceleration, variable passage handling, and custom system messages for generating prompts.
         RankListWiseOSLLM uses the default implementations for sliding_window

         Parameters:
         - model (str): Identifier for the language model to be used for ranking tasks.
         - context_size (int, optional): Maximum number of tokens that can be handled in a single prompt. Defaults to 4096.
        - prompt_mode (PromptMode, optional): Specifies the mode of prompt generation, with the default set to RANK_GPT,
         indicating that this class is designed primarily for listwise ranking tasks following the RANK_GPT methodology.
         - num_few_shot_examples (int, optional): Number of few-shot learning examples to include in the prompt, allowing for
         the integration of example-based learning to improve model performance. Defaults to 0, indicating no few-shot examples
         by default.
         - device (str, optional): Specifies the device for model computation ('cuda' for GPU or 'cpu'). Defaults to 'cuda'.
         - num_gpus (int, optional): Number of GPUs to use for model loading and inference. Defaults to 1.
         - variable_passages (bool, optional): Indicates whether the number of passages to rank can vary. Defaults to False.
         - window_size (int, optional): The window size for handling text inputs. Defaults to 20.
         - system_message (Optional[str], optional): Custom system message to be included in the prompt for additional
         instructions or context. Defaults to None.
         - vllm_batched (bool, optional): Indicates whether batched inference using VLLM is leveraged. Defaults to False.
         - sglang_batched (bool, optional): Indicates whether batched inference using SGLang is leveraged. Defaults to False.

         Raises:
         - AssertionError: If CUDA is specified as the device but is not available on the system.
         - ValueError: If an unsupported prompt mode is provided.

         Note:
         - This class is operates given scenarios where listwise ranking is required, with support for dynamic
         passage handling and customization of prompts through system messages and few-shot examples.
         - GPU acceleration is supported and recommended for faster computations.
        TODO: Make repetition_penalty configurable
        """
        super().__init__(
            model, context_size, prompt_mode, num_few_shot_examples, window_size
        )
        self._device = device
        self._vllm_batched = vllm_batched
        self._sglang_batched = sglang_batched
        self._name = name
        self._variable_passages = variable_passages
        self._system_message = system_message
        self._output_token_estimate = None

        if num_few_shot_examples > 0:
            with open("data/output_v2_aug_filtered.jsonl", "r") as json_file:
                self._examples = list(json_file)[1:-1]
        if self._device == "cuda":
            assert torch.cuda.is_available()
        #print(prompt_mode , PromptMode.RANK_GPT )
        if prompt_mode != PromptMode.RANK_GPT:
            raise ValueError(
                f"Unsupported prompt mode: {prompt_mode}. The only prompt mode currently supported is a slight variation of {PromptMode.RANK_GPT} prompt."
            )
        if vllm_batched and LLM is None:
            raise ImportError(
                "Please install rank-llm with `pip install rank-llm[vllm]` to use batch inference."
            )
        elif vllm_batched:
            # TODO: find max_model_len given gpu
            self._llm = LLM(
                model,
                download_dir=os.getenv("HF_HOME"),
                enforce_eager=False,
                max_logprobs=30,
                tensor_parallel_size=num_gpus
            )
            self._tokenizer = self._llm.get_tokenizer()
        elif sglang_batched and Engine is None:
            raise ImportError(
                "Please install rank-llm with `pip install rank-llm[sglang]` to use sglang batch inference."
            )
        elif sglang_batched:
            port = random.randint(30000, 35000)
            self._llm = Engine(model, port=port)
            self._tokenizer = self._llm.get_tokenizer()
        else:
            self._llm, self._tokenizer = load_model(
                model, device=device, num_gpus=num_gpus
            )

    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: Any,
    ) -> List[Result]:
        top_k_retrieve: int = kwargs.get("top_k_retrieve", 50)
        window_size: int = kwargs.get("window_size", 20)
        window_size = min(window_size, top_k_retrieve)
        step: int = kwargs.get("step", 10)
        populate_exec_summary: bool = kwargs.get("populate_exec_summary", False)

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

            return self.sliding_windows_batched(
                requests,
                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,
                use_logits=use_logits,
                use_alpha=use_alpha
            )
        else:
            if use_logits:
                raise TypeError("Reranking using logits of first identifier is currently only supported when vllm_batch=True")

            # 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,
                    use_logits=use_logits,
                    use_alpha=use_alpha
                )
                results.append(result)
            return results

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

        return result_string, evaluations

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

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

        if isinstance(self._llm, LLM):
            logger.info(f"VLLM Generating!")
            if current_window_size is None:
                current_window_size = self._window_size

            if use_logits:
                params = SamplingParams(
                    min_tokens=2,
                    max_tokens=2, 
                    temperature=0.0,
                    logprobs=30,
                )
                outputs = self._llm.generate(prompts, sampling_params=params)
                arr = [self._get_logits_single_digit(output, use_alpha=use_alpha) for output in outputs]
                return [(s, len(s)) for s, __ in arr]
            else:
                sampling_params = SamplingParams(
                    temperature=0.0,
                    max_tokens=self.num_output_tokens(current_window_size, use_alpha),
                    min_tokens=self.num_output_tokens(current_window_size, use_alpha),
                )
                outputs = self._llm.generate(prompts, sampling_params)
                return [
                    (output.outputs[0].text, len(output.outputs[0].token_ids))
                    for output in outputs
                ]
        else:
            logger.info(f"SGLang Generating!")
            sampling_params = {
                "temperature": 0.0,
                "max_new_tokens": self.num_output_tokens(current_window_size, use_alpha),
                "min_new_tokens": self.num_output_tokens(current_window_size, use_alpha),
            }
            outputs = self._llm.generate(prompts, sampling_params)
            return [
                # completion_tokens counts stop token
                (output["text"], output["meta_info"]["completion_tokens"] - 1)
                for output in outputs
            ]

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

        if use_logits:
            params = SamplingParams(min_tokens=1, max_tokens=1, temperature=0.0, logprobs=30)
            output = self._llm.generate([prompt+"["], sampling_params=params)[0]
            s, _ = self._get_logits_single_digit(output, effective_location=0, use_alpha=use_alpha)
            return s, len(s)
        else:
            inputs = self._tokenizer([prompt])
            inputs = {k: torch.tensor(v).to(self._device) for k, v in inputs.items()}
            gen_cfg = GenerationConfig.from_model_config(self._llm.config)
            gen_cfg.max_new_tokens = self.num_output_tokens(current_window_size, use_alpha)
            gen_cfg.min_new_tokens = self.num_output_tokens(current_window_size, use_alpha)
            # gen_cfg.temperature = 0
            gen_cfg.do_sample = False
            output_ids = self._llm.generate(**inputs, generation_config=gen_cfg)

            if self._llm.config.is_encoder_decoder:
                output_ids = output_ids[0]
            else:
                output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
            outputs = self._tokenizer.decode(
                output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
            )
            return outputs, output_ids.size(0)

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

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

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

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

        if self._output_token_estimate is None and self._window_size == current_window_size:
            self._output_token_estimate = _output_token_estimate

        return _output_token_estimate

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

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

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

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

    def create_prompt(
        self, result: Result, rank_start: int, rank_end: int, use_alpha: bool = False
    ) -> Tuple[str, int]:
        query = result.query["text"]
        query = self._replace_number(query)
        num = len(result.candidates[rank_start:rank_end])
        max_length = 300 * (20 / (rank_end - rank_start))
        while True:
            if self._vllm_batched:
                messages = list()
                if self._system_message:
                    messages.append({"role": "system", "content": self._system_message})
                messages = self._add_few_shot_examples_messages(messages)
                prefix = self._add_prefix_prompt(query, num, use_alpha=use_alpha)
                rank = 0
                input_context = f"{prefix}\n"
                for cand in result.candidates[rank_start:rank_end]:
                    rank += 1
                    content = self.convert_doc_to_prompt_content(cand["doc"], max_length)

                    identifier = chr(ALPH_START_IDX + rank) if use_alpha else str(rank)
                    input_context += f"[{identifier}] {self._replace_number(content)}\n"

                input_context += self._add_post_prompt(query, num, use_alpha=use_alpha)
                messages.append({"role": "user", "content": input_context})

                prompt = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
                prompt = fix_text(prompt)
                num_tokens = self.get_num_tokens(prompt)
                if num_tokens <= self.max_tokens() - self.num_output_tokens(
                    rank_end - rank_start, use_alpha
                ):
                    break
                else:
                    max_length -= max(
                        1,
                        (
                            num_tokens
                            - self.max_tokens()
                            + self.num_output_tokens(rank_end - rank_start, use_alpha)
                        )
                        // ((rank_end - rank_start) * 4),
                    )
            else:
                conv = get_conversation_template(self._model)
                if self._system_message:
                    conv.set_system_message(self._system_message)
                conv = self._add_few_shot_examples(conv)
                prefix = self._add_prefix_prompt(query, num, use_alpha=use_alpha)
                rank = 0
                input_context = f"{prefix}\n"
                for cand in result.candidates[rank_start:rank_end]:
                    rank += 1
                    # For Japanese should cut by character: content = content[:int(max_length)]
                    content = self.convert_doc_to_prompt_content(cand["doc"], max_length)

                    identifier = chr(ALPH_START_IDX + rank) if use_alpha else str(rank)
                    input_context += f"[{identifier}] {self._replace_number(content)}\n"

                input_context += self._add_post_prompt(query, num, use_alpha=use_alpha)
                conv.append_message(conv.roles[0], input_context)
                conv.append_message(conv.roles[1], None)
                prompt = conv.get_prompt()
                prompt = fix_text(prompt)
                num_tokens = self.get_num_tokens(prompt)
                if num_tokens <= self.max_tokens() - self.num_output_tokens(
                    rank_end - rank_start, use_alpha
                ):
                    break
                else:
                    max_length -= max(
                        1,
                        (
                            num_tokens
                            - self.max_tokens()
                            + self.num_output_tokens(rank_end - rank_start, use_alpha)
                        )
                        // ((rank_end - rank_start) * 4),
                    )
        return prompt, self.get_num_tokens(prompt)

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

        all_completed_prompts = []

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

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

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

    def get_name(self) -> str:
        return self._name

__init__(model, name='', context_size=4096, prompt_mode=PromptMode.RANK_GPT, num_few_shot_examples=0, device='cuda', num_gpus=1, variable_passages=False, window_size=20, system_message=None, vllm_batched=False, sglang_batched=False)

Creates instance of the RankListwiseOSLLM class, an extension of RankLLM designed for performing listwise ranking of passages using a specified language model. Advanced configurations are supported such as GPU acceleration, variable passage handling, and custom system messages for generating prompts. RankListWiseOSLLM uses the default implementations for sliding_window

Parameters: - model (str): Identifier for the language model to be used for ranking tasks. - context_size (int, optional): Maximum number of tokens that can be handled in a single prompt. Defaults to 4096. - prompt_mode (PromptMode, optional): Specifies the mode of prompt generation, with the default set to RANK_GPT, indicating that this class is designed primarily for listwise ranking tasks following the RANK_GPT methodology. - num_few_shot_examples (int, optional): Number of few-shot learning examples to include in the prompt, allowing for the integration of example-based learning to improve model performance. Defaults to 0, indicating no few-shot examples by default. - device (str, optional): Specifies the device for model computation ('cuda' for GPU or 'cpu'). Defaults to 'cuda'. - num_gpus (int, optional): Number of GPUs to use for model loading and inference. Defaults to 1. - variable_passages (bool, optional): Indicates whether the number of passages to rank can vary. Defaults to False. - window_size (int, optional): The window size for handling text inputs. Defaults to 20. - system_message (Optional[str], optional): Custom system message to be included in the prompt for additional instructions or context. Defaults to None. - vllm_batched (bool, optional): Indicates whether batched inference using VLLM is leveraged. Defaults to False. - sglang_batched (bool, optional): Indicates whether batched inference using SGLang is leveraged. Defaults to False.

Raises: - AssertionError: If CUDA is specified as the device but is not available on the system. - ValueError: If an unsupported prompt mode is provided.

Note: - This class is operates given scenarios where listwise ranking is required, with support for dynamic passage handling and customization of prompts through system messages and few-shot examples. - GPU acceleration is supported and recommended for faster computations. TODO: Make repetition_penalty configurable

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

def __init__(
    self,
    model: str,
    name: str = "",
    context_size: int = 4096,
    prompt_mode: PromptMode = PromptMode.RANK_GPT,
    num_few_shot_examples: int = 0,
    device: str = "cuda",
    num_gpus: int = 1,
    variable_passages: bool = False,
    window_size: int = 20,
    system_message: str = None,
    vllm_batched: bool = False,
    sglang_batched: bool = False,
) -> None:
    """
     Creates instance of the RankListwiseOSLLM class, an extension of RankLLM designed for performing listwise ranking of passages using a specified language model. Advanced configurations are supported such as GPU acceleration, variable passage handling, and custom system messages for generating prompts.
     RankListWiseOSLLM uses the default implementations for sliding_window

     Parameters:
     - model (str): Identifier for the language model to be used for ranking tasks.
     - context_size (int, optional): Maximum number of tokens that can be handled in a single prompt. Defaults to 4096.
    - prompt_mode (PromptMode, optional): Specifies the mode of prompt generation, with the default set to RANK_GPT,
     indicating that this class is designed primarily for listwise ranking tasks following the RANK_GPT methodology.
     - num_few_shot_examples (int, optional): Number of few-shot learning examples to include in the prompt, allowing for
     the integration of example-based learning to improve model performance. Defaults to 0, indicating no few-shot examples
     by default.
     - device (str, optional): Specifies the device for model computation ('cuda' for GPU or 'cpu'). Defaults to 'cuda'.
     - num_gpus (int, optional): Number of GPUs to use for model loading and inference. Defaults to 1.
     - variable_passages (bool, optional): Indicates whether the number of passages to rank can vary. Defaults to False.
     - window_size (int, optional): The window size for handling text inputs. Defaults to 20.
     - system_message (Optional[str], optional): Custom system message to be included in the prompt for additional
     instructions or context. Defaults to None.
     - vllm_batched (bool, optional): Indicates whether batched inference using VLLM is leveraged. Defaults to False.
     - sglang_batched (bool, optional): Indicates whether batched inference using SGLang is leveraged. Defaults to False.

     Raises:
     - AssertionError: If CUDA is specified as the device but is not available on the system.
     - ValueError: If an unsupported prompt mode is provided.

     Note:
     - This class is operates given scenarios where listwise ranking is required, with support for dynamic
     passage handling and customization of prompts through system messages and few-shot examples.
     - GPU acceleration is supported and recommended for faster computations.
    TODO: Make repetition_penalty configurable
    """
    super().__init__(
        model, context_size, prompt_mode, num_few_shot_examples, window_size
    )
    self._device = device
    self._vllm_batched = vllm_batched
    self._sglang_batched = sglang_batched
    self._name = name
    self._variable_passages = variable_passages
    self._system_message = system_message
    self._output_token_estimate = None

    if num_few_shot_examples > 0:
        with open("data/output_v2_aug_filtered.jsonl", "r") as json_file:
            self._examples = list(json_file)[1:-1]
    if self._device == "cuda":
        assert torch.cuda.is_available()
    #print(prompt_mode , PromptMode.RANK_GPT )
    if prompt_mode != PromptMode.RANK_GPT:
        raise ValueError(
            f"Unsupported prompt mode: {prompt_mode}. The only prompt mode currently supported is a slight variation of {PromptMode.RANK_GPT} prompt."
        )
    if vllm_batched and LLM is None:
        raise ImportError(
            "Please install rank-llm with `pip install rank-llm[vllm]` to use batch inference."
        )
    elif vllm_batched:
        # TODO: find max_model_len given gpu
        self._llm = LLM(
            model,
            download_dir=os.getenv("HF_HOME"),
            enforce_eager=False,
            max_logprobs=30,
            tensor_parallel_size=num_gpus
        )
        self._tokenizer = self._llm.get_tokenizer()
    elif sglang_batched and Engine is None:
        raise ImportError(
            "Please install rank-llm with `pip install rank-llm[sglang]` to use sglang batch inference."
        )
    elif sglang_batched:
        port = random.randint(30000, 35000)
        self._llm = Engine(model, port=port)
        self._tokenizer = self._llm.get_tokenizer()
    else:
        self._llm, self._tokenizer = load_model(
            model, device=device, num_gpus=num_gpus
        )

VicunaReranker

Bases: BaseRanking

Implements RankVicuna, a zero-shot listwise document reranking method using Vicuna.

RankVicuna is a listwise ranking method that leverages open-source large language models (LLMs) for document reranking without requiring fine-tuning.

References

• Pradeep et al. (2023): RankVicuna: Zero-Shot Listwise Document Reranking with Open-Source LLMs. Paper

Attributes:

Name	Type	Description
method	str	The reranking method name.
model_name	str	The Vicuna model used for reranking.
window_size	int	The window size for listwise ranking.
_reranker	RankListwiseOSLLM	The RankVicuna reranking model instance.

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 health benefits of green tea?")
contexts = [
    Context(text="Green tea is rich in antioxidants that improve heart health.", id=0),
    Context(text="Drinking green tea may boost brain function and alertness.", id=1),
    Context(text="Excessive sugar intake increases the risk of diabetes.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize Vicuna reranker
model = Reranking(method='vicuna_reranker', model_name='rank_vicuna_7b_v1')
model.rank([document])

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

Source code in rankify/models/vicuna_reranker.py

class VicunaReranker(BaseRanking):
    """
    Implements **RankVicuna**, a **zero-shot listwise document reranking** method using **Vicuna**.



    RankVicuna is a **listwise** ranking method that leverages open-source **large language models (LLMs)** 
    for **document reranking** without requiring fine-tuning.

    References:
        - **Pradeep et al. (2023)**: *RankVicuna: Zero-Shot Listwise Document Reranking with Open-Source LLMs*.
          [Paper](https://arxiv.org/abs/2309.15088)

    Attributes:
        method (str): The reranking method name.
        model_name (str): The **Vicuna** model used for reranking.
        window_size (int): The **window size** for listwise ranking.
        _reranker (RankListwiseOSLLM): The **RankVicuna** reranking model instance.

    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 health benefits of green tea?")
        contexts = [
            Context(text="Green tea is rich in antioxidants that improve heart health.", id=0),
            Context(text="Drinking green tea may boost brain function and alertness.", id=1),
            Context(text="Excessive sugar intake increases the risk of diabetes.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize Vicuna reranker
        model = Reranking(method='vicuna_reranker', model_name='rank_vicuna_7b_v1')
        model.rank([document])

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

    def __init__(
        self,
        method: str = None,
        model_name: str = "castorini/rank_vicuna_7b_v1", **kwargs
    ):
        """
        Initializes **RankVicuna** for listwise document reranking.

        Args:
            method (str, optional): The reranking method name.
            model_name (str, optional): The **Vicuna** model used for ranking.
                Defaults to `"castorini/rank_vicuna_7b_v1"`.
            **kwargs: Additional parameters:
                - `context_size` (int, default=4096): Maximum context size for the model.
                - `num_few_shot_examples` (int, default=0): Number of few-shot examples.
                - `device` (str, default="cuda"): Computation device (`"cpu"` or `"cuda"`).
                - `num_gpus` (int, default=1): Number of GPUs to use.
                - `variable_passages` (bool, default=False): Whether to handle variable passage lengths.
                - `window_size` (int, default=20): Sliding window size for ranking.
                - `system_message` (str, optional): System message for the model.
        """
        context_size: int =  kwargs.get("context_size", 4096)
        num_few_shot_examples: int = kwargs.get("num_few_shot_examples", 0) 
        device: str = kwargs.get("device", "cuda") 
        num_gpus: int = kwargs.get("num_gpus", 1) 
        variable_passages: bool = kwargs.get("variable_passages", False) 
        window_size: int = kwargs.get("window_size", 20)
        system_message: str = kwargs.get("system_message", None)

        self.window_size = window_size
        self._reranker = RankListwiseOSLLM(
            model=model_name,
            context_size=context_size,
            num_few_shot_examples=num_few_shot_examples,
            device=device,
            num_gpus=num_gpus,
            variable_passages=variable_passages,
            window_size=window_size,
            system_message=system_message,
        )

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks a list of **Document** instances using **RankVicuna**.

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

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

    def _rerank_document(self, document: Document):
        """
        Reranks a single document using **RankVicuna**.

        Args:
            document (Document): A **Document** instance to rerank.

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

        rank_start = 0
        rank_end = 100
        window_size = self.window_size
        step = 10
        shuffle_candidates = False
        logging = False

        # Rerank using the Vicuna model
        reranked_result = self._reranker.rerank_batch(
            requests=[request],
            rank_start=rank_start,
            rank_end=rank_end,
            window_size=window_size,
            step=step,
            shuffle_candidates=shuffle_candidates,
            logging=logging,
        )[0]

        contexts = copy.deepcopy(document.contexts)

        # Create a mapping from docid to the original context
        docid_to_context = {str(ctx.id): ctx for ctx in contexts}

        # Reorder contexts based on reranked_result
        reorder_contexts = []
        for candidate in reranked_result.candidates:
            d = docid_to_context[str(candidate["docid"])]
            d.score = candidate["score"]
            reorder_contexts.append(d)
        document.reorder_contexts = reorder_contexts
        return document

__init__(method=None, model_name='castorini/rank_vicuna_7b_v1', **kwargs)

Initializes RankVicuna for listwise document reranking.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	The Vicuna model used for ranking. Defaults to "castorini/rank_vicuna_7b_v1".	'castorini/rank_vicuna_7b_v1'
**kwargs		Additional parameters: - context_size (int, default=4096): Maximum context size for the model. - num_few_shot_examples (int, default=0): Number of few-shot examples. - device (str, default="cuda"): Computation device ("cpu" or "cuda"). - num_gpus (int, default=1): Number of GPUs to use. - variable_passages (bool, default=False): Whether to handle variable passage lengths. - window_size (int, default=20): Sliding window size for ranking. - system_message (str, optional): System message for the model.	{}

Source code in rankify/models/vicuna_reranker.py

def __init__(
    self,
    method: str = None,
    model_name: str = "castorini/rank_vicuna_7b_v1", **kwargs
):
    """
    Initializes **RankVicuna** for listwise document reranking.

    Args:
        method (str, optional): The reranking method name.
        model_name (str, optional): The **Vicuna** model used for ranking.
            Defaults to `"castorini/rank_vicuna_7b_v1"`.
        **kwargs: Additional parameters:
            - `context_size` (int, default=4096): Maximum context size for the model.
            - `num_few_shot_examples` (int, default=0): Number of few-shot examples.
            - `device` (str, default="cuda"): Computation device (`"cpu"` or `"cuda"`).
            - `num_gpus` (int, default=1): Number of GPUs to use.
            - `variable_passages` (bool, default=False): Whether to handle variable passage lengths.
            - `window_size` (int, default=20): Sliding window size for ranking.
            - `system_message` (str, optional): System message for the model.
    """
    context_size: int =  kwargs.get("context_size", 4096)
    num_few_shot_examples: int = kwargs.get("num_few_shot_examples", 0) 
    device: str = kwargs.get("device", "cuda") 
    num_gpus: int = kwargs.get("num_gpus", 1) 
    variable_passages: bool = kwargs.get("variable_passages", False) 
    window_size: int = kwargs.get("window_size", 20)
    system_message: str = kwargs.get("system_message", None)

    self.window_size = window_size
    self._reranker = RankListwiseOSLLM(
        model=model_name,
        context_size=context_size,
        num_few_shot_examples=num_few_shot_examples,
        device=device,
        num_gpus=num_gpus,
        variable_passages=variable_passages,
        window_size=window_size,
        system_message=system_message,
    )

rank(documents)

Reranks a list of Document instances using RankVicuna.

Parameters:

Name	Type	Description	Default
documents	List[Document]	A list of Document instances to rerank.	required

Returns:

Type	Description
List[Document]	List[Document]: The reranked list of Documents with updated reorder_contexts.

Source code in rankify/models/vicuna_reranker.py

def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks a list of **Document** instances using **RankVicuna**.

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

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