diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index 503fbeb4c3..9bbfee92ac 100644 --- a/tests/test_raw_text.py +++ b/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 = "" - - 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) \ No newline at end of file + sys.exit(0 if success else 1) diff --git a/unsloth-cli.py b/unsloth-cli.py index 044ad93f31..44232e19d7 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -106,7 +106,7 @@ def run(args): # Use raw text loader loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride) dataset = loader.load_from_file(args.raw_text_file) - elif args.dataset.endswith(('.txt', '.md', '.json', '.jsonl')): + elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")): # Auto-detect local raw text files loader = RawTextDataLoader(tokenizer) dataset = loader.load_from_file(args.dataset) @@ -409,35 +409,27 @@ if __name__ == "__main__": ) parser.add_argument( - "--raw_text_file", - type=str, - help="Path to raw text file for training" + "--raw_text_file", type = str, help = "Path to raw text file for training" ) parser.add_argument( - "--chunk_size", - type=int, - default=2048, - help="Size of text chunks for training" + "--chunk_size", type = int, default = 2048, help = "Size of text chunks for training" ) parser.add_argument( - "--stride", - type=int, - default=512, - help="Overlap between chunks" + "--stride", type = int, default = 512, help = "Overlap between chunks" ) TRAINING_MODES = { - 'instruction': 'Standard instruction-following', - 'causal': 'Causal language modeling (raw text)', - 'completion': 'Text completion tasks' + "instruction": "Standard instruction-following", + "causal": "Causal language modeling (raw text)", + "completion": "Text completion tasks", } parser.add_argument( "--training_mode", - type=str, - default="instruction", - choices=list(TRAINING_MODES.keys()), - help="Training mode for the model" + type = str, + default = "instruction", + choices = list(TRAINING_MODES.keys()), + help = "Training mode for the model", ) args = parser.parse_args() diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py index dd6d6b96d0..d809eb312a 100644 --- a/unsloth/dataprep/raw_text.py +++ b/unsloth/dataprep/raw_text.py @@ -26,23 +26,24 @@ __all__ = [ ] SUPPORTED_FORMATS = { - '.txt': 'plain_text', - '.md': 'markdown', - '.json': 'json_lines', - '.jsonl': 'json_lines', - '.csv': 'csv_text_column' + ".txt": "plain_text", + ".md": "markdown", + ".json": "json_lines", + ".jsonl": "json_lines", + ".csv": "csv_text_column", } + class RawTextDataLoader: - def __init__(self, tokenizer, chunk_size=2048, stride=512): + def __init__(self, tokenizer, chunk_size = 2048, stride = 512): self.tokenizer = tokenizer - self.chunk_size = chunk_size + self.chunk_size = chunk_size self.stride = stride - + def detect_format(self, file_path): """Auto-detect file format and parse accordingly""" extension = Path(file_path).suffix.lower() - return SUPPORTED_FORMATS.get(extension, 'plain_text') + return SUPPORTED_FORMATS.get(extension, "plain_text") def load_from_file(self, file_path): """Load raw text and convert to dataset""" @@ -50,7 +51,7 @@ class RawTextDataLoader: text_content = self._read_file_by_format(file_path, file_format) chunks = self.smart_chunk_text(text_content, self.chunk_size, self.stride) return self.create_causal_dataset(chunks) - + def load_from_files(self, file_paths): """Load multiple text files""" all_chunks = [] @@ -61,11 +62,10 @@ class RawTextDataLoader: all_chunks.extend(chunks) return self.create_causal_dataset(all_chunks) - def chunk_text(self, text): """Split text into overlapping chunks""" return self.smart_chunk_text(text, self.chunk_size, self.stride) - + def create_causal_dataset(self, chunks): """Create dataset for causal language modeling""" return Dataset.from_dict({"text": chunks}) @@ -79,51 +79,50 @@ class RawTextDataLoader: 4. Adds proper EOS tokens """ # First pass: tokenize the entire text to get accurate token counts - tokenized = self.tokenizer(text, return_tensors="pt", add_special_tokens=False) + tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False) tokens = tokenized["input_ids"] - + # Handle different tokenizer return formats - if hasattr(tokens, '__len__') and len(tokens) > 0: + if hasattr(tokens, "__len__") and len(tokens) > 0: # If it's a nested structure, get the first element - if hasattr(tokens[0], '__len__'): + if hasattr(tokens[0], "__len__"): tokens = tokens[0] elif isinstance(tokens, int): # If tokenizer returns just a count, create a simple range tokens = list(range(tokens)) - + if len(tokens) <= chunk_size: # Text is small enough to fit in one chunk eos_token = self.tokenizer.eos_token if self.tokenizer.eos_token else "" return [text + eos_token] - + chunks = [] start_idx = 0 - + while start_idx < len(tokens): # Calculate end index for this chunk end_idx = min(start_idx + chunk_size, len(tokens)) - + # Extract tokens for this chunk chunk_tokens = tokens[start_idx:end_idx] - + # Decode back to text - chunk_text = self.tokenizer.decode(chunk_tokens, skip_special_tokens=True) - + chunk_text = self.tokenizer.decode(chunk_tokens, skip_special_tokens = True) + # Add EOS token if it's the last chunk or chunk is complete if end_idx == len(tokens) or len(chunk_tokens) == chunk_size: eos_token = self.tokenizer.eos_token if self.tokenizer.eos_token else "" chunk_text += eos_token - + chunks.append(chunk_text) - + # Move to next chunk with stride overlap if end_idx == len(tokens): break start_idx += chunk_size - stride - + return chunks - def tokenize_and_chunk(self, text): """ Tokenize first, then chunk by token count: @@ -134,10 +133,10 @@ class RawTextDataLoader: def _read_file_by_format(self, file_path, file_format): """Read file content based on detected format.""" - with open(file_path, 'r', encoding='utf-8') as f: - if file_format == 'plain_text' or file_format == 'markdown': + with open(file_path, "r", encoding = "utf-8") as f: + if file_format == "plain_text" or file_format == "markdown": return f.read() - elif file_format == 'json_lines': + elif file_format == "json_lines": lines = [] for line in f: try: @@ -147,42 +146,43 @@ class RawTextDataLoader: lines.append(text) except json.JSONDecodeError: continue - return '\n\n'.join(lines) - elif file_format == 'csv_text_column': + return "\n\n".join(lines) + elif file_format == "csv_text_column": reader = csv.DictReader(f) texts = [] for row in reader: text = self._extract_text_from_csv_row(row) if text: texts.append(text) - return '\n\n'.join(texts) + return "\n\n".join(texts) return "" - + def _extract_text_from_json(self, data): """Extract text from JSON object using common field names.""" - text_fields = ['text', 'content', 'message', 'body', 'description', 'prompt'] + text_fields = ["text", "content", "message", "body", "description", "prompt"] for field in text_fields: if field in data and isinstance(data[field], str): return data[field] return "" - + def _extract_text_from_csv_row(self, row): """Extract text from CSV row using common column names.""" - text_columns = ['text', 'content', 'message', 'body', 'description', 'prompt'] + text_columns = ["text", "content", "message", "body", "description", "prompt"] for column in text_columns: if column in row and row[column]: return row[column] return "" + class TextPreprocessor: def clean_text(self, text): """Remove unwanted characters, normalize whitespace""" - text = re.sub(r'\s+', ' ', text) - text = re.sub(r'[^\x20-\x7E\n\t]', '', text) - text = text.replace('\r\n', '\n').replace('\r', '\n') - text = re.sub(r'\n{3,}', '\n\n', text) + text = re.sub(r"\s+", " ", text) + text = re.sub(r"[^\x20-\x7E\n\t]", "", text) + text = text.replace("\r\n", "\n").replace("\r", "\n") + text = re.sub(r"\n{3,}", "\n\n", text) return text.strip() - + def extract_sections(self, text, patterns): """Extract specific sections (e.g., code blocks, quotes)""" sections = [] @@ -193,12 +193,20 @@ class TextPreprocessor: def add_structure_tokens(self, text): """Add special tokens for structure (chapters, sections)""" - text = re.sub(r'^# (.+)$', r'<|chapter|>\1<|/chapter|>', text, flags=re.MULTILINE) - text = re.sub(r'^## (.+)$', r'<|section|>\1<|/section|>', text, flags=re.MULTILINE) - text = re.sub(r'^### (.+)$', r'<|subsection|>\1<|/subsection|>', text, flags=re.MULTILINE) - text = re.sub(r'```(\w*)\n(.*?)\n```', r'<|code|\1|>\2<|/code|>', text, flags=re.DOTALL) + text = re.sub( + r"^# (.+)$", r"<|chapter|>\1<|/chapter|>", text, flags = re.MULTILINE + ) + text = re.sub( + r"^## (.+)$", r"<|section|>\1<|/section|>", text, flags = re.MULTILINE + ) + text = re.sub( + r"^### (.+)$", r"<|subsection|>\1<|/subsection|>", text, flags = re.MULTILINE + ) + text = re.sub( + r"```(\w*)\n(.*?)\n```", r"<|code|\1|>\2<|/code|>", text, flags = re.DOTALL + ) return text - + def validate_dataset(self, dataset): """ Check for: @@ -208,61 +216,66 @@ class TextPreprocessor: - Empty chunks """ stats = { - 'total_samples': len(dataset), - 'empty_samples': 0, - 'min_length': float('inf'), - 'max_length': 0, - 'avg_length': 0, - 'repeated_content': 0, - 'encoding_issues': 0, - 'warnings': [] + "total_samples": len(dataset), + "empty_samples": 0, + "min_length": float("inf"), + "max_length": 0, + "avg_length": 0, + "repeated_content": 0, + "encoding_issues": 0, + "warnings": [], } - - texts = dataset['text'] + + texts = dataset["text"] text_lengths = [] seen_texts = set() - + for i, text in enumerate(texts): if not text or len(text.strip()) == 0: - stats['empty_samples'] += 1 + stats["empty_samples"] += 1 continue - + # Check for encoding issues try: - text.encode('utf-8') + text.encode("utf-8") except UnicodeEncodeError: - stats['encoding_issues'] += 1 - + stats["encoding_issues"] += 1 + # Calculate lengths length = len(text) text_lengths.append(length) - stats['min_length'] = min(stats['min_length'], length) - stats['max_length'] = max(stats['max_length'], length) - + stats["min_length"] = min(stats["min_length"], length) + stats["max_length"] = max(stats["max_length"], length) + # Check for repeated content text_hash = hash(text.strip()) if text_hash in seen_texts: - stats['repeated_content'] += 1 + stats["repeated_content"] += 1 else: seen_texts.add(text_hash) - + # Calculate average length if text_lengths: - stats['avg_length'] = sum(text_lengths) / len(text_lengths) - stats['min_length'] = stats['min_length'] if stats['min_length'] != float('inf') else 0 - - # Generate warnings - if stats['empty_samples'] > 0: - stats['warnings'].append(f"Found {stats['empty_samples']} empty samples") - - if stats['repeated_content'] > 0: - stats['warnings'].append(f"Found {stats['repeated_content']} repeated samples") - - if stats['encoding_issues'] > 0: - stats['warnings'].append(f"Found {stats['encoding_issues']} encoding issues") - - if stats['min_length'] < 10: - stats['warnings'].append("Some samples are very short (< 10 characters)") - - return stats + stats["avg_length"] = sum(text_lengths) / len(text_lengths) + stats["min_length"] = ( + stats["min_length"] if stats["min_length"] != float("inf") else 0 + ) + # Generate warnings + if stats["empty_samples"] > 0: + stats["warnings"].append(f"Found {stats['empty_samples']} empty samples") + + if stats["repeated_content"] > 0: + stats["warnings"].append( + f"Found {stats['repeated_content']} repeated samples" + ) + + if stats["encoding_issues"] > 0: + stats["warnings"].append( + f"Found {stats['encoding_issues']} encoding issues" + ) + + if stats["min_length"] < 10: + stats["warnings"].append("Some samples are very short (< 10 characters)") + + return stats