Sentence Transformer Reranker

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

SentenceTransformerReranker

Bases: BaseRanking

Implements SentenceTransformerReranker, a dense retrieval reranking approach using Sentence Transformers for encoding queries and passages.

This method leverages dual encoders, distilled self-attention, and corpus-aware pre-training to enhance retrieval quality. It supports Sentence-BERT (SBERT) embeddings, MiniLM, and Sentence-T5 for ranking.

References

• Ni et al. (2021): Large Dual Encoders are Generalizable Retrievers. Paper
• Wang et al. (2020): MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression. Paper
• Ni et al. (2021): Sentence-T5: Scalable Sentence Encoders from Pre-trained Text-to-Text Models. Paper
• Gao & Callan (2021): Unsupervised Corpus Aware Language Model Pre-training for Dense Passage Retrieval. Paper
• Reimers (2019): Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks. Paper

Attributes:

Name	Type	Description
method	str	The name of the reranking method.
model_name	str	The name or path of the Sentence Transformer model.
device	str	The device (CPU/GPU) used for inference.
query_prefix	str	Prefix to prepend to query texts.
document_prefix	str	Prefix to prepend to document texts.
normalize_embeddings	bool	Whether to normalize embeddings before computing similarity.
model	SentenceTransformer	The Sentence Transformer model used for reranking.

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 benefits of machine learning?")
contexts = [
    Context(text="Machine learning improves decision-making and automation.", id=0),
    Context(text="Quantum computing explores new paradigms in computation.", id=1),
    Context(text="Deep learning allows neural networks to learn from large data.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize SentenceTransformerReranker
model = Reranking(method='sentence_transformer_reranker', model_name='all-MiniLM-L6-v2')
model.rank([document])

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

Source code in rankify/models/sentence_transformer_reranker.py

class SentenceTransformerReranker(BaseRanking):
    """
    Implements **SentenceTransformerReranker**, 
    a **dense retrieval reranking approach** using **Sentence Transformers** for encoding queries and passages.

    This method **leverages dual encoders**, **distilled self-attention**, and **corpus-aware pre-training**
    to enhance retrieval quality. It supports **Sentence-BERT (SBERT) embeddings**, **MiniLM**, and **Sentence-T5** for ranking.

    References:
        - **Ni et al. (2021)**: *Large Dual Encoders are Generalizable Retrievers*. [Paper](https://arxiv.org/abs/2112.07899)
        - **Wang et al. (2020)**: *MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression*. [Paper](https://arxiv.org/abs/2002.10957)
        - **Ni et al. (2021)**: *Sentence-T5: Scalable Sentence Encoders from Pre-trained Text-to-Text Models*. [Paper](https://arxiv.org/abs/2108.08877)
        - **Gao & Callan (2021)**: *Unsupervised Corpus Aware Language Model Pre-training for Dense Passage Retrieval*.  [Paper](https://arxiv.org/abs/2108.05540)
        - **Reimers (2019)**: *Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks*.  [Paper](https://arxiv.org/abs/1908.10084)

    Attributes:
        method (str): The name of the reranking method.
        model_name (str): The name or path of the **Sentence Transformer** model.
        device (str): The device (CPU/GPU) used for inference.
        query_prefix (str): Prefix to prepend to query texts.
        document_prefix (str): Prefix to prepend to document texts.
        normalize_embeddings (bool): Whether to **normalize embeddings** before computing similarity.
        model (SentenceTransformer): The **Sentence Transformer** model used for reranking.

    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 benefits of machine learning?")
        contexts = [
            Context(text="Machine learning improves decision-making and automation.", id=0),
            Context(text="Quantum computing explores new paradigms in computation.", id=1),
            Context(text="Deep learning allows neural networks to learn from large data.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize SentenceTransformerReranker
        model = Reranking(method='sentence_transformer_reranker', model_name='all-MiniLM-L6-v2')
        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 = "all-MiniLM-L6-v2",
        **kwargs
    ):
        """
        Initializes **Sentence Transformer Reranker** for reranking tasks.

        Args:
            method (str): The name of the reranking method.
            model_name (str): The name or path to the **Sentence Transformer** model.
            **kwargs: Additional parameters:
                - device (str, optional): The computation device (`"auto"`, `"cuda"`, `"cpu"`). Default is `"auto"`.
                - query_prefix (str, optional): Prefix for query texts.
                - document_prefix (str, optional): Prefix for document texts.
                - use_fp16 (bool, optional): Whether to use **FP16** inference (default: `True`).
                - normalize_embeddings (bool, optional): Whether to **normalize embeddings** (default: `True`).
                - max_seq_length (int, optional): Maximum tokenization length (default: `512`).
        """
        super().__init__(method)
        self.device = self._detect_device(kwargs.get("device", "auto"))
        if model_name =="other":
            model_name =  kwargs.get("name", "all-MiniLM-L6-v2")
        self.model = SentenceTransformer(
            model_name, device=self.device, trust_remote_code=True
        )
        if kwargs.get("use_fp16", True) and "cuda" in self.device:
            self.model.half()

        self.query_prefix = kwargs.get("query_prefix", "")
        self.document_prefix = kwargs.get("document_prefix", "")
        self.normalize_embeddings = kwargs.get("normalize_embeddings", True)
        self.model.max_seq_length = kwargs.get("max_seq_length", 512)

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks a list of **Document** instances based on **Sentence Transformer** similarity.

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

        Returns:
            List[Document]: The 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) -> Document:
        """
        Reranks a single document's contexts using the **Sentence Transformer** model.

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

        Returns:
            Document: The reranked **Document** with updated `reorder_contexts`.
        """
        query = document.question.question  # Extract query text
        contexts = copy.deepcopy(document.contexts)

        # Extract context texts
        context_texts = [ctx.text for ctx in contexts]

        # Compute relevance scores
        scores = self._rerank(query, context_texts)

        # Map scores back to contexts and update the scores
        for ctx, score in zip(contexts, scores):
            ctx.score = score

        # Map scores back to contexts
        scored_contexts = list(zip(contexts, scores))
        scored_contexts.sort(key=lambda x: x[1], reverse=True)

        # Reorder contexts in the document
        document.reorder_contexts = [ctx for ctx, _ in scored_contexts]
        return document

    def _rerank(self, query: str, documents: List[str]) -> List[float]:
        """
        Computes similarity scores between a query and documents.

        Args:
            query (str): The **query** text.
            documents (List[str]): A list of **document texts**.

        Returns:
            List[float]: Similarity scores for each document.
        """
        # Add prefixes to query and documents
        query_text = self.query_prefix + query
        documents = [self.document_prefix + doc for doc in documents]

        # Encode query and documents
        embeddings = self.model.encode(
            [query_text] + documents,
            normalize_embeddings=self.normalize_embeddings,
            convert_to_tensor=True,
            show_progress_bar=False  # Disable tqdm progress bar
        )
        query_emb = embeddings[0]
        document_embs = embeddings[1:]

        # Compute cosine similarity scores
        scores = torch.nn.functional.cosine_similarity(
            query_emb.unsqueeze(0), document_embs
        ).tolist()
        return scores

    @staticmethod
    def _detect_device(device: str) -> str:
        """
        Detects the appropriate device for computation.

        Args:
            device (str): Desired device (`"auto"`, `"cuda"`, or `"cpu"`).

        Returns:
            str: The detected device.
        """
        if device == "auto":
            return "cuda" if torch.cuda.is_available() else "cpu"
        return device

__init__(method=None, model_name='all-MiniLM-L6-v2', **kwargs)

Initializes Sentence Transformer Reranker for reranking tasks.

Parameters:

Name	Type	Description	Default
method	str	The name of the reranking method.	None
model_name	str	The name or path to the Sentence Transformer model.	'all-MiniLM-L6-v2'
**kwargs		Additional parameters: - device (str, optional): The computation device ("auto", "cuda", "cpu"). Default is "auto". - query_prefix (str, optional): Prefix for query texts. - document_prefix (str, optional): Prefix for document texts. - use_fp16 (bool, optional): Whether to use FP16 inference (default: True). - normalize_embeddings (bool, optional): Whether to normalize embeddings (default: True). - max_seq_length (int, optional): Maximum tokenization length (default: 512).	{}

Source code in rankify/models/sentence_transformer_reranker.py

def __init__(
    self,
    method: str = None,
    model_name: str = "all-MiniLM-L6-v2",
    **kwargs
):
    """
    Initializes **Sentence Transformer Reranker** for reranking tasks.

    Args:
        method (str): The name of the reranking method.
        model_name (str): The name or path to the **Sentence Transformer** model.
        **kwargs: Additional parameters:
            - device (str, optional): The computation device (`"auto"`, `"cuda"`, `"cpu"`). Default is `"auto"`.
            - query_prefix (str, optional): Prefix for query texts.
            - document_prefix (str, optional): Prefix for document texts.
            - use_fp16 (bool, optional): Whether to use **FP16** inference (default: `True`).
            - normalize_embeddings (bool, optional): Whether to **normalize embeddings** (default: `True`).
            - max_seq_length (int, optional): Maximum tokenization length (default: `512`).
    """
    super().__init__(method)
    self.device = self._detect_device(kwargs.get("device", "auto"))
    if model_name =="other":
        model_name =  kwargs.get("name", "all-MiniLM-L6-v2")
    self.model = SentenceTransformer(
        model_name, device=self.device, trust_remote_code=True
    )
    if kwargs.get("use_fp16", True) and "cuda" in self.device:
        self.model.half()

    self.query_prefix = kwargs.get("query_prefix", "")
    self.document_prefix = kwargs.get("document_prefix", "")
    self.normalize_embeddings = kwargs.get("normalize_embeddings", True)
    self.model.max_seq_length = kwargs.get("max_seq_length", 512)

rank(documents)

Reranks a list of Document instances based on Sentence Transformer similarity.

Parameters:

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

Returns:

Type	Description
List[Document]	List[Document]: The documents with updated reorder_contexts.

Source code in rankify/models/sentence_transformer_reranker.py

def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks a list of **Document** instances based on **Sentence Transformer** similarity.

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

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