Problem 112 · Retrieval

Recursive Sentence-Aware Chunker with Token Overlap

HardAI engineering

01Problem

Write `chunk(text, max_tokens, overlap_tokens)` (tokens = whitespace words). Split recursively: paragraphs on blank lines, then sentences on `.`, `!` or `?` followed by whitespace, and only a sentence longer than `max_tokens` is cut into word windows. Greedily pack sentences into chunks of at most `max_tokens`; when starting a new chunk, carry over the trailing whole sentences of the previous chunk whose total is at most `overlap_tokens` (and that still leave room for the next sentence). Raise `ValueError` unless `0 <= overlap_tokens < max_tokens`.

02Key insight

Chunk boundaries decide what retrieval can ever return: a fact cut mid-sentence is unretrievable as a unit, so split on the largest natural boundary that fits and fall back to smaller ones only when forced. Overlap made of whole sentences preserves context across the cut without duplicating half-sentences that confuse the reader model. The progress guarantee matters: carried overlap must always leave room for the next unit, or a naive implementation loops forever when overlap is close to max.

03Solution

solution.pypython52 lines

import re

_SENT = re.compile(r"(?<=[.!?])\s+")
_PARA = re.compile(r"\n\s*\n")


def ntok(s: str) -> int:
    return len(s.split())


def split_units(text: str, max_tokens: int) -> list[str]:
    units: list[str] = []
    for para in _PARA.split(text.strip()):
        for sent in _SENT.split(para.strip()):
            words = sent.split()
            if not words:
                continue
            if len(words) <= max_tokens:
                units.append(" ".join(words))
            else:
                units.extend(" ".join(words[i:i + max_tokens]) for i in range(0, len(words), max_tokens))
    return units


def chunk(text: str, max_tokens: int, overlap_tokens: int) -> list[str]:
    if not 0 <= overlap_tokens < max_tokens:
        raise ValueError("need 0 <= overlap_tokens < max_tokens")
    units = split_units(text, max_tokens)
    chunks: list[str] = []
    cur: list[str] = []
    cur_len = 0
    i = 0
    while i < len(units):
        n = ntok(units[i])
        if cur and cur_len + n > max_tokens:
            chunks.append(" ".join(cur))
            carry: list[str] = []
            carry_len = 0
            for prev in reversed(cur):
                p = ntok(prev)
                if carry_len + p > overlap_tokens or carry_len + p + n > max_tokens:
                    break
                carry.insert(0, prev)
                carry_len += p
            cur, cur_len = carry, carry_len
            continue
        cur.append(units[i])
        cur_len += n
        i += 1
    if cur:
        chunks.append(" ".join(cur))
    return chunks

Try it first.

TimeO(n) tokens plus O(units * overlap units) for carry

SpaceO(n)

tests.py · 11 passing

tests.pypython24 lines

text = ("Alpha beta gamma. Delta epsilon zeta eta. Theta iota.\n\n"
        "Kappa lambda mu nu xi. Omicron pi rho sigma tau upsilon phi chi psi omega! Short one?")
chunks = chunk(text, 8, 3)
assert all(ntok(c) <= 8 for c in chunks)
assert chunks[0] == "Alpha beta gamma. Delta epsilon zeta eta."
assert chunks[1].startswith("Theta iota."), chunks
units = split_units(text, 8)
assert all(any(u in c for c in chunks) for u in units)
assert "Omicron pi rho sigma tau upsilon phi chi" in units
for a, b in zip(chunks, chunks[1:]):
    assert a != b
no_ov = chunk(text, 8, 0)
assert sum(ntok(c) for c in no_ov) == ntok(text)
assert chunk("", 5, 1) == []
assert chunk("one two three", 5, 4) == ["one two three"]
for mx, ov in [(5, 5), (5, -1)]:
    try:
        chunk(text, mx, ov)
        assert False
    except ValueError:
        pass
long = " ".join(f"S{i} w w w." for i in range(50))
out = chunk(long, 9, 8)
assert all(ntok(c) <= 9 for c in out) and len(out) == 49

04Staff-level follow-up

Swap whitespace counting for the embedding model's tokenizer (counts differ 20-40% for code and non-Latin scripts), and attach metadata so each chunk carries its heading path and source offsets for citations. Then argue for parent-document retrieval or late chunking when answers need more context than any single chunk, and measure chunking choices by retrieval recall on a labelled set, not by eye.