diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index 5306c68fa6..7c7272a551 100644 --- a/tests/test_raw_text.py +++ b/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" diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py index f7c1bf7856..da64565bbc 100644 --- a/unsloth/dataprep/raw_text.py +++ b/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"]