diff --git a/pageindex/page_index_classic.py b/pageindex/page_index_classic.py index 446e8893b..29d8117c4 100644 --- a/pageindex/page_index_classic.py +++ b/pageindex/page_index_classic.py @@ -528,7 +528,8 @@ def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, over average_tokens_per_part = math.ceil(((num_tokens / expected_parts_num) + max_tokens) / 2) for i, (page_content, page_tokens) in enumerate(zip(page_contents, token_lengths)): - if current_token_count + page_tokens > average_tokens_per_part: + # without the guard, an oversized first page flushes an empty chunk + if current_subset and current_token_count + page_tokens > average_tokens_per_part: subsets.append(''.join(current_subset)) # Start new subset from overlap if specified diff --git a/tests/test_page_index.py b/tests/test_page_index.py index da69a9a79..4422d5c1c 100644 --- a/tests/test_page_index.py +++ b/tests/test_page_index.py @@ -3,6 +3,7 @@ from pageindex.page_index_classic import ( _secure_doc_text, + page_list_to_group_text, process_no_toc, process_toc_no_page_numbers, ) @@ -65,5 +66,39 @@ def test_secure_doc_text_neutralizes_document_delimiters(self): self.assertIn("", wrapped) +class PageListToGroupTextTest(unittest.TestCase): + PAGES = ["PAGE-A", "PAGE-B", "PAGE-C"] + + def test_oversized_first_page_does_not_emit_an_empty_chunk(self): + chunks = page_list_to_group_text( + self.PAGES, [60000, 500, 500], max_tokens=20000 + ) + + self.assertNotIn("", chunks) + + def test_every_page_survives_an_oversized_first_page(self): + chunks = page_list_to_group_text( + self.PAGES, [60000, 500, 500], max_tokens=20000 + ) + + joined = "".join(chunks) + for page in self.PAGES: + self.assertIn(page, joined) + + def test_oversized_middle_page_is_unaffected(self): + chunks = page_list_to_group_text( + self.PAGES, [100, 60000, 100], max_tokens=20000 + ) + + self.assertNotIn("", chunks) + + def test_document_under_the_limit_stays_one_chunk(self): + chunks = page_list_to_group_text( + self.PAGES, [100, 100, 100], max_tokens=20000 + ) + + self.assertEqual(chunks, ["PAGE-APAGE-BPAGE-C"]) + + if __name__ == "__main__": unittest.main()