PyLate

// 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.

pip install pylate

from pylate import models

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

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()

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// 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

// How PyLate compares

PyLate alongside other open-source rerank, search & hybrid tools AI/TLDR tracks, ranked by GitHub stars.