Semantic Chunker (Token Budget)

Start from brute force

The version most people write first adds the sentence, then checks: current.append(i), total += sz, and flush if total > max_tokens. It looks like the same greedy idea, and it does keep sentences whole and in order. It costs O(n) time and O(n) space for the output, so speed is not the problem. The order of the two steps is the problem: the sentence that causes the overflow has already joined the full chunk when the check fires. On the first example it returns [[0, 1, 2], [3, 4]], and the first chunk holds 3 + 4 + 2 = 9 tokens against a budget of 7.

The approach

The insight is to decide before adding a sentence. A chunk is flushed exactly when it is non-empty and the next sentence would push it past the budget, so every chunk either fits max_tokens or is one oversized sentence. Greedy is also the best possible packing by count: each chunk extends as far as the budget allows, so its end is never earlier than the matching chunk of any other valid split, and no in-order split uses fewer chunks.

The reference keeps chunks (the output), current (indices in the open chunk) and total (its running token count, so nothing is re-summed). For each i, sz it tests current and total + sz > max_tokens; if true it appends current to chunks and resets both. Then it always appends i and adds sz, which is how an oversized sentence gets a chunk of its own: the current and guard stops an empty chunk being flushed in front of it. After the loop, if current: flushes the last chunk, so empty input returns [].

Walk through an example

The first example, sizes = [3, 4, 2, 5, 1] and max_tokens = 7, one row per sentence as the reference processes it.

Complexity

Reference: O(n) time, one pass with constant work per sentence because total is a running sum, and O(n) space for the output, since every index appears once. Add-then-check: the same O(n) time and O(n) space; it is wrong, not slow. Re-summing current before each add would be correct but O(n²) when everything fits in one chunk.

Edge cases to test

• sizes = []: []. The final flush is skipped because current is empty, so no empty chunk appears.
• sizes = [4, 4], max_tokens = 4: [[0], [1]]. An exact fit is allowed (the test is >, not >=), and the second sentence starts a new chunk.
• sizes = [9, 1, 1], max_tokens = 5: [[0], [1, 2]]. An oversized first sentence gets its own chunk, and without the current and guard the output would start with [].
• sizes = [5, 5, 5], max_tokens = 4: [[0], [1], [2]]. Every sentence is oversized, and none is dropped.
• sizes = [1, 1, 1], max_tokens = 10: [[0, 1, 2]]. Everything fits, and the only flush is the final one.

Common mistakes

• Checking after appending: the sentence that overflows stays in the chunk, so chunk([3, 4, 2, 5, 1], 7) yields [[0, 1, 2], [3, 4]], a 9-token chunk.
• Dropping the current and guard: an oversized first sentence flushes an empty chunk, and chunk([5, 5, 5], 4) returns [[], [0], [1], [2]].
• Testing >= instead of >: a chunk that exactly fills the budget is cut early, and the first example becomes [[0], [1, 2], [3, 4]].
• Forgetting the final flush, which silently drops the last chunk: [4] disappears from the first example.
• Resetting current but not total after a flush: the stale total makes every later sentence overflow, and the first example becomes [[0, 1], [2], [3], [4]].

Follow-ups an interviewer asks next

Add overlap: each chunk should repeat the end of the previous one.

Overlap keeps a fact that sits on a boundary retrievable from both sides. After a flush, seed the new current with the trailing sentences of the flushed chunk whose sizes sum to at most an overlap budget, and start total at that sum. Drop the seed whenever seed plus the next sentence would itself exceed max_tokens, or the overlap breaks the budget it exists to serve.

Sentences arrive as a stream, not a list.

The reference already suits a stream: it reads each size once and never looks back. Turn it into a generator that yields current at every flush and once more at the end. Memory drops from every chunk to the open chunk alone, O(L) for L sentences per chunk, and the first chunk can be embedded before the document has finished arriving.

Greedy leaves a tiny last chunk. Can you balance the sizes?

Keep the chunk count greedy found, then binary search the smallest budget that still achieves it, running this same packer at each candidate budget; the count never rises as the budget grows, so the search is valid. On the first example, a budget of 6 still gives three chunks, [[0], [1, 2], [3, 4]] with 3, 6 and 6 tokens instead of 7, 7 and 1, while a budget of 5 needs five.

The token counts are not given. How do you get them?

Tokenize each sentence with the tokenizer of the embedding model the chunks are for, because different tokenizers count the same text differently, and cache the counts so re-chunking never re-tokenizes. Tokenizing a joined chunk can differ by a token or so from the sum of its sentences, since whitespace merges into neighboring tokens, so keep a small margin below the model’s hard limit.

Where this shows up

Chunking strategy directly caps RAG retrieval quality: if the key context is split across two chunks, neither retrieves well and the model answers from the wrong context. Packing whole sentences up to a token budget is the usual first defense: the budget keeps each chunk inside the embedding model’s input limit, and cutting only between sentences keeps every chunk readable on its own. LangChain’s recursive text splitter, for one, has the same split-then-merge shape: it splits on separators such as paragraph breaks, then merges the pieces up to a chunk size.

Reference solution

def chunk(sizes, max_tokens):
    chunks = []
    current = []
    total = 0
    for i, sz in enumerate(sizes):
        if current and total + sz > max_tokens:   # would overflow -> flush
            chunks.append(current)
            current = []
            total = 0
        current.append(i)                          # add whole sentence
        total += sz
    if current:
        chunks.append(current)                     # flush the last chunk
    return chunks


print(chunk([3, 4, 2, 5, 1], 7))   # -> [[0, 1], [2, 3], [4]]
print(chunk([5, 5, 5], 4))         # -> [[0], [1], [2]]

function chunk(sizes: number[], maxTokens: number): number[][] {
  const chunks: number[][] = [];
  let current: number[] = [];
  let total = 0;
  for (let i = 0; i < sizes.length; i++) {
    if (current.length && total + sizes[i] > maxTokens) {  // would overflow
      chunks.push(current);
      current = [];
      total = 0;
    }
    current.push(i);                                        // add whole sentence
    total += sizes[i];
  }
  if (current.length) chunks.push(current);                 // flush the last
  return chunks;
}

console.log(JSON.stringify(chunk([3, 4, 2, 5, 1], 7)));  // -> [[0,1],[2,3],[4]]
console.log(JSON.stringify(chunk([5, 5, 5], 4)));        // -> [[0],[1],[2]]