TwoLAR Reranker

rankify.models.twolar

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

Score

Score offers a collection of static methods for extracting specific scores from a tensor of logits, typically obtained as an output from models like T5 and flan-T5. These scores are computed based on predefined token IDs corresponding to 'true' and 'false' representations in the T5 and flan-T5 vocabularies.

Methods include: - extra_id: Extracts the logit corresponding to token ID 32089. - difference: Calculates the difference between the logits for the 'true' (token ID 1176) and 'false' (token ID 6136) tokens. - softmax: Computes the softmax scores for the 'true' (token ID 1176) and 'false' (token ID 6136) logits, and returns the softmax value for the 'true' token.

Note: - Token ID 1176 corresponds to the 'true' token in the T5 and flan-T5 vocabularies. - Token ID 6136 corresponds to the 'false' token in the T5 and flan-T5 vocabularies. - Token ID 32089 corresponds to the '' token in the T5 and flan-T5 vocabularies.

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

class Score:
    """
    Score offers a collection of static methods for extracting specific scores from a tensor of logits,
    typically obtained as an output from models like T5 and flan-T5. These scores are computed based on
    predefined token IDs corresponding to 'true' and 'false' representations in the T5 and flan-T5 vocabularies.

    Methods include:
    - extra_id: Extracts the logit corresponding to token ID 32089.
    - difference: Calculates the difference between the logits for the 'true' (token ID 1176) and 'false' (token ID 6136) tokens.
    - softmax: Computes the softmax scores for the 'true' (token ID 1176) and 'false' (token ID 6136) logits, and returns the softmax value for the 'true' token.

    Note:
    - Token ID 1176 corresponds to the 'true' token in the T5 and flan-T5 vocabularies.
    - Token ID 6136 corresponds to the 'false' token in the T5 and flan-T5 vocabularies.
    - Token ID 32089 corresponds to the '<extra_id_10>' token in the T5 and flan-T5 vocabularies.
    """

    @staticmethod
    def extra_id(logits):
        return logits[:, 32089]

    @staticmethod
    def difference(logits):
        true_logits = logits[:, 1176]
        false_logits = logits[:, 6136]
        return true_logits - false_logits

    @staticmethod
    def softmax(logits):
        true_logits = logits[:, 1176]
        false_logits = logits[:, 6136]
        scores = [torch.nn.functional.softmax(torch.stack([true_logit, false_logit]), dim=0)[0]
                 for true_logit, false_logit in zip(true_logits, false_logits)]
        return torch.stack(scores) 

TWOLAR

Bases: BaseRanking

Implements TWOLAR, a two-step LLM-augmented distillation method for passage reranking.

TWOLAR enhances passage ranking by using a two-step distillation approach. It first generates an LLM-augmented score and then refines it using a trained ranking model.

References

• Baldelli et al. (2024): TWOLAR: A TWO-step LLM-Augmented Distillation Method for Passage Reranking. Paper

Attributes:

Name	Type	Description
method	str	The reranking method name.
model_name	str	The name or path of the pre-trained TWOLAR model.
device	device	The computation device (CPU/GPU).
tokenizer	AutoTokenizer	The tokenizer for encoding queries and passages.
model	AutoModelForSeq2SeqLM	The TWOLAR reranking model.
batch_size	int	The batch size for inference.
max_length	int	The maximum sequence length for encoding passages.
score_strategy	str	The scoring strategy for ranking.

Examples:

Basic Usage:

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

