diff --git a/crawl4ai/utils.py b/crawl4ai/utils.py index 279c27708..8e3556319 100644 --- a/crawl4ai/utils.py +++ b/crawl4ai/utils.py @@ -198,26 +198,25 @@ def merge_chunks( if not total_tokens: return [] - # Pre-allocate chunks - num_chunks = max(1, (total_tokens + target_size - 1) // target_size) - chunks: List[List[str]] = [[] for _ in range(num_chunks)] - - curr_chunk = 0 + # Chunks grow on demand instead of being pre-allocated from total_tokens // target_size: + # that estimate doesn't account for the tokens re-injected at every overlap boundary, so a + # fixed chunk count silently runs out and the last chunk absorbs all remaining tokens + # uncapped (growing without bound as the document gets longer). + chunks: List[List[str]] = [[]] curr_size = 0 - + # Distribute tokens for tokens in chain.from_iterable(all_tokens): - if curr_size >= target_size and curr_chunk < num_chunks - 1: + if curr_size >= target_size: if overlap > 0: - overlap_tokens = chunks[curr_chunk][-overlap:] - curr_chunk += 1 - chunks[curr_chunk].extend(overlap_tokens) + overlap_tokens = chunks[-1][-overlap:] + chunks.append(list(overlap_tokens)) curr_size = len(overlap_tokens) else: - curr_chunk += 1 + chunks.append([]) curr_size = 0 - - chunks[curr_chunk].append(tokens) + + chunks[-1].append(tokens) curr_size += 1 # Return only non-empty chunks diff --git a/tests/test_merge_chunks.py b/tests/test_merge_chunks.py new file mode 100644 index 000000000..847ea53b4 --- /dev/null +++ b/tests/test_merge_chunks.py @@ -0,0 +1,40 @@ +import unittest +from crawl4ai.utils import merge_chunks + +class TestMergeChunks(unittest.TestCase): + + def test_no_overlap_respects_target_size(self): + docs = [f"word{i} " * 50 for i in range(200)] # 10000 tokens total + chunks = merge_chunks(docs, target_size=2048, overlap=0) + for chunk in chunks: + self.assertLessEqual(len(chunk.split()), 2048) + + def test_overlap_does_not_overshoot_target_size(self): + # chunk_token_threshold=2048 / overlap_rate=0.1 are this repo's own defaults. + docs = [f"word{i} " * 50 for i in range(200)] # 10000 tokens total + chunks = merge_chunks(docs, target_size=2048, overlap=205) + for chunk in chunks: + self.assertLessEqual(len(chunk.split()), 2048) + + def test_overlap_does_not_overshoot_on_longer_documents(self): + # Regression guard: the pre-allocated-chunk-count bug made the last chunk grow + # without bound as more documents were merged in. + docs = [f"word{i} " * 50 for i in range(3200)] # 160000 tokens total + chunks = merge_chunks(docs, target_size=2048, overlap=205) + for chunk in chunks: + self.assertLessEqual(len(chunk.split()), 2048) + + def test_overlap_tokens_are_repeated_at_chunk_boundary(self): + docs = [f"word{i}" for i in range(10)] + chunks = merge_chunks(docs, target_size=6, overlap=2) + self.assertEqual(chunks[0].split()[-2:], chunks[1].split()[:2]) + + def test_all_tokens_are_preserved_without_overlap(self): + docs = [f"word{i}" for i in range(10)] + chunks = merge_chunks(docs, target_size=6, overlap=0) + merged = " ".join(chunks).split() + self.assertEqual(merged, [f"word{i}" for i in range(10)]) + + def test_empty_docs_returns_empty_list(self): + self.assertEqual(merge_chunks([], target_size=100), []) + self.assertEqual(merge_chunks(["", " "], target_size=100), [])