Chunking documents for retrieval

Step 4 of 8. Step 3 ended with the retrieval primitive at tag llm-ch03: a real model turns a sentence into 384 numbers, and a ten-line cosine function says which of two sentences a question points at more. This step points the same machinery at something retrieval actually consumes — a document too long to embed as one unit — and cuts it into chunks. By the end you have one document embedded twice under two different chunking strategies, plus one number for the mistake both strategies exist to avoid, and the scores make the trade-offs concrete: what overlap buys, what it costs, and what a chunk boundary quietly throws away. The step ends at tag llm-ch04.

Why the document cannot be embedded whole

The step’s document is 114 words about running a model locally. Embedding it is trivial — tokenizer.encode(DOC, true) runs — but the output line in this step’s run says something is wrong before any similarity is computed:

document: 114 words, 128 tokens

114 words produce more than 128 tokens, and this model’s tokenizer is configured to truncate at 128 and pad every input to exactly 128 — the count on that line is the ceiling, not the length. The whole-document vector never saw the last two sentences at all. A larger encoder would fit the text, and the problem would only hide: a document that spans several topics becomes one average of them, and an average of many directions points at none of them. The last block of output below makes that measurable — the whole document, embedded anyway, scores worse than any properly cut chunk.

So the document has to be cut, and the cut raises two questions with no universal answer: how big is a chunk, and do consecutive chunks share text? This step answers both by experiment rather than by doctrine — one chunker, two configurations, one question.

The chunker

The chunker is deliberately dumb: fixed windows of window words, advancing stride words each time. Setting stride equal to window makes adjacent chunks share nothing; setting it lower keeps the tail of each chunk at the start of the next.

/// Cut `doc` into word windows of `window` words, advancing `stride`
/// words each time. `stride == window` means adjacent chunks share
/// nothing; `stride < window` means each chunk keeps `window - stride`
/// words of the previous chunk as overlap.
fn chunk(doc: &str, window: usize, stride: usize) -> Vec<String> {
    let words: Vec<&str> = doc.split_whitespace().collect();
    let mut out = Vec::new();
    let mut start = 0;
    while start < words.len() {
        let end = (start + window).min(words.len());
        out.push(words[start..end].join(" "));
        if end == words.len() {
            break;
        }
        start += stride;
    }
    out
}

Chunk size is measured in words here, but the budget that matters is tokens — the same tokenizer from step 3 counts them, and the tokenizer has opinions. The model was trained on inputs of at most 128 word pieces, its tokenizer truncates there, and a 30-word chunk lands at 34–43 tokens (the [CLS] and [SEP] markers add two) — comfortably inside. That headroom is what lets overlap be nearly free at this size: a 6-word overlap costs 6 words out of a 128-token budget that 30-word chunks barely use.

The whole program

The document is embedded three ways: 30-word windows with no overlap, 30-word windows with a 6-word overlap, and the whole document as a single chunk. One question — about the document’s memory footprint — is scored against every chunk from each set. embed is the tokenize, forward, mean-pool pipeline from step 3, lifted into a function so it can run on any batch; cosine is unchanged.

[package]
name = "llm-apps"
version = "0.1.0"
edition = "2024"

[dependencies]
anyhow = "1.0.104"
# The runtime the whole track runs on. `default-features = false` keeps the
# CPU path only - no Metal, no CUDA - which is everything this track needs
# and the shortest dependency tree candle has.
candle-core = { version = "0.11.0", default-features = false }
candle-nn = "0.11.0"
candle-transformers = "0.11.0"
hf-hub = "0.3.2"
serde_json = "1.0.154"
tokenizers = "0.23.2"

//! Step 4: chunking documents for retrieval.
//!
//! Same model and same cosine arithmetic as step 3, but the input is a
//! whole document instead of three sentences. The document is cut into
//! fixed 30-word windows - once with no overlap, once with a 6-word
//! overlap - every chunk is embedded, and one question is scored against
//! each set. The chunk that wins changes, and that difference is what
//! overlap buys.

use anyhow::Result;
use candle_core::{Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config, DTYPE};
use hf_hub::api::sync::Api;
use tokenizers::Tokenizer;

const MODEL_ID: &str = "sentence-transformers/all-MiniLM-L6-v2";

const DOC: &str = "Running a sentence encoder on your own machine is mostly an exercise in memory budgeting. The runtime in this project is candle, which loads the safetensors weights through memory mapping, so nothing is copied until a layer actually runs. Startup cost is dominated entirely by paging the model in on the first forward pass. Its peak memory footprint is about 340 megabytes. That is with three sentences in the batch; a long document pushes the number up. Tokenization itself is cheap, microseconds per sentence, and it happens before anything reaches the model. Embedding a whole document at once is the mistake this step fixes: the vector averages every topic together and retrieval degrades quietly.";

