ListT5 Reranker

rankify.models.listt5

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

FiDT5

Bases: T5ForConditionalGeneration

Source code in rankify/utils/models/fidt5.py

class FiDT5(transformers.T5ForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        self.wrap_encoder()

    def forward_(self, **kwargs):
        if 'input_ids' in kwargs:
            kwargs['input_ids'] = kwargs['input_ids'].view(kwargs['input_ids'].size(0), -1)
        if 'attention_mask' in kwargs:
            kwargs['attention_mask'] = kwargs['attention_mask'].view(kwargs['attention_mask'].size(0), -1)

        return super(FiDT5, self).forward(
            **kwargs
        )

    # We need to resize as B x (N * L) instead of (B * N) x L here
    # because the T5 forward method uses the input tensors to infer
    # dimensions used in the decoder.
    # EncoderWrapper resizes the inputs as (B * N) x L.
    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        if input_ids != None:
            # inputs might have already be resized in the generate method
            if input_ids.dim() == 3:
                self.encoder.n_passages = input_ids.size(1)
            input_ids = input_ids.view(input_ids.size(0), -1)
        if attention_mask != None:
            attention_mask = attention_mask.view(attention_mask.size(0), -1)
        if kwargs.get('return_dict') is None:
            kwargs['return_dict'] = False
        return super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            **kwargs
        )

    # We need to resize the inputs here, as the generate method expect 2D tensors
    def generate(self, input_ids, attention_mask, max_length, **kwargs):
        self.encoder.n_passages = input_ids.size(1)  # Adjust for the actual number of passages
        return super().generate(
            input_ids=input_ids.reshape(input_ids.size(0), -1),
            attention_mask=attention_mask.reshape(attention_mask.size(0), -1),
            max_length=max_length,
            **kwargs
        )

    def wrap_encoder(self, use_checkpoint=False):
        """
        Wrap T5 encoder to obtain a Fusion-in-Decoder model.
        """
        self.encoder = EncoderWrapper(self.encoder, use_checkpoint=use_checkpoint)

    def unwrap_encoder(self):
        """
        Unwrap Fusion-in-Decoder encoder, useful to load T5 weights.
        """
        self.encoder = self.encoder.encoder
        block = []
        for mod in self.encoder.block:
            block.append(mod.module)
        block = nn.ModuleList(block)
        self.encoder.block = block

    def load_t5(self, state_dict):
        self.unwrap_encoder()
        self.load_state_dict(state_dict)
        self.wrap_encoder()

    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.score_storage = None

    def get_crossattention_scores(self, context_mask):
        """
        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.
        """
        scores = []
        n_passages = context_mask.size(1)
        for mod in self.decoder.block:
            scores.append(mod.layer[1].EncDecAttention.score_storage)
        scores = torch.cat(scores, dim=2)
        bsz, n_heads, n_layers, _ = scores.size()
        # batch_size, n_head, n_layers, n_passages, text_maxlength
        scores = scores.view(bsz, n_heads, n_layers, n_passages, -1)
        scores = scores.masked_fill(~context_mask[:, None, None], 0.)
        scores = scores.sum(dim=[1, 2, 4])
        ntokens = context_mask.sum(dim=[2]) * n_layers * n_heads
        scores = scores/ntokens
        return scores

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

wrap_encoder(use_checkpoint=False)

Wrap T5 encoder to obtain a Fusion-in-Decoder model.

Source code in rankify/utils/models/fidt5.py

def wrap_encoder(self, use_checkpoint=False):
    """
    Wrap T5 encoder to obtain a Fusion-in-Decoder model.
    """
    self.encoder = EncoderWrapper(self.encoder, use_checkpoint=use_checkpoint)

unwrap_encoder()

Unwrap Fusion-in-Decoder encoder, useful to load T5 weights.

Source code in rankify/utils/models/fidt5.py

def unwrap_encoder(self):
    """
    Unwrap Fusion-in-Decoder encoder, useful to load T5 weights.
    """
    self.encoder = self.encoder.encoder
    block = []
    for mod in self.encoder.block:
        block.append(mod.module)
    block = nn.ModuleList(block)
    self.encoder.block = block

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/fidt5.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/fidt5.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.score_storage = None

get_crossattention_scores(context_mask)

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/fidt5.py

