CODING CHALLENGE · N°53

Semantic Chunker (Token Budget)

Easy AIRAGChunking

Before you can embed a document for retrieval, you have to split it into chunks that fit a token budget — without slicing through a sentence. The greedy packer fills each chunk with whole sentences until the next one would overflow, then starts fresh. It’s the unglamorous step that makes or breaks a RAG pipeline. Solve it in Python or TypeScript, with hidden tests.

The problem

You are given sizes, the token count of each sentence in order, and a max_tokens budget per chunk. Greedily pack consecutive sentences into a chunk; when adding the next sentence would exceed the budget, start a new chunk. A single sentence larger than the budget becomes its own chunk. Implement chunk(sizes, max_tokens): return a list of chunks, each a list of the sentence indices it contains.

EXAMPLE 1
Input sizes = [3, 4, 2, 5, 1], max_tokens = 7
Output [[0, 1], [2, 3], [4]]
3+4=7 fits; +2 would overflow → new chunk; 2+5=7 fits; +1 overflows
EXAMPLE 2
Input sizes = [5, 5, 5], max_tokens = 4
Output [[0], [1], [2]]
every sentence exceeds the budget, so each is its own chunk
CONSTRAINTS
  • Sentences stay in order and are never split — chunks contain whole sentences only.
  • Start a new chunk only when the current chunk is non-empty and adding the next sentence would exceed max_tokens.
  • A sentence whose size alone exceeds the budget still forms a (one-sentence) chunk rather than being dropped.
SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement chunk(sizes, max_tokens): walk the sentences keeping a running total; before adding a sentence, if the current chunk is non-empty and the total would overflow, flush it and start a new one. Flush the last chunk at the end.

HINTS — 4 IDEAS
  1. Keep a current chunk (list of indices) and its running token total.
  2. For each sentence i: if the current chunk is non-empty and total + sizes[i] > max_tokens, push the current chunk and reset.
  3. Then append i to the current chunk and add its size to the total.
  4. After the loop, do not forget to push the final non-empty chunk.
CPython · WebAssembly

Explore the topic

See this challenge alongside everything else on the same subject — handbooks, system designs, algorithms and tools, in one place.