# Define a query and contexts
question = Question("What are the effects of climate change?")
contexts = [
    Context(text="Climate change leads to rising sea levels and extreme weather.", id=0),
    Context(text="Renewable energy helps reduce carbon emissions.", id=1),
    Context(text="Deforestation accelerates global warming.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize TWOLAR reranker
model = Reranking(method='twolar', model_name='twolar-xl')
model.rank([document])

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

Source code in rankify/models/twolar.py

class TWOLAR(BaseRanking):
    """
    Implements **TWOLAR**, a **two-step LLM-augmented distillation method** for passage reranking.



    TWOLAR enhances passage ranking by using a **two-step distillation approach**. It first generates 
    an **LLM-augmented score** and then **refines it** using a trained ranking model.

    References:
        - **Baldelli et al. (2024)**: *TWOLAR: A TWO-step LLM-Augmented Distillation Method for Passage Reranking*.
          [Paper](https://arxiv.org/abs/2403.17759)

    Attributes:
        method (str): The **reranking method name**.
        model_name (str): The **name or path** of the pre-trained **TWOLAR** model.
        device (torch.device): The computation device (**CPU/GPU**).
        tokenizer (AutoTokenizer): The tokenizer for encoding **queries and passages**.
        model (AutoModelForSeq2SeqLM): The **TWOLAR** reranking model.
        batch_size (int): The **batch size** for inference.
        max_length (int): The **maximum sequence length** for encoding passages.
        score_strategy (str): The **scoring strategy** for ranking.

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

        # Define a query and contexts
        question = Question("What are the effects of climate change?")
        contexts = [
            Context(text="Climate change leads to rising sea levels and extreme weather.", id=0),
            Context(text="Renewable energy helps reduce carbon emissions.", id=1),
            Context(text="Deforestation accelerates global warming.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize TWOLAR reranker
        model = Reranking(method='twolar', model_name='twolar-xl')
        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 = None, api_key: str = None, **kwargs):
        """
        Initializes **TWOLAR** for reranking tasks.

        Args:
            method (str, optional): The **reranking method name**.
            model_name (str, optional): The **name of the pre-trained TWOLAR model**.
            api_key (str, optional): API key if required (**default: None**).
            **kwargs: Additional parameters such as `batch_size`, `max_length`, and `score_strategy`.
        """
        self.method = method
        self.model_name = model_name or "Dundalia/TWOLAR-xl"
        self.api_key = api_key
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.batch_size = kwargs.get("batch_size", 8)
        self.max_length = kwargs.get("max_length", 500)
        self.score_strategy = kwargs.get("score_strategy", "difference")

        # Load model and tokenizer
        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.get_score = getattr(Score, self.score_strategy)

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

        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"):
            query = document.question.question
            contexts = [ctx.text for ctx in document.contexts]

            # Encode the inputs
            input_features = self._prepare_inputs(query, contexts)

            # Perform inference
            logits = self._inference(input_features)

            # Compute scores
            scores = self.get_score(logits).tolist()

            # Assign scores to contexts and sort
            copy_context = copy.deepcopy(document.contexts)
            for i, context in enumerate(copy_context):
                context.score = scores[i]

            ranked_contexts = sorted(copy_context, key=lambda ctx: ctx.score, reverse=True)

            document.reorder_contexts = ranked_contexts
            #reranked_documents.append(document)

        return documents

    def _prepare_inputs(self, query: str, contexts: List[str]):
        """
        Prepares input features for the **TWOLAR model**.

        Args:
            query (str): The **query text**.
            contexts (List[str]): A list of **context passages**.

        Returns:
            Dict[str, torch.Tensor]: Tokenized input features for the model.
        """
        inputs = [f"Query: {query} Document: {ctx} Relevant: " for ctx in contexts]
        features = self.tokenizer(
            inputs,
            truncation=True,
            return_tensors="pt",
            max_length=self.max_length,
            padding=True,
        )
        features["input_ids"] = features.input_ids.to(self.device)
        features["attention_mask"] = features.attention_mask.to(self.device)
        features["decoder_input_ids"] = torch.full(
            (features.input_ids.size(0), 1),
            self.model.config.decoder_start_token_id,
            dtype=torch.long,
            device=self.device,
        )
        return features

    def _inference(self, features):
        """
        Runs inference on the **TWOLAR model** to compute passage scores.

        Args:
            features (Dict[str, torch.Tensor]): Tokenized input features.

        Returns:
            torch.Tensor: Logits from the model.
        """
        with torch.no_grad():
            output = self.model(
                input_ids=features["input_ids"],
                attention_mask=features["attention_mask"],
                decoder_input_ids=features["decoder_input_ids"],
            )
        return output.logits[:, 0, :]

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

Initializes TWOLAR for reranking tasks.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	The name of the pre-trained TWOLAR model.	None
api_key	str	API key if required (default: None).	None
**kwargs		Additional parameters such as batch_size, max_length, and score_strategy.	{}

Source code in rankify/models/twolar.py

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

    Args:
        method (str, optional): The **reranking method name**.
        model_name (str, optional): The **name of the pre-trained TWOLAR model**.
        api_key (str, optional): API key if required (**default: None**).
        **kwargs: Additional parameters such as `batch_size`, `max_length`, and `score_strategy`.
    """
    self.method = method
    self.model_name = model_name or "Dundalia/TWOLAR-xl"
    self.api_key = api_key
    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    self.batch_size = kwargs.get("batch_size", 8)
    self.max_length = kwargs.get("max_length", 500)
    self.score_strategy = kwargs.get("score_strategy", "difference")

    # Load model and tokenizer
    self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name).to(self.device)
    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    self.get_score = getattr(Score, self.score_strategy)

rank(documents)

Reranks a list of Document instances using TWOLAR.

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

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

    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"):
        query = document.question.question
        contexts = [ctx.text for ctx in document.contexts]

        # Encode the inputs
        input_features = self._prepare_inputs(query, contexts)

        # Perform inference
        logits = self._inference(input_features)

        # Compute scores
        scores = self.get_score(logits).tolist()

        # Assign scores to contexts and sort
        copy_context = copy.deepcopy(document.contexts)
        for i, context in enumerate(copy_context):
            context.score = scores[i]

        ranked_contexts = sorted(copy_context, key=lambda ctx: ctx.score, reverse=True)

        document.reorder_contexts = ranked_contexts
        #reranked_documents.append(document)

    return documents