In-Context Reranker

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

DynamicCacheWithQuery

Bases: DynamicCache

Cache class used for In-context RAG

Source code in rankify/utils/models/incontext_reranker/custom_cache.py

class DynamicCacheWithQuery(DynamicCache):
    '''
    Cache class used for In-context RAG
    '''
    def __init__(self, query_indices=[]) -> None:
        super().__init__()
        self._query_indices = query_indices # indices for query vectors to save
        self.query_cache = []

    def update(
        self,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

        Parameters:
            query_states (`torch.Tensor`):
                The new query states to cache.
            key_states (`torch.Tensor`):
                The new key states to cache.
            value_states (`torch.Tensor`):
                The new value states to cache.
            layer_idx (`int`):
                The index of the layer to cache the states for.
            cache_kwargs (`Dict[str, Any]`, `optional`):
                Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`.

        Return:
            A tuple containing the updated key and value states.
        """
        # Update the number of seen tokens
        if layer_idx == 0:
            self._seen_tokens += key_states.shape[-2]

        # Update the cache
        if len(self.key_cache) <= layer_idx:
            self.key_cache.append(key_states)
            self.value_cache.append(value_states)
        else:
            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)

        # IC-RAG
        if query_states is not None:
            if len(self.query_cache) <= layer_idx:
                self.query_cache.append(query_states)
            else:
                self.query_cache[layer_idx] = torch.cat([self.query_cache[layer_idx], query_states], dim=-2)
        return self.key_cache[layer_idx], self.value_cache[layer_idx]

    @classmethod
    def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "DynamicCache":
        """Converts a cache in the legacy cache format into an equivalent `DynamicCache`."""
        cache = cls()
        if past_key_values is not None:
            for layer_idx in range(len(past_key_values)):
                key_states, value_states = past_key_values[layer_idx]
                cache.update(None, key_states, value_states, layer_idx)
        return cache

update(query_states, key_states, value_states, layer_idx, cache_kwargs=None)

Updates the cache with the new key_states and value_states for the layer layer_idx.

Parameters:

Name	Type	Description	Default
query_states	`torch.Tensor`	The new query states to cache.	required
key_states	`torch.Tensor`	The new key states to cache.	required
value_states	`torch.Tensor`	The new value states to cache.	required
layer_idx	`int`	The index of the layer to cache the states for.	required
cache_kwargs	`Dict[str, Any]`, `optional`	Additional arguments for the cache subclass. No additional arguments are used in DynamicCache.	None

Return

A tuple containing the updated key and value states.

Source code in rankify/utils/models/incontext_reranker/custom_cache.py

def update(
    self,
    query_states: torch.Tensor,
    key_states: torch.Tensor,
    value_states: torch.Tensor,
    layer_idx: int,
    cache_kwargs: Optional[Dict[str, Any]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

    Parameters:
        query_states (`torch.Tensor`):
            The new query states to cache.
        key_states (`torch.Tensor`):
            The new key states to cache.
        value_states (`torch.Tensor`):
            The new value states to cache.
        layer_idx (`int`):
            The index of the layer to cache the states for.
        cache_kwargs (`Dict[str, Any]`, `optional`):
            Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`.

    Return:
        A tuple containing the updated key and value states.
    """
    # Update the number of seen tokens
    if layer_idx == 0:
        self._seen_tokens += key_states.shape[-2]

    # Update the cache
    if len(self.key_cache) <= layer_idx:
        self.key_cache.append(key_states)
        self.value_cache.append(value_states)
    else:
        self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
        self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)

    # IC-RAG
    if query_states is not None:
        if len(self.query_cache) <= layer_idx:
            self.query_cache.append(query_states)
        else:
            self.query_cache[layer_idx] = torch.cat([self.query_cache[layer_idx], query_states], dim=-2)
    return self.key_cache[layer_idx], self.value_cache[layer_idx]

from_legacy_cache(past_key_values=None) classmethod

Converts a cache in the legacy cache format into an equivalent DynamicCache.

Source code in rankify/utils/models/incontext_reranker/custom_cache.py

@classmethod
def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "DynamicCache":
    """Converts a cache in the legacy cache format into an equivalent `DynamicCache`."""
    cache = cls()
    if past_key_values is not None:
        for layer_idx in range(len(past_key_values)):
            key_states, value_states = past_key_values[layer_idx]
            cache.update(None, key_states, value_states, layer_idx)
    return cache

MistralForCausalLM

Bases: MistralPreTrainedModel

Source code in rankify/utils/models/incontext_reranker/custom_modeling_mistral.py

class MistralForCausalLM(MistralPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.model = MistralModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        r"""
        Args:
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Returns:

        Example:

        ```python
        >>> from transformers import AutoTokenizer, MistralForCausalLM

        >>> model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
        >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Ensure tensors are on the same device
            shift_labels = shift_labels.to(shift_logits.device)
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        use_cache=True,
        **kwargs,
    ):
        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        # Exception 1: when passing input_embeds, input_ids may be missing entries
        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        if past_key_values is not None:
            if inputs_embeds is not None:  # Exception 1
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -input_ids.shape[1] :]

                # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s  `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture.
                position_ids = position_ids.clone(memory_format=torch.contiguous_format)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and cache_position[0] == 0:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases

        model_inputs.update(
            {
                "position_ids": position_ids,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
            }
        )
        return model_inputs

forward(input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, cache_position=None)

Parameters:

Name	Type	Description	Default
labels	`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*	Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].	None

Returns:

Example:

>>> from transformers import AutoTokenizer, MistralForCausalLM

>>> model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
>>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."

Source code in rankify/utils/models/incontext_reranker/custom_modeling_mistral.py

@add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
def forward(
    self,
    input_ids: torch.LongTensor = None,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
    inputs_embeds: Optional[torch.FloatTensor] = None,
    labels: Optional[torch.LongTensor] = None,
    use_cache: Optional[bool] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
    cache_position: Optional[torch.LongTensor] = None,
) -> Union[Tuple, CausalLMOutputWithPast]:
    r"""
    Args:
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

    Returns:

    Example:

    ```python
    >>> from transformers import AutoTokenizer, MistralForCausalLM

    >>> model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
    >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

    >>> prompt = "Hey, are you conscious? Can you talk to me?"
    >>> inputs = tokenizer(prompt, return_tensors="pt")

    >>> # Generate
    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
    ```"""

    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
    outputs = self.model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
        cache_position=cache_position,
    )

    hidden_states = outputs[0]
    logits = self.lm_head(hidden_states)
    logits = logits.float()

    loss = None
    if labels is not None:
        # Shift so that tokens < n predict n
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        # Flatten the tokens
        shift_logits = shift_logits.view(-1, self.config.vocab_size)
        shift_labels = shift_labels.view(-1)
        # Ensure tensors are on the same device
        shift_labels = shift_labels.to(shift_logits.device)
        loss_fct = CrossEntropyLoss()
        loss = loss_fct(shift_logits, shift_labels)

    if not return_dict:
        output = (logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output

    return CausalLMOutputWithPast(
        loss=loss,
        logits=logits,
        past_key_values=outputs.past_key_values,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

LlamaForCausalLM

Bases: LlamaPreTrainedModel

Source code in rankify/utils/models/incontext_reranker/custom_modeling_llama.py

class LlamaForCausalLM(LlamaPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.model = LlamaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        r"""
        Args:
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Returns:

        Example:

        ```python
        >>> from transformers import AutoTokenizer, LlamaForCausalLM

        >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
        )

        hidden_states = outputs[0]
        if self.config.pretraining_tp > 1:
            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
            logits = torch.cat(logits, dim=-1)
        else:
            logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        use_cache=True,
        **kwargs,
    ):
        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        # Exception 1: when passing input_embeds, input_ids may be missing entries
        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        if past_key_values is not None:
            if inputs_embeds is not None:  # Exception 1
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -input_ids.shape[1] :]

                # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s  `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture.
                position_ids = position_ids.clone(memory_format=torch.contiguous_format)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and cache_position[0] == 0:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
            if inputs_embeds is not None:
                batch_size, sequence_length = inputs_embeds.shape
                device = inputs_embeds.device
            else:
                batch_size, sequence_length = input_ids.shape
                device = input_ids.device

            dtype = self.lm_head.weight.dtype
            min_dtype = torch.finfo(dtype).min

            attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
                attention_mask,
                sequence_length=sequence_length,
                target_length=past_key_values.get_max_length(),
                dtype=dtype,
                device=device,
                min_dtype=min_dtype,
                cache_position=cache_position,
                batch_size=batch_size,
            )

        model_inputs.update(
            {
                "position_ids": position_ids,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
            }
        )
        return model_inputs