def get_crossattention_scores(self, context_mask):
    """
    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.
    """
    scores = []
    n_passages = context_mask.size(1)
    for mod in self.decoder.block:
        scores.append(mod.layer[1].EncDecAttention.score_storage)
    scores = torch.cat(scores, dim=2)
    bsz, n_heads, n_layers, _ = scores.size()
    # batch_size, n_head, n_layers, n_passages, text_maxlength
    scores = scores.view(bsz, n_heads, n_layers, n_passages, -1)
    scores = scores.masked_fill(~context_mask[:, None, None], 0.)
    scores = scores.sum(dim=[1, 2, 4])
    ntokens = context_mask.sum(dim=[2]) * n_layers * n_heads
    scores = scores/ntokens
    return scores

overwrite_forward_crossattention()

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

Source code in rankify/utils/models/fidt5.py

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

ListT5

Bases: BaseRanking

Implements ListT5, a listwise reranker based on Fusion-in-Decoder (FiD) models.

ListT5 ranks passages in chunks using a tournament-style sorting approach.
The model processes query-context pairs in batches, extracts ranking scores,
and produces reordered results.

References

• Yoon et al. (2024): ListT5: Listwise reranking with fusion-in-decoder improves zero-shot retrieval.
  Paper

Attributes:

Name	Type	Description
method	str	The reranking method name.
model_name	str	Name of the pre-trained ListT5 model.
api_key	str	API key for authentication (if needed).
use_gpu	bool	Whether to use GPU acceleration.
batch_size	int	Number of query-document pairs processed in a batch.
max_length	int	Maximum tokenized sequence length.
padding	str	Padding strategy for tokenization ("max_length" or "longest").
listwise_k	int	Number of contexts per chunk during listwise ranking.
out_k	int	Number of top contexts retained from each chunk.
model	FiDT5	The pre-trained Fusion-in-Decoder model for ranking.
tokenizer	T5Tokenizer	The tokenizer associated with ListT5.

See Also

• Fusion-in-Decoder: The FiD model architecture for improved document fusion.

Example

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

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

# Initialize Reranking with ListT5
model = Reranking(method='listt5', model_name='listt5-base')
model.rank([document])

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

Notes

• Uses listwise ranking instead of pointwise or pairwise methods.
• Processes passages in chunks to maintain model efficiency.
• Supports batch processing for large-scale reranking tasks.

Source code in rankify/models/listt5.py

