# How to Implement Reranking Functionality Using PyLate: A Complete Guide

> Learn how to implement reranking functionality using PyLate with our complete guide. Easily reorder documents using ColBERT scoring with query embeddings, doc embeddings, and doc IDs.

- Repository: [LightOn/pylate](https://github.com/lightonai/pylate)
- Tags: how-to-guide
- Published: 2026-03-06

---

**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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/pylate/rank/rank.py)** – Implements the main `rerank` function, `reshape_embeddings` for input normalization, and the `RerankResult` TypedDict structure.
- **[`pylate/scores/scores.py`](https://github.com/lightonai/pylate/blob/main/pylate/scores/scores.py)** – Contains `colbert_scores`, the similarity engine that computes the ColBERT scoring matrix.
- **[`pylate/utils/tensor.py`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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:

```python
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:

```python
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:

```python

# 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`](https://github.com/lightonai/pylate/blob/main/pylate/rank/rank.py). You can use it for type annotation in your own code:

```python
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:

```python

# 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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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.