diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index d7f6c317fe..c0f5d5f398 100644 --- a/tests/test_raw_text.py +++ b/tests/test_raw_text.py @@ -141,6 +141,26 @@ def test_raw_text_loader(): except ValueError as e: assert "stride" in str(e) and "chunk_size" in str(e) + # smart_chunk_text validation: called directly, chunk_size/stride are its own + # arguments and bypass the constructor guard, so it must guard itself or an + # invalid stride makes `start_idx += chunk_size - stride` non-positive and the + # chunking loop never terminates (hangs). + long_text = "This is a test file for raw text training. " * 10 + valid_chunks = loader.smart_chunk_text(long_text, chunk_size = 5, stride = 2) + assert len(valid_chunks) > 0, "Valid stride should produce chunks" + + try: + loader.smart_chunk_text(long_text, chunk_size = 5, stride = 5) + assert False, "Should raise ValueError for stride == chunk_size" + except ValueError as e: + assert "stride" in str(e) and "chunk_size" in str(e) + + try: + loader.smart_chunk_text(long_text, chunk_size = 5, stride = 10) + assert False, "Should raise ValueError for stride > chunk_size" + except ValueError as e: + assert "stride" in str(e) and "chunk_size" in str(e) + # Preprocessor. preprocessor = TextPreprocessor() clean_text = preprocessor.clean_text(" messy text \n\n\n ") diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py index 7d18b1ff29..0993f0b4e9 100644 --- a/unsloth/dataprep/raw_text.py +++ b/unsloth/dataprep/raw_text.py @@ -132,6 +132,13 @@ class RawTextDataLoader: 3. Maintains context with stride overlap 4. Returns tokenized chunks directly (more efficient) or text chunks """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if stride >= chunk_size: + raise ValueError( + f"stride ({stride}) must be smaller than chunk_size ({chunk_size}) to progress the chunking loop" + ) + # Tokenize the whole text once for accurate token counts tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False) tokens = tokenized["input_ids"]