class ListT5(BaseRanking):
    """
    Implements **ListT5**, a **listwise reranker** based on **Fusion-in-Decoder (FiD)** models.


    **ListT5** ranks passages in **chunks** using a **tournament-style sorting** approach.  
    The model **processes query-context pairs in batches**, extracts **ranking scores**,  
    and produces **reordered results**.

    References:
        - **Yoon et al. (2024)**: *ListT5: Listwise reranking with fusion-in-decoder improves zero-shot retrieval*.  
          [Paper](https://arxiv.org/abs/2402.15838)

    Attributes:
        method (str, optional): The reranking method name.
        model_name (str): Name of the **pre-trained ListT5 model**.
        api_key (str, optional): API key for authentication (if needed).
        use_gpu (bool): Whether to use **GPU acceleration**.
        batch_size (int): Number of query-document pairs processed in a batch.
        max_length (int): Maximum tokenized sequence length.
        padding (str): Padding strategy for tokenization (`"max_length"` or `"longest"`).
        listwise_k (int): Number of **contexts per chunk** during **listwise ranking**.
        out_k (int): Number of **top contexts retained** from each chunk.
        model (FiDT5): The **pre-trained Fusion-in-Decoder model** for ranking.
        tokenizer (T5Tokenizer): The tokenizer associated with **ListT5**.

    See Also:
        - `Fusion-in-Decoder`: The FiD model architecture for improved document fusion.

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

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

        # Initialize Reranking with ListT5
        model = Reranking(method='listt5', model_name='listt5-base')
        model.rank([document])

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

    Notes:
        - Uses **listwise ranking** instead of **pointwise** or **pairwise** methods.
        - Processes passages in **chunks** to maintain model efficiency.
        - Supports **batch processing** for large-scale reranking tasks.
    """

    def __init__(self, method: str = None, model_name: str = None, api_key: str = None,  **kwargs):
        """
        Initializes **ListT5** for **listwise document reranking**.

        Args:
            method (str, optional): The reranking method name.
            model_name (str): Name of the **pre-trained ListT5 model**.
            api_key (str, optional): API key for authentication (if needed).
            **kwargs: Additional parameters, including:
                - `batch_size` (int): Number of query-document pairs per batch (default: `20`).
                - `max_length` (int): Maximum sequence length for tokenization (default: `512`).
                - `padding` (str): Padding strategy (`"max_length"` or `"longest"`, default: `"max_length"`).
                - `listwise_k` (int): Number of **contexts per chunk** (default: `5`).
                - `out_k` (int): Number of **top contexts retained** per chunk (default: `2`).

        Example:
            ```python
            model = ListT5(method='listt5', model_name='listt5-base')
            ```
        """
        self.method = method
        self.model_name = model_name
        self.api_key = api_key
        self.use_gpu = torch.cuda.is_available()
        self.batch_size = kwargs.get("batch_size", 20)
        self.max_length = kwargs.get("max_length", 512)
        self.padding = kwargs.get("padding", "max_length")
        self.listwise_k = kwargs.get("listwise_k", 5)  # Max contexts per chunk
        self.out_k = kwargs.get("out_k", 2)  # Top contexts per chunk
        self.model, self.tokenizer = self._load_model()

    def _load_model(self):
        """
        Loads the **pre-trained ListT5 model** and tokenizer.

        Returns:
            Tuple[FiDT5, T5Tokenizer]: The **model** and its **tokenizer**.
        """

        tokenizer = T5Tokenizer.from_pretrained(self.model_name)
        model = FiDT5.from_pretrained(self.model_name)
        if self.use_gpu:
            model = model.cuda()
        model.eval()
        return model, tokenizer

    def rank(self, documents: list[Document]):
        """
        Reranks documents using **ListT5’s listwise tournament sorting approach**.

        Each document’s contexts are processed in **chunks**, ranked, and merged.

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

        Returns:
            List[Document]: The reranked list of `Document` instances with updated `reorder_contexts`.

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

    def _run_tournament_sort(self, document):
        """
        Applies **tournament-style ranking** on a document's contexts.

        Args:
            document (Document): The document whose contexts will be reranked.

        Returns:
            List[Context]: The reordered list of contexts.
        """
        question = document.question.question
        contexts = [ctx.text for ctx in document.contexts]
        scores = [ctx.score for ctx in document.contexts]
        full_list_idx = list(range(len(contexts)))
        saved_top_indices = []

        # Chunk processing based on listwise_k
        chunks = [full_list_idx[i:i + self.listwise_k] for i in range(0, len(full_list_idx), self.listwise_k)]
        for chunk in chunks:
            if len(chunk) < self.listwise_k:
                chunk = self._pad_chunk(chunk, full_list_idx)  # Ensure chunk has listwise_k items
            top_indices = self._get_out_k(question, contexts, chunk)
            saved_top_indices.extend(top_indices[:self.out_k])

        # Deduplicate and fill remaining slots
        saved_top_indices = self._remove_duplicates(saved_top_indices)
        if len(saved_top_indices) < len(full_list_idx):
            remaining_indices = [idx for idx in full_list_idx if idx not in saved_top_indices]
            saved_top_indices.extend(remaining_indices)
        reranked_contexts= []
        copy_document= copy.deepcopy(document)
        for idx , score in zip(saved_top_indices,scores):
            copy_document.contexts[idx].score= score 
            reranked_contexts.append(copy_document.contexts[idx])
        return reranked_contexts #[document.contexts[idx] for idx in saved_top_indices]

    def _get_out_k(self, question, contexts, chunk):
        """
        Extracts **top-k relevant passages** from a chunk.

        Args:
            question (str): The query text.
            contexts (List[str]): List of passage texts.
            chunk (List[int]): Indices of passages in the chunk.

        Returns:
            List[int]: Indices of **top-ranked** passages.
        """
        chunk_texts = self._prepare_texts(question, [contexts[i] for i in chunk])
        input_tensors = self._prepare_inputs(chunk_texts)
        output = self._run_inference(input_tensors)
        ranked_indices = self._get_rel_index(output, k=self.listwise_k)

        # Map relative indices to original chunk indices
        return [chunk[i - 1] for i in ranked_indices]

    def _prepare_texts(self, question, contexts):
        """
        Formats input texts for **ListT5 processing**.

        Args:
            question (str): The query text.
            contexts (List[str]): List of passage texts.

        Returns:
            List[str]: Formatted query-context strings.
        """
        return [f"Query: {question}, Index: {i + 1}, Context: {ctx}" for i, ctx in enumerate(contexts)]

    def _prepare_inputs(self, texts):
        """
        Tokenizes and prepares input tensors for the model.

        Args:
            texts (List[str]): The input texts for each context.

        Returns:
            Dict[str, torch.Tensor]: A dictionary with `input_ids` and `attention_mask`.
        """
        encoding = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=self.padding,
            max_length=self.max_length,
            truncation=True
        )
        input_tensors = {
            "input_ids": encoding["input_ids"].unsqueeze(0),
            "attention_mask": encoding["attention_mask"].unsqueeze(0)
        }
        if self.use_gpu:
            input_tensors = {key: tensor.cuda() for key, tensor in input_tensors.items()}
        return input_tensors

    def _run_inference(self, input_tensors):
        """
        Runs inference using **ListT5’s decoder**.

        Args:
            input_tensors (Dict[str, torch.Tensor]): Tokenized inputs for the model.

        Returns:
            torch.Tensor: Output tensor containing rankings.
        """
        output = self.model.generate(
            **input_tensors,
            max_length=self.listwise_k + 2,
            return_dict_in_generate=True,
            output_scores=True
        )
        return output

    def _get_rel_index(self, output, k):
        """
        Extracts **top-k ranked indices** from the model output.

        Args:
            output (torch.Tensor): Model output tensor.
            k (int): Number of top indices to extract.

        Returns:
            List[int]: Extracted ranked indices.
        """
        generated_text = self.tokenizer.batch_decode(output.sequences, skip_special_tokens=True)[0]
        return [int(idx) for idx in generated_text.split()[-k:]]

    def _pad_chunk(self, chunk, full_list):
        """
        Ensures a chunk has `listwise_k` items by **padding with existing elements**.

        Args:
            chunk (List[int]): The chunk to pad.
            full_list (List[int]): Full list of passage indices.

        Returns:
            List[int]: Padded chunk of indices.
        """
        padded_chunk = chunk[:]
        i = 0
        while len(padded_chunk) < self.listwise_k:
            padded_chunk.append(full_list[i % len(full_list)])
            i += 1
        return padded_chunk

    def _remove_duplicates(self, indices):
        """
        Removes **duplicate indices** while maintaining order.

        Args:
            indices (List[int]): List of ranked indices.

        Returns:
            List[int]: Deduplicated list of indices.
        """
        seen = set()
        unique_indices = []
        for idx in indices:
            if idx not in seen:
                seen.add(idx)
                unique_indices.append(idx)
        return unique_indices

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

Initializes ListT5 for listwise document reranking.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	Name of the pre-trained ListT5 model.	None
api_key	str	API key for authentication (if needed).	None
**kwargs		Additional parameters, including: - batch_size (int): Number of query-document pairs per batch (default: 20). - max_length (int): Maximum sequence length for tokenization (default: 512). - padding (str): Padding strategy ("max_length" or "longest", default: "max_length"). - listwise_k (int): Number of contexts per chunk (default: 5). - out_k (int): Number of top contexts retained per chunk (default: 2).	{}

Example

model = ListT5(method='listt5', model_name='listt5-base')

Source code in rankify/models/listt5.py

def __init__(self, method: str = None, model_name: str = None, api_key: str = None,  **kwargs):
    """
    Initializes **ListT5** for **listwise document reranking**.

    Args:
        method (str, optional): The reranking method name.
        model_name (str): Name of the **pre-trained ListT5 model**.
        api_key (str, optional): API key for authentication (if needed).
        **kwargs: Additional parameters, including:
            - `batch_size` (int): Number of query-document pairs per batch (default: `20`).
            - `max_length` (int): Maximum sequence length for tokenization (default: `512`).
            - `padding` (str): Padding strategy (`"max_length"` or `"longest"`, default: `"max_length"`).
            - `listwise_k` (int): Number of **contexts per chunk** (default: `5`).
            - `out_k` (int): Number of **top contexts retained** per chunk (default: `2`).

    Example:
        ```python
        model = ListT5(method='listt5', model_name='listt5-base')
        ```
    """
    self.method = method
    self.model_name = model_name
    self.api_key = api_key
    self.use_gpu = torch.cuda.is_available()
    self.batch_size = kwargs.get("batch_size", 20)
    self.max_length = kwargs.get("max_length", 512)
    self.padding = kwargs.get("padding", "max_length")
    self.listwise_k = kwargs.get("listwise_k", 5)  # Max contexts per chunk
    self.out_k = kwargs.get("out_k", 2)  # Top contexts per chunk
    self.model, self.tokenizer = self._load_model()

rank(documents)

Reranks documents using ListT5’s listwise tournament sorting approach.

Each document’s contexts are processed in chunks, ranked, and merged.

Parameters:

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

Returns:

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

Example

reranked_docs = model.rank(documents)

Source code in rankify/models/listt5.py

def rank(self, documents: list[Document]):
    """
    Reranks documents using **ListT5’s listwise tournament sorting approach**.

    Each document’s contexts are processed in **chunks**, ranked, and merged.

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

    Returns:
        List[Document]: The reranked list of `Document` instances with updated `reorder_contexts`.

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