MonoBERT

rankify.models.monobert

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

MonoBERT

Bases: BaseRanking

Implements MonoBERT Reranking, a BERT-based multi-stage ranking approach for improving document retrieval in information retrieval tasks.

MonoBERT re-ranks retrieved documents based on query-passage relevance scores using a pretrained BERT model.

References

• Nogueira et al. (2019): Multi-stage Document Ranking with BERT. Paper

Attributes:

Name	Type	Description
method	str	The name of the reranking method.
model_name	str	The name of the pre-trained MonoBERT model used for reranking.
device	device	The device (CPU/GPU) on which the model runs.
use_amp	bool	Whether to use Automatic Mixed Precision (AMP) for faster inference.
model	AutoModelForSequenceClassification	The pretrained MonoBERT model for reranking.
tokenizer	AutoTokenizer	The tokenizer for MonoBERT.

Example

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

# Define a query and contexts
question = Question("What are the health benefits of green tea?")
contexts = [
    Context(text="Green tea contains antioxidants that promote heart health.", id=0),
    Context(text="Excessive caffeine intake can cause insomnia.", id=1),
    Context(text="Green tea consumption is linked to improved metabolism.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize MonoBERT Reranker
model = Reranking(method='monobert', model_name='monobert-large')
model.rank([document])

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

Source code in rankify/models/monobert.py

class MonoBERT(BaseRanking):
    """
    Implements **MonoBERT Reranking**, a **BERT-based multi-stage ranking approach**
    for **improving document retrieval** in information retrieval tasks.


    MonoBERT **re-ranks retrieved documents** based on query-passage **relevance scores** using a **pretrained BERT model**.

    References:
        - **Nogueira et al. (2019)**: *Multi-stage Document Ranking with BERT*.
          [Paper](https://arxiv.org/abs/1910.14424)

    Attributes:
        method (str): The **name of the reranking method**.
        model_name (str): The **name of the pre-trained MonoBERT model** used for reranking.
        device (torch.device): The **device (CPU/GPU)** on which the model runs.
        use_amp (bool): Whether to use **Automatic Mixed Precision (AMP)** for **faster inference**.
        model (AutoModelForSequenceClassification): The **pretrained MonoBERT model** for reranking.
        tokenizer (AutoTokenizer): The **tokenizer** for MonoBERT.

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

        # Define a query and contexts
        question = Question("What are the health benefits of green tea?")
        contexts = [
            Context(text="Green tea contains antioxidants that promote heart health.", id=0),
            Context(text="Excessive caffeine intake can cause insomnia.", id=1),
            Context(text="Green tea consumption is linked to improved metabolism.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize MonoBERT Reranker
        model = Reranking(method='monobert', model_name='monobert-large')
        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 **MonoBERT** for reranking tasks.

        Args:
            method (str, optional): The **reranking method name**.
            model_name (str, optional): The **name of the pretrained MonoBERT model** 
                (default: `"castorini/monobert-large-msmarco"`).
            api_key (str, optional): **Not used**, but included for framework consistency.
            kwargs (dict): Additional parameters such as `use_amp` for **mixed precision inference**.
        """
        self.method = method
        self.model_name = model_name or "castorini/monobert-large-msmarco"
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.use_amp = kwargs.get("use_amp", False)
        self.model = self.get_model(self.model_name)
        self.tokenizer = self.get_tokenizer()

    @staticmethod
    def get_model(pretrained_model_name_or_path: str) -> AutoModelForSequenceClassification:
        """
        Loads the **MonoBERT model**.

        Args:
            pretrained_model_name_or_path (str): Path to the **pretrained MonoBERT model**.

        Returns:
            AutoModelForSequenceClassification: The **MonoBERT model**.
        """
        return AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path).to(
            torch.device("cuda" if torch.cuda.is_available() else "cpu")
        ).eval()

    @staticmethod
    def get_tokenizer(pretrained_model_name_or_path: str = "bert-large-uncased") -> AutoTokenizer:
        """
        Loads the **tokenizer** for MonoBERT.

        Args:
            pretrained_model_name_or_path (str): Path to the **pretrained tokenizer**.

        Returns:
            AutoTokenizer: The **MonoBERT tokenizer**.
        """
        return AutoTokenizer.from_pretrained(pretrained_model_name_or_path, use_fast=False)

    @torch.no_grad()
    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks each document's **contexts** using **MonoBERT** and updates `reorder_contexts`.

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

        Returns:
            List[Document]: Documents with updated **`reorder_contexts`** after reranking.
        """
        for document in tqdm(documents, desc="Reranking Documents"):
            query = document.question.question
            contexts = deepcopy(document.contexts)

            # Rescore contexts using MonoBERT
            for context in contexts:
                inputs = self.tokenizer.encode_plus(
                    query,
                    context.text,
                    max_length=512,
                    truncation=True,
                    return_token_type_ids=True,
                    return_tensors="pt"
                )
                with autocast(enabled=self.use_amp):
                    input_ids = inputs["input_ids"].to(self.device)
                    token_type_ids = inputs["token_type_ids"].to(self.device)
                    outputs = self.model(input_ids, token_type_ids=token_type_ids, return_dict=False)
                    logits = outputs[0]

                    # Handle binary and multi-class classification
                    if logits.size(1) > 1:
                        context.score = torch.nn.functional.log_softmax(logits, dim=1)[0, -1].item()
                    else:
                        context.score = logits.item()

            # Sort contexts by score in descending order
            ranked_contexts = sorted(contexts, key=lambda ctx: ctx.score, reverse=True)

            # Update `reorder_contexts` in the document
            document.reorder_contexts = ranked_contexts

        return documents

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

Initializes MonoBERT for reranking tasks.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	The name of the pretrained MonoBERT model (default: "castorini/monobert-large-msmarco").	None
api_key	str	Not used, but included for framework consistency.	None
kwargs	dict	Additional parameters such as use_amp for mixed precision inference.	{}

Source code in rankify/models/monobert.py

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

    Args:
        method (str, optional): The **reranking method name**.
        model_name (str, optional): The **name of the pretrained MonoBERT model** 
            (default: `"castorini/monobert-large-msmarco"`).
        api_key (str, optional): **Not used**, but included for framework consistency.
        kwargs (dict): Additional parameters such as `use_amp` for **mixed precision inference**.
    """
    self.method = method
    self.model_name = model_name or "castorini/monobert-large-msmarco"
    self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.use_amp = kwargs.get("use_amp", False)
    self.model = self.get_model(self.model_name)
    self.tokenizer = self.get_tokenizer()

get_model(pretrained_model_name_or_path) staticmethod

Loads the MonoBERT model.

Parameters:

Name	Type	Description	Default
pretrained_model_name_or_path	str	Path to the pretrained MonoBERT model.	required

Returns:

Name	Type	Description
AutoModelForSequenceClassification	AutoModelForSequenceClassification	The MonoBERT model.

Source code in rankify/models/monobert.py

@staticmethod
def get_model(pretrained_model_name_or_path: str) -> AutoModelForSequenceClassification:
    """
    Loads the **MonoBERT model**.

    Args:
        pretrained_model_name_or_path (str): Path to the **pretrained MonoBERT model**.

    Returns:
        AutoModelForSequenceClassification: The **MonoBERT model**.
    """
    return AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path).to(
        torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ).eval()

