ColBERT Reranker

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

ColBERTModel

Bases: BertPreTrainedModel

Source code in rankify/utils/models/colbert.py

class ColBERTModel(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config)

        # TODO: Load from artifact.metadata
        if "small" in config._name_or_path:
            linear_dim = 96
        else:
            linear_dim = 128
        #print(f"Linear Dim set to: {linear_dim} for downcasting")
        self.linear = nn.Linear(config.hidden_size, linear_dim, bias=False)
        self.init_weights()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=None,
        output_hidden_states=None,
    ):
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=True,  # Always output hidden states
        )

        sequence_output = outputs[0]

        return self.linear(sequence_output)

    def _encode(self, texts: list[str], insert_token_id: int, is_query: bool = False):
        encoding = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            max_length=self.max_length - 1,  # for insert token
            truncation=True,
        )
        encoding = _insert_token(encoding, insert_token_id)  # type: ignore

        if is_query:
            mask_token_id = self.tokenizer.mask_token_id

            new_encodings = {"input_ids": [], "attention_mask": []}

            for i, input_ids in enumerate(encoding["input_ids"]):
                original_length = (
                    (input_ids != self.tokenizer.pad_token_id).sum().item()
                )

                # Calculate QLEN dynamically for each query
                if original_length % 32 <= 8:
                    QLEN = original_length + 8
                else:
                    QLEN = ceil(original_length / 32) * 32

                if original_length < QLEN:
                    pad_length = QLEN - original_length
                    padded_input_ids = input_ids.tolist() + [mask_token_id] * pad_length
                    padded_attention_mask = (
                        encoding["attention_mask"][i].tolist() + [0] * pad_length
                    )
                else:
                    padded_input_ids = input_ids[:QLEN].tolist()
                    padded_attention_mask = encoding["attention_mask"][i][
                        :QLEN
                    ].tolist()

                new_encodings["input_ids"].append(padded_input_ids)
                new_encodings["attention_mask"].append(padded_attention_mask)

            for key in new_encodings:
                new_encodings[key] = torch.tensor(
                    new_encodings[key], device=self.device
                )

            encoding = new_encodings

        encoding = {key: value.to(self.device) for key, value in encoding.items()}
        return encoding

    def _query_encode(self, query: list[str]):
        return self._encode(query, self.query_token_id, is_query=True)

    def _document_encode(self, documents: list[str]):
        return self._encode(documents, self.document_token_id)

    def _to_embs(self, encoding) -> torch.Tensor:
        with torch.inference_mode():
            # embs = self.model(**encoding).last_hidden_state.squeeze(1)
            embs = self.model(**encoding)
        if self.normalize:
            embs = embs / embs.norm(dim=-1, keepdim=True)
        return embs

    def _rerank(self, query: str, documents: list[str]) -> list[float]:
        query_encoding = self._query_encode([query])
        documents_encoding = self._document_encode(documents)
        query_embeddings = self._to_embs(query_encoding)
        document_embeddings = self._to_embs(documents_encoding)
        scores = (
            _colbert_score(
                query_embeddings,
                document_embeddings,
                query_encoding["attention_mask"],
                documents_encoding["attention_mask"],
            )
            .cpu()
            .tolist()[0]
        )
        return scores

ColBERTReranker

Bases: BaseRanking

A reranking model that leverages ColBERT (Contextualized Late Interaction over BERT) to reorder document contexts based on semantic relevance.

ColBERT introduces a late interaction mechanism that efficiently compares query and document token embeddings to determine ranking scores.

Attributes:

Name	Type	Description
method	str	The reranking method name.
model_name	str	The pretrained ColBERT model to be used for reranking.
api_key	str	API key for external services (not used in this model).
device	str	The device on which the model runs (cuda if available, otherwise cpu).
batch_size	int	The batch size for encoding documents and queries (default: 32).
normalize	bool	Whether to normalize embeddings for cosine similarity computation (default: True).
query_token	str	Special token used to distinguish query embeddings (default: [unused0]).
document_token	str	Special token used to distinguish document embeddings (default: [unused1]).
tokenizer	PreTrainedTokenizer	The tokenizer used to process queries and contexts.
model	ColBERTModel	The ColBERT model for generating contextual embeddings.
query_max_length	int	The maximum token length for queries (default: 32).
doc_max_length	int	The maximum token length for documents (determined by model constraints).