forward(input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, cache_position=None)

Parameters:

Name	Type	Description	Default
labels	`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*	Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].	None

Returns:

Example:

>>> from transformers import AutoTokenizer, LlamaForCausalLM

>>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."

Source code in rankify/utils/models/incontext_reranker/custom_modeling_llama.py

@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
def forward(
    self,
    input_ids: torch.LongTensor = None,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
    inputs_embeds: Optional[torch.FloatTensor] = None,
    labels: Optional[torch.LongTensor] = None,
    use_cache: Optional[bool] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
    cache_position: Optional[torch.LongTensor] = None,
) -> Union[Tuple, CausalLMOutputWithPast]:
    r"""
    Args:
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

    Returns:

    Example:

    ```python
    >>> from transformers import AutoTokenizer, LlamaForCausalLM

    >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
    >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

    >>> prompt = "Hey, are you conscious? Can you talk to me?"
    >>> inputs = tokenizer(prompt, return_tensors="pt")

    >>> # Generate
    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
    ```"""
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
    outputs = self.model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
        cache_position=cache_position,
    )

    hidden_states = outputs[0]
    if self.config.pretraining_tp > 1:
        lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
        logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
        logits = torch.cat(logits, dim=-1)
    else:
        logits = self.lm_head(hidden_states)
    logits = logits.float()

    loss = None
    if labels is not None:
        # Shift so that tokens < n predict n
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        # Flatten the tokens
        loss_fct = CrossEntropyLoss()
        shift_logits = shift_logits.view(-1, self.config.vocab_size)
        shift_labels = shift_labels.view(-1)
        # Enable model parallelism
        shift_labels = shift_labels.to(shift_logits.device)
        loss = loss_fct(shift_logits, shift_labels)

    if not return_dict:
        output = (logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output

    return CausalLMOutputWithPast(
        loss=loss,
        logits=logits,
        past_key_values=outputs.past_key_values,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

InContextReranker

Bases: BaseRanking

Implements In-Context Reranking (ICR) using Large Language Models (LLMs), specifically Mistral and Llama.

ICR reranks passages using query-aware attention scores computed across multiple LLM layers.
It employs a sliding window strategy for efficient ranking and supports different retrieval types:

• QA Mode: Answers a given question using retrieved contexts.
• IE Mode: Extracts information relevant to a query.

The available scoring strategies include: - query_last: Default method, uses query token attention scores. - attention_sorting: Ranks contexts based on attention weights. - NA_only: Estimates the model's intrinsic bias. - NA_calibration_no_agg: Uses a neutral "N/A" query for bias correction. - masked_NA_calibration: Default ICR method with token-level masking.

References

• Chen et al. (2024) Attention in Large Language Models Yields Efficient Zero-Shot Re-Rankers.
  Paper

Attributes:

Name	Type	Description
method	str	The reranking method name.
model_name	str	Name of the LLM used for reranking.
tokenizer	AutoTokenizer	The tokenizer for text processing.
llm	Module	The pre-trained LLM model (Mistral or Llama).
prompt_template	str	Template for constructing query-context prompts.
scoring_strategy	str	The attention-based scoring strategy used for reranking.
retrieval_type	str	The retrieval type, either "QA" (question answering) or "IE" (information extraction).
sliding_window_size	int	Size of the sliding window for reranking.
sliding_window_stride	int	Step size for moving the sliding window.
reverse_doc_order	bool	Whether to reverse document order before reranking.

See Also

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

Example

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

question = Question("What are the symptoms of COVID-19?")
answers = Answer(["COVID-19 symptoms include fever and cough."])
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, answers=answers, contexts=contexts)

# Initialize Reranking with InContextReranker
model = Reranking(method='incontext_reranker', model_name='llamav3.1-8b')
model.rank([document])

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

Notes

• Requires a Mistral or Llama model.
• Uses sliding windows for efficient ranking.
• Supports attention-based scoring for zero-shot ranking.

Source code in rankify/models/incontext_reranker.py

class InContextReranker(BaseRanking):
    """
    Implements **In-Context Reranking (ICR)** using **Large Language Models (LLMs)**, specifically **Mistral** and **Llama**.


    **ICR** reranks passages using query-aware **attention scores** computed across multiple LLM layers.  
    It employs a **sliding window strategy** for efficient ranking and supports different retrieval types:

    - **QA Mode**: Answers a given question using retrieved contexts.
    - **IE Mode**: Extracts information relevant to a query.

    The available **scoring strategies** include:
    - **`query_last`**: Default method, uses query token attention scores.
    - **`attention_sorting`**: Ranks contexts based on attention weights.
    - **`NA_only`**: Estimates the model's intrinsic bias.
    - **`NA_calibration_no_agg`**: Uses a neutral `"N/A"` query for bias correction.
    - **`masked_NA_calibration`**: Default ICR method with token-level masking.

    References:
        - **Chen et al. (2024)** *Attention in Large Language Models Yields Efficient Zero-Shot Re-Rankers.*  
          [Paper](https://dl.acm.org/doi/10.1145/3626772.3657855)

    Attributes:
        method (str, optional): The reranking method name.
        model_name (str): Name of the LLM used for reranking.
        tokenizer (AutoTokenizer): The tokenizer for text processing.
        llm (torch.nn.Module): The pre-trained LLM model (Mistral or Llama).
        prompt_template (str): Template for constructing query-context prompts.
        scoring_strategy (str): The attention-based scoring strategy used for reranking.
        retrieval_type (str): The retrieval type, either `"QA"` (question answering) or `"IE"` (information extraction).
        sliding_window_size (int): Size of the **sliding window** for reranking.
        sliding_window_stride (int): Step size for moving the **sliding window**.
        reverse_doc_order (bool): Whether to reverse document order before reranking.

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

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

        question = Question("What are the symptoms of COVID-19?")
        answers = Answer(["COVID-19 symptoms include fever and cough."])
        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, answers=answers, contexts=contexts)

        # Initialize Reranking with InContextReranker
        model = Reranking(method='incontext_reranker', model_name='llamav3.1-8b')
        model.rank([document])

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

    Notes:
        - Requires a **Mistral** or **Llama** model.
        - Uses **sliding windows** for efficient ranking.
        - Supports **attention-based scoring** for **zero-shot ranking**.
    """
    def __init__(self, method=None, model_name=None, **kwargs):
        """
        Initializes the In-Context Reranker.

        Args:
            method (str, optional): The reranking method name.
            model_name (str): Name of the LLM model (e.g., `"meta-llama/Meta-Llama-3.1-8B-Instruct"`).
            **kwargs: Additional customization parameters, including:
                - `scoring_strategy` (str): Strategy for computing attention scores (`"query_last"`, `"attention_sorting"`, etc.).
                - `retrieval_type` (str): Type of retrieval (`"QA"` or `"IE"`).
                - `sliding_window_size` (int): Number of documents considered per window.

        Example:
            ```python
            model = InContextReranker(method='incontext_reranker', model_name='llamav3.1-8b')
            ```
        """
        # Setup the base LLM
        self._base_llm_name =model_name
        tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
        self.tokenizer = tokenizer
        print(f"initialized tokenizer for [{model_name}]")

        if any([x in model_name.lower() for x in ['mistralai/mistral', ]]):
            BaseLLMClass = MistralForCausalLM
        elif any([x in model_name.lower() for x in ['llama']]):
            BaseLLMClass = LlamaForCausalLM
        else:
            print(f"Warning: The model family for [{model_name}] is not supported by InContextRAGModel!")
            raise NotImplementedError

        prompt_template = kwargs.get("prompt_template", "instruct")

        prompt_template, prompt_prefix, prompt_suffix, = self._setup_llm_prompts(prompt_template, model_name)

        self.use_fa2 = kwargs.get("use_fa2", True)
        if self.use_fa2:
            _attn_implementation = "flash_attention_2"
        else:
            _attn_implementation = "eager"

        llm = BaseLLMClass.from_pretrained(
                model_name, 
                torch_dtype=torch.float16, 
                attn_implementation=_attn_implementation,
                device_map='auto'
            )
        self.llm = llm
        self.llm.config.pad_token_id = self.llm.config.eos_token_id

        # Setup prompts for ICR
        assert prompt_template in ['instruct', 'simple', 'simple_instruct'], "Invalid prompt template!"

        self.prompt_template = prompt_template
        self.prompt_prefix = prompt_prefix


        self.prompt_suffix = prompt_suffix
        self.scoring_strategy  = kwargs.get("scoring_strategy", "query_last")
        retrieval_type = kwargs.get("retrieval_type", "QA")

        if retrieval_type == 'QA':
            print('[ICR is using QA prompt type]')
            self.retrieval_instruction = ' Here are some paragraphs:'
            self.retrieval_instruction_late = 'Please answer the following question based on the information in the paragraphs above.'
        elif retrieval_type == 'IE':
            print('[ICR is using IE prompt type]')
            self.retrieval_instruction = ' Here are some paragraphs:'
            self.retrieval_instruction_late = 'Please find information that are relevant to the following query in the paragraphs above.'
        else:
            raise NotImplementedError('Invalid retrieval type! Should be one of [QA, IE]')

        assert self.scoring_strategy in ['query_last', 'attention_sorting', 'NA_only', 'NA_calibration_no_agg', 'masked_NA_calibration'], "Invalid scoring strategy!"

        self._use_fa2 = self.use_fa2
        if self.use_fa2:
            print('Using FA2 for retrieval score computation.')
        else:
            print('Using eager attention weights for retrieval score computation.')
        self.num_layers = self.llm.config.num_hidden_layers


        self.start_layer = 0
        self.end_layer = self.num_layers - 1

        print('[ICR is using layers from {} to {}.]'.format(self.start_layer, self.end_layer))


        # The following settings are for constructing the input prompt.
        self.prompt_bos_length=1
        if any(x in self._base_llm_name.lower() for x in ['mistral-']):
            self.additional_prompt_offset = 1 # for models that adds a ' ' at the beginning when tokenizing the prompt. e.g. '\n\n' -> [<s>, ' ', '\n\n']
            self.prompt_separator = '\n\n'
        elif any([x in self._base_llm_name.lower() for x in ['llama']]):
            self.additional_prompt_offset = 0
            self.prompt_separator = ' \n\n'
        else:
            self.additional_prompt_offset = 0
            self.prompt_separator = '\n\n'

        kwargs.get("reverse_doc_order", False)
        # Setup sliding window.
        # ICR typically works worse with sliding window, especially with smaller window sizes. Try to fit all documents to be re-ranked in the context as much as possible. 
        self.reverse_doc_order = kwargs.get("reverse_doc_order", False)
        self.sliding_window_size = kwargs.get("sliding_window_size", 10)
        if self.sliding_window_size is None:
            self.sliding_window_stride = self.sliding_window_size//2
        else:
            self.sliding_window_stride = self.sliding_window_size

    def _setup_llm_prompts(self, prompt_template, base_llm_name):
        """
        Configures the LLM prompt format based on the model type.

        Args:
            prompt_template (str): The prompt template type (`"instruct"`, `"simple"`, `"simple_instruct"`).
            base_llm_name (str): The base model name.

        Returns:
            tuple: The `prompt_template`, `prompt_prefix`, and `prompt_suffix`.

        Example:
            ```python
            prompt_template, prompt_prefix, prompt_suffix = model._setup_llm_prompts("instruct", "llamav3.1-8b")
            ```
        """
        if prompt_template == '':
            prompt_template='instruct' if any(x in base_llm_name.lower() for x in ['instruct']) else 'simple'
        else:
            assert prompt_template in ['instruct', 'simple', 'simple_instruct']
        print('ICR is using prompt template [{}] for in-context retrieval'.format(prompt_template))

        if  'mistral' in base_llm_name.lower():
            prompt_prefix = '[INST]'
            prompt_suffix = '[/INST]'
        elif 'llama-3' in base_llm_name.lower():
            prompt_prefix = '<|start_header_id|>user<|end_header_id|>'
            prompt_suffix = '<|eot_id|><|start_header_id|>assistant<|end_header_id|>'
        else:
            raise NotImplementedError("Prompt prefix and suffix not defined for the model family of {}.".format(base_llm_name))

        return prompt_template, prompt_prefix, prompt_suffix

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks documents using **In-Context Reranking (ICR)**.

        Args:
            documents (List[Document]): A list of `Document` instances containing query and contexts.

        Returns:
            List[Document]: Documents with reordered contexts based on reranking.

        Example:
            ```python
            reranked_docs = model.rank(documents)
            ```
        """
        max_length = 300
        for document in tqdm(documents, desc="Reranking Documents"):
            query = document.question.question
            contexts = [ctx.text.split()[:int(max_length)] for ctx in document.contexts]

            # Perform reranking with sliding windows
            #print(contexts)
            sorted_doc_ids, sorted_doc_scores = self.rerank(query, contexts, order="desc")[0]
            #print(sorted_doc_ids, sorted_doc_scores)
            # Assign scores and reorder contexts
            copy_context = copy.deepcopy(document.contexts)
            for idx, ctx_id in enumerate(sorted_doc_ids):
                copy_context[ctx_id].score = sorted_doc_scores[idx]

            ranked_contexts = sorted(copy_context, key=lambda x: x.score, reverse=True)
            document.reorder_contexts = ranked_contexts


        return documents

    def rerank(self, query, documents,  return_per_doc_results=False, order="desc"):
        """
        Applies **In-Context Reranking** using a **sliding window** approach.

        Args:
            query (str): The query for which documents need to be ranked.
            documents (List[str]): List of document texts.
            return_per_doc_results (bool, optional): Whether to return detailed per-document scores.
            order (str, optional): Sorting order (`"desc"` or `"asc"`).

        Returns:
            tuple: A tuple containing sorted document IDs and their scores.

        Example:
            ```python
            sorted_ids, scores = model.rerank("What is quantum computing?", documents)
            ```
        """

        N_docs = len(documents)

        if self.sliding_window_size < 0:
            self.sliding_window_size = N_docs

        sorted_doc_ids = list(range(N_docs))
        sorted_doc_ids.reverse()

        sorted_doc_scores = []
        if return_per_doc_results == 'tok':
            per_doc_results = []
        else:
            per_doc_results = None

        _i = 0
        _j = min(self.sliding_window_size, N_docs)
        while True:

            ids = [sorted_doc_ids[i] for i in range(_i, _j)]
            if not self.reverse_doc_order:
                # Put the most relevant documents at the front of document list.
                ids.reverse()

            docs = [documents[i] for i in ids]
            (_sorted_doc_ids, _sorted_doc_scores), _per_doc_results = self.get_sorted_docs(query, docs, return_per_doc_results=return_per_doc_results, order='asc')

            __sorted_doc_ids = [ids[i] for i in _sorted_doc_ids]
            for i in range(_i, _j):
                sorted_doc_ids[i] = __sorted_doc_ids[i-_i]

            if _j < N_docs:
                sorted_doc_scores.extend(_sorted_doc_scores[:self.sliding_window_stride])
                if return_per_doc_results == 'tok':
                    per_doc_results.extend(_per_doc_results[:self.sliding_window_stride])
            else:
                sorted_doc_scores.extend(_sorted_doc_scores)
                if return_per_doc_results == 'tok':
                    per_doc_results.extend(_per_doc_results)
                break

            _i += self.sliding_window_stride
            _j += self.sliding_window_stride
            _j = min(_j, N_docs)


        if order == 'desc':
            sorted_doc_ids.reverse()
            sorted_doc_scores.reverse()
            if return_per_doc_results == 'tok':
                per_doc_results.reverse()

        assert len(sorted_doc_ids) == len(sorted_doc_scores), "Length mismatch between sorted doc ids ({}) and scores({})!".format(len(sorted_doc_ids), len(sorted_doc_scores))
        return (sorted_doc_ids, sorted_doc_scores), per_doc_results

    def get_sorted_docs(self, query, retrieval_doc_pool, return_per_doc_results=False, prompt_prefix='', order='desc'):
        """
        Scores and sorts documents using **attention-based scoring**.

        Args:
            query (str): The query text.
            retrieval_doc_pool (List[str]): List of documents to be ranked.
            order (str, optional): Sort order (`"desc"` or `"asc"`).

        Returns:
            tuple: Sorted document IDs and their scores.

        Example:
            ```python
            sorted_ids, scores = model.get_sorted_docs("What is climate change?", documents)
            ```
        """
        kv_cache = None

        if self.scoring_strategy == 'query_last':
            # ICR without calibration.
            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
            doc_scores, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results)

        elif self.scoring_strategy == 'attention_sorting':
            # ICR without both calibration and attention aggregation.
            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
            query_start_idx = query_end_idx # Only using last query token (i.e. attention sorting).
            doc_scores, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results)

        elif self.scoring_strategy == 'NA_only':
            # For analyzing the intrinsic bias captured by calibration scores.
            query = 'N/A'

            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
            doc_scores, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results)

        elif self.scoring_strategy == 'NA_calibration_no_agg':
            # ICR without attention aggregation.

            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
            query_start_idx = query_end_idx
            doc_scores_query, perdoc_result, kv_cache = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results, return_cache=True)


            calibration_query = 'N/A'
            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(calibration_query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')

            # Use kv_cache from first query to speed up forward() for the calibration query.
            # query_start_idx should be the same for both queries.
            for i in range(len(kv_cache.key_cache)):
                kv_cache.key_cache[i] = kv_cache.key_cache[i][:,:,:query_start_idx,:]
                kv_cache.value_cache[i] = kv_cache.value_cache[i][:,:,:query_start_idx,:]
            kv_cache._seen_tokens = query_start_idx


            if kv_cache is not None:
                context_start_idx=query_start_idx
            else:
                context_start_idx=0

            query_start_idx = query_end_idx
            doc_scores_calib, doc_tok_scores_calib_na,_ = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results,  kv_cache=kv_cache, context_start_idx=context_start_idx)

            doc_scores = doc_scores_query - doc_scores_calib

            if return_per_doc_results != 'none':
                for i in range(len(perdoc_result)):
                    perdoc_result[i][1] -= doc_tok_scores_calib_na[i][1]

        elif self.scoring_strategy == 'masked_NA_calibration':
            return_per_doc_results = 'tok'
            # The default ICR method

            # FP with calibration query
            calibration_query = 'N/A'
            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(calibration_query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')

            doc_scores_calib, doc_tok_scores_calib_na, kv_cache = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results,  return_cache=True)

            # Use kv_cache from first query to speed up forward() for the calibration query.
            # query_start_idx should be the same for both queries.
            for i in range(len(kv_cache.key_cache)):
                kv_cache.key_cache[i] = kv_cache.key_cache[i][:,:,:query_start_idx,:]
                kv_cache.value_cache[i] = kv_cache.value_cache[i][:,:,:query_start_idx,:]
            kv_cache._seen_tokens = query_start_idx

            if kv_cache is not None:
                context_start_idx=query_start_idx
            else:
                context_start_idx=0

            # FP with the actual query            
            llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')

            doc_scores_query, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results,  kv_cache=kv_cache, context_start_idx=context_start_idx)


            _i = 0
            doc_scores = torch.zeros(len(retrieval_doc_pool))

            for doc_tok_score, doc_tok_score_na in zip(perdoc_result, doc_tok_scores_calib_na):
                doc_tok_score[1] = doc_tok_score[1].to(doc_tok_score_na[1].device)
                calibrated_scores = doc_tok_score[1] - doc_tok_score_na[1]

                mean_bias = calibrated_scores.mean()
                std_bias = calibrated_scores.std()
                threshold = mean_bias - 2*std_bias
                tok_mask = (calibrated_scores>threshold)

                doc_tok_score[1] = doc_tok_score[1] * tok_mask
                doc_tok_score_na[1] = doc_tok_score_na[1] * tok_mask
                doc_tok_score[1] = doc_tok_score[1] - doc_tok_score_na[1]
                doc_scores[_i] = doc_tok_score[1].sum()
                _i+=1

        per_doc_result = None
        if order in ['desc', 'asc']:
            sorted_results = torch.sort(doc_scores, descending=(order=='desc'))
            if return_per_doc_results != 'none':
                per_doc_result = [(perdoc_result[i][0], perdoc_result[i][1]) for i in sorted_results.indices]

            return (sorted_results.indices.tolist(), sorted_results.values.tolist()), per_doc_result
        elif order=='none':
            # Only return the scores and the per-doc results for documents in the input order.
            # Used during development.
            return list(range(len(retrieval_doc_pool))), doc_scores, per_doc_result
        else:
            print(f"Invalid order: {order}. Please use 'desc', 'asc' or 'none")
            raise NotImplementedError
    def score_documents(
            self,
            llm_input,
            doc_tok_idx_spans,
            query_start_tok_idx,
            query_end_tok_idx,
            context_start_idx=0,
            return_per_doc_results=False,
            long_prompt=False,
            return_cache=False,
            kv_cache=None,
        ):
        """
        Computes **attention-based document scores**.

        Args:
            llm_input (str): The full LLM input prompt.
            doc_tok_idx_spans (List[tuple]): Token spans for each document.
            query_start_tok_idx (int): Start index of query tokens.
            query_end_tok_idx (int): End index of query tokens.

        Returns:
            List[float]: Document relevance scores.

        Example:
            ```python
            scores = model.score_documents(prompt, spans, 10, 20)
            ```
        """
        tokenized_input = self.tokenizer(llm_input,return_tensors='pt').to(self.llm.device)
        _input_ids = tokenized_input.input_ids[:, context_start_idx:]
        _query_indices = list(range(query_start_tok_idx-context_start_idx, query_end_tok_idx-context_start_idx+1))

        if kv_cache is None:
            if self._use_fa2:
                kv_cache=DynamicCacheWithQuery(query_indices=_query_indices)
            else:
                kv_cache=DynamicCache()
        else:
            kv_cache.query_cache = []
            _query_indices = _query_indices
            kv_cache._query_indices = _query_indices

        with torch.no_grad():
            output = self.llm(
                input_ids=_input_ids,
                use_cache=True,
                past_key_values=kv_cache,
                output_attentions=True
                )

        if self._use_fa2:
            # Extract key and query vectors from FA2. Then recompute attention scores for re-ranking.
            kv_cache = output.past_key_values

            long_prompt = False
            if len(_input_ids[0]) > 40000:
                # For sequences that are too long, compute scores on CPU to void GPU OOM.
                # Adjust the limit here depending on your system configuration.
                print('Long sequence of more than 40K tokens detected. Computing attention scores on CPU.')
                long_prompt = True

            attention_weights = []
            doc_tok_weights = []

            if long_prompt:
                _device = 'cpu'
            else:
                _device = 'cuda:0'

            # loop through all layers and compute attention scores
            for i in range(self.start_layer, self.end_layer+1):                     
                attn_weights = self._get_attn_weights(kv_cache.key_cache[i][:,:,:query_end_tok_idx+1], kv_cache.query_cache[i],  use_cpu=long_prompt).to(_device).squeeze(0)
                attn_weights = attn_weights.mean(1) # average over query tokens
                attention_weights.append(attn_weights.squeeze(0))

        else:
            # Directly extract attention weights from the attention layers of the LLM.
            attention_weights = [attn[0][:,query_start_tok_idx:query_end_tok_idx+1,:].mean(1) for attn in output.attentions]

        attention_weights = torch.stack(attention_weights, dim=0)

        if return_per_doc_results != 'none':
            per_doc_results = [[None, None] for _ in range(len(doc_tok_idx_spans))]
        else:
            per_doc_results = None

        attention_weights = attention_weights.sum(0) # sum attention scores across layers            
        attention_weights = attention_weights.sum(0) # sum attention scores across attention heads
        doc_scores = []

        for i, doc_span in enumerate(doc_tok_idx_spans): 
            _tok_score = attention_weights[doc_span[0]:doc_span[1]]
            doc_scores.append(_tok_score.sum())

            if return_per_doc_results != 'none':
                _doc_tok_ids = tokenized_input.input_ids[0][doc_span[0]:doc_span[1]]
                _doc_toks = self.tokenizer.convert_ids_to_tokens(_doc_tok_ids)
                per_doc_results[i][0] = _doc_toks
                per_doc_results[i][1] = _tok_score.clone().detach() # sum over layers

        doc_scores = torch.tensor(doc_scores)
        gc.collect()
        torch.cuda.empty_cache()

        if return_cache:
            return doc_scores, per_doc_results, kv_cache
        else:
            return doc_scores, per_doc_results
    def _prepare_input_for_document_retrieval(self, query, documents, system_prompt='', query_position='last'):
        """
        Constructs the **ICR prompt** with query and document contexts.

        Args:
            query (str): Query text.
            documents (List[str]): List of documents.
            system_prompt (str, optional): Custom system prompt.
            query_position (str, optional): Query placement (`"first"` or `"last"`).

        Returns:
            tuple: The formatted prompt and document spans.

        Example:
            ```python
            prompt, spans, query_start, query_end = model._prepare_input_for_document_retrieval("What is AI?", docs)
            ```
        """
        llm_prompt = ''
        document_span_intervals = []


        if self.prompt_template == 'simple':
            system_prompt = ''
        elif self.prompt_template == 'simple_instruct':
            system_prompt = system_prompt
        elif self.prompt_template == 'instruct':
            if system_prompt != '':
                system_prompt = self.retrieval_instruction.format(len(documents), query) + self.prompt_separator + system_prompt
            else:
                system_prompt = self.retrieval_instruction.format(len(documents), query)

        system_prompt = self.prompt_prefix + system_prompt

        query_start_idx = None
        query_end_idx = None


        separator_length = self.tokenizer(self.prompt_separator, return_tensors='pt').input_ids.size(1) - self.prompt_bos_length - self.additional_prompt_offset # remove the leading ['<s>', '_'] tokens

        llm_prompt = system_prompt


        prompt_length = self.tokenizer(llm_prompt+self.prompt_separator, return_tensors='pt').input_ids.size(1)-separator_length # add and subtract separator tokens for accurate prefix length

        if query_position == 'first':
            if self.prompt_template in ['simple', 'instruct']:
                instruction_prompt = f'Query:'

                llm_prompt += self.prompt_separator + instruction_prompt 
                prompt_length += separator_length
                prompt_length += self.tokenizer(self.prompt_separator + instruction_prompt, return_tensors='pt').input_ids.size(1) - self.prompt_bos_length - separator_length - self.additional_prompt_offset
                query_start_idx = prompt_length - 1 # The ':' after 'Query'    
            else:
                llm_prompt += self.prompt_separator
                prompt_length += separator_length
                query_start_idx = prompt_length # The start of the query context

            if self.prompt_template == 'simple':
                query_prompt = f' {query.strip()}{self.prompt_separator}Answer:'
            elif self.prompt_template in ['instruct', 'simple_instruct']:
                query_prompt = f' {query.strip()}'

            llm_prompt += query_prompt
            prompt_length += self.tokenizer(query_prompt, return_tensors='pt').input_ids.size(1) - self.prompt_bos_length - self.additional_prompt_offset
            query_end_idx = prompt_length - 1 


        if prompt_length != self.tokenizer(llm_prompt, return_tensors='pt').input_ids.size(1):
                print('Prompt length mismatch!')
                print(prompt_length, ' vs ', self.tokenizer(llm_prompt, return_tensors='pt').input_ids.size(1))
                print('-'*30)
                self.__show_tokens(llm_prompt)
                raise Exception('ICR prompt length mismatch before adding docs.')

        _doc_separator_length = separator_length

        for i, doc in enumerate(documents):

            doc = f'[{i+1}] {doc}'
            prompt_length += _doc_separator_length
            llm_prompt += self.prompt_separator + doc
            doc_length = self.tokenizer(self.prompt_separator + doc, return_tensors='pt').input_ids.size(1) - self.prompt_bos_length - _doc_separator_length - self.additional_prompt_offset # - bos_length for the leading ['<s>'] token, -additional for the potential extra tokens, e.g. the '_' token between <s> and <0x0A> when <0x0A> is the first token for mistral models.

            document_span_intervals.append((prompt_length, prompt_length + doc_length))
            prompt_length += doc_length

            if prompt_length != self.tokenizer(llm_prompt, return_tensors='pt').input_ids.size(1):
                print('Prompt length mismatch @ doc {}!'.format(i))
                print(prompt_length, ' vs ', self.tokenizer(llm_prompt, return_tensors='pt').input_ids.size(1))
                print('-'*30)
                self.__show_tokens(llm_prompt)
                print('-'*30)
                print('doc length:', doc_length)
                self.__show_tokens(self.prompt_separator+doc)
                raise Exception('ICR prompt length mismatch after adding docs.')


        if query_position == 'last':
            query_start_idx = prompt_length + separator_length
            if self.prompt_template in ['simple', 'instruct']:
                instruction_prompt = self.retrieval_instruction_late + self.prompt_separator + 'Query:'
                llm_prompt += self.prompt_separator + instruction_prompt 
                prompt_length += separator_length
                prompt_length += self.tokenizer(self.prompt_separator + instruction_prompt, return_tensors='pt').input_ids.size(1) - self.prompt_bos_length - separator_length - self.additional_prompt_offset

            else:
                llm_prompt += self.prompt_separator
                prompt_length += separator_length

        if self.prompt_template == 'simple':
            query_prompt = f' {query.strip()}{self.prompt_separator}Answer:'
        elif self.prompt_template in ['instruct', 'simple_instruct']:
            query_prompt = f' {query.strip()}'
            if query_position == 'last':
                query_prompt += self.prompt_suffix.format(len(documents))

        llm_prompt += query_prompt
        prompt_length += self.tokenizer(query_prompt, return_tensors='pt').input_ids.size(1) - self.prompt_bos_length - self.additional_prompt_offset
        if query_position == 'last':
            query_end_idx = prompt_length - 1
        return llm_prompt, document_span_intervals, query_start_idx, query_end_idx

    @classmethod
    def __show_tokens(self, string):
        # Shows tokenized string.
        # Mainly used for debugging prompt construction for document retrieval.
        tokenized_string_ids = self.tokenizer(string, return_tensors='pt').input_ids[0]
        print(self.tokenizer.convert_ids_to_tokens(tokenized_string_ids), tokenized_string_ids.size(0))

    @classmethod
    def _get_attn_weights(cls, key_states, query_states, use_cpu=False):

        bsz, num_heads, q_len, head_dim = query_states.size()
        num_key_value_heads = key_states.size(1)
        num_key_value_groups = num_heads // num_key_value_heads
        kv_seq_len = key_states.size(-2)

        if use_cpu:
            query_states = query_states.cpu()
            key_states = key_states.cpu()

        key_states = repeat_kv(key_states, num_key_value_groups)


        attn_weights = torch.matmul(query_states, key_states.transpose(2,3)) / math.sqrt(head_dim)

        if attn_weights.size() != (bsz, num_heads, q_len, kv_seq_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz, num_heads, q_len, kv_seq_len)}, but is"
                f" {attn_weights.size()}"
            )

        # Make causal mask and add it to attention weights.
        causal_mask = cls._get_causal_mask(attn_weights).to(attn_weights.device)
        attn_weights += causal_mask.unsqueeze(0)
        attn_lses = torch.logsumexp(attn_weights, dim=-1, keepdim=True) # Log-sum-exp of attention weights for numerical stability in softmax.
        attn_weights = torch.exp(attn_weights - attn_lses) # softmax

        return attn_weights

    @classmethod
    def _get_causal_mask(cls, attn_weights):
        # Make causal mask for attention weights.
        query_len, seq_len = attn_weights.size(-2), attn_weights.size(-1)
        causal_mask = torch.ones_like(attn_weights.transpose(-1,-2).squeeze(0))
        causal_mask = torch.triu(causal_mask, diagonal=-(seq_len-query_len))
        causal_mask = causal_mask.transpose(-1,-2)
        causal_mask = (1-causal_mask) * torch.finfo(causal_mask.dtype).min
        return causal_mask

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

Initializes the In-Context Reranker.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	Name of the LLM model (e.g., "meta-llama/Meta-Llama-3.1-8B-Instruct").	None
**kwargs		Additional customization parameters, including: - scoring_strategy (str): Strategy for computing attention scores ("query_last", "attention_sorting", etc.). - retrieval_type (str): Type of retrieval ("QA" or "IE"). - sliding_window_size (int): Number of documents considered per window.	{}

Example

model = InContextReranker(method='incontext_reranker', model_name='llamav3.1-8b')

Source code in rankify/models/incontext_reranker.py

def __init__(self, method=None, model_name=None, **kwargs):
    """
    Initializes the In-Context Reranker.

    Args:
        method (str, optional): The reranking method name.
        model_name (str): Name of the LLM model (e.g., `"meta-llama/Meta-Llama-3.1-8B-Instruct"`).
        **kwargs: Additional customization parameters, including:
            - `scoring_strategy` (str): Strategy for computing attention scores (`"query_last"`, `"attention_sorting"`, etc.).
            - `retrieval_type` (str): Type of retrieval (`"QA"` or `"IE"`).
            - `sliding_window_size` (int): Number of documents considered per window.

    Example:
        ```python
        model = InContextReranker(method='incontext_reranker', model_name='llamav3.1-8b')
        ```
    """
    # Setup the base LLM
    self._base_llm_name =model_name
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
    self.tokenizer = tokenizer
    print(f"initialized tokenizer for [{model_name}]")

    if any([x in model_name.lower() for x in ['mistralai/mistral', ]]):
        BaseLLMClass = MistralForCausalLM
    elif any([x in model_name.lower() for x in ['llama']]):
        BaseLLMClass = LlamaForCausalLM
    else:
        print(f"Warning: The model family for [{model_name}] is not supported by InContextRAGModel!")
        raise NotImplementedError

    prompt_template = kwargs.get("prompt_template", "instruct")

    prompt_template, prompt_prefix, prompt_suffix, = self._setup_llm_prompts(prompt_template, model_name)

    self.use_fa2 = kwargs.get("use_fa2", True)
    if self.use_fa2:
        _attn_implementation = "flash_attention_2"
    else:
        _attn_implementation = "eager"

    llm = BaseLLMClass.from_pretrained(
            model_name, 
            torch_dtype=torch.float16, 
            attn_implementation=_attn_implementation,
            device_map='auto'
        )
    self.llm = llm
    self.llm.config.pad_token_id = self.llm.config.eos_token_id

    # Setup prompts for ICR
    assert prompt_template in ['instruct', 'simple', 'simple_instruct'], "Invalid prompt template!"

    self.prompt_template = prompt_template
    self.prompt_prefix = prompt_prefix


    self.prompt_suffix = prompt_suffix
    self.scoring_strategy  = kwargs.get("scoring_strategy", "query_last")
    retrieval_type = kwargs.get("retrieval_type", "QA")

    if retrieval_type == 'QA':
        print('[ICR is using QA prompt type]')
        self.retrieval_instruction = ' Here are some paragraphs:'
        self.retrieval_instruction_late = 'Please answer the following question based on the information in the paragraphs above.'
    elif retrieval_type == 'IE':
        print('[ICR is using IE prompt type]')
        self.retrieval_instruction = ' Here are some paragraphs:'
        self.retrieval_instruction_late = 'Please find information that are relevant to the following query in the paragraphs above.'
    else:
        raise NotImplementedError('Invalid retrieval type! Should be one of [QA, IE]')

    assert self.scoring_strategy in ['query_last', 'attention_sorting', 'NA_only', 'NA_calibration_no_agg', 'masked_NA_calibration'], "Invalid scoring strategy!"

    self._use_fa2 = self.use_fa2
    if self.use_fa2:
        print('Using FA2 for retrieval score computation.')
    else:
        print('Using eager attention weights for retrieval score computation.')
    self.num_layers = self.llm.config.num_hidden_layers


    self.start_layer = 0
    self.end_layer = self.num_layers - 1

    print('[ICR is using layers from {} to {}.]'.format(self.start_layer, self.end_layer))


    # The following settings are for constructing the input prompt.
    self.prompt_bos_length=1
    if any(x in self._base_llm_name.lower() for x in ['mistral-']):
        self.additional_prompt_offset = 1 # for models that adds a ' ' at the beginning when tokenizing the prompt. e.g. '\n\n' -> [<s>, ' ', '\n\n']
        self.prompt_separator = '\n\n'
    elif any([x in self._base_llm_name.lower() for x in ['llama']]):
        self.additional_prompt_offset = 0
        self.prompt_separator = ' \n\n'
    else:
        self.additional_prompt_offset = 0
        self.prompt_separator = '\n\n'

    kwargs.get("reverse_doc_order", False)
    # Setup sliding window.
    # ICR typically works worse with sliding window, especially with smaller window sizes. Try to fit all documents to be re-ranked in the context as much as possible. 
    self.reverse_doc_order = kwargs.get("reverse_doc_order", False)
    self.sliding_window_size = kwargs.get("sliding_window_size", 10)
    if self.sliding_window_size is None:
        self.sliding_window_stride = self.sliding_window_size//2
    else:
        self.sliding_window_stride = self.sliding_window_size

rank(documents)

Reranks documents using In-Context Reranking (ICR).

Parameters:

Name	Type	Description	Default
documents	List[Document]	A list of Document instances containing query and contexts.	required

Returns:

Type	Description
List[Document]	List[Document]: Documents with reordered contexts based on reranking.

Example

reranked_docs = model.rank(documents)

Source code in rankify/models/incontext_reranker.py

def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks documents using **In-Context Reranking (ICR)**.

    Args:
        documents (List[Document]): A list of `Document` instances containing query and contexts.

    Returns:
        List[Document]: Documents with reordered contexts based on reranking.

    Example:
        ```python
        reranked_docs = model.rank(documents)
        ```
    """
    max_length = 300
    for document in tqdm(documents, desc="Reranking Documents"):
        query = document.question.question
        contexts = [ctx.text.split()[:int(max_length)] for ctx in document.contexts]

        # Perform reranking with sliding windows
        #print(contexts)
        sorted_doc_ids, sorted_doc_scores = self.rerank(query, contexts, order="desc")[0]
        #print(sorted_doc_ids, sorted_doc_scores)
        # Assign scores and reorder contexts
        copy_context = copy.deepcopy(document.contexts)
        for idx, ctx_id in enumerate(sorted_doc_ids):
            copy_context[ctx_id].score = sorted_doc_scores[idx]

        ranked_contexts = sorted(copy_context, key=lambda x: x.score, reverse=True)
        document.reorder_contexts = ranked_contexts


    return documents

rerank(query, documents, return_per_doc_results=False, order='desc')

Applies In-Context Reranking using a sliding window approach.

Parameters:

Name	Type	Description	Default
query	str	The query for which documents need to be ranked.	required
documents	List[str]	List of document texts.	required
return_per_doc_results	bool	Whether to return detailed per-document scores.	False
order	str	Sorting order ("desc" or "asc").	'desc'

Returns:

Name	Type	Description
tuple		A tuple containing sorted document IDs and their scores.

Example

sorted_ids, scores = model.rerank("What is quantum computing?", documents)

Source code in rankify/models/incontext_reranker.py

def rerank(self, query, documents,  return_per_doc_results=False, order="desc"):
    """
    Applies **In-Context Reranking** using a **sliding window** approach.

    Args:
        query (str): The query for which documents need to be ranked.
        documents (List[str]): List of document texts.
        return_per_doc_results (bool, optional): Whether to return detailed per-document scores.
        order (str, optional): Sorting order (`"desc"` or `"asc"`).

    Returns:
        tuple: A tuple containing sorted document IDs and their scores.

    Example:
        ```python
        sorted_ids, scores = model.rerank("What is quantum computing?", documents)
        ```
    """

    N_docs = len(documents)

    if self.sliding_window_size < 0:
        self.sliding_window_size = N_docs

    sorted_doc_ids = list(range(N_docs))
    sorted_doc_ids.reverse()

    sorted_doc_scores = []
    if return_per_doc_results == 'tok':
        per_doc_results = []
    else:
        per_doc_results = None

    _i = 0
    _j = min(self.sliding_window_size, N_docs)
    while True:

        ids = [sorted_doc_ids[i] for i in range(_i, _j)]
        if not self.reverse_doc_order:
            # Put the most relevant documents at the front of document list.
            ids.reverse()

        docs = [documents[i] for i in ids]
        (_sorted_doc_ids, _sorted_doc_scores), _per_doc_results = self.get_sorted_docs(query, docs, return_per_doc_results=return_per_doc_results, order='asc')

        __sorted_doc_ids = [ids[i] for i in _sorted_doc_ids]
        for i in range(_i, _j):
            sorted_doc_ids[i] = __sorted_doc_ids[i-_i]

        if _j < N_docs:
            sorted_doc_scores.extend(_sorted_doc_scores[:self.sliding_window_stride])
            if return_per_doc_results == 'tok':
                per_doc_results.extend(_per_doc_results[:self.sliding_window_stride])
        else:
            sorted_doc_scores.extend(_sorted_doc_scores)
            if return_per_doc_results == 'tok':
                per_doc_results.extend(_per_doc_results)
            break

        _i += self.sliding_window_stride
        _j += self.sliding_window_stride
        _j = min(_j, N_docs)


    if order == 'desc':
        sorted_doc_ids.reverse()
        sorted_doc_scores.reverse()
        if return_per_doc_results == 'tok':
            per_doc_results.reverse()

    assert len(sorted_doc_ids) == len(sorted_doc_scores), "Length mismatch between sorted doc ids ({}) and scores({})!".format(len(sorted_doc_ids), len(sorted_doc_scores))
    return (sorted_doc_ids, sorted_doc_scores), per_doc_results

get_sorted_docs(query, retrieval_doc_pool, return_per_doc_results=False, prompt_prefix='', order='desc')

Scores and sorts documents using attention-based scoring.

Parameters:

Name	Type	Description	Default
query	str	The query text.	required
retrieval_doc_pool	List[str]	List of documents to be ranked.	required
order	str	Sort order ("desc" or "asc").	'desc'

Returns:

Name	Type	Description
tuple		Sorted document IDs and their scores.

Example

sorted_ids, scores = model.get_sorted_docs("What is climate change?", documents)

Source code in rankify/models/incontext_reranker.py

def get_sorted_docs(self, query, retrieval_doc_pool, return_per_doc_results=False, prompt_prefix='', order='desc'):
    """
    Scores and sorts documents using **attention-based scoring**.

    Args:
        query (str): The query text.
        retrieval_doc_pool (List[str]): List of documents to be ranked.
        order (str, optional): Sort order (`"desc"` or `"asc"`).

    Returns:
        tuple: Sorted document IDs and their scores.

    Example:
        ```python
        sorted_ids, scores = model.get_sorted_docs("What is climate change?", documents)
        ```
    """
    kv_cache = None

    if self.scoring_strategy == 'query_last':
        # ICR without calibration.
        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
        doc_scores, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results)

    elif self.scoring_strategy == 'attention_sorting':
        # ICR without both calibration and attention aggregation.
        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
        query_start_idx = query_end_idx # Only using last query token (i.e. attention sorting).
        doc_scores, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results)

    elif self.scoring_strategy == 'NA_only':
        # For analyzing the intrinsic bias captured by calibration scores.
        query = 'N/A'

        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
        doc_scores, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results)

    elif self.scoring_strategy == 'NA_calibration_no_agg':
        # ICR without attention aggregation.

        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')
        query_start_idx = query_end_idx
        doc_scores_query, perdoc_result, kv_cache = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results, return_cache=True)


        calibration_query = 'N/A'
        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(calibration_query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')

        # Use kv_cache from first query to speed up forward() for the calibration query.
        # query_start_idx should be the same for both queries.
        for i in range(len(kv_cache.key_cache)):
            kv_cache.key_cache[i] = kv_cache.key_cache[i][:,:,:query_start_idx,:]
            kv_cache.value_cache[i] = kv_cache.value_cache[i][:,:,:query_start_idx,:]
        kv_cache._seen_tokens = query_start_idx


        if kv_cache is not None:
            context_start_idx=query_start_idx
        else:
            context_start_idx=0

        query_start_idx = query_end_idx
        doc_scores_calib, doc_tok_scores_calib_na,_ = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results,  kv_cache=kv_cache, context_start_idx=context_start_idx)

        doc_scores = doc_scores_query - doc_scores_calib

        if return_per_doc_results != 'none':
            for i in range(len(perdoc_result)):
                perdoc_result[i][1] -= doc_tok_scores_calib_na[i][1]

    elif self.scoring_strategy == 'masked_NA_calibration':
        return_per_doc_results = 'tok'
        # The default ICR method

        # FP with calibration query
        calibration_query = 'N/A'
        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(calibration_query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')

        doc_scores_calib, doc_tok_scores_calib_na, kv_cache = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results,  return_cache=True)

        # Use kv_cache from first query to speed up forward() for the calibration query.
        # query_start_idx should be the same for both queries.
        for i in range(len(kv_cache.key_cache)):
            kv_cache.key_cache[i] = kv_cache.key_cache[i][:,:,:query_start_idx,:]
            kv_cache.value_cache[i] = kv_cache.value_cache[i][:,:,:query_start_idx,:]
        kv_cache._seen_tokens = query_start_idx

        if kv_cache is not None:
            context_start_idx=query_start_idx
        else:
            context_start_idx=0

        # FP with the actual query            
        llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx = self._prepare_input_for_document_retrieval(query, retrieval_doc_pool, system_prompt=prompt_prefix, query_position='last')

        doc_scores_query, perdoc_result = self.score_documents(llm_prompt, doc_tok_idx_spans, query_start_idx, query_end_idx, return_per_doc_results=return_per_doc_results,  kv_cache=kv_cache, context_start_idx=context_start_idx)


        _i = 0
        doc_scores = torch.zeros(len(retrieval_doc_pool))

        for doc_tok_score, doc_tok_score_na in zip(perdoc_result, doc_tok_scores_calib_na):
            doc_tok_score[1] = doc_tok_score[1].to(doc_tok_score_na[1].device)
            calibrated_scores = doc_tok_score[1] - doc_tok_score_na[1]

            mean_bias = calibrated_scores.mean()
            std_bias = calibrated_scores.std()
            threshold = mean_bias - 2*std_bias
            tok_mask = (calibrated_scores>threshold)

            doc_tok_score[1] = doc_tok_score[1] * tok_mask
            doc_tok_score_na[1] = doc_tok_score_na[1] * tok_mask
            doc_tok_score[1] = doc_tok_score[1] - doc_tok_score_na[1]
            doc_scores[_i] = doc_tok_score[1].sum()
            _i+=1

    per_doc_result = None
    if order in ['desc', 'asc']:
        sorted_results = torch.sort(doc_scores, descending=(order=='desc'))
        if return_per_doc_results != 'none':
            per_doc_result = [(perdoc_result[i][0], perdoc_result[i][1]) for i in sorted_results.indices]

        return (sorted_results.indices.tolist(), sorted_results.values.tolist()), per_doc_result
    elif order=='none':
        # Only return the scores and the per-doc results for documents in the input order.
        # Used during development.
        return list(range(len(retrieval_doc_pool))), doc_scores, per_doc_result
    else:
        print(f"Invalid order: {order}. Please use 'desc', 'asc' or 'none")
        raise NotImplementedError

score_documents(llm_input, doc_tok_idx_spans, query_start_tok_idx, query_end_tok_idx, context_start_idx=0, return_per_doc_results=False, long_prompt=False, return_cache=False, kv_cache=None)

Computes attention-based document scores.

Parameters:

Name	Type	Description	Default
llm_input	str	The full LLM input prompt.	required
doc_tok_idx_spans	List[tuple]	Token spans for each document.	required
query_start_tok_idx	int	Start index of query tokens.	required
query_end_tok_idx	int	End index of query tokens.	required

Returns:

Type	Description
	List[float]: Document relevance scores.

Example

scores = model.score_documents(prompt, spans, 10, 20)

Source code in rankify/models/incontext_reranker.py

def score_documents(
        self,
        llm_input,
        doc_tok_idx_spans,
        query_start_tok_idx,
        query_end_tok_idx,
        context_start_idx=0,
        return_per_doc_results=False,
        long_prompt=False,
        return_cache=False,
        kv_cache=None,
    ):
    """
    Computes **attention-based document scores**.

    Args:
        llm_input (str): The full LLM input prompt.
        doc_tok_idx_spans (List[tuple]): Token spans for each document.
        query_start_tok_idx (int): Start index of query tokens.
        query_end_tok_idx (int): End index of query tokens.

    Returns:
        List[float]: Document relevance scores.

    Example:
        ```python
        scores = model.score_documents(prompt, spans, 10, 20)
        ```
    """
    tokenized_input = self.tokenizer(llm_input,return_tensors='pt').to(self.llm.device)
    _input_ids = tokenized_input.input_ids[:, context_start_idx:]
    _query_indices = list(range(query_start_tok_idx-context_start_idx, query_end_tok_idx-context_start_idx+1))

    if kv_cache is None:
        if self._use_fa2:
            kv_cache=DynamicCacheWithQuery(query_indices=_query_indices)
        else:
            kv_cache=DynamicCache()
    else:
        kv_cache.query_cache = []
        _query_indices = _query_indices
        kv_cache._query_indices = _query_indices

    with torch.no_grad():
        output = self.llm(
            input_ids=_input_ids,
            use_cache=True,
            past_key_values=kv_cache,
            output_attentions=True
            )

    if self._use_fa2:
        # Extract key and query vectors from FA2. Then recompute attention scores for re-ranking.
        kv_cache = output.past_key_values

        long_prompt = False
        if len(_input_ids[0]) > 40000:
            # For sequences that are too long, compute scores on CPU to void GPU OOM.
            # Adjust the limit here depending on your system configuration.
            print('Long sequence of more than 40K tokens detected. Computing attention scores on CPU.')
            long_prompt = True

        attention_weights = []
        doc_tok_weights = []

        if long_prompt:
            _device = 'cpu'
        else:
            _device = 'cuda:0'

        # loop through all layers and compute attention scores
        for i in range(self.start_layer, self.end_layer+1):                     
            attn_weights = self._get_attn_weights(kv_cache.key_cache[i][:,:,:query_end_tok_idx+1], kv_cache.query_cache[i],  use_cpu=long_prompt).to(_device).squeeze(0)
            attn_weights = attn_weights.mean(1) # average over query tokens
            attention_weights.append(attn_weights.squeeze(0))

    else:
        # Directly extract attention weights from the attention layers of the LLM.
        attention_weights = [attn[0][:,query_start_tok_idx:query_end_tok_idx+1,:].mean(1) for attn in output.attentions]

    attention_weights = torch.stack(attention_weights, dim=0)

    if return_per_doc_results != 'none':
        per_doc_results = [[None, None] for _ in range(len(doc_tok_idx_spans))]
    else:
        per_doc_results = None

    attention_weights = attention_weights.sum(0) # sum attention scores across layers            
    attention_weights = attention_weights.sum(0) # sum attention scores across attention heads
    doc_scores = []

    for i, doc_span in enumerate(doc_tok_idx_spans): 
        _tok_score = attention_weights[doc_span[0]:doc_span[1]]
        doc_scores.append(_tok_score.sum())

        if return_per_doc_results != 'none':
            _doc_tok_ids = tokenized_input.input_ids[0][doc_span[0]:doc_span[1]]
            _doc_toks = self.tokenizer.convert_ids_to_tokens(_doc_tok_ids)
            per_doc_results[i][0] = _doc_toks
            per_doc_results[i][1] = _tok_score.clone().detach() # sum over layers

    doc_scores = torch.tensor(doc_scores)
    gc.collect()
    torch.cuda.empty_cache()

    if return_cache:
        return doc_scores, per_doc_results, kv_cache
    else:
        return doc_scores, per_doc_results