get_tokenizer(pretrained_model_name_or_path='bert-large-uncased') staticmethod

Loads the tokenizer for MonoBERT.

Parameters:

Name	Type	Description	Default
pretrained_model_name_or_path	str	Path to the pretrained tokenizer.	'bert-large-uncased'

Returns:

Name	Type	Description
AutoTokenizer	AutoTokenizer	The MonoBERT tokenizer.

Source code in rankify/models/monobert.py

@staticmethod
def get_tokenizer(pretrained_model_name_or_path: str = "bert-large-uncased") -> AutoTokenizer:
    """
    Loads the **tokenizer** for MonoBERT.

    Args:
        pretrained_model_name_or_path (str): Path to the **pretrained tokenizer**.

    Returns:
        AutoTokenizer: The **MonoBERT tokenizer**.
    """
    return AutoTokenizer.from_pretrained(pretrained_model_name_or_path, use_fast=False)

rank(documents)

Reranks each document's contexts using MonoBERT and updates reorder_contexts.

Parameters:

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

Returns:

Type	Description
List[Document]	List[Document]: Documents with updated reorder_contexts after reranking.

Source code in rankify/models/monobert.py

@torch.no_grad()
def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks each document's **contexts** using **MonoBERT** and updates `reorder_contexts`.

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

    Returns:
        List[Document]: Documents with updated **`reorder_contexts`** after reranking.
    """
    for document in tqdm(documents, desc="Reranking Documents"):
        query = document.question.question
        contexts = deepcopy(document.contexts)

        # Rescore contexts using MonoBERT
        for context in contexts:
            inputs = self.tokenizer.encode_plus(
                query,
                context.text,
                max_length=512,
                truncation=True,
                return_token_type_ids=True,
                return_tensors="pt"
            )
            with autocast(enabled=self.use_amp):
                input_ids = inputs["input_ids"].to(self.device)
                token_type_ids = inputs["token_type_ids"].to(self.device)
                outputs = self.model(input_ids, token_type_ids=token_type_ids, return_dict=False)
                logits = outputs[0]

                # Handle binary and multi-class classification
                if logits.size(1) > 1:
                    context.score = torch.nn.functional.log_softmax(logits, dim=1)[0, -1].item()
                else:
                    context.score = logits.item()

        # Sort contexts by score in descending order
        ranked_contexts = sorted(contexts, key=lambda ctx: ctx.score, reverse=True)

        # Update `reorder_contexts` in the document
        document.reorder_contexts = ranked_contexts

    return documents