References

• Khattab, Omar, and Matei Zaharia. ColBERT: Efficient and effective passage search via contextualized late interaction over BERT.
  Paper
• Santhanam, Keshav, et al. Colbertv2: Effective and efficient retrieval via lightweight late interaction.
  Paper

See Also

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

Example

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

question = Question("When did Thomas Edison invent the light bulb?")
answers = Answer(["1879"])
contexts = [
    Context(text="Lightning strike at Seoul National University", id=1),
    Context(text="Thomas Edison invented the light bulb in 1879", id=4),
]
document = Document(question=question, answers=answers, contexts=contexts)

# Initialize Reranking with ColBERTReranker
model = Reranking(method='colbert_ranker', model_name='Colbert')
model.rank([document])

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

Source code in rankify/models/colbert_ranker.py

class ColBERTReranker(BaseRanking):
    """
    A reranking model that leverages **ColBERT (Contextualized Late Interaction over BERT)** to reorder document contexts based on semantic relevance.

    ColBERT introduces a **late interaction mechanism** that efficiently compares query and document token embeddings to determine ranking scores.

    Attributes:
        method (str, optional): The reranking method name.
        model_name (str): The pretrained ColBERT model to be used for reranking.
        api_key (str, optional): API key for external services (not used in this model).
        device (str): The device on which the model runs (`cuda` if available, otherwise `cpu`).
        batch_size (int): The batch size for encoding documents and queries (default: 32).
        normalize (bool): Whether to normalize embeddings for cosine similarity computation (default: True).
        query_token (str): Special token used to distinguish query embeddings (default: `[unused0]`).
        document_token (str): Special token used to distinguish document embeddings (default: `[unused1]`).
        tokenizer (transformers.PreTrainedTokenizer): The tokenizer used to process queries and contexts.
        model (ColBERTModel): The ColBERT model for generating contextual embeddings.
        query_max_length (int): The maximum token length for queries (default: 32).
        doc_max_length (int): The maximum token length for documents (determined by model constraints).

    References:
        - **Khattab, Omar, and Matei Zaharia.** *ColBERT: Efficient and effective passage search via contextualized late interaction over BERT.*  
          [Paper](https://arxiv.org/abs/2004.12832)
        - **Santhanam, Keshav, et al.** *Colbertv2: Effective and efficient retrieval via lightweight late interaction.*  
          [Paper](https://aclanthology.org/2022.naacl-main.272/)

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

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

        question = Question("When did Thomas Edison invent the light bulb?")
        answers = Answer(["1879"])
        contexts = [
            Context(text="Lightning strike at Seoul National University", id=1),
            Context(text="Thomas Edison invented the light bulb in 1879", id=4),
        ]
        document = Document(question=question, answers=answers, contexts=contexts)

        # Initialize Reranking with ColBERTReranker
        model = Reranking(method='colbert_ranker', model_name='Colbert')
        model.rank([document])

        # Print reordered contexts
        print("Reordered Contexts:")
        for context in document.reorder_contexts:
            print(context.text)
        ```
    """
    def __init__(
        self,
        method: str = None,
        model_name: str = None,
        api_key: str = None,
        **kwargs,
    ):        
        """
        Initializes a **ColBERTReranker** instance.

        Args:
            method (str, optional): The reranking method name.
            model_name (str): The pretrained ColBERT model.
            api_key (str, optional): API key (not used).
            **kwargs: Additional parameters.

        Example:
            ```python
            model = ColBERTReranker(method='colbert_ranker', model_name='Colbert')
            ```
        """
        self.method = method
        self.model_name = model_name
        self.api_key = api_key
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.batch_size = kwargs.get("batch_size", 32)
        self.normalize = kwargs.get("normalize", True)
        self.query_token = kwargs.get("query_token", "[unused0]")
        self.document_token = kwargs.get("document_token", "[unused1]")

        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = ColBERTModel.from_pretrained(self.model_name).to(self.device)
        self.query_token_id = self.tokenizer.convert_tokens_to_ids(self.query_token)
        self.document_token_id = self.tokenizer.convert_tokens_to_ids(self.document_token)
        self.query_max_length = 32  # Lower bound
        self.doc_max_length = (
            self.model.config.max_position_embeddings - 2
        )  # Upper bound

        self.model.eval()

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

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

        Returns:
            List[Document]: Documents with updated `reorder_contexts`.

        Example:
            ```python
            reranked_documents = model.rank(documents)
            ```
        """
        for document in tqdm(documents, desc="Reranking Documents"):
            query = document.question.question
            if not document.contexts:
                # Skip document if no valid contexts
                print(f"[SKIP] Document skipped — no valid contexts. Query: '{query}'")
                continue
            contexts = [copy.deepcopy(ctx.text) for ctx in document.contexts]

            # Compute scores
            scores = self._colbert_rank(query, contexts)

            # Reorder contexts based on scores
            document.reorder_contexts = []
            ranked_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
            for index in ranked_indices:
                d =copy.deepcopy(document.contexts[index]) 
                d.score= scores[index]
                document.reorder_contexts.append(d)
        return documents

    @torch.inference_mode()
    def _colbert_rank(self, query: str, contexts: List[str]) -> List[float]:
        """
        Computes **relevance scores** for a query and its associated contexts.

        Args:
            query (str): The query string.
            contexts (List[str]): A list of context strings.

        Returns:
            List[float]: Relevance scores for each context.
        """
        query_encoding = self._query_encode([query])
        documents_encoding = self._document_encode(contexts)
        query_embeddings = self._to_embs(query_encoding)
        document_embeddings = self._to_embs(documents_encoding)
        scores = (
            _colbert_score(
                query_embeddings,
                document_embeddings,
                query_encoding["attention_mask"],
                documents_encoding["attention_mask"],
            )
            .cpu()
            .tolist()[0]
        )
        return scores

    def _query_encode(self, query: list[str]):
        """
        Encodes the query using **ColBERT tokenization**.

        Args:
            query (List[str]): List containing a single query.

        Returns:
            Dict[str, torch.Tensor]: Tokenized query encoding.
        """
        return self._encode(
            query, self.query_token_id, max_length=self.doc_max_length, is_query=True
        )

    def _document_encode(self, documents: list[str]):
        """
        Encodes document passages while ensuring alignment with **ColBERT constraints**.

        Args:
            documents (List[str]): List of document texts.

        Returns:
            Dict[str, torch.Tensor]: Tokenized document encoding.
        """
        tokenized_doc_lengths = [
            len(
                self.tokenizer.encode(
                    doc, max_length=self.doc_max_length, truncation=True
                )
            )
            for doc in documents
        ]
        max_length = max(tokenized_doc_lengths)
        max_length = (
            ceil(max_length / 32) * 32
        )  # Round up to the nearest multiple of 32
        max_length = max(
            max_length, self.query_max_length
        )  # Ensure not smaller than query_max_length
        max_length = int(
            min(max_length, self.doc_max_length)
        )  # Ensure not larger than doc_max_length
        return self._encode(documents, self.document_token_id, max_length)

    def _encode(
        self,
        texts: list[str],
        insert_token_id: int,
        max_length: int,
        is_query: bool = False,
    ):
        """
        Tokenizes and encodes text while inserting **ColBERT-specific tokens**.

        Args:
            texts (List[str]): List of input texts.
            insert_token_id (int): Special token ID for ColBERT.
            max_length (int): Maximum sequence length.
            is_query (bool, optional): Whether encoding is for a query.

        Returns:
            Dict[str, torch.Tensor]: Tokenized and padded encoding.
        """
        encoding = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            max_length=max_length - 1,  # for insert token
            truncation=True,
        )
        encoding = _insert_token(encoding, insert_token_id)  # type: ignore

        if is_query:
            mask_token_id = self.tokenizer.mask_token_id

            new_encodings = {"input_ids": [], "attention_mask": []}

            for i, input_ids in enumerate(encoding["input_ids"]):
                original_length = (
                    (input_ids != self.tokenizer.pad_token_id).sum().item()
                )

                # Calculate QLEN dynamically for each query
                if original_length % 16 <= 8:
                    QLEN = original_length + 8
                else:
                    QLEN = ceil(original_length / 16) * 16

                if original_length < QLEN:
                    pad_length = QLEN - original_length
                    padded_input_ids = input_ids.tolist() + [mask_token_id] * pad_length
                    padded_attention_mask = (
                        encoding["attention_mask"][i].tolist() + [0] *  pad_length
                    )
                else:
                    padded_input_ids = input_ids[:QLEN].tolist()
                    padded_attention_mask = encoding["attention_mask"][i][
                        :QLEN
                    ].tolist()

                new_encodings["input_ids"].append(padded_input_ids)
                new_encodings["attention_mask"].append(padded_attention_mask)

            for key in new_encodings:
                new_encodings[key] = torch.tensor(
                    new_encodings[key], device=self.device
                )

            encoding = new_encodings

        encoding = {key: value.to(self.device) for key, value in encoding.items()}
        return encoding

    def _to_embs(self, encoding) -> torch.Tensor:
        """
        Converts tokenized inputs into **ColBERT embeddings**.

        Args:
            encoding (Dict[str, torch.Tensor]): Tokenized input encoding.

        Returns:
            torch.Tensor: ColBERT embeddings.
        """
        with torch.inference_mode():
            batched_embs = []
            for i in range(0, encoding["input_ids"].size(0), self.batch_size):
                batch_encoding = {
                    key: val[i : i + self.batch_size] for key, val in encoding.items()
                }
                batch_embs = self.model(**batch_encoding)
                batched_embs.append(batch_embs)
            embs = torch.cat(batched_embs, dim=0)
        if self.normalize:
            embs = embs / embs.norm(dim=-1, keepdim=True)
        return embs

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

Initializes a ColBERTReranker instance.

Parameters:

Name	Type	Description	Default
method	str	The reranking method name.	None
model_name	str	The pretrained ColBERT model.	None
api_key	str	API key (not used).	None
**kwargs		Additional parameters.	{}

Example

model = ColBERTReranker(method='colbert_ranker', model_name='Colbert')

Source code in rankify/models/colbert_ranker.py

def __init__(
    self,
    method: str = None,
    model_name: str = None,
    api_key: str = None,
    **kwargs,
):        
    """
    Initializes a **ColBERTReranker** instance.

    Args:
        method (str, optional): The reranking method name.
        model_name (str): The pretrained ColBERT model.
        api_key (str, optional): API key (not used).
        **kwargs: Additional parameters.

    Example:
        ```python
        model = ColBERTReranker(method='colbert_ranker', model_name='Colbert')
        ```
    """
    self.method = method
    self.model_name = model_name
    self.api_key = api_key
    self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.batch_size = kwargs.get("batch_size", 32)
    self.normalize = kwargs.get("normalize", True)
    self.query_token = kwargs.get("query_token", "[unused0]")
    self.document_token = kwargs.get("document_token", "[unused1]")

    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    self.model = ColBERTModel.from_pretrained(self.model_name).to(self.device)
    self.query_token_id = self.tokenizer.convert_tokens_to_ids(self.query_token)
    self.document_token_id = self.tokenizer.convert_tokens_to_ids(self.document_token)
    self.query_max_length = 32  # Lower bound
    self.doc_max_length = (
        self.model.config.max_position_embeddings - 2
    )  # Upper bound

    self.model.eval()

rank(documents)

Reranks each document's contexts using ColBERT scoring and updates reorder_contexts.

Parameters:

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

Returns:

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

Example

reranked_documents = model.rank(documents)

Source code in rankify/models/colbert_ranker.py

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

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

    Returns:
        List[Document]: Documents with updated `reorder_contexts`.

    Example:
        ```python
        reranked_documents = model.rank(documents)
        ```
    """
    for document in tqdm(documents, desc="Reranking Documents"):
        query = document.question.question
        if not document.contexts:
            # Skip document if no valid contexts
            print(f"[SKIP] Document skipped — no valid contexts. Query: '{query}'")
            continue
        contexts = [copy.deepcopy(ctx.text) for ctx in document.contexts]

        # Compute scores
        scores = self._colbert_rank(query, contexts)

        # Reorder contexts based on scores
        document.reorder_contexts = []
        ranked_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
        for index in ranked_indices:
            d =copy.deepcopy(document.contexts[index]) 
            d.score= scores[index]
            document.reorder_contexts.append(d)
    return documents

