Upload hf_scripts/prepare_data.py with huggingface_hub

hf_scripts/prepare_data.py

@@ -0,0 +1,429 @@
+"""
+CogNet Data Preparation Script
+===============================
+Prepares and tokenizes multiple datasets for training:
+- Wikipedia (multilingual)
+- Code datasets (The Stack, CodeParrot)
+- Books (BookCorpus)
+- Common Crawl subsets
+- Custom local files
+
+Outputs pre-tokenized .pt files for maximum training throughput.
+
+Usage:
+    python prepare_data.py --output-dir ./data_cache --vocab-size 32000
+    python prepare_data.py --output-dir ./data_cache --datasets wiki code books
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+# ─── Dataset Configs ─────────────────────────────────────────────────────────
+
+DATASET_CONFIGS = {
+    'wiki': {
+        'path': 'wikimedia/wikipedia',
+        'subset': '20231101.en',
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 5_000_000_000,  # 5B chars
+        'description': 'Wikipedia English',
+    },
+    'wiki_fr': {
+        'path': 'wikimedia/wikipedia',
+        'subset': '20231101.fr',
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 2_000_000_000,
+        'description': 'Wikipedia French',
+    },
+    'code': {
+        'path': 'bigcode/the-stack',
+        'subset': 'data',
+        'split': 'train',
+        'text_field': 'content',
+        'max_docs': None,
+        'max_chars': 5_000_000_000,
+        'description': 'The Stack (multi-language code)',
+        'languages': ['python', 'javascript', 'java', 'cpp', 'c', 'rust', 'go', 'typescript'],
+    },
+    'code_python': {
+        'path': 'bigcode/the-stack',
+        'subset': 'data',
+        'split': 'train',
+        'text_field': 'content',
+        'max_docs': None,
+        'max_chars': 3_000_000_000,
+        'description': 'Python code from The Stack',
+        'languages': ['python'],
+    },
+    'books': {
+        'path': 'bookcorpus/bookcorpus',
+        'subset': None,
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 3_000_000_000,
+        'description': 'BookCorpus',
+    },
+    'c4': {
+        'path': 'allenai/c4',
+        'subset': 'en',
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 10_000_000_000,
+        'description': 'C4 (Colossal Clean Crawled Corpus)',
+    },
+    'openwebtext': {
+        'path': 'openwebtext',
+        'subset': None,
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 5_000_000_000,
+        'description': 'OpenWebText',
+    },
+    'alpaca': {
+        'path': 'tatsu-lab/alpaca',
+        'subset': None,
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 500_000_000,
+        'description': 'Alpaca instruction data',
+        'format_fn': 'alpaca_format',
+    },
+    'redpajama': {
+        'path': 'togethercomputer/RedPajama-Data-1T',
+        'subset': None,
+        'split': 'train',
+        'text_field': 'text',
+        'max_docs': None,
+        'max_chars': 10_000_000_000,
+        'description': 'RedPajama 1T',
+    },
+}
+
+
+def alpaca_format(example: Dict) -> str:
+    """Format Alpaca data into text."""
+    instruction = example.get('instruction', '')
+    input_text = example.get('input', '')
+    output = example.get('output', '')
+    if input_text:
+        return f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n{output}"
+    return f"### Instruction:\n{instruction}\n\n### Response:\n{output}"
+
+
+# ─── Tokenizer Training ──────────────────────────────────────────────────────
+
+def train_bpe_tokenizer(output_dir: str, vocab_size: int = 32000,
+                         sample_files: Optional[List[str]] = None) -> str:
+    """
+    Train a BPE tokenizer on sample text data.
+    Returns the path to the saved tokenizer.
+    """
+    try:
+        from tokenizers import Tokenizer
+        from tokenizers.models import BPE
+        from tokenizers.trainers import BpeTrainer
+        from tokenizers.pre_tokenizers import Metaspace, ByteLevel
+        from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+    except ImportError:
+        print("ERROR: 'tokenizers' library not installed.")
+        print("Install with: pip install tokenizers")
+        sys.exit(1)
+
+    tokenizer_path = os.path.join(output_dir, f"bpe_tokenizer_{vocab_size}.json")
+    if os.path.exists(tokenizer_path):
+        print(f"Tokenizer already exists at {tokenizer_path}")
+        return tokenizer_path
+
+    print(f"\nTraining BPE tokenizer (vocab_size={vocab_size})...")
+
+    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = ByteLevelDecoder()
+
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=[
+            "[PAD]",    # 0
+            "[UNK]",    # 1
+            "[BOS]",    # 2
+            "[EOS]",    # 3
+        ],
+        show_progress=True,
+        initial_alphabet=ByteLevel.alphabet(),
+    )
+
+    if sample_files and len(sample_files) > 0:
+        print(f"Training on {len(sample_files)} files...")
+        tokenizer.train(sample_files, trainer)
+    else:
+        print("No sample files provided. Training on built-in data...")
+        # Generate diverse sample text for tokenizer training
+        sample_texts = []
+        # English
+        sample_texts.extend([
+            "The quick brown fox jumps over the lazy dog. " * 500,
+            "Science and technology have transformed our understanding of the universe. " * 500,
+            "In the field of artificial intelligence, neural networks learn from data. " * 500,
+        ])
+        # French
+        sample_texts.extend([
+            "Le renard brun rapide saute par-dessus le chien paresseux. " * 500,
+            "La science et la technologie ont transforme notre comprehension de l'univers. " * 500,
+        ])
+        # Code
+        sample_texts.extend([
+            "def hello_world():\n    print('Hello, World!')\n    return True\n" * 500,
+            "class NeuralNetwork:\n    def __init__(self, layers):\n        self.layers = layers\n" * 500,
+            "import torch\nimport torch.nn as nn\nmodel = nn.Sequential(nn.Linear(768, 768))\n" * 500,
+            "function fibonacci(n) {\n  if (n <= 1) return n;\n  return fibonacci(n-1) + fibonacci(n-2);\n}\n" * 500,
+        ])
+        tokenizer.train_from_iterator(sample_texts, trainer)
+
+    os.makedirs(output_dir, exist_ok=True)
+    tokenizer.save(tokenizer_path)
+    print(f"Saved tokenizer to {tokenizer_path}")
+    print(f"Vocabulary size: {tokenizer.get_vocab_size()}")
+
+    return tokenizer_path
+
+
+# ─── Data Processing ─────────────────────────────────────────────────────────
+
+def process_dataset(name: str, config: Dict, tokenizer, output_dir: str,
+                    seq_len: int = 4096) -> Optional[str]:
+    """
+    Process a single dataset and save as pre-tokenized .pt file.
+    Returns the output path or None if failed.
+    """
+    print(f"\n{'='*60}")
+    print(f"Processing: {name} — {config.get('description', '')}")
+    print(f"{'='*60}")
+
+    output_path = os.path.join(output_dir, f"{name}_packed_seq{seq_len}.pt")
+    if os.path.exists(output_path):
+        print(f"Already exists: {output_path}")
+        return output_path
+
+    try:
+        from datasets import load_dataset
+    except ImportError:
+        print("ERROR: 'datasets' library not installed.")
+        print("Install with: pip install datasets")
+        return None
+
+    # Load dataset
+    print(f"Loading {config['path']}...")
+    try:
+        if config.get('subset'):
+            ds = load_dataset(
+                config['path'],
+                config['subset'],
+                split=config['split'],
+                streaming=True,
+                trust_remote_code=True,
+            )
+        else:
+            ds = load_dataset(
+                config['path'],
+                split=config['split'],
+                streaming=True,
+                trust_remote_code=True,
+            )
+    except Exception as e:
+        print(f"Failed to load dataset: {e}")
+        return None
+
+    # Filter by language if specified (for code datasets)
+    if config.get('languages'):
+        languages = set(config['languages'])
+        def lang_filter(example):
+            return example.get('language', '') in languages
+        ds = ds.filter(lang_filter)
+
+    # Tokenize
+    all_ids = []
+    doc_count = 0
+    total_chars = 0
+    max_chars = config.get('max_chars', 5_000_000_000)
+    text_field = config.get('text_field', 'text')
+    format_fn_name = config.get('format_fn')
+
+    t0 = time.time()
+
+    for example in ds:
+        # Get text
+        if format_fn_name == 'alpaca_format':
+            text = alpaca_format(example)
+        else:
+            text = example.get(text_field, '')
+
+        if not text or len(text.strip()) < 20:
+            continue
+
+        # Tokenize
+        ids = tokenizer.encode(text)
+        if isinstance(ids, list):
+            all_ids.extend(ids)
+        elif hasattr(ids, 'ids'):
+            all_ids.extend(ids.ids)
+        else:
+            all_ids.extend(list(ids))
+
+        # Add EOS between documents
+        all_ids.append(3)  # [EOS] token id
+
+        doc_count += 1
+        total_chars += len(text)
+
+        if doc_count % 10000 == 0:
+            elapsed = time.time() - t0
+            print(f"  {doc_count:,} docs | {len(all_ids):,} tokens | "
+                  f"{total_chars/1e9:.2f}B chars | {elapsed:.0f}s")
+
+        if total_chars >= max_chars:
+            print(f"  Reached char limit ({max_chars/1e9:.1f}B)")
+            break
+
+        if config.get('max_docs') and doc_count >= config['max_docs']:
+            print(f"  Reached doc limit ({config['max_docs']:,})")
+            break
+
+    if len(all_ids) == 0:
+        print("  No tokens collected!")
+        return None
+
+    # Save
+    elapsed = time.time() - t0
+    print(f"\n  Final: {doc_count:,} docs, {len(all_ids):,} tokens, {total_chars/1e9:.2f}B chars")
+    print(f"  Time: {elapsed:.0f}s ({doc_count/max(elapsed,1):,.0f} docs/s)")
+
+    # Pack into sequences and save
+    import torch
+    tensor_data = torch.tensor(all_ids, dtype=torch.long)
+    torch.save(tensor_data, output_path)
+    size_gb = os.path.getsize(output_path) / 1e9
+    print(f"  Saved to {output_path} ({size_gb:.2f} GB)")
+
+    return output_path
+
+
+# ─── Merge Datasets ──────────────────────────────────────────────────────────
+
+def merge_datasets(paths: List[str], output_path: str):
+    """Merge multiple pre-tokenized datasets into one."""
+    print(f"\nMerging {len(paths)} datasets...")
+    all_data = []
+
+    for path in paths:
+        if not os.path.exists(path):
+            print(f"  Skipping (not found): {path}")
+            continue
+        data = torch.load(path, map_location='cpu', weights_only=True)
+        all_data.append(data)
+        print(f"  {path}: {len(data):,} tokens")
+
+    if not all_data:
+        print("  No data to merge!")
+        return
+
+    merged = torch.cat(all_data, dim=0)
+    print(f"  Total: {len(merged):,} tokens")
+
+    torch.save(merged, output_path)
+    size_gb = os.path.getsize(output_path) / 1e9
+    print(f"  Saved to {output_path} ({size_gb:.2f} GB)")
+
+
+# ─── Main ────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description='CogNet Data Preparation')
+    parser.add_argument('--output-dir', type=str, default='./data_cache',
+                        help='Output directory for processed data')
+    parser.add_argument('--vocab-size', type=int, default=32000,
+                        help='BPE vocabulary size')
+    parser.add_argument('--seq-len', type=int, default=4096,
+                        help='Sequence length for packing')
+    parser.add_argument('--datasets', nargs='+',
+                        default=['wiki', 'code'],
+                        choices=list(DATASET_CONFIGS.keys()) + ['all'],
+                        help='Datasets to process')
+    parser.add_argument('--merge', action='store_true',
+                        help='Merge all datasets into one file')
+    parser.add_argument('--local-data', type=str, default=None,
+                        help='Path to local data directory with .txt/.py files')
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Train tokenizer
+    tokenizer_path = train_bpe_tokenizer(args.output_dir, args.vocab_size)
+
+    # Load tokenizer
+    from tokenizers import Tokenizer
+    tokenizer = Tokenizer.from_file(tokenizer_path)
+    print(f"\nTokenizer loaded: {tokenizer.get_vocab_size()} vocab")
+
+    # Process datasets
+    if 'all' in args.datasets:
+        datasets_to_process = list(DATASET_CONFIGS.keys())
+    else:
+        datasets_to_process = args.datasets
+
+    output_paths = []
+    for name in datasets_to_process:
+        config = DATASET_CONFIGS[name]
+        path = process_dataset(name, config, tokenizer, args.output_dir, args.seq_len)
+        if path:
+            output_paths.append(path)
+
+    # Process local data
+    if args.local_data and os.path.exists(args.local_data):
+        print(f"\nProcessing local data from {args.local_data}...")
+        local_ids = []
+        for ext in ['*.txt', '*.md', '*.py', '*.js', '*.java', '*.c', '*.cpp', '*.rs', '*.go']:
+            for fpath in Path(args.local_data).rglob(ext):
+                try:
+                    with open(fpath, 'r', encoding='utf-8', errors='ignore') as f:
+                        text = f.read()
+                    ids = tokenizer.encode(text)
+                    if isinstance(ids, list):
+                        local_ids.extend(ids)
+                    elif hasattr(ids, 'ids'):
+                        local_ids.extend(ids.ids)
+                    local_ids.append(3)  # EOS
+                except Exception as e:
+                    print(f"  Skipping {fpath}: {e}")
+
+        if local_ids:
+            local_path = os.path.join(args.output_dir, "local_packed_seq{args.seq_len}.pt")
+            torch.save(torch.tensor(local_ids, dtype=torch.long), local_path)
+            output_paths.append(local_path)
+            print(f"  Local data: {len(local_ids):,} tokens")
+
+    # Merge
+    if args.merge and len(output_paths) > 1:
+        merge_path = os.path.join(args.output_dir, f"train_packed_seq{args.seq_len}.pt")
+        merge_datasets(output_paths, merge_path)
+
+    print("\n" + "=" * 60)
+    print("Data preparation complete!")
+    print(f"Output directory: {args.output_dir}")
+    print("=" * 60)
+
+
+if __name__ == '__main__':
+    main()