LLM Layerwise Reranker

rankify.models.llm_layerwise_ranker

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

LLMLayerWiseRanker

Bases: BaseRanking

Implements LLM Layer-Wise Reranking, a zero-shot ranking approach using large language models (LLMs).

This method performs layer-wise scoring of query-passage pairs, leveraging self-knowledge distillation to enhance ranking performance across multiple granularities.

References

• Li et al. (2023): Making Large Language Models A Better Foundation For Dense Retrieval. Paper
• Chen et al. (2024): BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation. Paper

Attributes:

Name	Type	Description
model_name	str	The name of the pre-trained model used for reranking.
max_sequence_length	int	The maximum token length of the input sequence (default: 512).
device	device	The device (CPU/GPU) on which the model runs.
dtype	dtype	The tensor data type for model inference.
batch_size	int	The batch size for processing query-passage pairs.
tokenizer	AutoTokenizer	Tokenizer for the model.
model	AutoModelForCausalLM	The transformer-based LLM model for reranking.
prompt	str	The default prompt template used for ranking.
params	dict	Model-specific configuration parameters.

Example

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

# Define a query and contexts
question = Question("What are the effects of climate change?")
contexts = [
    Context(text="Rising temperatures are causing ice caps to melt.", id=0),
    Context(text="Many species face extinction due to habitat loss.", id=1),
    Context(text="Ocean acidification is increasing.", id=2),
]
document = Document(question=question, contexts=contexts)

# Initialize LLM Layer-Wise Ranker
model = Reranking(method='llm_layerwise_ranker', model_name='bge-multilingual-gemma2')
model.rank([document])

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

Source code in rankify/models/llm_layerwise_ranker.py

