# PyLate

> Train and serve ColBERT late-interaction retrieval models on top of Sentence Transformers

- Category: Retrieval & Knowledge (RAG) / Rerank, Search & Hybrid
- Language: Python
- Licence: MIT
- Repository: https://github.com/lightonai/pylate
- Website: https://lightonai.github.io/pylate/
- Canonical page: https://ai-tldr.dev/tools/pylate/

## In short

PyLate is a Python library for working with ColBERT-style models, where queries and documents are compared token by token (late interaction) instead of being squeezed into a single embedding. It sits in the Rerank, Search & Hybrid category of the Retrieval & Knowledge (RAG) tools and is open-source, written in Python and licensed MIT.

## Overview

PyLate is a Python library for working with ColBERT-style models, where queries and documents are compared token by token (late interaction) instead of being squeezed into a single embedding. It is built on top of Sentence Transformers and reuses its trainer and training arguments, so the workflow will feel familiar if you have used that library before.

It is aimed at developers and researchers who want to fine-tune their own retrieval or reranking models and run document retrieval over their data. You can start from most pre-trained language models (for example bert-base-uncased) and PyLate will turn them into ColBERT models, adding a linear projection layer when the base encoder is not already a ColBERT model.

Within the RAG and retrieval space, PyLate sits in the rerankers and hybrid-search area: late-interaction scoring tends to be more accurate than a plain single-vector search, which makes it useful as a stronger retriever or as a reranking step in a RAG pipeline.

## What it does

- Builds ColBERT models from most pre-trained language models, adding a linear layer when the base encoder is not already ColBERT
- Built on Sentence Transformers, reusing SentenceTransformerTrainer and SentenceTransformerTrainingArguments
- Contrastive training with options like a tunable temperature and CachedContrastive (GradCache) to emulate larger batches without more memory
- Knowledge distillation training from a strong teacher model's scores for higher quality
- Single- and multi-GPU training, including gathering examples across devices for larger effective batch sizes
- Built-in evaluation with ColBERTTripletEvaluator for held-out triplet eval sets

## Getting started

Install PyLate with pip, then load or train a ColBERT model. The example below trains a model with contrastive loss on MS MARCO triplets.

1. Install PyLate

   Install the package from PyPI. For evaluation dependencies, install the eval extra instead.

```bash
pip install pylate
```

2. Load or create a ColBERT model

   Pass any base model name; if it is not already a ColBERT model, PyLate adds a linear layer to the base encoder.

```python
from pylate import models

model = models.ColBERT(model_name_or_path="bert-base-uncased")
```

3. Train with contrastive loss

   Use the Sentence Transformers trainer with PyLate's loss, evaluator, and ColBERT data collator on a triplet dataset.

```python
import torch
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)

from pylate import evaluation, losses, models, utils

model = models.ColBERT(model_name_or_path="bert-base-uncased")

dataset = load_dataset("sentence-transformers/msmarco-bm25", "triplet", split="train")
splits = dataset.train_test_split(test_size=0.01)
train_dataset = splits["train"]
eval_dataset = splits["test"]

train_loss = losses.Contrastive(model=model)

dev_evaluator = evaluation.ColBERTTripletEvaluator(
    anchors=eval_dataset["query"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
)

args = SentenceTransformerTrainingArguments(
    output_dir="output/contrastive-bert-base-uncased",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    learning_rate=3e-6,
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=train_loss,
    evaluator=dev_evaluator,
    data_collator=utils.ColBERTCollator(model.tokenize),
)
trainer.train()
```

## When to use it

- Fine-tune a ColBERT retrieval model on your own data using contrastive training or knowledge distillation
- Add a late-interaction reranking step to a RAG pipeline for more accurate query-document scoring
- Turn an existing pre-trained language model into a ColBERT model without writing the architecture by hand
- Train on single or multiple GPUs, using cached contrastive loss or cross-device gathering to scale batch size