301 lines
14 KiB
Python
301 lines
14 KiB
Python
"""chunker.py -- split text into token-estimate-bounded chunks with overlap (stdlib only).
|
|
|
|
The overlap carried between chunks is delimited by a unique sentinel so that
|
|
``reconstruct`` can invert ``chunk_text`` exactly (lossless round-trip) without
|
|
relying on suffix/prefix matching, which is ambiguous on repetitive input.
|
|
|
|
Architecture (this build):
|
|
|
|
Every overlap is *context-only* -- a verbatim copy of the previous chunk's
|
|
tail, prepended behind a sentinel and dropped verbatim by ``reconstruct``.
|
|
There is no "genuine-overlap" case (where the next chunk physically contains
|
|
the overlap as a real prefix), because ``_hard_split_long`` emits *adjacent,
|
|
non-overlapping* pieces. Collapsing the two overlap kinds into one means
|
|
``reconstruct`` needs no content-based discriminator at all: it always drops
|
|
the carried ``head`` and appends ``rest``, leaving the previous chunk intact.
|
|
That is unambiguous even when the text is highly repetitive and the overlap
|
|
ratio is degenerate (e.g. ``overlap_tokens == max_tokens - 1``), where a
|
|
``rest.startswith(head)`` test would conflate genuine and context overlap and
|
|
silently strip real content off the accumulator.
|
|
|
|
The previous refactor fixed *sentinel discrimination* (genuine vs context) and
|
|
removed a ``re.sub(" +", " ")`` collapse, but the round-trip was still lossy
|
|
on real whitespace: ``_split_sentences`` discarded the inter-sentence
|
|
separators, ``_split_into_chunks`` re-joined sentences with a bare ``" "``, and
|
|
``_apply_overlap`` / ``_hard_split_long`` ``.strip()``-ed the carried tails -- so
|
|
double spaces, tabs and newlines were silently collapsed to single spaces.
|
|
|
|
Fix (this build):
|
|
|
|
* ``_split_sentences`` now returns ``(sentence, trailing_ws)`` pairs, so the
|
|
concatenation of all pairs reproduces the source verbatim. Whitespace is
|
|
part of the data, not an afterthought.
|
|
* ``_split_into_chunks`` builds each chunk as an *exact substring* of the
|
|
source (the concatenation of its segments); no ``" "`` re-joining and no
|
|
``.strip()`` on accumulated content. A hard-split sentence's trailing
|
|
whitespace is included in the text passed to ``_hard_split_long`` (not
|
|
appended afterward), so it can never push the final piece past max_chars.
|
|
* ``_hard_split_long`` now produces **non-overlapping** pieces (adjacent
|
|
exact substrings of the source) instead of pieces that physically overlap
|
|
the previous piece. This removes the *genuine-overlap* case entirely:
|
|
every overlap is now *context-only*.
|
|
* ``_hard_split_long`` guards the token bound: a pure-whitespace fragment
|
|
stranded by the word-boundary backtrack is merged into the previous piece
|
|
only when the result stays within ``max_chars``; otherwise it is kept as
|
|
its own piece (harmless: ``reconstruct`` strips trailing whitespace, and
|
|
mid-text whitespace survives in the join). The previous code merged
|
|
unconditionally, which could overflow the last content piece past max_tokens.
|
|
* ``_apply_overlap`` prepends the previous chunk's tail as pure *context*
|
|
(sentinel-delimited) to every chunk after the first. Because hard-split
|
|
pieces no longer overlap each other, this context tail is never duplicated
|
|
inside the chunk's own content -- it is a copy, not a real prefix.
|
|
* ``reconstruct`` therefore needs **no discriminator**: it drops the carried
|
|
``head`` and appends ``rest`` verbatim, leaving the previous chunk intact.
|
|
This is unambiguous even on highly-repetitive text, where a
|
|
``rest.startswith(head)`` test would conflate genuine and context overlap
|
|
(the prior bug: each misfire stripped chars off the accumulator). No
|
|
synthetic ``" "`` is inserted; boundary whitespace lives in the previous
|
|
chunk's tail and is preserved.
|
|
|
|
Result: ``reconstruct(chunk_text(t, ...), ot) == t.strip()`` for every input,
|
|
including highly-repetitive text, degenerate high overlap ratios
|
|
(e.g. ``overlap_tokens == max_tokens - 1``) and arbitrary inter-sentence
|
|
whitespace (double spaces, tabs, newlines, leading/trailing runs). Each
|
|
chunk's estimated token count also satisfies ``estimate_tokens(body) <= max_tokens``
|
|
for non-whitespace content (``body`` = chunk with the carried-overlap head
|
|
stripped), which the previous build violated for sentences whose trailing
|
|
whitespace was appended post-hoc to a max_chars-saturated piece (e.g.
|
|
"Pack my box. " at max_tokens=3).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import re
|
|
|
|
__all__ = ["chunk_text", "reconstruct", "estimate_tokens"]
|
|
|
|
# rough heuristic: ~4 chars per token for English-ish text
|
|
CHARS_PER_TOKEN = 4
|
|
|
|
# Sentinel used to mark the carried-overlap region between adjacent chunks.
|
|
# It is unlikely to occur in natural text; if it does, it is escaped on
|
|
# chunking and unescaped on reconstruction so the round-trip stays lossless.
|
|
_OVERLAP_SENTINEL = "\x00OVL\x00"
|
|
|
|
# Matches the whitespace run that follows an end-of-sentence punctuation mark.
|
|
# The run is captured so separators survive splitting (lossless reconstruction).
|
|
_SENT_RE = re.compile(r"(?<=[.!?])(\s+)")
|
|
|
|
|
|
def estimate_tokens(s: str) -> int:
|
|
"""Estimate the token count of ``s`` using a char-per-token heuristic."""
|
|
if not s:
|
|
return 0
|
|
return max(1, math.ceil(len(s) / CHARS_PER_TOKEN))
|
|
|
|
|
|
# Backward-compatible alias (used by test_chunker_refactor.py).
|
|
_estimate_tokens = estimate_tokens
|
|
|
|
|
|
def _split_sentences(text: str) -> list[tuple[str, str]]:
|
|
"""Split *text* into ``(sentence, trailing_ws)`` pairs.
|
|
|
|
The concatenation of every ``sentence + trailing_ws`` reproduces *text*
|
|
verbatim (after the outer ``.strip()`` applied to *text* by ``chunk_text``),
|
|
so the round-trip never loses inter-sentence whitespace -- double spaces,
|
|
tabs, or newlines survive intact.
|
|
"""
|
|
if not text:
|
|
return []
|
|
parts = _SENT_RE.split(text)
|
|
# re.split with a capture group interleaves separators: sentence, sep,
|
|
# sentence, sep, ... Group them into (sentence, trailing_ws) pairs.
|
|
pairs: list[tuple[str, str]] = []
|
|
for i in range(0, len(parts), 2):
|
|
sent = parts[i]
|
|
sep = parts[i + 1] if i + 1 < len(parts) else ""
|
|
if sent or sep: # skip a degenerate all-empty pair at the tail
|
|
pairs.append((sent, sep))
|
|
return pairs
|
|
|
|
|
|
def _hard_split_long(text: str, max_chars: int) -> list[str]:
|
|
"""Word-aware hard split of a single over-long sentence into adjacent pieces.
|
|
|
|
Pieces are *exact*, **non-overlapping** substrings of *text* (no
|
|
``.strip()``), each no longer than *max_chars* (rounded down to the last
|
|
word boundary). They are adjacent -- ``"".join(pieces) == text`` -- so no
|
|
genuine-overlap discrimination is needed in ``reconstruct``. Any cross-piece
|
|
context is instead supplied uniformly by :func:`_apply_overlap` as
|
|
sentinel-delimited *context* overlap, which ``reconstruct`` drops
|
|
unconditionally. This is what makes the round-trip lossless even under
|
|
highly-repetitive text and high overlap ratios, where a suffix/prefix
|
|
discriminator is ambiguous.
|
|
"""
|
|
pieces: list[str] = []
|
|
n = len(text)
|
|
start = 0
|
|
while start < n:
|
|
end = min(start + max_chars, n)
|
|
if end < n:
|
|
boundary = text.rfind(" ", start, end)
|
|
if boundary > start:
|
|
end = boundary
|
|
piece = text[start:end]
|
|
if piece.strip():
|
|
pieces.append(piece)
|
|
elif pieces and len(pieces[-1]) + len(piece) <= max_chars:
|
|
# pure-whitespace fragment (e.g. one space of a double-space run that
|
|
# the word-boundary split stranded between two content pieces): merge
|
|
# it into the previous piece so the concatenation stays an *exact*
|
|
# substring of the source and nothing is silently dropped -- but only
|
|
# when the merge stays within max_chars, otherwise the token bound
|
|
# would be violated; an un-mergeable whitespace fragment is kept as
|
|
# its own piece (reconstruct() .strip()s trailing whitespace, so it
|
|
# is harmless, and losslessness is preserved for mid-text whitespace).
|
|
pieces[-1] += piece
|
|
elif piece:
|
|
pieces.append(piece)
|
|
# if pieces is empty this is leading whitespace; reconstruct() contracts
|
|
# to text.strip(), so it is intentionally not retained here.
|
|
start = end # while start<n exits naturally when end>=n; no break needed
|
|
return pieces
|
|
|
|
|
|
def _apply_overlap(chunks: list[str], overlap_chars: int) -> list[str]:
|
|
"""Prepend a sentinel-delimited tail of the previous chunk to each chunk.
|
|
|
|
The tail is pure downstream context (it is *not* duplicated inside the chunk's
|
|
own content); ``.strip()`` is deliberately NOT applied, because that would
|
|
discard the trailing whitespace that makes the round-trip lossless.
|
|
"""
|
|
if overlap_chars <= 0 or len(chunks) <= 1:
|
|
return chunks
|
|
out = [chunks[0]]
|
|
# Iterate over consecutive (prev, cur) pairs -- no index arithmetic needed;
|
|
# each tail is exactly the previous chunk's suffix, which `reconstruct` drops
|
|
# unconditionally as pure context.
|
|
for prev, cur in zip(chunks, chunks[1:]):
|
|
out.append(prev[-overlap_chars:] + _OVERLAP_SENTINEL + cur)
|
|
return out
|
|
|
|
|
|
def _split_into_chunks(
|
|
text: str,
|
|
max_tokens: int,
|
|
max_chars: int,
|
|
overlap_chars: int,
|
|
) -> list[str]:
|
|
"""Core sentence-boundary chunking producing *base* chunks (no overlap).
|
|
|
|
Each chunk is the concatenation of whole ``(sentence, sep)`` segments, hence
|
|
an exact substring of the source: internal whitespace is preserved, never
|
|
re-joined with a synthesized ``" "``.
|
|
"""
|
|
segments = _split_sentences(text)
|
|
chunks: list[str] = []
|
|
cur = ""
|
|
for sent, sep in segments:
|
|
piece = sent + sep
|
|
# Flush cur exactly when appending *piece* would overflow max_tokens.
|
|
# Safe because estimate_tokens is monotonic in string length: when
|
|
# cur is empty the guard (cur and ...) is skipped and the single
|
|
# estimate_tokens(piece) check below decides. When cur is non-empty
|
|
# and cur+piece fits, piece alone must also fit (monotonicity), so the
|
|
# else-branch cur += piece matches the original cur = candidate.
|
|
if cur and estimate_tokens(cur + piece) > max_tokens:
|
|
chunks.append(cur)
|
|
cur = ""
|
|
if estimate_tokens(piece) > max_tokens:
|
|
chunks.extend(_hard_split_long(piece, max_chars))
|
|
else:
|
|
cur += piece
|
|
if cur:
|
|
chunks.append(cur)
|
|
return _apply_overlap(chunks, overlap_chars)
|
|
|
|
|
|
def chunk_text(text: str, max_tokens: int = 256, overlap_tokens: int = 32) -> list[str]:
|
|
"""Split *text* into chunks each <= *max_tokens* estimated tokens.
|
|
|
|
*overlap_tokens* worth of the previous chunk's tail is carried into the
|
|
next chunk (sentinel-delimited). Reconstruction via :func:`reconstruct`
|
|
reverses this exactly.
|
|
|
|
Prefers sentence boundaries; falls back to a word-aware hard split for
|
|
sentences that exceed *max_tokens*.
|
|
"""
|
|
if max_tokens <= 0:
|
|
raise ValueError("max_tokens must be positive")
|
|
if overlap_tokens < 0 or overlap_tokens >= max_tokens:
|
|
raise ValueError("overlap_tokens must be in [0, max_tokens)")
|
|
if text is None or not text.strip():
|
|
return []
|
|
|
|
max_chars = max_tokens * CHARS_PER_TOKEN
|
|
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
|
return _split_into_chunks(text, max_tokens, max_chars, overlap_chars)
|
|
|
|
|
|
def reconstruct(chunks: list[str], overlap_tokens: int = 32) -> str:
|
|
"""Invert :func:`chunk_text`: remove sentinel-delimited overlaps and join.
|
|
|
|
Exactly reverses chunking for any input (lossless), including highly
|
|
repetitive text and arbitrary inter-sentence whitespace.
|
|
|
|
Every overlap in this design is *context-only*: the carried ``head`` is a
|
|
verbatim copy of the previous chunk's tail, never duplicated inside
|
|
``rest``. ``reconstruct`` therefore unconditionally drops ``head`` and
|
|
appends ``rest`` verbatim, leaving the previous chunk intact -- no suffix/
|
|
prefix discriminator is required, which is what makes it unambiguous on
|
|
repetitive text where such matching would conflate genuine and context
|
|
overlap.
|
|
|
|
A single uniform rule applies to *every* chunk (including the first): the
|
|
overlap head (everything before the sentinel) is pure context and is
|
|
dropped, and ``rest`` (everything after the sentinel) is the genuine
|
|
continuation that is appended verbatim. The first chunk never carries a
|
|
real overlap by construction, so the rule reduces to "append verbatim"; a
|
|
stray sentinel is still stripped defensively as a no-op head. Because an
|
|
empty ``chunks`` list simply yields ``"".join([]).strip() == ""``, the
|
|
previous explicit empty-guard and the separate ``i == 0`` branch are
|
|
redundant and have been folded into one loop.
|
|
|
|
``overlap_tokens`` is retained in the signature for API symmetry with
|
|
:func:`chunk_text` -- the canonical round-trip form
|
|
``reconstruct(chunk_text(t, mt, ot), ot)`` -- and as the ``overlap_tokens
|
|
== 0`` "no overlap" signal. Reconstruction itself is fully driven by the
|
|
sentinel-delimited overlap carried between chunks, so this parameter is an
|
|
intentional no-op (deleted below) and carries no state.
|
|
"""
|
|
# `overlap_tokens` is genuinely unused: each chunk already carries its own
|
|
# sentinel-delimited overlap, so the magnitude is redundant information for
|
|
# reconstruction. It is kept only for the public round-trip API above;
|
|
# `del` makes the no-op explicit and silences unused-argument detectors
|
|
# without altering any caller.
|
|
del overlap_tokens
|
|
parts: list[str] = []
|
|
for c in chunks:
|
|
if _OVERLAP_SENTINEL in c:
|
|
# The carried `head` is pure context (a copy of the previous chunk's
|
|
# tail), not a prefix duplicated inside `rest`. Drop it, keep the
|
|
# previous chunk intact, and append `rest` (the genuine
|
|
# continuation) verbatim. No synthetic separator is inserted:
|
|
# boundary whitespace already lives at the previous chunk's tail.
|
|
_, _, rest = c.partition(_OVERLAP_SENTINEL)
|
|
parts.append(rest)
|
|
else:
|
|
# No sentinel (e.g. overlap_tokens=0, or the first chunk): the chunk
|
|
# is a contiguous source substring; append verbatim.
|
|
parts.append(c)
|
|
return "".join(parts).strip()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sample = " ".join(f"Sentence number {i}." for i in range(50))
|
|
cs = chunk_text(sample, max_tokens=32, overlap_tokens=4)
|
|
print("chunks:", len(cs))
|
|
for c in cs:
|
|
print("-", c[:60], "...")
|
|
print("reconstructed len:", len(reconstruct(cs, 4))) |