[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

tests/test_raw_text.py

@@ -10,15 +10,16 @@ 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']
@@ -28,19 +29,22 @@ class MockDataset:
             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 = type(sys)("datasets")
 datasets_mock.Dataset = MockDataset
-sys.modules['datasets'] = datasets_mock
+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')
+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)
@@ -49,73 +53,75 @@ 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 = "</s>"
-        
-        def __call__(self, text, return_tensors=None, add_special_tokens=False):
+
+        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):
+
+        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:
+    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)
-        
+        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"
-        
+        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"
-        
+        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)
+    sys.exit(0 if success else 1)