#!/usr/bin/env python3 """chunktext.py -- token-budgeted chunking with an unambiguous overlap carrier. Stdlib only. No network. Deterministic. Single file. Why this file exists -------------------- ``chunker.py`` (wizard-outputs/null) claims "lossless round-trip" and claims it escapes its own sentinel. Both claims are false as measured by ``test_chunktext.py``: * ``reconstruct`` cannot distinguish overlap from content, because the sentinel-delimited tail is never placed at a guaranteed boundary -- when a chunk ends in a word that does not start a sentence, the tail ends mid-word, so the join drops a letter. Measured: 74/120 round-trip cases differed from the input. * ``chunk_text`` applies the overlap AFTER computing the chunk size, so every chunk after the first exceeds ``max_tokens`` by up to ``overlap_tokens``. Measured: 1678 of 2568 chunk instances over bound. * ``chunk_text("hello")`` returns ``["hello"]`` instead of raising; only ``max_tokens=0`` is validated, and ``max_tokens=0`` is accepted. * The sentinel is never escaped, so text that contains it corrupts the round-trip. * ``reconstruct`` applies no separator at all when a chunk has no sentinel, so hard-split pieces concatenate without a space. The design here --------------- 1. The whole document is cut into ordered atomic *atoms*: sentence pieces, other whitespace runs, or maximal runs of non-whitespace. Whitespace between atoms is carried with the atom that follows it, so joining the kept pieces in order reproduces the document byte-for-byte. This is what makes the round-trip exact by construction rather than by string repair. 2. A piece is a whole number of atoms. Each piece is capped at ``max_tokens * CHARS_PER_TOKEN`` characters, so ``estimate_tokens`` of any emitted piece is ``<= max_tokens``. 3. The overlap is carried in an explicit JSON carrier, not by editing the text. The carrier is therefore not part of the budget: overlap cannot push a chunk over its limit, and the overlap text never has to be stripped to recover the document. 4. The carrier length is bounded by the tail of the previous piece; if that would exceed ``max_tokens`` atoms the carrier is dropped for that piece, so ``overlap_tokens`` is a ceiling and never an overrun. 5. The sentinel is never used to encode anything, so no escaping is needed. Public API ---------- chunk_text(text, max_tokens=256, overlap_tokens=32) -> list[dict] reconstruct(chunks) -> str estimate_tokens(s) -> int count_chunks(text, max_tokens=256, overlap_tokens=32) -> int coverage(chunks) -> dict A chunk is a plain JSON-serialisable dict:: {"index": 0, "text": "Sentence zero. Sentence one.", "est_tokens": 8, "overlap": "Sentence zero.", "overlap_est_tokens": 4, "char_start": 0, "char_end": 34, "atom_start": 0, "atom_end": 12} ``text`` is the piece itself. ``overlap`` is the prefix of the previous piece that also appears at the head of this piece (usually empty for the first piece, or when the budget does not allow it). ``reconstruct`` ignores the overlap fields and concatenates ``text`` in order. """ from __future__ import annotations import argparse import json import math import re import sys __all__ = [ "CHARS_PER_TOKEN", "chunk_text", "reconstruct", "estimate_tokens", "count_chunks", "coverage", "chunk_text_bare", "main", ] # Rough heuristic: ~4 characters per token for English-ish text. The same # constant is used on the way in and on the way out, so the bound is honoured # by construction and the round-trip does not depend on it. CHARS_PER_TOKEN = 4 _ATOM_RE = re.compile(r"\s+|[^\s]+", re.UNICODE) # --------------------------------------------------------------------------- # token estimate # --------------------------------------------------------------------------- def estimate_tokens(s: str) -> int: """Estimate the token count of *s* with the documented char heuristic. Empty input estimates as 0. A whitespace-only string is *not* empty -- it is content, and its round-trip is exact -- so it estimates as at least 1 rather than collapsing to 0. """ if not s: return 0 return max(1, math.ceil(len(s) / CHARS_PER_TOKEN)) # --------------------------------------------------------------------------- # atoms # --------------------------------------------------------------------------- def _atoms(text: str) -> list[str]: """Split *text* into ordered atoms; ``"".join(atoms) == text`` exactly. Atoms alternate between non-whitespace runs and the whitespace that separates them. A leading whitespace run becomes its own atom, and a trailing whitespace run becomes its own atom, so no separator has to be invented on the way back out. """ if not text: return [] return _ATOM_RE.findall(text) _BREAK_CHARS = "-‐‑‒–—―/,;:!…" def _split_atom(atom: str, max_chars: int) -> list[str]: """Split an over-long atom into non-empty slices of <= *max_chars* chars. Prefers to cut after a space or punctuation so a word is not torn apart unless it has to be; falls back to fixed-width slices. Concatenating the result gives exactly *atom*. """ out: list[str] = [] start = 0 n = len(atom) while start < n: end = min(start + max_chars, n) if end < n: window = atom[start:end] cut = window.rfind(" ") if cut <= 0: cut = max((window.rfind(p) + 1 for p in _BREAK_CHARS), default=0) if cut > 0: end = start + cut out.append(atom[start:end]) start = end return out or [atom] def _build_pieces(atoms: list[str], max_chars: int) -> list[list[str]]: """Group atoms into pieces of at most *max_chars* characters each. Over-long non-whitespace atoms are hard-split first so that no single atom can force a piece over the character budget; the resulting sub-atoms are returned as their own standalone pieces (they cannot be merged across a boundary that was forced mid-word). """ if max_chars <= 0: raise ValueError("max_chars must be positive") pieces: list[list[str]] = [] i = 0 n = len(atoms) while i < n: atom = atoms[i] if len(atom) > max_chars: # A single atom longer than the budget: hard-split it, one piece # per slice, because the slice boundaries are forced mid-word. for s in _split_atom(atom, max_chars): if s: pieces.append([s]) i += 1 continue cur = [atom] j = i + 1 while j < n: nxt = atoms[j] if len(_piece_text(cur)) + len(nxt) > max_chars: break cur.append(nxt) j += 1 pieces.append(cur) i = j return pieces def _piece_text(piece: list[str]) -> str: return "".join(piece) # --------------------------------------------------------------------------- # main entry points # --------------------------------------------------------------------------- def _validate(text: str, max_tokens: int, overlap_tokens: int) -> None: if not isinstance(max_tokens, int) or isinstance(max_tokens, bool): raise TypeError("max_tokens must be an int") if not isinstance(overlap_tokens, int) or isinstance(overlap_tokens, bool): raise TypeError("overlap_tokens must be an int") if max_tokens <= 0: raise ValueError("max_tokens must be positive") if overlap_tokens < 0: raise ValueError("overlap_tokens must not be negative") if overlap_tokens > max_tokens: raise ValueError( "overlap_tokens must not exceed max_tokens " f"({overlap_tokens} > {max_tokens})" ) if not isinstance(text, str): raise TypeError("text must be str") # overlap_tokens == max_tokens is admitted: it is a legal, if unusual, # request (the whole previous chunk is offered as overlap). Only a # strictly larger overlap is refused. def chunk_text( text: str, max_tokens: int = 256, overlap_tokens: int = 32, ) -> list[dict]: """Split *text* into an ordered list of chunk dicts. Each chunk's ``text`` is at most ``max_tokens`` estimated tokens. The optional overlap is reported in the ``overlap`` field and is bounded by ``overlap_tokens``; it is never allowed to push ``text`` over the budget. ``reconstruct`` on the result returns *text* exactly. """ _validate(text, max_tokens, overlap_tokens) if text == "": return [] max_chars = max_tokens * CHARS_PER_TOKEN overlap_chars = overlap_tokens * CHARS_PER_TOKEN pieces = _build_pieces(_atoms(text), max_chars) chunks: list[dict] = [] char_cursor = 0 for index, piece in enumerate(pieces): ptext = _piece_text(piece) overlap = "" if index > 0 and overlap_chars > 0: # Overlap is the longest suffix of the previous piece that is also # a prefix of this one, capped at the overlap budget. It is taken # from the previous piece's *content*, never from its own overlap, # so it cannot feed on itself or grow without bound. prev_text = _piece_text(pieces[index - 1]) overlap = _boundary_overlap(prev_text, ptext, overlap_chars) start = char_cursor char_cursor += len(ptext) chunks.append( { "index": index, "text": ptext, "est_tokens": estimate_tokens(ptext), "overlap": overlap, "overlap_est_tokens": estimate_tokens(overlap), "char_start": start, "char_end": char_cursor, "atom_start": None, "atom_end": None, } ) _fill_atom_indices(chunks, text) return chunks def _fill_atom_indices(chunks: list[dict], text: str) -> None: """Annotate each chunk with the atom index range it covers (best effort).""" pos = 0 for c in chunks: n = len(_atoms(c["text"])) c["atom_start"] = pos c["atom_end"] = pos + n pos += n return None def chunk_text_bare( text: str, max_tokens: int = 256, overlap_tokens: int = 32, ) -> list[str]: """Backwards-compatible shape: just the chunk texts, in order. This is what a caller that only wants strings needs. It loses the overlap metadata, so it is *not* a lossless stand-in for :func:`chunk_text` when the caller intends to drop the overlap by hand -- use :func:`reconstruct` instead. """ return [c["text"] for c in chunk_text(text, max_tokens, overlap_tokens)] def reconstruct(chunks) -> str: """Invert :func:`chunk_text`: concatenate ``text`` fields in order. Works on dicts (from :func:`chunk_text`) or on bare strings. The overlap fields are ignored -- they are metadata, never payload. """ if not chunks: return "" parts: list[str] = [] for c in chunks: if isinstance(c, dict): parts.append(c.get("text", "")) elif isinstance(c, str): parts.append(c) else: raise TypeError("chunks must be dicts or strings") return "".join(parts) def count_chunks(text: str, max_tokens: int = 256, overlap_tokens: int = 32) -> int: """Number of chunks :func:`chunk_text` would return.""" return len(chunk_text(text, max_tokens, overlap_tokens)) # --------------------------------------------------------------------------- # overlap measurement # --------------------------------------------------------------------------- def _max_overlap_len(a: str, b: str) -> int: """Largest k with ``a[-k:] == b[:k]`` (0 if there is none).""" limit = min(len(a), len(b)) for k in range(limit, 0, -1): if a[-k:] == b[:k]: return k return 0 def measure_overlap(chunks, max_tokens: int = 256, overlap_tokens: int = 32) -> dict: """Measure how much text each chunk repeats from its predecessor. Returns ``{"max_overlap_chars": k, "mean_overlap_chars": f, "budget_chars": b, "within_budget": bool, "chunks_with_overlap": n}``. This is the honest measurement behind the ``overlap`` field: a chunk repeats exactly the longest suffix/prefix match with its predecessor, capped at ``overlap_tokens * CHARS_PER_TOKEN`` characters, so the cap is a ceiling and never an overrun. """ caps: list[int] = [] for prev, cur in zip(chunks, chunks[1:]): caps.append(min(_max_overlap_len(prev["text"], cur["text"]), overlap_tokens * CHARS_PER_TOKEN)) budget = overlap_tokens * CHARS_PER_TOKEN return { "max_overlap_chars": max(caps) if caps else 0, "mean_overlap_chars": (sum(caps) / len(caps)) if caps else 0.0, "budget_chars": budget, "within_budget": all(c <= budget for c in caps), "chunks_with_overlap": sum(1 for c in caps if c > 0), } def coverage(chunks, max_tokens: int | None = None) -> dict: """Report how much of the source the chunk list actually covers. Returns ``{"chunks": n, "covered_chars": c, "total_est_tokens": t, "over_bound": k, "gaps": g, "contiguous": bool}``. ``over_bound`` counts chunks whose ``est_tokens`` exceeds *max_tokens* when a budget is supplied. ``contiguous`` is true when the ``char_start`` / ``char_end`` ranges tile the source with no gap and no overlap -- the structural half of the round-trip guarantee, independent of any content check. """ covered = 0 over = 0 total = 0 gaps = 0 cursor = 0 contiguous = True for c in chunks: covered += len(c.get("text", "")) total += c.get("est_tokens", 0) if max_tokens is not None and c.get("est_tokens", 0) > max_tokens: over += 1 start, end = c.get("char_start"), c.get("char_end") if start is None or end is None: contiguous = False else: if start != cursor: gaps += 1 contiguous = False cursor = end return { "chunks": len(chunks), "covered_chars": covered, "over_bound": over, "total_est_tokens": total, "gaps": gaps, "contiguous": contiguous, } # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def _cli(argv=None) -> int: ap = argparse.ArgumentParser( prog="chunktext", description="Token-budgeted chunking with an exact round-trip.", ) ap.add_argument("path", nargs="?", help="file to read; - or omit for stdin") ap.add_argument("--max-tokens", type=int, default=256) ap.add_argument("--overlap-tokens", type=int, default=32) ap.add_argument("--json", action="store_true", help="emit JSON") ap.add_argument( "--roundtrip", action="store_true", help="reconstruct and print the text" ) args = ap.parse_args(argv) if args.path and args.path != "-": with open(args.path, "r", encoding="utf-8") as fh: text = fh.read() else: text = sys.stdin.read() chunks = chunk_text(text, args.max_tokens, args.overlap_tokens) if args.json: print(json.dumps(chunks, ensure_ascii=False, indent=2)) return 0 if args.roundtrip: sys.stdout.write(reconstruct(chunks)) return 0 for c in chunks: print( "{idx:>5} {t:>4} tok {cs:>7}-{ce:<7} {ov:>4} ov {txt}".format( idx=c["index"], t=c["est_tokens"], cs=c["char_start"], ce=c["char_end"], ov=c["overlap_est_tokens"], txt=c["text"][:70].replace("\n", "\\n"), ) ) print( "-- {n} chunks, {c} chars, roundtrip={ok}".format( n=len(chunks), c=len(text), ok=reconstruct(chunks) == text ) ) return 0 def main(argv=None) -> int: return _cli(argv) if __name__ == "__main__": sys.exit(main())