const QUESTION: &str = "How much memory does inference use at peak?";

/// Cut `doc` into word windows of `window` words, advancing `stride`
/// words each time. `stride == window` means adjacent chunks share
/// nothing; `stride < window` means each chunk keeps `window - stride`
/// words of the previous chunk as overlap.
fn chunk(doc: &str, window: usize, stride: usize) -> Vec<String> {
    let words: Vec<&str> = doc.split_whitespace().collect();
    let mut out = Vec::new();
    let mut start = 0;
    while start < words.len() {
        let end = (start + window).min(words.len());
        out.push(words[start..end].join(" "));
        if end == words.len() {
            break;
        }
        start += stride;
    }
    out
}

fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let len_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let len_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    dot / (len_a * len_b)
}

/// Tokenize, forward, mean-pool: one Vec<f32> of 384 numbers per input.
fn embed(
    model: &BertModel,
    device: &Device,
    tokenizer: &Tokenizer,
    texts: &[String],
) -> Result<Vec<Vec<f32>>> {
    let encodings = tokenizer
        .encode_batch(texts.to_vec(), true)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let len = encodings.iter().map(|e| e.get_ids().len()).max().unwrap();
    let token_ids: Vec<Vec<u32>> = encodings
        .iter()
        .map(|e| {
            let mut ids = e.get_ids().to_vec();
            ids.resize(len, 0);
            ids
        })
        .collect();
    let attention: Vec<Vec<u32>> = encodings
        .iter()
        .map(|e| {
            let mut mask = e.get_attention_mask().to_vec();
            mask.resize(len, 0);
            mask
        })
        .collect();
    let types = vec![vec![0u32; len]; texts.len()];

    let ids = Tensor::new(token_ids, device)?;
    let mask = Tensor::new(attention.clone(), device)?;
    let types = Tensor::new(types, device)?;

    let hidden = model.forward(&ids, &types, Some(&mask))?.to_vec3::<f32>()?;
    Ok(hidden
        .iter()
        .zip(attention.iter())
        .map(|(tokens, mask)| {
            let mut v = vec![0f32; tokens[0].len()];
            let mut n = 0f32;
            for (t, &keep) in tokens.iter().zip(mask) {
                if keep == 1 {
                    n += 1.0;
                    for (v, &x) in v.iter_mut().zip(t) {
                        *v += x;
                    }
                }
            }
            v.iter_mut().for_each(|x| *x /= n);
            v
        })
        .collect())
}

/// The first 10 words of a chunk, for the printed table.
fn preview(s: &str) -> String {
    let words: Vec<&str> = s.split_whitespace().take(10).collect();
    let mut out = words.join(" ");
    if s.split_whitespace().count() > 10 {
        out.push_str(" ...");
    }
    out
}

fn main() -> Result<()> {
    let device = Device::Cpu;
    let started = std::time::Instant::now();

    let api = Api::new()?;
    let repo = api.model(MODEL_ID.to_string());
    let config: Config = serde_json::from_str(&std::fs::read_to_string(repo.get("config.json")?)?)?;
    let tokenizer = Tokenizer::from_file(repo.get("tokenizer.json")?)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let weights = repo.get("model.safetensors")?;
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], DTYPE, &device)? };
    let model = BertModel::load(vb, &config)?;
    println!("model loaded in {:.1?} on {device:?}", started.elapsed());

    // The model's tokenizer pads every input to 128 tokens, so the real
    // length is the number of positions the attention mask marks as
    // actual tokens - not the number of ids in the encoding.
    let real_tokens = |e: &tokenizers::Encoding| e.get_attention_mask().iter().sum::<u32>();
    let doc_tokens = real_tokens(&tokenizer.encode(DOC, true).map_err(|e| anyhow::anyhow!("{e}"))?);
    println!(
        "document: {} words, {doc_tokens} tokens",
        DOC.split_whitespace().count()
    );

    for (label, stride) in [("no overlap", 30), ("6-word overlap", 24)] {
        let chunks = chunk(DOC, 30, stride);
        println!(
            "\n{label}: {} chunks of 30 words, stride {stride}",
            chunks.len()
        );
        for (i, c) in chunks.iter().enumerate() {
            let tokens = real_tokens(
                &tokenizer.encode(c.as_str(), true).map_err(|e| anyhow::anyhow!("{e}"))?,
            );
            println!("  chunk {i}: {tokens} tokens, starts {:?}", preview(c));
        }
        let mut texts = chunks.clone();
        texts.push(QUESTION.to_string());
        let vectors = embed(&model, &device, &tokenizer, &texts)?;
        let q = vectors.last().unwrap();
        let mut scored: Vec<(usize, f32)> = vectors[..vectors.len() - 1]
            .iter()
            .enumerate()
            .map(|(i, v)| (i, cosine(v, q)))
            .collect();
        scored.sort_by(|a, b| b.1.total_cmp(&a.1));
        println!("  question: {QUESTION:?}");
        println!(
            "    best chunk {}: {:.4}, then chunk {}: {:.4}",
            scored[0].0, scored[0].1, scored[1].0, scored[1].1
        );
    }

    let whole = vec![DOC.to_string(), QUESTION.to_string()];
    let vectors = embed(&model, &device, &tokenizer, &whole)?;
    println!(
        "\nwhole document as one chunk: {:.4}",
        cosine(&vectors[0], &vectors[1])
    );

    Ok(())
}