class LLMLayerWiseRanker(BaseRanking):
    """
    Implements **LLM Layer-Wise Reranking**, a **zero-shot ranking approach** using large language models (LLMs).


    This method performs **layer-wise scoring** of query-passage pairs, leveraging **self-knowledge distillation** to 
    enhance ranking performance across **multiple granularities**.

    References:
        - **Li et al. (2023)**: *Making Large Language Models A Better Foundation For Dense Retrieval*.
          [Paper](https://arxiv.org/abs/2312.15503)
        - **Chen et al. (2024)**: *BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation*.
          [Paper](https://arxiv.org/abs/2402.03216)

    Attributes:
        model_name (str): The **name of the pre-trained model** used for reranking.
        max_sequence_length (int): The **maximum token length** of the input sequence (default: `512`).
        device (torch.device): The **device (CPU/GPU)** on which the model runs.
        dtype (torch.dtype): The **tensor data type** for model inference.
        batch_size (int): The **batch size** for processing query-passage pairs.
        tokenizer (AutoTokenizer): **Tokenizer** for the model.
        model (AutoModelForCausalLM): The **transformer-based LLM model** for reranking.
        prompt (str): The **default prompt template** used for ranking.
        params (dict): **Model-specific configuration parameters**.

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

        # Define a query and contexts
        question = Question("What are the effects of climate change?")
        contexts = [
            Context(text="Rising temperatures are causing ice caps to melt.", id=0),
            Context(text="Many species face extinction due to habitat loss.", id=1),
            Context(text="Ocean acidification is increasing.", id=2),
        ]
        document = Document(question=question, contexts=contexts)

        # Initialize LLM Layer-Wise Ranker
        model = Reranking(method='llm_layerwise_ranker', model_name='bge-multilingual-gemma2')
        model.rank([document])

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

    # Class-level constants
    PROMPTS = {
        "BAAI/bge-reranker-v2.5-gemma2-lightweight": "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.",
        "default": "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.",
    }

    DEFAULT_PARAMS = {
        "default": {},
        "BAAI/bge-multilingual-gemma2": {},
        "BAAI/bge-reranker-v2-gemma": {},
        "BAAI/bge-reranker-v2-minicpm-layerwise": {"cutoff_layers": [28]},
        "BAAI/bge-reranker-v2.5-gemma2-lightweight": {
            "cutoff_layers": [28],
            "compress_ratio": 2,
            "compress_layer": [24, 40],
        },
    }

    def __init__(
        self, method = None, model_name = None, api_key = None, **kwargs):
        """
        Initializes the **LLM Layer-Wise Ranker**.

        Args:
            method (str, optional): The **reranking method name**.
            model_name (str, optional): The **name of the pre-trained model** (default: `"BAAI/bge-reranker-v2.5-gemma2-lightweight"`).
            api_key (str, optional): API key for **authentication** (if needed).
            kwargs (dict): Additional configuration arguments (e.g., `max_sequence_length`, `batch_size`, `device`, `dtype`).
        """

        max_sequence_length= kwargs.get("max_sequence_length", 1024)
        device = kwargs.get("device", "cuda")
        self.device = get_device(device)
        self.batch_size = kwargs.get("batch_size", 4)
        self.cutoff_layer =  kwargs.get("cutoff_layer", None)
        self.compress_ratio = kwargs.get("compress_ratio", None)  
        self.compress_layer =  kwargs.get("compress_layer", None)  
        self.prompt =  kwargs.get("prompt", None)  

        self.model_name = model_name

        self.dtype = get_dtype( kwargs.get("dtype", torch.float16), self.device)

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=True,
        )
        self.max_sequence_length = max_sequence_length
        self.tokenizer.model_max_length = self.max_sequence_length
        self.tokenizer.padding_side = "right"


        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            trust_remote_code=True,
            torch_dtype=self.dtype,
        ).to(self.device)
        self.model.eval()

         # Create params dict based on specified values or defaults
        params = {}
        if self.cutoff_layer is not None:
            params["cutoff_layers"] = self.cutoff_layer
        if self.compress_ratio is not None:
            params["compress_ratio"] = self.compress_ratio
        if self.compress_layer is not None:
            params["compress_layer"] = self.compress_layer
        if not params:
            params = self.DEFAULT_PARAMS.get(model_name, self.DEFAULT_PARAMS["default"])
        self.params = params

        self.prompt = self.prompt
        if self.prompt is None:
            self.prompt = self.PROMPTS.get(model_name, self.PROMPTS["default"])

    def _get_inputs(self, pairs, max_sequence_length: int):
        prompt = self.prompt
        sep = "\n"
        prompt_inputs = self.tokenizer(
            prompt, return_tensors=None, add_special_tokens=False
        )["input_ids"]
        sep_inputs = self.tokenizer(sep, return_tensors=None, add_special_tokens=False)[
            "input_ids"
        ]
        inputs = []
        for query, passage in pairs:
            query_inputs = self.tokenizer(
                f"A: {query}",
                return_tensors=None,
                add_special_tokens=False,
                max_length=max_sequence_length * 3 // 4,
                truncation=True,
            )
            passage_inputs = self.tokenizer(
                f"B: {passage}",
                return_tensors=None,
                add_special_tokens=False,
                max_length=max_sequence_length,
                truncation=True,
            )
            item = self.tokenizer.prepare_for_model(
                [self.tokenizer.bos_token_id] + query_inputs["input_ids"],
                sep_inputs + passage_inputs["input_ids"],
                truncation="only_second",
                max_length=max_sequence_length,
                padding=False,
                return_attention_mask=False,
                return_token_type_ids=False,
                add_special_tokens=False,
            )
            item["input_ids"] = item["input_ids"] + sep_inputs + prompt_inputs
            item["attention_mask"] = [1] * len(item["input_ids"])
            inputs.append(item)

        return self.tokenizer.pad(
            inputs,
            padding=True,
            max_length=max_sequence_length + len(sep_inputs) + len(prompt_inputs),
            pad_to_multiple_of=8,
            return_tensors="pt",
        )

    @torch.no_grad()
    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks each document's **contexts** using the **LLM Layer-Wise Ranking approach**.

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

        Returns:
            List[Document]: The reranked list of **Document** instances with updated `reorder_contexts`.
        """
        for doc in tqdm(documents, desc="Reranking Documents"):
            query = doc.question.question
            pairs = [(query, context.text) for context in doc.contexts]
            batched_pairs = [
                pairs[i : i + self.batch_size] for i in range(0, len(pairs), self.batch_size)
            ]
            scores = []

            for batch in batched_pairs:
                inputs = self._get_inputs(batch, max_sequence_length=self.max_sequence_length)
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                print(self.params)
                #aaaaaaa
                outputs = self.model(**inputs, return_dict=True, **self.params)
                #print(f"[DEBUG] Hidden states returned: {len(outputs.hidden_states)}")

                all_scores = [
                    scores[:, -1]
                    .view(
                        -1,
                    )
                    .float()
                    for scores in outputs[0]
                ]
                batch_scores = all_scores[-1].cpu().numpy().tolist()

                scores.extend(batch_scores)

            # Assign scores and reorder
            contexts = copy.deepcopy(doc.contexts)
            for context, score in zip(contexts, scores):
                context.score = score

            doc.reorder_contexts = sorted(contexts, key=lambda x: x.score, reverse=True)

        return documents


    @torch.inference_mode()
    def score(self, query: str, doc: str) -> float:
        inputs = self._get_inputs(
            [(query, doc)], max_sequence_length=self.max_sequence_length
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        outputs = self.model(**inputs, return_dict=True, **self.params)
        all_scores = [
            scores[:, -1]
            .view(
                -1,
            )
            .float()
            for scores in outputs[0]
        ]
        score = all_scores[-1].item()

        return score

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

Initializes the LLM Layer-Wise Ranker.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	The name of the pre-trained model (default: "BAAI/bge-reranker-v2.5-gemma2-lightweight").	None
api_key	str	API key for authentication (if needed).	None
kwargs	dict	Additional configuration arguments (e.g., max_sequence_length, batch_size, device, dtype).	{}

Source code in rankify/models/llm_layerwise_ranker.py

def __init__(
    self, method = None, model_name = None, api_key = None, **kwargs):
    """
    Initializes the **LLM Layer-Wise Ranker**.

    Args:
        method (str, optional): The **reranking method name**.
        model_name (str, optional): The **name of the pre-trained model** (default: `"BAAI/bge-reranker-v2.5-gemma2-lightweight"`).
        api_key (str, optional): API key for **authentication** (if needed).
        kwargs (dict): Additional configuration arguments (e.g., `max_sequence_length`, `batch_size`, `device`, `dtype`).
    """

    max_sequence_length= kwargs.get("max_sequence_length", 1024)
    device = kwargs.get("device", "cuda")
    self.device = get_device(device)
    self.batch_size = kwargs.get("batch_size", 4)
    self.cutoff_layer =  kwargs.get("cutoff_layer", None)
    self.compress_ratio = kwargs.get("compress_ratio", None)  
    self.compress_layer =  kwargs.get("compress_layer", None)  
    self.prompt =  kwargs.get("prompt", None)  

    self.model_name = model_name

    self.dtype = get_dtype( kwargs.get("dtype", torch.float16), self.device)

    self.tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True,
    )
    self.max_sequence_length = max_sequence_length
    self.tokenizer.model_max_length = self.max_sequence_length
    self.tokenizer.padding_side = "right"


    self.model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        torch_dtype=self.dtype,
    ).to(self.device)
    self.model.eval()

     # Create params dict based on specified values or defaults
    params = {}
    if self.cutoff_layer is not None:
        params["cutoff_layers"] = self.cutoff_layer
    if self.compress_ratio is not None:
        params["compress_ratio"] = self.compress_ratio
    if self.compress_layer is not None:
        params["compress_layer"] = self.compress_layer
    if not params:
        params = self.DEFAULT_PARAMS.get(model_name, self.DEFAULT_PARAMS["default"])
    self.params = params

    self.prompt = self.prompt
    if self.prompt is None:
        self.prompt = self.PROMPTS.get(model_name, self.PROMPTS["default"])

