Fix SyntheticDataKit.chunk_data dropping single-chunk documents (#6595)

* Fix SyntheticDataKit.chunk_data dropping single-chunk documents

chunk_data turns the n boundary points from np.linspace into n-1 ranges
via the boundaries[:-1] / [1:] pairing. When a document fits in a single
chunk (n_chunks == 1) that produces zero ranges, so the loop writes no
files and the whole document is silently dropped. Emit the full
[0, length] range when n_chunks <= 1; the multi-chunk path is unchanged.

Added a regression test covering the single-chunk and multi-chunk cases.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* chunk_data: emit nothing for an empty document (no empty chunk file)

Addresses review feedback: when the input document is empty (length == 0),
return no chunks instead of writing a single empty chunk file. Added a
regression test for the empty-document case.

* chunk_data: reject overlap >= chunk size (non-positive stride)

Per review feedback: when overlap >= max_tokens the chunk stride is
non-positive, which would divide by zero or silently emit one oversized
chunk. Raise a clear RuntimeError for that unusable configuration. Added
a regression test.

* Broaden single-chunk guard to length <= max_tokens (also fixes sub-overlap docs); expand tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>

tests/test_synthetic_chunk_data.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""Regression tests for SyntheticDataKit.chunk_data: short-document handling and overlap validation."""
+
+import os
+import tempfile
+from types import SimpleNamespace
+
+from unsloth.dataprep.synthetic import SyntheticDataKit
+
+
+class _MockTokenizer:
+    def __call__(
+        self,
+        text,
+        add_special_tokens = False,
+    ):
+        return SimpleNamespace(input_ids = list(range(len(text.split()))))
+
+    def decode(self, token_ids):
+        return " ".join(f"w{i}" for i in token_ids)
+
+
+def _make_kit(
+    max_seq_length = 2048,
+    max_generation_tokens = 512,
+    overlap = 64,
+):
+    kit = SyntheticDataKit.__new__(SyntheticDataKit)
+    kit.tokenizer = _MockTokenizer()
+    kit.max_seq_length = max_seq_length
+    kit.max_generation_tokens = max_generation_tokens
+    kit.overlap = overlap
+    return kit
+
+
+def _chunk(text, kit = None):
+    """Returns (chunk_filenames, chunk_contents); reads content before cleanup."""
+    if kit is None:
+        kit = _make_kit()
+    with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f:
+        f.write(text)
+        path = f.name
+    created = []
+    try:
+        created = kit.chunk_data(filename = path)
+        contents = []
+        for fn in created:
+            with open(fn, encoding = "utf-8") as fh:
+                contents.append(fh.read())
+        return list(created), contents
+    finally:
+        os.unlink(path)
+        for fn in created:
+            if os.path.exists(fn):
+                os.unlink(fn)
+
+
+def test_chunk_data_keeps_single_chunk_document():
+    # A short document fits in one chunk (n_chunks == 1) and must still produce
+    # one output file rather than silently vanishing.
+    out, contents = _chunk("word " * 50)
+    assert len(out) == 1, f"single-chunk doc should yield 1 file, got {len(out)}"
+    assert contents[0] != "", "the chunk file must contain the document text"
+
+
+def test_chunk_data_still_splits_long_document():
+    # A long document (n_chunks > 1) must still produce multiple chunks.
+    out, _ = _chunk("word " * 5000)
+    assert len(out) > 1, f"long doc should yield multiple chunks, got {len(out)}"
+
+
+def test_chunk_data_empty_document_yields_no_chunks():
+    # An empty document must not produce an (empty) chunk file.
+    out, _ = _chunk("")
+    assert out == [], f"empty doc should yield no files, got {len(out)}"
+
+
+def test_chunk_data_short_document_is_not_split_into_fragments():
+    # A document shorter than the overlap previously reached the multi-chunk path
+    # (n_chunks >= 3) where linspace produced negative start indices, slicing the
+    # wrong tail tokens. It must be emitted as one chunk covering the whole document.
+    kit = _make_kit(max_seq_length = 2048, max_generation_tokens = 920, overlap = 64)  # max_tokens = 80
+    out, contents = _chunk("word " * 50, kit = kit)  # 50 tokens < overlap (would be 4 chunks)
+    assert len(out) == 1, f"sub-overlap doc should yield 1 chunk, got {len(out)}"
+    assert contents[0].split() == [
+        f"w{i}" for i in range(50)
+    ], f"chunk must cover the whole document, not a fragment; got: {contents[0]!r}"
+
+
+def test_chunk_data_rejects_overlap_not_smaller_than_chunk():
+    # If overlap >= chunk size the stride is non-positive, which would divide by zero
+    # or emit one oversized chunk. The config must be rejected with a clear error.
+    kit = _make_kit(max_seq_length = 2048, max_generation_tokens = 950, overlap = 64)  # max_tokens = 20
+    with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f:
+        f.write("word " * 50)
+        path = f.name
+    try:
+        try:
+            kit.chunk_data(filename = path)
+            raise AssertionError("expected RuntimeError when overlap >= chunk size")
+        except RuntimeError as e:
+            assert "overlap" in str(e), f"error should mention overlap, got: {e}"
+    finally:
+        os.unlink(path)
+
+
+if __name__ == "__main__":
+    test_chunk_data_keeps_single_chunk_document()
+    test_chunk_data_still_splits_long_document()
+    test_chunk_data_empty_document_yields_no_chunks()
+    test_chunk_data_short_document_is_not_split_into_fragments()
+    test_chunk_data_rejects_overlap_not_smaller_than_chunk()
+    print("OK")

unsloth/dataprep/synthetic.py

@@ -403,14 +403,32 @@ class SyntheticDataKit:
         )  # -128 to reduce errors
         if max_tokens <= 5:
             raise RuntimeError("Generation length is way too long!")
+        if max_tokens <= self.overlap:
+            # A non-positive stride (max_tokens - overlap) makes the n_chunks
+            # computation below divide by zero or go negative, so reject it.
+            raise RuntimeError(
+                f"The chunk size (max_seq_length - 2 * max_generation_tokens - 128 = "
+                f"{max_tokens}) must be larger than the overlap ({self.overlap}). "
+                f"Reduce overlap or max_generation_tokens."
+            )
         input_ids = self.tokenizer(text, add_special_tokens = False).input_ids
 
         # Get left and right boundaries
         length = len(input_ids)
-        n_chunks = int(np.ceil(length / (max_tokens - self.overlap)))
-        boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks)).astype(int)
-        boundaries = np.stack((boundaries[:-1], (boundaries + self.overlap)[1:])).T
-        boundaries = np.minimum(boundaries, length).tolist()
+        if length <= max_tokens:
+            # The whole document fits in one chunk window, so emit it as a single
+            # chunk. Routing it through the multi-chunk path below would drop it
+            # (the linspace/stack pairing emits one fewer range than boundary
+            # points) or, for a document shorter than the overlap, slice the wrong
+            # tokens via negative start indices. Empty doc -> no chunk.
+            boundaries = [[0, length]] if length > 0 else []
+        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)
+            boundaries = np.stack((boundaries[:-1], (boundaries + self.overlap)[1:])).T
+            boundaries = np.minimum(boundaries, length).tolist()
 
         filename, extension = os.path.splitext(filename)
         if filename.endswith("/"):