InRanker

rankify.models.inranker

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

Context

Represents a context with metadata such as score and title.

Attributes:

Name	Type	Description
score	float	The relevance score of the context.
has_answer	bool	Whether the context contains an answer.
id	int	The identifier of the context.
title	str	The title of the context.
text	str	The text of the context.

Source code in rankify/dataset/dataset.py

class Context:
    """
    Represents a context with metadata such as score and title.

    Attributes:
        score (float, optional): The relevance score of the context.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.
    """
    def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
        """
        Initializes a Context instance.

        Args:
            score (float, optional): The relevance score.
            has_answer (bool, optional): Whether the context contains an answer.
            id (int, optional): The identifier of the context.
            title (str, optional): The title of the context.
            text (str, optional): The text of the context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c)
            ```
        """
        self.score: Optional[float] = score
        self.has_answer: Optional[bool] = has_answer
        self.id: Optional[str] = id
        self.title: Optional[str] = title
        self.text: Optional[str] = text

    def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

        """
        Converts the Context instance to a dictionary.

        Args:
            save_text (bool): Whether to include text in the output dictionary.

        Returns:
            dict: The context data.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c.to_dict())
            ```
        """
        context_dict = {
            "score": float(self.score) if self.score is not None else None,
            "has_answer": self.has_answer,
            "id": self.id,
            }

        # Include 'text' only if save_text is True
        if save_text:
            context_dict["text"] = self.text
            context_dict["title"] =  self.title

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

        Returns:
            str: The formatted context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(str(c))
            ```
        """
        return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

__init__(score=None, has_answer=None, id=None, title=None, text=None)

Initializes a Context instance.

Parameters:

Name	Type	Description	Default
score	float	The relevance score.	None
has_answer	bool	Whether the context contains an answer.	None
id	int	The identifier of the context.	None
title	str	The title of the context.	None
text	str	The text of the context.	None

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c)

Source code in rankify/dataset/dataset.py

def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
    """
    Initializes a Context instance.

    Args:
        score (float, optional): The relevance score.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c)
        ```
    """
    self.score: Optional[float] = score
    self.has_answer: Optional[bool] = has_answer
    self.id: Optional[str] = id
    self.title: Optional[str] = title
    self.text: Optional[str] = text

to_dict(save_text=False)

Converts the Context instance to a dictionary.

Parameters:

Name	Type	Description	Default
save_text	bool	Whether to include text in the output dictionary.	False

Returns:

Name	Type	Description
dict	Dict[str, Optional[object]]	The context data.

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c.to_dict())

Source code in rankify/dataset/dataset.py

def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

    """
    Converts the Context instance to a dictionary.

    Args:
        save_text (bool): Whether to include text in the output dictionary.

    Returns:
        dict: The context data.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c.to_dict())
        ```
    """
    context_dict = {
        "score": float(self.score) if self.score is not None else None,
        "has_answer": self.has_answer,
        "id": self.id,
        }

    # Include 'text' only if save_text is True
    if save_text:
        context_dict["text"] = self.text
        context_dict["title"] =  self.title

    return context_dict

__str__()

Returns a string representation of the Context instance.

Returns:

Name	Type	Description
str	str	The formatted context.

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(str(c))

Source code in rankify/dataset/dataset.py

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

    Returns:
        str: The formatted context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(str(c))
        ```
    """
    return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

InRanker

Bases: BaseRanking

Implements InRanker, a distilled ranker for zero-shot information retrieval based on Seq2Seq models.

InRanker ranks passages by estimating their relevance probability using a pre-trained language model.
It tokenizes query-document pairs and predicts relevance via binary classification.

The model assigns softmax scores between "false" and "true" tokens,
where the "true" probability determines document relevance.

References

• Laitz et al. (2024): InRanker: Distilled Rankers for Zero-shot Information Retrieval.
  Paper

Attributes:

Name	Type	Description
method	str	The reranking method name.
model_name	str	Name of the pre-trained Seq2Seq model.
api_key	str	API key for authentication (if needed).
tokenizer	AutoTokenizer	The tokenizer for processing queries and documents.
model	AutoModelForSeq2SeqLM	The sequence-to-sequence model used for reranking.
precision	str	Model precision ("bf16", "fp16", or "fp32").
device	str	The device used for computation ("cuda" or "cpu").
batch_size	int	Number of query-document pairs processed in a batch.
max_length	int	Maximum length of tokenized sequences.
token_false_id	int	Token ID for "false" (indicating irrelevance).
token_true_id	int	Token ID for "true" (indicating relevance).

See Also

• Reranking: Main interface for reranking models, including InRanker.

Example

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

question = Question("What are the symptoms of COVID-19?")
contexts = [
    Context(text="Fever and cough are common symptoms of COVID-19.", id=0),
    Context(text="Headache is a rare symptom.", id=1),
    Context(text="Fatigue and loss of taste are also common.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize Reranking with InRanker
model = Reranking(method='inranker', model_name='inranker-small')
model.rank([document])

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

Notes

• Uses a Seq2Seq binary classification approach ("false" vs "true").
• Supports batch processing for efficiency.
• Works in zero-shot retrieval scenarios without fine-tuning.

Source code in rankify/models/inranker.py

class InRanker(BaseRanking):
    """
    Implements **InRanker**, a **distilled ranker** for **zero-shot information retrieval** based on **Seq2Seq models**.



    **InRanker** ranks passages by estimating their **relevance probability** using a **pre-trained language model**.  
    It **tokenizes** query-document pairs and **predicts relevance** via **binary classification**.

    The model assigns **softmax scores** between `"false"` and `"true"` tokens,  
    where the `"true"` probability determines **document relevance**.

    References:
        - **Laitz et al. (2024)**: *InRanker: Distilled Rankers for Zero-shot Information Retrieval*.  
          [Paper](https://arxiv.org/abs/2401.06910)

    Attributes:
        method (str, optional): The reranking method name.
        model_name (str): Name of the pre-trained **Seq2Seq model**.
        api_key (str, optional): API key for authentication (if needed).
        tokenizer (AutoTokenizer): The tokenizer for processing queries and documents.
        model (AutoModelForSeq2SeqLM): The **sequence-to-sequence model** used for reranking.
        precision (str): Model precision (`"bf16"`, `"fp16"`, or `"fp32"`).
        device (str): The device used for computation (`"cuda"` or `"cpu"`).
        batch_size (int): Number of query-document pairs processed in a batch.
        max_length (int): Maximum length of tokenized sequences.
        token_false_id (int): Token ID for `"false"` (indicating irrelevance).
        token_true_id (int): Token ID for `"true"` (indicating relevance).

    See Also:
        - `Reranking`: Main interface for reranking models, including `InRanker`.

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

        question = Question("What are the symptoms of COVID-19?")
        contexts = [
            Context(text="Fever and cough are common symptoms of COVID-19.", id=0),
            Context(text="Headache is a rare symptom.", id=1),
            Context(text="Fatigue and loss of taste are also common.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize Reranking with InRanker
        model = Reranking(method='inranker', model_name='inranker-small')
        model.rank([document])

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

    Notes:
        - Uses a **Seq2Seq binary classification** approach (`"false"` vs `"true"`).
        - Supports **batch processing** for efficiency.
        - Works in **zero-shot retrieval** scenarios without fine-tuning.
    """
    def __init__(self, method: str= None, model_name: str= None, api_key: str= None, **kwargs) ->None:
        """
        Initializes **InRanker** for **zero-shot document reranking**.

        Args:
            method (str, optional): The reranking method name.
            model_name (str): Name of the pre-trained **Seq2Seq model**.
            api_key (str, optional): API key for authentication (if needed).
            **kwargs: Additional parameters, including:
                - `precision` (str): `"bf16"`, `"fp16"`, or `"fp32"` for model precision.
                - `device` (str): `"cuda"` or `"cpu"` (default: `"cuda"`).
                - `batch_size` (int): Number of query-document pairs per batch (default: `32`).
                - `max_length` (int): Maximum sequence length for tokenization (default: `512`).

        Example:
            ```python
            model = InRanker(method='inranker', model_name='inranker-small')
            ```
        """
        model_args = {}

        self.precision = kwargs.get("precision", "bf16")
        self.slient = kwargs.get("slient", True)
        self.batch_size = kwargs.get("batch_size", 32)
        self.device = kwargs.get("device" , "cuda")
        self.max_length = kwargs.get("max_length",512)
        if self.precision == "bf16":
            model_args["torch_dtype"] =  torch.bfloat16
        elif self.precision == "fp16":
            model_args["torch_dtype"] = torch.float16
        else:
            model_args["torch_dtype"] = torch.float32

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name,**model_args)
        self.model.to(self.device)
        self.model.eval()
        #token_false, token_true = ["▁false", "▁true"]
        token_false, token_true = ["▁false","▁true"]
        self.token_false_id = self.tokenizer.get_vocab()[token_false]
        self.token_true_id = self.tokenizer.get_vocab()[token_true]


    def rank(self, documents: list[Document])-> List[Document]:
        """
        Reranks documents using **InRanker's binary classification** approach.

        Each document's **contexts** are scored based on the probability of being **relevant** (`"true"` token).

        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"):
            scores =[]
            logits = []
            context_copy= copy.deepcopy(document.contexts)
            for batch in self._chunks(document.contexts,self.batch_size):


                queries_documents = [ 
                    f"Query: {document.question.question} Document: {context.text} Relevant:" for context in batch
                ]
                tokenized = self.tokenizer(queries_documents,
                                           padding=True,
                                           truncation="longest_first",
                                           return_tensors="pt",
                                           max_length=self.max_length).to(self.device)


                input_ids = tokenized["input_ids"].to(self.device)
                attention_mask = tokenized["attention_mask"].to(self.device)
                _ , batch_scores = self._greedy_decode(
                    model = self.model,
                    input_ids = input_ids,
                    length= 1,
                    attention_mask= attention_mask,
                    return_last_logits=True
                )

                batch_scores = batch_scores[:,[self.token_false_id, self.token_true_id]]
                logits.extend(batch_scores.tolist())

                batch_scores = torch.log_softmax(batch_scores,dim=-1)
                batch_scores = torch.exp(batch_scores[:,1])
                batch_scores = batch_scores.tolist()

                scores.extend(batch_scores)

            for score, context in zip(scores, context_copy):
                 context.score = score

            context_copy.sort(key=lambda x:x.score, reverse=True)
            document.reorder_contexts = context_copy

        return documents


    @torch.no_grad()
    def _greedy_decode(self,model,
        input_ids: torch.Tensor,
        length: int,
        attention_mask: torch.Tensor = None,
        return_last_logits: bool = True):
        """
        Performs **greedy decoding** to generate the next token logits.

        Args:
            model (AutoModelForSeq2SeqLM): The **pre-trained Seq2Seq model**.
            input_ids (torch.Tensor): The **input token IDs**.
            length (int): The number of decoding steps.
            attention_mask (torch.Tensor, optional): The **attention mask** for input sequences.
            return_last_logits (bool, optional): Whether to return the last token's **logits**.

        Returns:
            tuple: A tuple containing:
                - **decoded token IDs** (torch.Tensor)
                - **last-step logits** (torch.Tensor)

        Example:
            ```python
            decoded_ids, logits = model._greedy_decode(model, input_ids, length=1)
            ```
        """
        decode_ids = torch.full(
            (input_ids.size(0), 1),
            model.config.decoder_start_token_id,
            dtype=torch.long,
        ).to(input_ids.device)
        encoder_outputs = model.get_encoder()(input_ids, attention_mask=attention_mask)
        next_token_logits = None
        for _ in range(length):
            model_inputs = model.prepare_inputs_for_generation(
                decode_ids,
                encoder_outputs=encoder_outputs,
                past=None,
                attention_mask=attention_mask,
                use_cache=True,
            )
            outputs = model(**model_inputs)  # (batch_size, cur_len, vocab_size)
            next_token_logits = outputs[0][:, -1, :]  # (batch_size, vocab_size)
            decode_ids = torch.cat(
                [decode_ids, next_token_logits.max(1)[1].unsqueeze(-1)], dim=-1
            )
        if return_last_logits:
            return decode_ids, next_token_logits
        return decode_ids
    @staticmethod
    def _chunks(contexts: list[Context],batch_size: int):
        """
        Splits a **list of contexts** into **smaller batches**.

        Args:
            contexts (List[Context]): The list of **contexts** to split.
            batch_size (int): The **batch size**.

        Yields:
            List[Context]: A **batch** of contexts.

        Example:
            ```python
            for batch in model._chunks(contexts, batch_size=8):
                process_batch(batch)
            ```
        """
        for i in range(0, len(contexts), batch_size):
            yield contexts[i:i+batch_size]

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

Initializes InRanker for zero-shot document reranking.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	Name of the pre-trained Seq2Seq model.	None
api_key	str	API key for authentication (if needed).	None
**kwargs		Additional parameters, including: - precision (str): "bf16", "fp16", or "fp32" for model precision. - device (str): "cuda" or "cpu" (default: "cuda"). - batch_size (int): Number of query-document pairs per batch (default: 32). - max_length (int): Maximum sequence length for tokenization (default: 512).	{}

Example

model = InRanker(method='inranker', model_name='inranker-small')

Source code in rankify/models/inranker.py

def __init__(self, method: str= None, model_name: str= None, api_key: str= None, **kwargs) ->None:
    """
    Initializes **InRanker** for **zero-shot document reranking**.

    Args:
        method (str, optional): The reranking method name.
        model_name (str): Name of the pre-trained **Seq2Seq model**.
        api_key (str, optional): API key for authentication (if needed).
        **kwargs: Additional parameters, including:
            - `precision` (str): `"bf16"`, `"fp16"`, or `"fp32"` for model precision.
            - `device` (str): `"cuda"` or `"cpu"` (default: `"cuda"`).
            - `batch_size` (int): Number of query-document pairs per batch (default: `32`).
            - `max_length` (int): Maximum sequence length for tokenization (default: `512`).

    Example:
        ```python
        model = InRanker(method='inranker', model_name='inranker-small')
        ```
    """
    model_args = {}

    self.precision = kwargs.get("precision", "bf16")
    self.slient = kwargs.get("slient", True)
    self.batch_size = kwargs.get("batch_size", 32)
    self.device = kwargs.get("device" , "cuda")
    self.max_length = kwargs.get("max_length",512)
    if self.precision == "bf16":
        model_args["torch_dtype"] =  torch.bfloat16
    elif self.precision == "fp16":
        model_args["torch_dtype"] = torch.float16
    else:
        model_args["torch_dtype"] = torch.float32

    self.tokenizer = AutoTokenizer.from_pretrained(model_name)
    self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name,**model_args)
    self.model.to(self.device)
    self.model.eval()
    #token_false, token_true = ["▁false", "▁true"]
    token_false, token_true = ["▁false","▁true"]
    self.token_false_id = self.tokenizer.get_vocab()[token_false]
    self.token_true_id = self.tokenizer.get_vocab()[token_true]

rank(documents)

Reranks documents using InRanker's binary classification approach.

Each document's contexts are scored based on the probability of being relevant ("true" token).

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 Document instances with updated reorder_contexts.

Example

reranked_docs = model.rank(documents)

Source code in rankify/models/inranker.py

def rank(self, documents: list[Document])-> List[Document]:
    """
    Reranks documents using **InRanker's binary classification** approach.

    Each document's **contexts** are scored based on the probability of being **relevant** (`"true"` token).

    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"):
        scores =[]
        logits = []
        context_copy= copy.deepcopy(document.contexts)
        for batch in self._chunks(document.contexts,self.batch_size):


            queries_documents = [ 
                f"Query: {document.question.question} Document: {context.text} Relevant:" for context in batch
            ]
            tokenized = self.tokenizer(queries_documents,
                                       padding=True,
                                       truncation="longest_first",
                                       return_tensors="pt",
                                       max_length=self.max_length).to(self.device)


            input_ids = tokenized["input_ids"].to(self.device)
            attention_mask = tokenized["attention_mask"].to(self.device)
            _ , batch_scores = self._greedy_decode(
                model = self.model,
                input_ids = input_ids,
                length= 1,
                attention_mask= attention_mask,
                return_last_logits=True
            )

            batch_scores = batch_scores[:,[self.token_false_id, self.token_true_id]]
            logits.extend(batch_scores.tolist())

            batch_scores = torch.log_softmax(batch_scores,dim=-1)
            batch_scores = torch.exp(batch_scores[:,1])
            batch_scores = batch_scores.tolist()

            scores.extend(batch_scores)

        for score, context in zip(scores, context_copy):
             context.score = score

        context_copy.sort(key=lambda x:x.score, reverse=True)
        document.reorder_contexts = context_copy

    return documents