LateChunker

Late-interaction embeddings: whole-document embeddings, then mean-pooled per chunk with contextual token-count rebalancing.

Tip

When to use: When you want chunk-level embeddings that capture document-wide context (colbert-style late interaction). Each chunk carries an .embedding attribute — the mean of all token embeddings in that chunk.

Initialization

from sentence_transformers import SentenceTransformer
from blazechunk import LateChunker

model = SentenceTransformer("all-MiniLM-L6-v2")  # or any model with token-level output
chunker = LateChunker(
    embed_batch=model.encode,        # embedding function (returns token-level embeddings)
    chunk_size=2048,                 # target max tokens per chunk
    tokenizer="character",           # size unit
)

Parameters

Parameter	Type	Default	Description
embed_batch	Callable[[list[str]], NDArray]	—	Function that returns token-level embeddings: shape (batch_size, num_tokens, embedding_dim).
chunk_size	int	2048	Target maximum tokens per chunk after rebalancing.
tokenizer	str	"character"	Token counter name, or path to a tokenizer.json.

Usage

from sentence_transformers import SentenceTransformer
from blazechunk import LateChunker

model = SentenceTransformer("all-MiniLM-L6-v2")
chunker = LateChunker(model.encode, chunk_size=2048)

chunks = chunker.chunk(document)
for c in chunks:
    print(c.text, c.start_index, c.end_index, c.token_count)
    print("Embedding shape:", c.embedding.shape)  # (embedding_dim,)

How it works

1.Splits the document into initial chunks using RecursiveChunker.
2.Embeds the whole document with your model to get token-level embeddings.
3.For each chunk, extracts its token embeddings and computes the mean (late interaction).
4.Rebalances chunk boundaries to respect chunk_size while preserving semantic flow.
5.Attaches the mean embedding to each chunk's .embedding attribute.

Chunk object with embedding

class Chunk:
    text: str                    # The chunk text
    start_index: int             # Byte offset in original document
    end_index: int               # Byte offset in original document
    token_count: int             # Token count in this chunk
    embedding: NDArray           # Mean of token embeddings (shape: embedding_dim,)
    # Invariant: text == original_document[start_index:end_index]

Note

Late-interaction semantics: Unlike SemanticChunker (which chunks first, then embeds), LateChunker embeds the whole document first. This means each chunk carries contextual information from the entire document — useful for ColBERT-style retrieval.

Warning

Token-level embeddings required: Your embedding function must return token-level embeddings (shape: batch_size × num_tokens × embedding_dim), not sentence/document-level embeddings. Most modern models (transformers, sentence-transformers) support this.