Fix bugs and add improvements to RawTextDataLoader

- Fix test file: use return_tokenized instead of return_tensors
- Fix test file: use text_dataset instead of undefined dataset variable
- Move parameter validation to constructor (fail fast on invalid params)
- Add labels field in tokenized output for causal LM training
- Add empty file handling with clear error message
- Add tests for constructor validation and labels field

tests/test_raw_text.py

@@ -100,12 +100,12 @@ def test_raw_text_loader():
         loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 2)
 
         # Test loading with text output (legacy mode)
-        text_dataset = loader.load_from_file(test_file, return_tensors = False)
+        text_dataset = loader.load_from_file(test_file, return_tokenized = False)
         assert len(text_dataset) > 0, "Should create at least one chunk"
         assert "text" in text_dataset.column_names, "Dataset should have 'text' column"
 
         # Test loading with tokenized output (new efficient mode)
-        tokenized_dataset = loader.load_from_file(test_file, return_tensors = True)
+        tokenized_dataset = loader.load_from_file(test_file, return_tokenized = True)
         assert len(tokenized_dataset) > 0, "Should create at least one tokenized chunk"
         assert (
             "input_ids" in tokenized_dataset.column_names
@@ -124,13 +124,30 @@ def test_raw_text_loader():
             first_sample["attention_mask"]
         ), "input_ids and attention_mask should have same length"
 
+        # Verify labels field exists (for causal LM training)
+        assert "labels" in tokenized_dataset.column_names, "Dataset should have 'labels' column"
+        assert first_sample["labels"] == first_sample["input_ids"], "labels should match input_ids"
+
+        # Test constructor validation
+        try:
+            bad_loader = RawTextDataLoader(tokenizer, chunk_size = 0, stride = 2)
+            assert False, "Should raise ValueError for chunk_size=0"
+        except ValueError as e:
+            assert "chunk_size must be positive" in str(e)
+
+        try:
+            bad_loader = RawTextDataLoader(tokenizer, 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)
+
         # Test preprocessor
         preprocessor = TextPreprocessor()
         clean_text = preprocessor.clean_text("  messy   text  \n\n\n  ")
         assert "messy text" in clean_text, "Should clean text properly"
 
         # Test validation
-        stats = preprocessor.validate_dataset(dataset)
+        stats = preprocessor.validate_dataset(text_dataset)
         assert stats["total_samples"] > 0, "Should count samples"
         assert "warnings" in stats, "Should include warnings"
 

unsloth/dataprep/raw_text.py

@@ -36,6 +36,12 @@ SUPPORTED_FORMATS = {
 
 class RawTextDataLoader:
     def __init__(self, tokenizer, chunk_size = 2048, stride = 512, return_tokenized = True):
+        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})"
+            )
         self.tokenizer = tokenizer
         self.chunk_size = chunk_size
         self.stride = stride
@@ -52,6 +58,8 @@ class RawTextDataLoader:
             return_tokenized = self.return_tokenized
         file_format = self.detect_format(file_path)
         text_content = self._read_file_by_format(file_path, file_format)
+        if not text_content or not text_content.strip():
+            raise ValueError(f"File '{file_path}' is empty or contains only whitespace")
         chunks = self.smart_chunk_text(
             text_content, self.chunk_size, self.stride, return_tokenized
         )
@@ -86,8 +94,10 @@ class RawTextDataLoader:
             # Reorganize the data structure for Dataset.from_dict
             input_ids = [chunk["input_ids"] for chunk in chunks]
             attention_mask = [chunk["attention_mask"] for chunk in chunks]
+            # Labels are same as input_ids for causal LM training
+            labels = [list(ids) for ids in input_ids]
             return Dataset.from_dict(
-                {"input_ids": input_ids, "attention_mask": attention_mask}
+                {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
             )
         else:
             # If chunks are text strings (backward compatibility)
@@ -101,13 +111,6 @@ 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"
-            )
-
         # First pass: tokenize the entire text to get accurate token counts
         tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False)
         tokens = tokenized["input_ids"]