#!/usr/bin/env python3 """Minimal test for raw text training, without heavy dependencies.""" import sys import os import tempfile from pathlib import Path import importlib.util # Mock the datasets module (not installed). class MockDataset: def __init__(self, data_dict): self.data = data_dict self.column_names = list(data_dict.keys()) def __len__(self): return len(next(iter(self.data.values()))) def __getitem__(self, idx): if isinstance(idx, str): # Column access, e.g. dataset['text']. return self.data[idx] elif isinstance(idx, int): # Row access by index. return {key: values[idx] for key, values in self.data.items()} else: raise TypeError(f"Invalid index type: {type(idx)}") @classmethod def from_dict(cls, data_dict): return cls(data_dict) # __spec__ must be set so importlib.util.find_spec doesn't raise ValueError when # transformers' import_utils later probes for the real `datasets` package. datasets_mock = type(sys)("datasets") datasets_mock.__spec__ = importlib.util.spec_from_loader("datasets", loader = None) datasets_mock.Dataset = MockDataset sys.modules["datasets"] = datasets_mock # Import raw_text directly to avoid unsloth/__init__.py dependencies. current_dir = os.path.dirname(__file__) raw_text_path = os.path.join(os.path.dirname(current_dir), "unsloth", "dataprep", "raw_text.py") spec = importlib.util.spec_from_file_location("raw_text", raw_text_path) raw_text_module = importlib.util.module_from_spec(spec) spec.loader.exec_module(raw_text_module) RawTextDataLoader = raw_text_module.RawTextDataLoader TextPreprocessor = raw_text_module.TextPreprocessor def test_raw_text_loader(): """Test basic RawTextDataLoader functionality.""" class MockTokenizer: def __init__(self): self.eos_token = "" self.eos_token_id = 2 def __call__( self, text, return_tensors = None, add_special_tokens = False, ): words = text.split() token_ids = list(range(len(words))) if return_tensors == "pt": class MockTensor: def __init__(self, data): self.data = data def __getitem__(self, idx): return self.data def __len__(self): return len(self.data) def tolist(self): return self.data return {"input_ids": [MockTensor(token_ids)]} return {"input_ids": token_ids} def decode( self, token_ids, skip_special_tokens = False, ): return " ".join([f"word_{i}" for i in token_ids]) test_content = "This is a test file for raw text training. " * 10 with tempfile.NamedTemporaryFile(mode = "w", suffix = ".txt", delete = False) as f: f.write(test_content) test_file = f.name try: tokenizer = MockTokenizer() loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 2) # Text output (legacy mode). 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" # Tokenized output (new efficient mode). 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 ), "Dataset should have 'input_ids' column" assert ( "attention_mask" in tokenized_dataset.column_names ), "Dataset should have 'attention_mask' column" first_sample = tokenized_dataset[0] assert isinstance(first_sample["input_ids"], list), "input_ids should be a list" assert isinstance(first_sample["attention_mask"], list), "attention_mask should be a list" assert len(first_sample["input_ids"]) == len( first_sample["attention_mask"] ), "input_ids and attention_mask should have same length" # labels field (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" # 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) # Preprocessor. preprocessor = TextPreprocessor() clean_text = preprocessor.clean_text(" messy text \n\n\n ") assert "messy text" in clean_text, "Should clean text properly" paragraph_text = preprocessor.clean_text("Line 1\r\n\r\n\r\nLine 2") assert ( paragraph_text == "Line 1\n\nLine 2" ), "Should preserve paragraph breaks while normalizing newlines" # Non-ASCII horizontal whitespace (NBSP, thin/em/ideographic space, VT, FF) must # normalize to one ASCII space, not be deleted, or adjacent words fuse on HTML/PDF/OCR input. unicode_whitespace_cases = [ ("hello\u00a0world", "hello world"), ("hello\u202fworld", "hello world"), ("hello\u2009world", "hello world"), ("hello\u3000world", "hello world"), ("hello\u2002world", "hello world"), ("hello\x0bworld", "hello world"), ("hello\x0cworld", "hello world"), ] for raw, expected in unicode_whitespace_cases: assert preprocessor.clean_text(raw) == expected, ( f"Should normalize Unicode/control whitespace to a single space " f"for {raw!r}" ) # Mixed paragraph + Unicode whitespace. mixed = preprocessor.clean_text("Section\u00a01\r\n\r\nBody\ftext\u202fhere") assert ( mixed == "Section 1\n\nBody text here" ), "Should preserve paragraph breaks and normalize Unicode whitespace simultaneously" # Tabs collapse to a single space. assert preprocessor.clean_text("a\tb") == "a b" assert preprocessor.clean_text("a\t\tb") == "a b" # Spaces around newlines trimmed on both sides, even across multiple newlines. assert preprocessor.clean_text("foo \n\n bar") == "foo\n\nbar" # Stripping a non-ASCII char between spaces must not leave a double space # (also guards idempotence: otherwise "word1 (c) word2" needs a second pass). assert preprocessor.clean_text("word1 \u00a9 word2") == "word1 word2" assert preprocessor.clean_text("a \u00e9 b") == "a b" assert preprocessor.clean_text("prefix \U0001f600 suffix") == "prefix suffix" # Stripping a non-ASCII char adjacent to a newline must not leave a stray space. assert preprocessor.clean_text("foo \u00e9\nbar") == "foo\nbar" assert preprocessor.clean_text("foo\n\u00e9 bar") == "foo\nbar" # The double-space collapse must not swallow a paragraph break near a non-ASCII char. assert preprocessor.clean_text("a \u00a9\n\nb") == "a\n\nb" # Idempotence: clean_text twice == once. idempotent_inputs = [ " messy text \n\n\n ", "Line 1\r\n\r\n\r\nLine 2", "hello\u00a0world", "Section\u00a01\r\n\r\nBody\ftext\u202fhere", "word1 \u00a9 word2", "a \u00e9 b", ] for raw in idempotent_inputs: once = preprocessor.clean_text(raw) twice = preprocessor.clean_text(once) assert once == twice, f"clean_text should be idempotent for {raw!r}" # Validation. stats = preprocessor.validate_dataset(text_dataset) assert stats["total_samples"] > 0, "Should count samples" assert "warnings" in stats, "Should include warnings" print("✅ All tests passed!") return True except Exception as e: print(f"❌ Test failed: {e}") return False finally: os.unlink(test_file) if __name__ == "__main__": success = test_raw_text_loader() sys.exit(0 if success else 1)