rank(documents)

Reranks each document's contexts using the LLM Layer-Wise Ranking approach.

Parameters:

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

Returns:

Type	Description
List[Document]	List[Document]: The reranked list of Document instances with updated reorder_contexts.

Source code in rankify/models/llm_layerwise_ranker.py

@torch.no_grad()
def rank(self, documents: List[Document]) -> List[Document]:
    """
    Reranks each document's **contexts** using the **LLM Layer-Wise Ranking approach**.

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

    Returns:
        List[Document]: The reranked list of **Document** instances with updated `reorder_contexts`.
    """
    for doc in tqdm(documents, desc="Reranking Documents"):
        query = doc.question.question
        pairs = [(query, context.text) for context in doc.contexts]
        batched_pairs = [
            pairs[i : i + self.batch_size] for i in range(0, len(pairs), self.batch_size)
        ]
        scores = []

        for batch in batched_pairs:
            inputs = self._get_inputs(batch, max_sequence_length=self.max_sequence_length)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            print(self.params)
            #aaaaaaa
            outputs = self.model(**inputs, return_dict=True, **self.params)
            #print(f"[DEBUG] Hidden states returned: {len(outputs.hidden_states)}")

            all_scores = [
                scores[:, -1]
                .view(
                    -1,
                )
                .float()
                for scores in outputs[0]
            ]
            batch_scores = all_scores[-1].cpu().numpy().tolist()

            scores.extend(batch_scores)

        # Assign scores and reorder
        contexts = copy.deepcopy(doc.contexts)
        for context, score in zip(contexts, scores):
            context.score = score

        doc.reorder_contexts = sorted(contexts, key=lambda x: x.score, reverse=True)

    return documents

get_device(device, no_mps=False)

Source code in rankify/utils/helper.py

def get_device(
        device: Optional[Union[str, torch.device]],
        no_mps: bool = False,
    ) -> Union[str, torch.device]:
        if not device:
            if torch.cuda.is_available():
                device = "cuda"
            elif torch.backends.mps.is_available() and not no_mps:
                device = "mps"
            else:
                device = "cpu"
        return device

get_dtype(dtype, device, verbose=1)

Source code in rankify/utils/helper.py

def get_dtype(
        dtype: Optional[Union[str, torch.dtype]],
        device: Optional[Union[str, torch.device]],
        verbose: int = 1,
    ) -> torch.dtype:
        if dtype is None:
            print("No dtype set")
        if device == "cpu":
            dtype = torch.float32
        if not isinstance(dtype, torch.dtype):
            if dtype == "fp16" or "float16":
                dtype = torch.float16
            elif dtype == "bf16" or "bfloat16":
                dtype = torch.bfloat16
            else:
                dtype = torch.float32
        return dtype