# Integration Patterns Between PyLate and Sentence Transformers: A Technical Deep Dive

> Explore PyLate integration patterns with Sentence Transformers. Leverage tokenization, training, and advanced embeddings for efficient retrieval.

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

---

**PyLate implements ColBERT retrieval by subclassing `SentenceTransformer`, enabling seamless reuse of tokenization, training utilities, and multi-process encoding while adding token-level embeddings and hierarchical pooling.**

PyLate is a high-performance ColBERT implementation built directly on top of the Sentence Transformers ecosystem. Understanding the integration patterns between PyLate and Sentence Transformers is essential for leveraging existing training pipelines, model loading mechanisms, and distributed encoding capabilities while implementing late-interaction retrieval architectures.

## Core Integration Patterns Between PyLate and Sentence Transformers

PyLate follows three primary architectural patterns to extend Sentence Transformers functionality without breaking API compatibility.

### Model Inheritance and Subclassing

The fundamental integration pattern is direct inheritance. In [`pylate/models/colbert.py`](https://github.com/lightonai/pylate/blob/main/pylate/models/colbert.py), the `ColBERT` class extends `SentenceTransformer` (lines 34-50), immediately gaining access to model loading, tokenizer initialization, and the underlying `modules` list management.

After the base initialization, PyLate injects a linear projection layer via the `Dense` class if the checkpoint lacks one (lines 26-38). For Stanford-NLP-style checkpoints, weights are imported using `Dense.from_stanford_weights` (lines 58-71), ensuring compatibility with existing ColBERT pretrained models while maintaining the Sentence Transformers module structure.

### Prompt and Prefix Token Handling

PyLate extends the tokenization pipeline to support query and document prefixes without modifying the base Sentence Transformers tokenizer API. In [`colbert.py`](https://github.com/lightonai/pylate/blob/main/colbert.py), prefixes are added as dedicated tokens to the tokenizer vocabulary (lines 81-89).

The `insert_prefix_token` method (lines 70-82) creates a tensor that inserts the prefix ID at position 1 of each sequence. During `tokenize` (lines 104-125), PyLate sets maximum lengths (`query_length` or `document_length`), handles padding, tokenizes raw texts, then inserts the appropriate prefix token into `input_ids`, `attention_mask`, and `token_type_ids`.

When `attend_to_expansion_tokens` is enabled, the attention mask is forced to 1 for all tokens (line 122), ensuring expanded queries maintain full attention across all token positions.

### Multi-Process and Multi-GPU Encoding

PyLate delegates distributed encoding to the native `SentenceTransformer.encode_multi_process` API while wrapping data-chunking logic for ColBERT-specific arguments. The `encode_multi_process` method (lines 260-300 in [`colbert.py`](https://github.com/lightonai/pylate/blob/main/colbert.py)) splits input lists into chunks and pushes work items onto a `multiprocessing.Queue`, collecting results in order.

The chunk size defaults to `min(⌈len/len(pool["processes"])/10⌉, 5000)` (line 256), a heuristic borrowed directly from Sentence Transformers to balance memory usage and process communication overhead. The underlying pool is created by `start_multi_process_pool`, which forwards to `pylate.utils._start_multi_process_pool` (line 91), maintaining compatibility with the Sentence Transformers multiprocessing utilities in [`pylate/utils/multi_process.py`](https://github.com/lightonai/pylate/blob/main/pylate/utils/multi_process.py).

## Architectural Implementation Details

Beyond the high-level patterns, PyLate implements specific modifications to the Sentence Transformers encoding flow to support ColBERT's late-interaction architecture.

### Constructor and Dense Layer Injection

When initializing a `ColBERT` instance, the constructor first calls `SentenceTransformer.__init__` to load the base encoder and build the tokenizer. PyLate then checks for the existence of a dense projection layer. If missing, it initializes a `Dense` layer that projects token embeddings to the desired ColBERT dimensionality.

For checkpoints trained with the original Stanford ColBERT implementation, `Dense.from_stanford_weights` handles weight loading, mapping the original tensor names to the Sentence Transformers module structure. This ensures that `model.save_pretrained()` and `model.from_pretrained()` work identically to standard Sentence Transformers models.

### Tokenization Pipeline Modifications

The tokenization pipeline in PyLate extends `SentenceTransformer.tokenize` with prefix injection and length management. When `is_query=True`, the tokenizer applies `query_length` and prepends the `query_prefix` token; when `is_query=False`, it uses `document_length` and the `document_prefix`.

The `insert_prefix_token` method manipulates the token tensors directly, creating a new tensor with the prefix ID inserted at position 1 and shifting subsequent tokens right. This occurs after the standard tokenizer call but before returning the `BatchEncoding`, ensuring compatibility with the `encode` method's expectation of tokenized inputs.

### ColBERT-Specific Encoding Flow

The `encode` method (lines 84-93) mirrors the Sentence Transformers API but inserts ColBERT-specific post-processing steps. After the forward pass through the underlying transformer, PyLate applies:

- **Skip-list masking**: Punctuation tokens identified by `self.skiplist_mask` (lines 311-322) are masked out to prevent them from contributing to similarity scores.
- **Query expansion**: When enabled, the attention mask logic either keeps all tokens (`torch.ones_like`) or prunes padding based on the `pool_factor` setting (lines 220-228).
- **Hierarchical pooling**: If `pool_factor > 1`, `pool_embeddings_hierarchical` (lines 194-215) groups adjacent token embeddings using max pooling to reduce sequence length while preserving local structure.

Finally, token-level embeddings are optionally normalized, padded to a fixed length for batching, and quantized if specified by the user.

## Training Integration with Sentence Transformers

Because `ColBERT` inherits from `SentenceTransformer`, it integrates seamlessly with the `SentenceTransformerTrainer` and `SentenceTransformerTrainingArguments` classes. This compatibility allows users to leverage existing training loops, optimizers, mixed-precision handling, and distributed training strategies without modification.

The training examples in `examples/train/*.py` demonstrate standard fine-tuning workflows where a `ColBERT` model is passed directly to the trainer. Similarly, the test suite (`tests/test_*.py`) validates that the model behaves identically to native Sentence Transformers during training, including gradient accumulation and checkpoint saving.

## Code Examples

### Basic Inference

```python
from pylate import models

# Load a pre-trained ColBERT model (inherits from SentenceTransformer)

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

    query_prefix="[Q] ",              # optional custom prefix

    document_prefix="[D] "
)

# Encode a single query and a list of documents

query_emb = model.encode(
    "What is the capital of France?",
    is_query=True,                    # adds query prefix

    normalize_embeddings=True
)

doc_embs = model.encode(
    ["Paris is the capital of France.", "Berlin is the capital of Germany."],
    is_query=False,                   # adds document prefix

    normalize_embeddings=True,
    pooling_factor=2,                  # hierarchical pooling

)

```

*Source*: `ColBERT.encode` – see lines 84-93 in [`pylate/models/colbert.py`](https://github.com/lightonai/pylate/blob/main/pylate/models/colbert.py).

### Multi-Process Encoding

```python
from pylate import models

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

# Spin up a pool (4 processes by default)

pool = model.start_multi_process_pool()

sentences = [
    "The quick brown fox jumps over the lazy dog.",
    "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
    # … thousands of sentences …

]

# Encode the whole list efficiently

embeddings = model.encode_multi_process(
    sentences=sentences,
    pool=pool,
    batch_size=64,
    is_query=False,
    pool_factor=1,
)

model.stop_multi_process_pool(pool)   # clean up

```

*Source*: `ColBERT.encode_multi_process` – lines 260-300 in [`pylate/models/colbert.py`](https://github.com/lightonai/pylate/blob/main/pylate/models/colbert.py) and pool creation in [`pylate/utils/multi_process.py`](https://github.com/lightonai/pylate/blob/main/pylate/utils/multi_process.py).

### Training Setup

```python
from pylate import models
from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments

train_dataset = ...          # a list of {"query": ..., "positive": ..., "negative": ...}

validation_dataset = ...

model = models.ColBERT(
    "sentence-transformers/all-MiniLM-L6-v2",
    device="cuda",
    query_prefix="[Q] ",
    document_prefix="[D] ",
)

training_args = SentenceTransformerTrainingArguments(
    output_dir="./colbert-output",
    num_train_epochs=2,
    per_device_train_batch_size=32,
    learning_rate=2e-5,
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=validation_dataset,
)

trainer.train()

```

*Source*: Training utilities are part of Sentence-Transformers; PyLate’s model behaves like any other `SentenceTransformer`. See example scripts such as [`examples/train/contrastive.py`](https://github.com/lightonai/pylate/blob/main/examples/train/contrastive.py).

## Summary

- **PyLate** implements **ColBERT** as a direct subclass of `SentenceTransformer` in [`pylate/models/colbert.py`](https://github.com/lightonai/pylate/blob/main/pylate/models/colbert.py), inheriting model loading, tokenization, and training utilities.
- The integration supports **custom prefix injection** for queries and documents through modified `tokenize` and `insert_prefix_token` methods while preserving the original Sentence Transformers pipeline.
- **Multi-process encoding** leverages `encode_multi_process` and `start_multi_process_pool`, delegating to Sentence Transformers utilities while handling ColBERT-specific arguments like `pool_factor` and `is_query`.
- Because `ColBERT` is a `SentenceTransformer`, it works out-of-the-box with `SentenceTransformerTrainer` and `SentenceTransformerTrainingArguments`, enabling standard fine-tuning workflows without custom training loops.

## Frequently Asked Questions

### How does PyLate extend Sentence Transformers without breaking existing code?

PyLate extends Sentence Transformers through inheritance rather than composition. The `ColBERT` class in [`pylate/models/colbert.py`](https://github.com/lightonai/pylate/blob/main/pylate/models/colbert.py) calls `SentenceTransformer.__init__` directly (lines 34-50), ensuring all base attributes like `tokenizer` and `modules` are initialized exactly as existing code expects. Method overrides like `encode` and `tokenize` maintain the same signatures as the parent class but add ColBERT-specific post-processing, allowing drop-in replacement of `SentenceTransformer` with `ColBERT` in existing training scripts.

### Can I use standard Sentence Transformers training utilities with PyLate models?

Yes, because `ColBERT` inherits from `SentenceTransformer`, it is fully compatible with `SentenceTransformerTrainer` and `SentenceTransformerTrainingArguments`. You can pass a `ColBERT` instance directly to the trainer constructor exactly as you would with any other Sentence Transformers model. The training examples in [`examples/train/contrastive.py`](https://github.com/lightonai/pylate/blob/main/examples/train/contrastive.py) demonstrate this pattern, using standard optimizers and mixed-precision handling from the Sentence Transformers ecosystem without requiring custom training loops.

### What are the specific ColBERT arguments added to the encode method?

PyLate extends the standard `encode` signature with ColBERT-specific parameters including `is_query` (boolean to trigger query vs document prefix insertion), `pool_factor` (integer controlling hierarchical pooling of token embeddings), and `pooling_strategy` (specifying how to aggregate token vectors). The method also handles `query_length` and `document_length` parameters that override the standard `max_seq_length` depending on the `is_query` flag. These additions are processed within the `encode` method (lines 84-93 in [`colbert.py`](https://github.com/lightonai/pylate/blob/main/colbert.py)) after calling the base tokenization logic but before returning the final embeddings.

### How does multi-process encoding work with ColBERT-specific parameters?

Multi-process encoding in PyLate delegates to `SentenceTransformer.encode_multi_process` but wraps the data chunking logic to handle ColBERT arguments like `is_query` and `pool_factor`. The `encode_multi_process` method (lines 260-300 in [`colbert.py`](https://github.com/lightonai/pylate/blob/main/colbert.py)) splits input lists into chunks using the heuristic `min(⌈len/len(pool["processes"])/10⌉, 5000)` (line 256), identical to Sentence Transformers, but ensures each worker receives the ColBERT-specific kwargs. The pool is created via `start_multi_process_pool`, which forwards to `pylate.utils._start_multi_process_pool` (line 91), maintaining full compatibility with the Sentence Transformers multiprocessing utilities while supporting token-level embedding generation across multiple GPUs or CPU cores.