458 lines
16 KiB
Python
458 lines
16 KiB
Python
#!/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())
|