_insert_token(output, insert_token_id, insert_position=1, token_type_id=0, attention_value=1)

Inserts a new token at a specified position into the sequences of a tokenized representation.

This function takes a dictionary containing tokenized representations (e.g., 'input_ids', 'token_type_ids', 'attention_mask') as PyTorch tensors, and inserts a specified token into each sequence at the given position. This can be used to add special tokens or other modifications to tokenized inputs.

• output (dict): A dictionary containing the tokenized representations. Expected keys are 'input_ids', 'token_type_ids', and 'attention_mask'. Each key is associated with a PyTorch tensor.
• insert_token_id (int): The token ID to be inserted into each sequence.
• insert_position (int, optional): The position in the sequence where the new token should be inserted. Defaults to 1, which typically follows a special starting token like '[CLS]' or '[BOS]'.
• token_type_id (int, optional): The token type ID to assign to the inserted token. Defaults to 0.
• attention_value (int, optional): The attention mask value to assign to the inserted token. Defaults to 1.

• updated_output (dict): A dictionary containing the updated tokenized representations, with the new token inserted at the specified position in each sequence. The structure and keys of the output dictionary are the same as the input.

Source code in rankify/utils/models/colbert.py

def _insert_token(
    output: dict,
    insert_token_id: int,
    insert_position: int = 1,
    token_type_id: int = 0,
    attention_value: int = 1,
):
    """
    Inserts a new token at a specified position into the sequences of a tokenized representation.

    This function takes a dictionary containing tokenized representations
    (e.g., 'input_ids', 'token_type_ids', 'attention_mask') as PyTorch tensors,
    and inserts a specified token into each sequence at the given position.
    This can be used to add special tokens or other modifications to tokenized inputs.

    Parameters:
    - output (dict): A dictionary containing the tokenized representations. Expected keys
                     are 'input_ids', 'token_type_ids', and 'attention_mask'. Each key
                     is associated with a PyTorch tensor.
    - insert_token_id (int): The token ID to be inserted into each sequence.
    - insert_position (int, optional): The position in the sequence where the new token
                                       should be inserted. Defaults to 1, which typically
                                       follows a special starting token like '[CLS]' or '[BOS]'.
    - token_type_id (int, optional): The token type ID to assign to the inserted token.
                                     Defaults to 0.
    - attention_value (int, optional): The attention mask value to assign to the inserted token.
                                        Defaults to 1.

    Returns:
    - updated_output (dict): A dictionary containing the updated tokenized representations,
                             with the new token inserted at the specified position in each sequence.
                             The structure and keys of the output dictionary are the same as the input.
    """
    updated_output = {}
    for key in output:
        updated_tensor_list = []
        for seqs in output[key]:
            if len(seqs.shape) == 1:
                seqs = seqs.unsqueeze(0)
            for seq in seqs:
                first_part = seq[:insert_position]
                second_part = seq[insert_position:]
                new_element = (
                    torch.tensor([insert_token_id])
                    if key == "input_ids"
                    else torch.tensor([token_type_id])
                )
                if key == "attention_mask":
                    new_element = torch.tensor([attention_value])
                updated_seq = torch.cat((first_part, new_element, second_part), dim=0)
                updated_tensor_list.append(updated_seq)
        updated_output[key] = torch.stack(updated_tensor_list)
    return updated_output

_colbert_score(q_reps, p_reps, q_mask, p_mask)

Source code in rankify/utils/models/colbert.py

def _colbert_score(q_reps, p_reps, q_mask: torch.Tensor, p_mask: torch.Tensor):
    # calc max sim
    # base code from: https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/BGE_M3/modeling.py

    # Assert that all q_reps are at least as long as the query length
    assert (
        q_reps.shape[1] >= q_mask.shape[1]
    ), f"q_reps should have at least {q_mask.shape[1]} tokens, but has {q_reps.shape[1]}"

    token_scores = torch.einsum("qin,pjn->qipj", q_reps, p_reps)
    token_scores = token_scores.masked_fill(p_mask.unsqueeze(0).unsqueeze(0) == 0, -1e4)
    scores, _ = token_scores.max(-1)
    scores = scores.sum(1) / q_mask.sum(-1, keepdim=True)
    return scores