diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py new file mode 100644 index 0000000000..503fbeb4c3 --- /dev/null +++ b/tests/test_raw_text.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +""" +Minimal test for raw text training implementation. +Tests basic functionality without heavy dependencies. +""" + +import sys +import os +import tempfile +from pathlib import Path +import importlib.util + +# Mock the datasets module since it's 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): + # Allow accessing columns by name like dataset['text'] + return self.data[idx] + elif isinstance(idx, int): + # Allow accessing individual rows 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) + +# Mock datasets module +datasets_mock = type(sys)('datasets') +datasets_mock.Dataset = MockDataset +sys.modules['datasets'] = datasets_mock + +# Import the raw_text module 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.""" + + # Mock tokenizer for testing + class MockTokenizer: + def __init__(self): + self.eos_token = "" + + def __call__(self, text, return_tensors=None, add_special_tokens=False): + words = text.split() + token_ids = list(range(len(words))) + + if return_tensors == "pt": + # Mock tensor-like object + class MockTensor: + def __init__(self, data): + self.data = data + def __getitem__(self, idx): + return self.data + def __len__(self): + return len(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]) + + + + + # Create test file + 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: + # Test loader + tokenizer = MockTokenizer() + loader = RawTextDataLoader(tokenizer, chunk_size=5, stride=2) + + # Test loading + dataset = loader.load_from_file(test_file) + assert len(dataset) > 0, "Should create at least one chunk" + assert 'text' in dataset.column_names, "Dataset should have 'text' column" + + # 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) + 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: + # Cleanup + os.unlink(test_file) + +if __name__ == "__main__": + success = test_raw_text_loader() + sys.exit(0 if success else 1) \ No newline at end of file