diff --git a/tests/test_synthetic_chunk_data.py b/tests/test_synthetic_chunk_data.py index abc2c01443..17e5228e32 100644 --- a/tests/test_synthetic_chunk_data.py +++ b/tests/test_synthetic_chunk_data.py @@ -129,6 +129,39 @@ def test_chunk_data_uninitialized_error_names_real_class(): os.unlink(path) +def test_chunk_data_chunks_do_not_exceed_max_tokens(): + # Every chunk must fit within max_tokens. The old multi-chunk path emitted one + # fewer, oversized chunk, and a doc just over the threshold came back unsplit. + kit = _make_kit(max_seq_length = 2048, max_generation_tokens = 760, overlap = 64) + max_tokens = 2048 - 760 * 2 - 128 # 400 + + for n_words in (500, 2000): + out, contents = _chunk("word " * n_words, kit = kit) + assert ( + len(out) >= 2 + ), f"a {n_words}-token doc (> max_tokens={max_tokens}) must be split, got {len(out)}" + for content in contents: + n_tokens = len(content.split()) + assert ( + n_tokens <= max_tokens + ), f"chunk has {n_tokens} tokens, exceeding max_tokens={max_tokens}" + + +def test_chunk_data_does_not_over_split(): + # n_chunks must be the minimum count: ceil((length - overlap) / stride), not + # ceil(length / stride) which over-splits just past a stride multiple. At 673 + # tokens (max_tokens=400, overlap=64) the tight count gives 2 chunks (~369+368). + kit = _make_kit(max_seq_length = 2048, max_generation_tokens = 760, overlap = 64) + max_tokens = 2048 - 760 * 2 - 128 # 400 + out, contents = _chunk("word " * 673, kit = kit) + assert len(out) == 2, f"673-token doc should yield the minimal 2 chunks, got {len(out)}" + for content in contents: + n_tokens = len(content.split()) + assert ( + n_tokens <= max_tokens + ), f"chunk has {n_tokens} tokens, exceeding max_tokens={max_tokens}" + + if __name__ == "__main__": test_chunk_data_keeps_single_chunk_document() test_chunk_data_still_splits_long_document() @@ -136,4 +169,6 @@ if __name__ == "__main__": test_chunk_data_short_document_is_not_split_into_fragments() test_chunk_data_rejects_overlap_not_smaller_than_chunk() test_chunk_data_uninitialized_error_names_real_class() + test_chunk_data_chunks_do_not_exceed_max_tokens() + test_chunk_data_does_not_over_split() print("OK") diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 6f025343f5..9c529d5a03 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -425,8 +425,13 @@ class SyntheticDataKit: else: # length > max_tokens > overlap here, so length - overlap > 0 and the # linspace boundaries below are always non-negative. - n_chunks = int(np.ceil(length / (max_tokens - self.overlap))) - boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks)).astype(int) + # Minimal count: overlapping chunks cover `length` in + # ceil((length - overlap) / stride) chunks, not ceil(length / stride) + # which over-splits just past a stride multiple. + n_chunks = int(np.ceil((length - self.overlap) / (max_tokens - self.overlap))) + # n_chunks + 1 points: [:-1]/[1:] pairing yields n_chunks ranges; using + # n_chunks points gave one fewer, oversized chunk (over max_tokens). + boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks + 1)).astype(int) boundaries = np.stack((boundaries[:-1], (boundaries + self.overlap)[1:])).T boundaries = np.minimum(boundaries, length).tolist()