Running it (cargo run -q, the model weights come from the local Hugging Face cache) prints:

model loaded in 93.4ms on Cpu
document: 114 words, 128 tokens

no overlap: 4 chunks of 30 words, stride 30
  chunk 0: 41 tokens, starts "Running a sentence encoder on your own machine is mostly ..."
  chunk 1: 35 tokens, starts "so nothing is copied until a layer actually runs. Startup ..."
  chunk 2: 43 tokens, starts "340 megabytes. That is with three sentences in the batch; ..."
  chunk 3: 34 tokens, starts "the model. Embedding a whole document at once is the ..."
  question: "How much memory does inference use at peak?"
    best chunk 2: 0.4711, then chunk 1: 0.3486

6-word overlap: 5 chunks of 30 words, stride 24
  chunk 0: 41 tokens, starts "Running a sentence encoder on your own machine is mostly ..."
  chunk 1: 38 tokens, starts "the safetensors weights through memory mapping, so nothing is copied ..."
  chunk 2: 39 tokens, starts "in on the first forward pass. Its peak memory footprint ..."
  chunk 3: 42 tokens, starts "document pushes the number up. Tokenization itself is cheap, microseconds ..."
  chunk 4: 25 tokens, starts "at once is the mistake this step fixes: the vector ..."
  question: "How much memory does inference use at peak?"
    best chunk 2: 0.5836, then chunk 0: 0.2620

whole document as one chunk: 0.3613

What the code does, move by move

chunk is the whole policy. Two numbers, window and stride, completely determine the segmentation — there is no sentence detection, no semantic boundary, nothing clever. Real chunkers add sentence-aware splitting on top, but the fixed window is the baseline every such refinement has to beat, and its failure modes are the reason the refinements exist.

The batch grows by one. The question is appended to the chunk list before embed, so one forward pass produces the question’s vector and every chunk’s vector together. Padding is to the longest text in this batch, which is why chunk token counts differ from each other — each run batches five or six short texts, not 128-token slabs.

Sorting, not argmax, on purpose. The scores are collected into a Vec<(usize, f32)> and sorted descending, and the print takes the top two. Retrieval rarely wants only the winner: the runner-up is what a re-ranking stage (a step this track skips) would reorder, and seeing the gap between first and second — 0.32 with overlap, 0.12 without — is how confidence gets calibrated later.

Reading the numbers

The document says “Its peak memory footprint is about 340 megabytes.” With no overlap, that sentence does not exist anywhere. Chunk 1 ends with “Its peak memory footprint is about” — a subject with no number — and scores 0.3486. Chunk 2 begins “340 megabytes.” — a number with no subject — and scores 0.4711, the higher of the two because “340 megabytes” is the lexically distinctive half. Neither chunk is wrong; the fact was cut in half by a boundary neither chunk can see across.

With a 6-word overlap, one chunk contains the whole sentence, and it scores 0.5836 — the best match in the run by a wide margin. That is what overlap buys: sentences that straddle an arbitrary boundary get to exist intact somewhere in the index.

What it costs is on the page too. Four chunks became five, so indexing this document means one extra forward pass and one extra 384-number vector to store — at 30-word windows, roughly 25% more embeddings per document, and every overlapping chunk repeats 6 words of its neighbor. Overlap also does not repair grammar: the no-overlap chunk that starts “340 megabytes.” is a fragment, and fragments embed anyway. Whether a fragment’s vector hurts depends on what the reader asks; nothing in this pipeline notices.

And the third number is the floor. The whole document as one chunk scores 0.3613 — below every cut version’s best chunk, and that is with the vector’s other handicap: truncated at the tokenizer’s 128-token ceiling, it never even saw the document’s closing sentences. Chunking is not an optimization here. It is the difference between the fact being retrievable and not.

The tag

The step’s end state is tagged in the examples repository:

git checkout llm-ch04

At that tag the project cuts a document into chunks two different ways, embeds every chunk, and scores a question against each set — the index-side half of RAG, working end to end. Step 5 gives those vectors a home: a vector store in SQLite, brute force first, then an index. The track overview keeps the full list current as articles land.

Embeddings and vector similarity — the previous step, whose tokenizer, mean-pooling and cosine this step reuses
What we’re building: a local RAG service — the map, if this is your first step in the track