PyLate

When backend="auto" is used, PyLate attempts to dispatch to the fastest available kernel based on your hardware and installed dependencies. The fallback logic is:

1. Try flash: Requires CUDA tensors, flash-maxsim to be installed, and the input shape to be supported.
2. Try lik: If flash fails, it tries late-interaction-kernels. Requires CUDA or MPS tensors, late-interaction-kernels to be installed, and supported dtype/head-dim.
3. Fallback to torch: If no kernel backend is compatible, it silently falls back to the standard PyTorch implementation.

Note on Error Handling:

• If you use backend="auto", failures in kernel compatibility are caught silently to allow the fallback.
• If you explicitly request backend="flash" or backend="lik", a precondition failure will raise an error instead of falling back, ensuring you are aware if your hardware/inputs do not support the requested optimization.

To create a searchable index, you must load a ColBERT model, initialize a PLAID index, encode your documents into embeddings, and then add those embeddings to the index.

Key steps:

1. Load Model: Use models.ColBERT(model_name_or_path=...).
2. Initialize Index: Use indexes.PLAID(index_folder=..., index_name=..., override=True). Setting override=True will replace any existing index with the same name in that folder.
3. Encode Documents: Use model.encode(documents, is_query=False, ...) to generate embeddings. For memory optimization, you can use pool_factor to compress embeddings.
4. Add to Index: Use index.add_documents(documents_ids=..., documents_embeddings=...) to populate the index.

from pylate import indexes, models

# 1. Load Model
model = models.ColBERT(model_name_or_path="lightonai/GTE-ModernColBERT-v1")

# 2. Initialize Index
index = indexes.PLAID(index_folder="pylate-colbert-index", index_name="my_documents", override=True)

# 3. Encode Documents
doc_ids = ["doc_001", "doc_002"]
docs = ["The Eiffel Tower...", "The Louvre..."]
doc_embeddings = model.encode(docs, batch_size=32, is_query=False, show_progress_bar=True)

# 4. Add to Index
index.add_documents(documents_ids=doc_ids, documents_embeddings=doc_embeddings)

To perform efficient similarity search, you can index documents using a ColBERT model and a PLAID index. This process involves loading the model, initializing the index, encoding the documents, and adding the embeddings to the index.

Note: When encoding documents, set is_query=False. When initializing the indexes.PLAID object, use override=True to overwrite any existing index in that folder.

from pylate import indexes, models, retrieve

# Step 1: Load the ColBERT model
model = models.ColBERT(
    model_name_or_path="pylate_model_id",
)

# Step 2: Initialize the PLAID index
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # This overwrites the existing index if any
)

# Step 3: Encode the documents
documents_ids = ["1", "2", "3"]
documents = ["document 1 text", "document 2 text", "document 3 text"]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # Ensure that it is set to False to indicate that these are documents, not queries
    show_progress_bar=True,
)

# Step 4: Add document embeddings to the index
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

Knowledge distillation trains ColBERT models to replicate the outputs of a teacher model (e.g., a cross-encoder). This requires a dataset containing queries, documents, and relevance scores.

Key steps:

1. Load the training, queries, and documents datasets.
2. Use utils.KDProcessing to transform the training dataset so it loads text on the fly using the provided query/document IDs.
3. Use losses.Distillation as the loss function.
4. Use utils.ColBERTCollator(tokenize_fn=model.tokenize) in the trainer.

import torch
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from pylate import losses, models, utils

# 1. Load datasets
train = load_dataset(path="lightonai/ms-marco-en-bge", name="train")
queries = load_dataset(path="lightonai/ms-marco-en-bge", name="queries")
documents = load_dataset(path="lightonai/ms-marco-en-bge", name="documents")

# 2. Set transformation for on-the-fly text loading
train.set_transform(utils.KDProcessing(queries=queries, documents=documents).transform)

# 3. Model and training setup
model = models.ColBERT(model_name_or_path="bert-base-uncased")
model = torch.compile(model)

args = SentenceTransformerTrainingArguments(
    output_dir="output/knowledge-distillation-bert-base",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    fp16=True,
    run_name="knowledge-distillation-bert-base",
    learning_rate=1e-5,
)

# 4. Loss and Trainer
train_loss = losses.Distillation(model=model)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train,
    loss=train_loss,
    data_collator=utils.ColBERTCollator(tokenize_fn=model.tokenize),
)

trainer.train()

To perform efficient retrieval, use a ColBERT model with a PLAID index. The process involves loading the model, initializing the index, and using a retrieve.ColBERT object to execute queries against the index.

Workflow

1. Initialize Model: Load a ColBERT model using models.ColBERT.
2. Initialize Index: Create an indexes.PLAID instance specifying an index_folder and index_name.
3. Setup Retriever: Create a retrieve.ColBERT instance passing the index.
4. Encode Documents: Use model.encode with is_query=False to generate embeddings for your documents.
5. Add to Index: Use index.add_documents to store document IDs and their embeddings.
6. Retrieve: Encode queries with is_query=True and call retriever.retrieve with the desired k value.

from pylate import indexes, models, retrieve

# 1. Load model
model = models.ColBERT(
    model_name_or_path="lightonai/GTE-ModernColBERT-v1",
)

# 2. Initialize PLAID index
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,
)

# 3. Setup retriever
retriever = retrieve.ColBERT(index=index)

# 4. Encode and add documents
doc_ids = ["1", "2"]
docs = ["Doc text 1", "Doc text 2"]

doc_embeddings = model.encode(
    docs,
    batch_size=32,
    is_query=False,
    show_progress_bar=True,
)

index.add_documents(
    documents_ids=doc_ids,
    documents_embeddings=doc_embeddings,
)

# 5. Retrieve
query_embeddings = model.encode(
    ["query text"],
    batch_size=32,
    is_query=True,
    show_progress_bar=True,
)

scores = retriever.retrieve(
    queries_embeddings=query_embeddings,
    k=10,
)
print(scores)