How to Implement Reranking Functionality Using PyLate: A Complete Guide

PyLate provides a dedicated reranking API via pylate.rank.rerank that reorders retrieved documents using the ColBERT scoring function, requiring only query embeddings, document embeddings, and document IDs to produce ranked results.

To implement reranking functionality using PyLate, you leverage the stateless reranking module located in pylate/rank/rank.py. This module integrates seamlessly with the library's ColBERT model implementation, allowing you to refine initial retrieval results by computing token-level similarity scores between queries and candidate documents.

Understanding the PyLate Reranking Architecture

Core Implementation Files

The reranking system in PyLate is distributed across several specialized modules:

• pylate/rank/rank.py – Implements the main rerank function, reshape_embeddings for input normalization, and the RerankResult TypedDict structure.
• pylate/scores/scores.py – Contains colbert_scores, the similarity engine that computes the ColBERT scoring matrix.
• pylate/utils/tensor.py – Provides convert_to_tensor (aliased as func_convert_to_tensor in the rank module) for flexible input conversion.

The Reranking Pipeline Steps

When you call rank.rerank(), the implementation executes the following sequence as defined in pylate/rank/rank.py:

1. Input normalization – reshape_embeddings expands 2-D embeddings to 3-D tensors (batch, tokens, dim) to handle single queries or document lists uniformly.
2. Tensor conversion – Embeddings are converted to torch.Tensor objects using func_convert_to_tensor.
3. Padding – Document embeddings per query are padded to equal length using torch.nn.utils.rnn.pad_sequence.
4. Device handling – Tensors move to the specified device (GPU/CPU) or default to the query tensor's device.
5. Scoring – colbert_scores computes the similarity matrix, extracts maximal token-wise similarities, and sums across tokens.
6. Sorting – Scores sort in descending order and pair with original document IDs.
7. Result aggregation – Returns lists of RerankResult dictionaries containing id and score keys.

How to Implement Reranking Functionality Using PyLate

Basic Reranking After Retrieval

The standard workflow involves encoding queries and candidate documents, then calling the rerank function:

from pylate import models, rank

# Initialize ColBERT model

model = models.ColBERT(
    model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
    device="cpu",
)

# Candidate documents from initial retrieval (e.g., FAISS)

candidate_texts = [
    ["doc A", "doc B", "doc C"],
    ["doc X", "doc Y"]
]

candidate_ids = [
    [101, 102, 103],
    [201, 202]
]

# Encode queries and documents

queries = ["first query", "second query"]
q_emb = model.encode(queries, is_query=True, batch_size=1)
d_emb = model.encode(candidate_texts, is_query=False, batch_size=1)

# Rerank

reranked = rank.rerank(
    documents_ids=candidate_ids,
    queries_embeddings=q_emb,
    documents_embeddings=d_emb,
    device="cpu"
)

# Process results

for i, results in enumerate(reranked):
    print(f"\nQuery {i+1} ranking:")
    for rank_pos, res in enumerate(results, start=1):
        print(f"  {rank_pos}. doc_id={res['id']}  score={res['score']:.4f}")

GPU Acceleration for Reranking

To implement reranking functionality using PyLate on GPU, specify the device parameter:

reranked_gpu = rank.rerank(
    documents_ids=candidate_ids,
    queries_embeddings=q_emb,
    documents_embeddings=d_emb,
    device="cuda:0"
)

Automatic Device Detection

If you omit the device parameter, PyLate automatically uses the device of the query embeddings:

# If q_emb is on GPU, reranking executes on GPU automatically

reranked_auto = rank.rerank(
    documents_ids=candidate_ids,
    queries_embeddings=q_emb,
    documents_embeddings=d_emb
)

Advanced Reranking Techniques

Working with RerankResult Types

The RerankResult type is a TypedDict defined in pylate/rank/rank.py. You can use it for type annotation in your own code:

from typing import List
from pylate.rank.rank import RerankResult

def process_rankings(rankings: List[List[RerankResult]]) -> None:
    for query_results in rankings:
        for result in query_results:
            doc_id = result["id"]
            score = result["score"]
            # Custom processing logic here

Integrating with Existing Retrieval Pipelines

Because the PyLate reranking module is stateless, you can integrate it with any initial retrieval system:

# Example: Using with FAISS initial retrieval

import faiss
import numpy as np

# Assume index and model already initialized

def retrieve_and_rerank(query, k=100):
    # Initial FAISS retrieval

    q_emb = model.encode([query], is_query=True)
    distances, indices = index.search(q_emb, k)
    
    # Fetch candidate texts and embeddings

    candidate_texts = [[doc_store[idx] for idx in indices[0]]]
    candidate_ids = [indices[0].tolist()]
    
    d_emb = model.encode(candidate_texts, is_query=False)
    
    # Rerank with PyLate

    return rank.rerank(
        documents_ids=candidate_ids,
        queries_embeddings=q_emb,
        documents_embeddings=d_emb
    )

Summary

• PyLate reranking is implemented in pylate/rank/rank.py and provides a stateless API for reordering documents using ColBERT scoring.
• The rerank function requires three inputs: documents_ids (nested list), queries_embeddings, and documents_embeddings.
• ColBERT scoring occurs in pylate/scores/scores.py via colbert_scores, computing token-level maximum similarities summed across dimensions.
• The module handles automatic tensor conversion, padding, device placement, and result sorting internally.
• GPU acceleration is available via the device parameter or automatic detection from query embeddings.
• Results return as RerankResult TypedDict objects containing id and score keys.

Frequently Asked Questions

What is the difference between retrieval and reranking in PyLate?

Retrieval in PyLate typically involves fast approximate nearest neighbor (ANN) search using indices like FAISS or ScaNN to fetch candidate documents. Reranking occurs after retrieval, using the pylate.rank.rerank function to reorder those candidates by computing precise ColBERT similarity scores between the query and each document. The reranking step is more computationally expensive but provides higher accuracy than initial retrieval alone.

Can I use PyLate reranking with FAISS or ScaNN indices?

Yes, the PyLate reranking module is stateless and designed to integrate with any initial retrieval system, including FAISS, ScaNN, or PyLate's own index wrappers. You perform initial retrieval using your preferred index to obtain candidate document IDs and texts, encode these candidates with the ColBERT model, then pass the embeddings to rank.rerank for final ordering.

How does the ColBERT scoring function work in reranking?

The colbert_scores function in pylate/scores/scores.py computes a similarity matrix between query tokens and document tokens, extracts the maximum similarity for each query token across all document tokens, and sums these maximum values across the token dimension. This "late interaction" approach captures fine-grained token-level relevance while maintaining computational efficiency during the reranking phase.

Is the PyLate reranking module stateful or stateless?

The PyLate reranking module is stateless. It does not maintain internal state between calls; each invocation of rank.rerank processes the provided embeddings and IDs independently. This design makes it straightforward to deploy in distributed systems, serverless functions, or multi-threaded applications where you need to rerank results on demand without managing persistent model state.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}