diff --git a/src/semchunk/semchunk.py b/src/semchunk/semchunk.py index 6538383..de39e50 100644 --- a/src/semchunk/semchunk.py +++ b/src/semchunk/semchunk.py @@ -487,13 +487,26 @@ def chunk( unoverlapped_chunk_size / subchunk_size ) # NOTE `math.ceil` would cause overlaps to be missed. - offsets = [ - ( - suboffsets[(start := i * subchunk_stride)][0], - suboffsets[min(start + subchunks_per_chunk, num_subchunks) - 1][1], - ) - for i in range(max(1, math.ceil((num_subchunks - subchunks_per_chunk) / subchunk_stride) + 1)) - ] + offsets = [] + i = 0 + + while True: + # Grow the window up to `subchunks_per_chunk` subchunks, then shrink it if the + # whitespace reintroduced between subchunks (stripped when the subchunks were + # formed but re-spanned here) would push the merged chunk over `chunk_size`. + end_i = min(i + subchunks_per_chunk, num_subchunks) - 1 + + while end_i > i and token_counter(text[suboffsets[i][0] : suboffsets[end_i][1]]) > chunk_size: + end_i -= 1 + + offsets.append((suboffsets[i][0], suboffsets[end_i][1])) + + if end_i >= num_subchunks - 1: + break + + # Advance by the stride but never past the last included subchunk so no + # subchunk is ever skipped (which would drop content). + i = min(i + subchunk_stride, end_i + 1) # Materialize chunks from offsets. chunks = [text[start:end] for start, end in offsets] @@ -611,13 +624,26 @@ def chunk( unoverlapped_chunk_size / subchunk_size ) # NOTE `math.ceil` would cause overlaps to be missed. - offsets = [ - ( - suboffsets[(start := i * subchunk_stride)][0], - suboffsets[min(start + subchunks_per_chunk, num_subchunks) - 1][1], - ) - for i in range(max(1, math.ceil((num_subchunks - subchunks_per_chunk) / subchunk_stride) + 1)) - ] + offsets = [] + i = 0 + + while True: + # Grow the window up to `subchunks_per_chunk` subchunks, then shrink it if the + # whitespace reintroduced between subchunks (stripped when the subchunks were + # formed but re-spanned here) would push the merged chunk over `chunk_size`. + end_i = min(i + subchunks_per_chunk, num_subchunks) - 1 + + while end_i > i and token_counter(text[suboffsets[i][0] : suboffsets[end_i][1]]) > chunk_size: + end_i -= 1 + + offsets.append((suboffsets[i][0], suboffsets[end_i][1])) + + if end_i >= num_subchunks - 1: + break + + # Advance by the stride but never past the last included subchunk so no + # subchunk is ever skipped (which would drop content). + i = min(i + subchunk_stride, end_i + 1) chunks = [text[start:end] for start, end in offsets] diff --git a/tests/test_semchunk.py b/tests/test_semchunk.py index 209de8f..0943171 100644 --- a/tests/test_semchunk.py +++ b/tests/test_semchunk.py @@ -200,4 +200,19 @@ def test_semchunk() -> None: semchunk.chunk('\n\n', 512, lambda *args: 0) if __name__ == '__main__': - test_semchunk() \ No newline at end of file + test_semchunk() + +def test_overlap_never_exceeds_chunk_size() -> None: + """Overlapped chunks must not exceed ``chunk_size`` when separators count as tokens. + + Merging subchunks by offset re-includes the whitespace stripped between them; with a + whitespace-significant counter (``len``) that can push a chunk over ``chunk_size``. + """ + text, chunk_size = "aa bb cc dd", 4 + chunks, offsets = semchunk.chunk(text, chunk_size, len, offsets=True, overlap=0.5, memoize=False) + + # Every chunk stays within the documented maximum... + for chunk in chunks: + assert len(chunk) <= chunk_size, f"{chunk!r} has {len(chunk)} tokens > chunk_size {chunk_size}" + # ...and the offsets still reconstruct the chunks (no content dropped). + assert chunks == [text[start:end] for start, end in offsets]