Upload cognet_data_prep.py with huggingface_hub

cognet_data_prep.py

@@ -0,0 +1,1252 @@
+#!/usr/bin/env python3
+"""
+CogNet 1B Data Preparation Script
+===================================
+Downloads, preprocesses, and tokenizes datasets from HuggingFace
+for training CogNet-1B on syntax, math, and code.
+
+Target: Character-level tokenizer (136 vocab: ASCII + French accents)
+Output: Pre-tokenized .pt files ready for training
+
+Usage:
+    python cognet_data_prep.py --output_dir /root/CogNet/data_1b [--max_gb 50] [--dry_run]
+
+Datasets covered:
+    CODE:   the-stack-smol, codeparrot-clean, CodeAlpaca, CodeSearchNet, python_code_instructions
+    MATH:   MathPile, OpenMathInstruct-1, MetaMathQA, GSM8K, HendrycksMath
+    SYNTAX: WikiText-103, C4 subset, Penn Treebank, Universal Dependencies
+    GENERAL: The Pile subset (for code+math+prose mix)
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+import unicodedata
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+# ─── Character-level Tokenizer ───────────────────────────────────────────────
+
+class CharTokenizer:
+    """Character-level tokenizer: printable ASCII + French accents + newline/tab."""
+
+    def __init__(self):
+        self.chars = sorted(set(
+            [chr(i) for i in range(32, 127)]
+            + list('àâäéèêëïîôùûüÿçœæÀÂÄÉÈÊËÏÎÔÙÛÜŸÇŒÆ')
+            + list('ëßñ¿«»')
+            + ['\t', '\n']
+        ))
+        self.char_to_id = {c: i for i, c in enumerate(self.chars)}
+        self.id_to_char = {i: c for i, c in enumerate(self.chars)}
+        self.vocab_size = len(self.chars)
+        self._allowed = set(self.chars)
+
+    def is_compatible(self, text: str) -> float:
+        """Return fraction of characters that are in vocab (1.0 = all compatible)."""
+        if not text:
+            return 1.0
+        compatible = sum(1 for c in text if c in self._allowed)
+        return compatible / len(text)
+
+    def clean_text(self, text: str) -> str:
+        """Clean text to be compatible with our 136-char vocab.
+        
+        Strategy:
+        - Map common Unicode to ASCII equivalents
+        - Replace smart quotes, em-dashes, etc.
+        - Strip characters that can't be mapped
+        - Preserve newlines and tabs
+        """
+        if not text:
+            return text
+
+        # Phase 1: Common Unicode replacements
+        replacements = {
+            '\u2018': "'",   # left single quote
+            '\u2019': "'",   # right single quote
+            '\u201c': '"',   # left double quote
+            '\u201d': '"',   # right double quote
+            '\u2013': '-',   # en dash
+            '\u2014': '--',  # em dash
+            '\u2026': '...',  # ellipsis
+            '\u00a0': ' ',   # non-breaking space
+            '\u2028': '\n',  # line separator
+            '\u2029': '\n',  # paragraph separator
+            '\u200b': '',    # zero-width space
+            '\u200c': '',    # zero-width non-joiner
+            '\u200d': '',    # zero-width joiner
+            '\ufeff': '',    # BOM
+            '\u2192': '->',  # right arrow
+            '\u2190': '<-',  # left arrow
+            '\u2194': '<->', # left-right arrow
+            '\u2264': '<=',  # less than or equal
+            '\u2265': '>=',  # greater than or equal
+            '\u2260': '!=',  # not equal
+            '\u00d7': '*',   # multiplication sign
+            '\u00f7': '/',   # division sign
+            '\u00b1': '+-',  # plus-minus
+            '\u2212': '-',   # minus sign
+            '\u2248': '~=',  # approximately equal
+            '\u221e': 'inf', # infinity
+            '\u03c0': 'pi',  # pi
+            '\u03b1': 'alpha',  # alpha
+            '\u03b2': 'beta',   # beta
+            '\u03b3': 'gamma',  # gamma
+            '\u03b4': 'delta',  # delta
+            '\u03b5': 'epsilon', # epsilon
+            '\u03b8': 'theta',  # theta
+            '\u03bb': 'lambda', # lambda
+            '\u03c3': 'sigma',  # sigma
+            '\u03c9': 'omega',  # omega
+            '\u2211': 'sum',    # summation
+            '\u220f': 'prod',   # product
+            '\u222b': 'int',    # integral
+            '\u221a': 'sqrt',   # square root
+            '\u2202': 'partial', # partial derivative
+            '\u2208': 'in',     # element of
+            '\u2282': 'subset', # subset
+            '\u2229': 'intersect', # intersection
+            '\u222a': 'union',  # union
+            '\u00b2': '^2',    # superscript 2
+            '\u00b3': '^3',    # superscript 3
+            '\u2082': '_2',    # subscript 2
+        }
+        
+        for old, new in replacements.items():
+            text = text.replace(old, new)
+        
+        # Phase 2: NFKC normalization for remaining Unicode
+        # This handles accented chars decomposition, etc.
+        normalized = []
+        for ch in text:
+            if ch in self._allowed:
+                normalized.append(ch)
+            else:
+                # Try NFKC normalization
+                nfkc = unicodedata.normalize('NFKC', ch)
+                if all(c in self._allowed for c in nfkc):
+                    normalized.append(nfkc)
+                else:
+                    # Try stripping diacritics
+                    stripped = unicodedata.normalize('NFD', ch)
+                    stripped = ''.join(
+                        c for c in stripped 
+                        if unicodedata.category(c) != 'Mn'
+                    )
+                    if all(c in self._allowed for c in stripped):
+                        normalized.append(stripped)
+                    # else: skip this character entirely
+        
+        return ''.join(normalized)
+
+    def encode(self, text: str) -> List[int]:
+        return [self.char_to_id.get(c, self.char_to_id.get(' ', 0)) for c in text]
+
+    def decode(self, ids: List[int]) -> str:
+        return ''.join(self.id_to_char.get(i, ' ') for i in ids)
+
+    def save(self, path: str):
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump({
+                'chars': self.chars,
+                'vocab_size': self.vocab_size,
+            }, f, ensure_ascii=False, indent=2)
+
+    @classmethod
+    def load(cls, path: str) -> 'CharTokenizer':
+        tok = cls.__new__(cls)
+        with open(path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        tok.chars = data['chars']
+        tok.char_to_id = {c: i for i, c in enumerate(tok.chars)}
+        tok.id_to_char = {i: c for i, c in enumerate(tok.chars)}
+        tok.vocab_size = data['vocab_size']
+        tok._allowed = set(tok.chars)
+        return tok
+
+
+# ─── Dataset Processors ──────────────────────────────────────────────────────
+
+class DatasetProcessor:
+    """Base class for dataset processors."""
+    
+    def __init__(self, name: str, category: str, tokenizer: CharTokenizer,
+                 output_dir: str, max_gb: float = 50):
+        self.name = name
+        self.category = category
+        self.tokenizer = tokenizer
+        self.output_dir = output_dir
+        self.max_gb = max_gb
+        self.stats = {
+            'name': name,
+            'category': category,
+            'raw_chars': 0,
+            'clean_chars': 0,
+            'tokens': 0,
+            'files': 0,
+            'skipped_chars': 0,
+        }
+    
+    def _should_stop(self, total_bytes: int) -> bool:
+        """Check if we've exceeded our storage budget."""
+        gb = total_bytes / (1024**3)
+        return gb >= self.max_gb
+
+    def _save_tokens(self, token_ids: List[int], split_name: str = 'train'):
+        """Save token IDs to a .pt file."""
+        import torch
+        out_path = os.path.join(self.output_dir, f'{self.name}_{split_name}.pt')
+        os.makedirs(self.output_dir, exist_ok=True)
+        torch.save(torch.tensor(token_ids, dtype=torch.long), out_path)
+        self.stats['files'] += 1
+        print(f"    Saved {len(token_ids):,} tokens to {out_path} "
+              f"({len(token_ids) * 8 / 1024**2:.1f} MB)")
+        return out_path
+
+    def process(self, dry_run: bool = False) -> Dict:
+        """Process the dataset. Override in subclasses."""
+        raise NotImplementedError
+
+
+class TheStackSmolProcessor(DatasetProcessor):
+    """bigcode/the-stack-smol — Multi-language code, manageable size."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('the_stack_smol', 'code', tokenizer, output_dir, max_gb)
+        self.languages = ['python', 'javascript', 'c', 'cpp', 'java', 'rust', 'go', 'typescript']
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (CODE)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for lang in self.languages:
+            print(f"  Loading language: {lang}")
+            try:
+                ds = load_dataset(
+                    "bigcode/the-stack-smol",
+                    data_dir=f"data/{lang}",
+                    split="train",
+                    streaming=True,
+                    trust_remote_code=True,
+                )
+                
+                for i, item in enumerate(ds):
+                    content = item.get('content', '')
+                    if not content or len(content) < 10:
+                        continue
+                    
+                    # Clean for char-level
+                    clean = self.tokenizer.clean_text(content)
+                    self.stats['raw_chars'] += len(content)
+                    self.stats['skipped_chars'] += len(content) - len(clean)
+                    
+                    # Add separator between files
+                    text = clean + '\n\n'
+                    self.stats['clean_chars'] += len(text)
+                    
+                    # Encode
+                    tokens = self.tokenizer.encode(text)
+                    all_tokens.extend(tokens)
+                    self.stats['tokens'] += len(tokens)
+                    total_bytes += len(text)
+                    
+                    if i % 5000 == 0 and i > 0:
+                        print(f"    {lang}: {i:,} files, {self.stats['tokens']:,} tokens total")
+                    
+                    if self._should_stop(total_bytes):
+                        print(f"    Storage limit reached at {total_bytes/1024**3:.1f} GB")
+                        break
+                    
+                    if i >= 8000:  # Cap per language for smol
+                        break
+                        
+            except Exception as e:
+                print(f"    Error loading {lang}: {e}")
+                continue
+            
+            if self._should_stop(total_bytes):
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class CodeParrotProcessor(DatasetProcessor):
+    """codeparrot/codeparrot-clean — Clean Python code."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('codeparrot_clean', 'code', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (CODE - Python)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        ds = load_dataset("codeparrot/codeparrot-clean", split="train", streaming=True)
+        
+        for i, item in enumerate(ds):
+            content = item.get('content', '')
+            if not content or len(content) < 20:
+                continue
+            
+            clean = self.tokenizer.clean_text(content)
+            self.stats['raw_chars'] += len(content)
+            self.stats['skipped_chars'] += len(content) - len(clean)
+            
+            text = clean + '\n\n'
+            self.stats['clean_chars'] += len(text)
+            
+            tokens = self.tokenizer.encode(text)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(text)
+            
+            if i % 10000 == 0 and i > 0:
+                print(f"    {i:,} files, {self.stats['tokens']:,} tokens, "
+                      f"{total_bytes/1024**3:.1f} GB")
+            
+            if self._should_stop(total_bytes) or i >= 200000:
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class CodeAlpacaProcessor(DatasetProcessor):
+    """sahil2801/CodeAlpaca-20k — Instruction-code pairs."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('code_alpaca', 'code', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (CODE - Instruction)")
+        print(f"{'='*60}")
+        
+        ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for i, item in enumerate(ds):
+            instruction = item.get('instruction', '')
+            inp = item.get('input', '')
+            output = item.get('output', '')
+            
+            # Flatten into linear text
+            parts = [f"### Instruction:\n{instruction}"]
+            if inp:
+                parts.append(f"### Input:\n{inp}")
+            parts.append(f"### Output:\n{output}")
+            text = '\n\n'.join(parts) + '\n\n'
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            tokens = self.tokenizer.encode(clean)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(clean)
+        
+        print(f"    {len(ds):,} samples, {self.stats['tokens']:,} tokens")
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class CodeSearchNetProcessor(DatasetProcessor):
+    """code_search_net — Code + documentation in 6 languages."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('codesearchnet', 'code', tokenizer, output_dir, max_gb)
+        self.languages = ['python', 'javascript', 'java', 'go', 'ruby', 'php']
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (CODE - Multi-language)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for lang in self.languages:
+            print(f"  Loading language: {lang}")
+            try:
+                ds = load_dataset("code_search_net", languages=[lang], 
+                                  split="train", streaming=True, trust_remote_code=True)
+                
+                for i, item in enumerate(ds):
+                    code = item.get('func_code_string', '')
+                    doc = item.get('func_documentation_string', '')
+                    
+                    text = ''
+                    if doc:
+                        text += f"# {doc}\n"
+                    text += code + '\n\n'
+                    
+                    clean = self.tokenizer.clean_text(text)
+                    self.stats['raw_chars'] += len(text)
+                    self.stats['clean_chars'] += len(clean)
+                    self.stats['skipped_chars'] += len(text) - len(clean)
+                    
+                    tokens = self.tokenizer.encode(clean)
+                    all_tokens.extend(tokens)
+                    self.stats['tokens'] += len(tokens)
+                    total_bytes += len(clean)
+                    
+                    if i % 10000 == 0 and i > 0:
+                        print(f"    {lang}: {i:,} funcs, {self.stats['tokens']:,} tokens")
+                    
+                    if self._should_stop(total_bytes) or i >= 50000:
+                        break
+            except Exception as e:
+                print(f"    Error with {lang}: {e}")
+                continue
+            
+            if self._should_stop(total_bytes):
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class PythonCodeInstructionsProcessor(DatasetProcessor):
+    """iamtarun/python_code_instructions_18k_alpaca — Python instruction pairs."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('python_code_instructions', 'code', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (CODE - Python Instructions)")
+        print(f"{'='*60}")
+        
+        ds = load_dataset("iamtarun/python_code_instructions_18k_alpaca", split="train")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for i, item in enumerate(ds):
+            instruction = item.get('instruction', '')
+            inp = item.get('input', '')
+            output_text = item.get('output', '')
+            
+            parts = [f"### Instruction:\n{instruction}"]
+            if inp:
+                parts.append(f"### Input:\n{inp}")
+            parts.append(f"### Output:\n{output_text}")
+            text = '\n\n'.join(parts) + '\n\n'
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            tokens = self.tokenizer.encode(clean)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(clean)
+        
+        print(f"    {len(ds):,} samples, {self.stats['tokens']:,} tokens")
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class MathPileProcessor(DatasetProcessor):
+    """GAIR/MathPile — Large-scale math pretraining corpus."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('mathpile', 'math', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (MATH)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        try:
+            ds = load_dataset("zwhe99/mathpile-text", split="train", streaming=True)
+        except Exception:
+            try:
+                ds = load_dataset("GAIR/MathPile", split="train", streaming=True)
+            except Exception as e:
+                print(f"    Could not load MathPile: {e}")
+                return self.stats
+        
+        for i, item in enumerate(ds):
+            text = item.get('text', '') or item.get('content', '')
+            if not text or len(text) < 20:
+                continue
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            text_out = clean + '\n\n'
+            tokens = self.tokenizer.encode(text_out)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(text_out)
+            
+            if i % 5000 == 0 and i > 0:
+                print(f"    {i:,} docs, {self.stats['tokens']:,} tokens, "
+                      f"{total_bytes/1024**3:.1f} GB")
+            
+            if self._should_stop(total_bytes) or i >= 100000:
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class OpenMathInstructProcessor(DatasetProcessor):
+    """nvidia/OpenMathInstruct-1 — Math reasoning with step-by-step solutions."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('openmath_instruct1', 'math', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (MATH - Reasoning)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        try:
+            ds = load_dataset("nvidia/OpenMathInstruct-1", split="train", streaming=True)
+        except Exception as e:
+            print(f"    Could not load OpenMathInstruct-1: {e}")
+            return self.stats
+        
+        for i, item in enumerate(ds):
+            problem = item.get('problem', '')
+            solution = item.get('generated_solution', '')
+            
+            text = f"Problem:\n{problem}\n\nSolution:\n{solution}\n\n"
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            tokens = self.tokenizer.encode(clean)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(clean)
+            
+            if i % 50000 == 0 and i > 0:
+                print(f"    {i:,} problems, {self.stats['tokens']:,} tokens, "
+                      f"{total_bytes/1024**3:.1f} GB")
+            
+            if self._should_stop(total_bytes) or i >= 500000:
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class MetaMathQAProcessor(DatasetProcessor):
+    """meta-math/MetaMathQA — Augmented math Q&A."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('metamath_qa', 'math', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (MATH - Q&A)")
+        print(f"{'='*60}")
+        
+        ds = load_dataset("meta-math/MetaMathQA", split="train")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for i, item in enumerate(ds):
+            query = item.get('query', '')
+            response = item.get('response', '')
+            
+            text = f"Question:\n{query}\n\nAnswer:\n{response}\n\n"
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            tokens = self.tokenizer.encode(clean)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(clean)
+        
+        print(f"    {len(ds):,} samples, {self.stats['tokens']:,} tokens")
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class GSM8KProcessor(DatasetProcessor):
+    """openai/gsm8k — Grade-school math word problems."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('gsm8k', 'math', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (MATH - Grade School)")
+        print(f"{'='*60}")
+        
+        ds = load_dataset("openai/gsm8k", "main", split="train")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for i, item in enumerate(ds):
+            question = item.get('question', '')
+            answer = item.get('answer', '')
+            
+            text = f"Question:\n{question}\n\nAnswer:\n{answer}\n\n"
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            tokens = self.tokenizer.encode(clean)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(clean)
+        
+        print(f"    {len(ds):,} problems, {self.stats['tokens']:,} tokens")
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class HendrycksMathProcessor(DatasetProcessor):
+    """EleutherAI/hendrycks_math — Competition-level math with LaTeX."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('hendrycks_math', 'math', tokenizer, output_dir, max_gb)
+        self.subjects = [
+            'algebra', 'counting_and_probability', 'geometry',
+            'intermediate_algebra', 'number_theory', 'prealgebra', 'precalculus'
+        ]
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (MATH - Competition)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for subject in self.subjects:
+            print(f"  Loading subject: {subject}")
+            try:
+                ds = load_dataset("EleutherAI/hendrycks_math", subject, split="train")
+                
+                for i, item in enumerate(ds):
+                    problem = item.get('problem', '')
+                    solution = item.get('solution', '')
+                    
+                    # LaTeX is ASCII — great for char-level!
+                    text = f"Problem:\n{problem}\n\nSolution:\n{solution}\n\n"
+                    
+                    clean = self.tokenizer.clean_text(text)
+                    self.stats['raw_chars'] += len(text)
+                    self.stats['clean_chars'] += len(clean)
+                    self.stats['skipped_chars'] += len(text) - len(clean)
+                    
+                    tokens = self.tokenizer.encode(clean)
+                    all_tokens.extend(tokens)
+                    self.stats['tokens'] += len(tokens)
+                    total_bytes += len(clean)
+            except Exception as e:
+                print(f"    Error with {subject}: {e}")
+                continue
+        
+        print(f"    {self.stats['tokens']:,} tokens total")
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class WikiText103Processor(DatasetProcessor):
+    """Salesforce/wikitext-103-raw-v1 — High-quality English prose."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('wikitext103', 'syntax', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (SYNTAX - English Prose)")
+        print(f"{'='*60}")
+        
+        ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1")
+        
+        for split_name in ['train', 'validation', 'test']:
+            split_ds = ds[split_name]
+            all_tokens = []
+            
+            for item in split_ds:
+                text = item.get('text', '')
+                if not text or text.strip() == '':
+                    continue
+                
+                clean = self.tokenizer.clean_text(text)
+                self.stats['raw_chars'] += len(text)
+                self.stats['clean_chars'] += len(clean)
+                self.stats['skipped_chars'] += len(text) - len(clean)
+                
+                # Add newline between paragraphs
+                text_out = clean + '\n'
+                tokens = self.tokenizer.encode(text_out)
+                all_tokens.extend(tokens)
+                self.stats['tokens'] += len(tokens)
+            
+            if not dry_run and all_tokens:
+                self._save_tokens(all_tokens, split_name)
+                print(f"    {split_name}: {len(all_tokens):,} tokens")
+        
+        return self.stats
+
+
+class C4SubsetProcessor(DatasetProcessor):
+    """allenai/c4 — Massive English web text (subset for syntax)."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('c4_subset', 'syntax', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (SYNTAX - Web Text)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        ds = load_dataset("allenai/c4", "en", split="train", streaming=True)
+        
+        for i, item in enumerate(ds):
+            text = item.get('text', '')
+            if not text or len(text) < 100:
+                continue
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            text_out = clean + '\n\n'
+            tokens = self.tokenizer.encode(text_out)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(text_out)
+            
+            if i % 10000 == 0 and i > 0:
+                print(f"    {i:,} docs, {self.stats['tokens']:,} tokens, "
+                      f"{total_bytes/1024**3:.1f} GB")
+            
+            # Cap at ~5GB of raw text for C4
+            if self._should_stop(total_bytes) or i >= 100000:
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class PennTreebankProcessor(DatasetProcessor):
+    """ptb_text_only — Gold-standard syntactic English."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('ptb', 'syntax', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (SYNTAX - Penn Treebank)")
+        print(f"{'='*60}")
+        
+        ds = load_dataset("ptb_text_only")
+        
+        for split_name in ['train', 'validation', 'test']:
+            if split_name not in ds:
+                continue
+            split_ds = ds[split_name]
+            all_tokens = []
+            
+            for item in split_ds:
+                text = item.get('sentence', '')
+                if not text:
+                    continue
+                
+                clean = self.tokenizer.clean_text(text)
+                self.stats['raw_chars'] += len(text)
+                self.stats['clean_chars'] += len(clean)
+                self.stats['skipped_chars'] += len(text) - len(clean)
+                
+                text_out = clean + '\n'
+                tokens = self.tokenizer.encode(text_out)
+                all_tokens.extend(tokens)
+                self.stats['tokens'] += len(tokens)
+            
+            if not dry_run and all_tokens:
+                self._save_tokens(all_tokens, split_name)
+                print(f"    {split_name}: {len(all_tokens):,} tokens")
+        
+        return self.stats
+
+
+class UniversalDependenciesProcessor(DatasetProcessor):
+    """universal-dependencies — Syntax annotations (CoNLL-U format)."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('universal_deps', 'syntax', tokenizer, output_dir, max_gb)
+        # English + French treebanks for accent coverage
+        self.treebanks = ['en_gum', 'en_ewt', 'fr_gsd', 'fr_sequoia']
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (SYNTAX - UD)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        for tb in self.treebanks:
+            print(f"  Loading treebank: {tb}")
+            try:
+                ds = load_dataset(
+                    "universal-dependencies/universal_dependencies", tb,
+                    split="train", trust_remote_code=True
+                )
+                
+                for i, item in enumerate(ds):
+                    # Build CoNLL-U style text from tokens
+                    tokens_list = item.get('tokens', [])
+                    lemmas = item.get('lemmas', [])
+                    upos = item.get('upos_tags', [])
+                    
+                    # Create a linear text: word/lemma/UPOS per line
+                    lines = []
+                    for j, (tok, lem, pos) in enumerate(zip(tokens_list, lemmas, upos)):
+                        lines.append(f"{tok}\t{lem}\t{pos}")
+                    
+                    text = '\n'.join(lines) + '\n\n'
+                    
+                    clean = self.tokenizer.clean_text(text)
+                    self.stats['raw_chars'] += len(text)
+                    self.stats['clean_chars'] += len(clean)
+                    self.stats['skipped_chars'] += len(text) - len(clean)
+                    
+                    tokens = self.tokenizer.encode(clean)
+                    all_tokens.extend(tokens)
+                    self.stats['tokens'] += len(tokens)
+                    total_bytes += len(clean)
+                    
+                    if self._should_stop(total_bytes):
+                        break
+            except Exception as e:
+                print(f"    Error with {tb}: {e}")
+                continue
+            
+            if self._should_stop(total_bytes):
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class ThePileSubsetProcessor(DatasetProcessor):
+    """EleutherAI/pile — All-in-one: code + math + prose (subset)."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb):
+        super().__init__('pile_subset', 'general', tokenizer, output_dir, max_gb)
+    
+    def process(self, dry_run=False):
+        from datasets import load_dataset
+        
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (GENERAL - The Pile)")
+        print(f"{'='*60}")
+        
+        all_tokens = []
+        total_bytes = 0
+        
+        try:
+            ds = load_dataset("EleutherAI/pile", split="train", streaming=True)
+        except Exception as e:
+            print(f"    Could not load The Pile: {e}")
+            return self.stats
+        
+        for i, item in enumerate(ds):
+            text = item.get('text', '')
+            if not text or len(text) < 50:
+                continue
+            
+            clean = self.tokenizer.clean_text(text)
+            self.stats['raw_chars'] += len(text)
+            self.stats['clean_chars'] += len(clean)
+            self.stats['skipped_chars'] += len(text) - len(clean)
+            
+            text_out = clean + '\n\n'
+            tokens = self.tokenizer.encode(text_out)
+            all_tokens.extend(tokens)
+            self.stats['tokens'] += len(tokens)
+            total_bytes += len(text_out)
+            
+            if i % 10000 == 0 and i > 0:
+                print(f"    {i:,} docs, {self.stats['tokens']:,} tokens, "
+                      f"{total_bytes/1024**3:.1f} GB")
+            
+            # Cap at ~10GB for Pile subset
+            if self._should_stop(total_bytes) or i >= 200000:
+                break
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+class CustomCodeProcessor(DatasetProcessor):
+    """Process user-provided code files from a directory."""
+    
+    def __init__(self, tokenizer, output_dir, max_gb, code_dir: str):
+        super().__init__('custom_code', 'code', tokenizer, output_dir, max_gb)
+        self.code_dir = code_dir
+    
+    def process(self, dry_run=False):
+        print(f"\n{'='*60}")
+        print(f"Processing: {self.name} (CODE - Custom Files)")
+        print(f"{'='*60}")
+        
+        if not os.path.exists(self.code_dir):
+            print(f"    Directory not found: {self.code_dir}")
+            return self.stats
+        
+        all_tokens = []
+        total_bytes = 0
+        supported_exts = {'.py', '.js', '.ts', '.c', '.cpp', '.h', '.hpp', 
+                         '.java', '.rs', '.go', '.rb', '.php', '.sh', '.sql',
+                         '.html', '.css', '.json', '.yaml', '.yml', '.toml',
+                         '.md', '.txt', '.tex', '.cfg', '.ini'}
+        
+        for root, dirs, files in os.walk(self.code_dir):
+            for fname in files:
+                ext = Path(fname).suffix.lower()
+                if ext not in supported_exts:
+                    continue
+                
+                fpath = os.path.join(root, fname)
+                try:
+                    with open(fpath, 'r', encoding='utf-8', errors='replace') as f:
+                        content = f.read()
+                except Exception:
+                    continue
+                
+                if not content or len(content) < 10:
+                    continue
+                
+                clean = self.tokenizer.clean_text(content)
+                self.stats['raw_chars'] += len(content)
+                self.stats['clean_chars'] += len(clean)
+                self.stats['skipped_chars'] += len(content) - len(clean)
+                
+                text = clean + '\n\n'
+                tokens = self.tokenizer.encode(text)
+                all_tokens.extend(tokens)
+                self.stats['tokens'] += len(tokens)
+                self.stats['files'] += 1
+                total_bytes += len(text)
+                
+                if self.stats['files'] % 100 == 0:
+                    print(f"    {self.stats['files']} files, {self.stats['tokens']:,} tokens")
+        
+        print(f"    Total: {self.stats['files']} files, {self.stats['tokens']:,} tokens")
+        
+        if not dry_run and all_tokens:
+            self._save_tokens(all_tokens)
+        
+        return self.stats
+
+
+# ─── Merge All Tokenized Data ────────────────────────────────────────────────
+
+def merge_all_data(output_dir: str, tokenizer_path: str, val_ratio: float = 0.02):
+    """Merge all .pt token files into unified train/val splits."""
+    import torch
+    
+    print(f"\n{'='*60}")
+    print(f"Merging all tokenized data")
+    print(f"{'='*60}")
+    
+    # Find all .pt files
+    pt_files = sorted(Path(output_dir).glob('*.pt'))
+    print(f"Found {len(pt_files)} tokenized files:")
+    for f in pt_files:
+        size_mb = f.stat().st_size / 1024**2
+        print(f"  {f.name} ({size_mb:.1f} MB)")
+    
+    if not pt_files:
+        print("No tokenized files found!")
+        return
+    
+    # Load and concatenate
+    all_tokens = []
+    for f in pt_files:
+        print(f"  Loading {f.name}...")
+        t = torch.load(f, weights_only=True)
+        all_tokens.append(t)
+    
+    all_tokens = torch.cat(all_tokens, dim=0)
+    total_tokens = len(all_tokens)
+    print(f"\nTotal tokens: {total_tokens:,}")
+    print(f"Total size: {total_tokens * 8 / 1024**3:.2f} GB (as int64)")
+    
+    # Convert to int16 to save space (vocab_size=136 fits in uint8 but int16 is safer)
+    all_tokens = all_tokens.to(torch.int16)
+    print(f"Compressed size: {total_tokens * 2 / 1024**3:.2f} GB (as int16)")
+    
+    # Shuffle
+    print("Shuffling tokens...")
+    perm = torch.randperm(total_tokens)
+    all_tokens = all_tokens[perm]
+    
+    # Split train/val
+    val_size = int(total_tokens * val_ratio)
+    train_tokens = all_tokens[val_size:]
+    val_tokens = all_tokens[:val_size]
+    
+    print(f"Train tokens: {len(train_tokens):,}")
+    print(f"Val tokens: {len(val_tokens):,}")
+    
+    # Save
+    train_path = os.path.join(output_dir, 'train_merged.pt')
+    val_path = os.path.join(output_dir, 'val_merged.pt')
+    
+    torch.save(train_tokens, train_path)
+    torch.save(val_tokens, val_path)
+    
+    print(f"\nSaved: {train_path} ({len(train_tokens)*2/1024**3:.2f} GB)")
+    print(f"Saved: {val_path} ({len(val_tokens)*2/1024**3:.2f} GB)")
+    
+    # Save dataset manifest
+    manifest = {
+        'total_tokens': total_tokens,
+        'train_tokens': len(train_tokens),
+        'val_tokens': len(val_tokens),
+        'vocab_size': 136,
+        'tokenizer': tokenizer_path,
+        'source_files': [f.name for f in pt_files],
+    }
+    manifest_path = os.path.join(output_dir, 'manifest.json')
+    with open(manifest_path, 'w') as f:
+        json.dump(manifest, f, indent=2)
+    print(f"Saved manifest: {manifest_path}")
+
+
+# ─── Main ────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description='CogNet 1B Data Preparation')
+    parser.add_argument('--output_dir', type=str, default='/root/CogNet/data_1b',
+                        help='Output directory for tokenized data')
+    parser.add_argument('--max_gb', type=float, default=50,
+                        help='Maximum GB of text data to download (default: 50)')
+    parser.add_argument('--dry_run', action='store_true',
+                        help='Only show what would be downloaded, dont save')
+    parser.add_argument('--tokenizer', type=str, default=None,
+                        help='Path to tokenizer JSON (default: create new)')
+    parser.add_argument('--skip_merge', action='store_true',
+                        help='Skip the final merge step')
+    parser.add_argument('--only', type=str, nargs='*', default=None,
+                        help='Only process these datasets (e.g., --only code_alpaca gsm8k)')
+    parser.add_argument('--custom_code_dir', type=str, default=None,
+                        help='Directory of custom code files to include')
+    args = parser.parse_args()
+    
+    # Tokenizer
+    if args.tokenizer and os.path.exists(args.tokenizer):
+        print(f"Loading tokenizer from {args.tokenizer}")
+        tokenizer = CharTokenizer.load(args.tokenizer)
+    else:
+        tokenizer = CharTokenizer()
+    print(f"Tokenizer: vocab_size={tokenizer.vocab_size}")
+    
+    # Save tokenizer
+    os.makedirs(args.output_dir, exist_ok=True)
+    tok_path = os.path.join(args.output_dir, 'tokenizer_v3.json')
+    tokenizer.save(tok_path)
+    
+    # Dataset processors (ordered by priority)
+    all_processors = [
+        # ── CODE ──
+        ('the_stack_smol',       TheStackSmolProcessor),
+        ('codeparrot_clean',     CodeParrotProcessor),
+        ('code_alpaca',          CodeAlpacaProcessor),
+        ('codesearchnet',        CodeSearchNetProcessor),
+        ('python_code_instructions', PythonCodeInstructionsProcessor),
+        # ── MATH ──
+        ('mathpile',             MathPileProcessor),
+        ('openmath_instruct1',   OpenMathInstructProcessor),
+        ('metamath_qa',          MetaMathQAProcessor),
+        ('gsm8k',                GSM8KProcessor),
+        ('hendrycks_math',       HendrycksMathProcessor),
+        # ── SYNTAX ──
+        ('wikitext103',          WikiText103Processor),
+        ('c4_subset',            C4SubsetProcessor),
+        ('ptb',                  PennTreebankProcessor),
+        ('universal_deps',       UniversalDependenciesProcessor),
+        # ── GENERAL ──
+        ('pile_subset',          ThePileSubsetProcessor),
+    ]
+    
+    # Filter if --only specified
+    if args.only:
+        all_processors = [(n, p) for n, p in all_processors if n in args.only]
+        print(f"Processing only: {args.only}")
+    
+    # Add custom code if specified
+    if args.custom_code_dir:
+        all_processors.append((
+            'custom_code',
+            lambda t, o, m: CustomCodeProcessor(t, o, m, args.custom_code_dir)
+        ))
+    
+    # Process all datasets
+    all_stats = []
+    start_time = time.time()
+    
+    for name, processor_cls in all_processors:
+        try:
+            processor = processor_cls(tokenizer, args.output_dir, args.max_gb)
+            stats = processor.process(dry_run=args.dry_run)
+            all_stats.append(stats)
+        except Exception as e:
+            print(f"\n  ERROR processing {name}: {e}")
+            import traceback
+            traceback.print_exc()
+            continue
+    
+    elapsed = time.time() - start_time
+    
+    # Summary
+    print(f"\n{'='*60}")
+    print(f"DATA PREPARATION SUMMARY")
+    print(f"{'='*60}")
+    print(f"Elapsed: {elapsed/60:.1f} minutes")
+    print(f"")
+    
+    total_tokens = 0
+    total_chars = 0
+    by_category = {}
+    
+    for s in all_stats:
+        total_tokens += s['tokens']
+        total_chars += s['clean_chars']
+        cat = s['category']
+        if cat not in by_category:
+            by_category[cat] = {'tokens': 0, 'chars': 0, 'datasets': 0}
+        by_category[cat]['tokens'] += s['tokens']
+        by_category[cat]['chars'] += s['clean_chars']
+        by_category[cat]['datasets'] += 1
+        
+        print(f"  {s['name']:30s} [{s['category']:8s}] "
+              f"{s['tokens']:>12,} tokens  "
+              f"{s['clean_chars']:>12,} chars  "
+              f"skipped: {s['skipped_chars']:,}")
+    
+    print(f"\n  {'TOTAL':30s} {'':8s} {total_tokens:>12,} tokens  {total_chars:>12,} chars")
+    print(f"\n  By category:")
+    for cat, info in by_category.items():
+        print(f"    {cat:10s}: {info['tokens']:>12,} tokens ({info['datasets']} datasets)")
+    
+    # Save stats
+    stats_path = os.path.join(args.output_dir, 'prep_stats.json')
+    with open(stats_path, 'w') as f:
+        json.dump({
+            'elapsed_seconds': elapsed,
+            'total_tokens': total_tokens,
+            'total_chars': total_chars,
+            'by_category': by_category,
+            'datasets': all_stats,
+        }, f, indent=2)
+    print(f"\nStats saved to: {stats_path}")
+    
+    # Merge
+    if not args.skip_merge and not args.dry_run:
+        merge_all_data(args.output_dir, tok_path)
+    
+    print(f"\nDone! Tokenized data ready in: {args.output_dir}")
+
+
+if __name__ == '__main__':
+    main()