Upload python/prepare_data.py with huggingface_hub

python/prepare_data.py

@@ -0,0 +1,270 @@
+"""
+Data download and tokenization pipeline for H4 Polytopic Attention experiments.
+
+Supports multiple datasets with automatic download and caching:
+  - synthetic: Fibonacci-structured phrases (no download needed)
+  - shakespeare: Tiny Shakespeare (~1MB character-level text)
+  - tinystories: TinyStories from HuggingFace (real children's stories)
+
+All datasets return the same interface:
+    (train_data, val_data, vocab_size, stoi, itos)
+"""
+
+import os
+import sys
+import json
+import torch
+import urllib.request
+
+DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'data')
+
+DATASETS = {
+    'synthetic': {
+        'source': 'synthetic',
+        'description': 'Fibonacci-structured phrases (built-in)',
+    },
+    'shakespeare': {
+        'source': 'url',
+        'url': 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt',
+        'filename': 'shakespeare.txt',
+        'description': 'Tiny Shakespeare (~1MB, character-level)',
+    },
+    'tinystories': {
+        'source': 'huggingface',
+        'path': 'roneneldan/TinyStories',
+        'split': 'train',
+        'val_split': 'validation',
+        'filename': 'tinystories.txt',
+        'val_filename': 'tinystories_val.txt',
+        'description': 'TinyStories (HuggingFace, real children\'s stories)',
+        # Fallback URL if HF datasets library is not installed
+        'fallback_url': None,  # Too large for raw URL fallback
+    },
+}
+
+
+def _ensure_data_dir():
+    """Create data/ directory if it doesn't exist."""
+    os.makedirs(DATA_DIR, exist_ok=True)
+
+
+def _download_url(url, filepath):
+    """Download a file from URL using urllib (stdlib)."""
+    print(f"Downloading {url} ...")
+    try:
+        urllib.request.urlretrieve(url, filepath)
+        print(f"  Saved to {filepath} ({os.path.getsize(filepath)} bytes)")
+        return True
+    except Exception as e:
+        print(f"  Download failed: {e}")
+        return False
+
+
+def _generate_synthetic_text():
+    """Generate synthetic text with Fibonacci-structured repetitions."""
+    base_phrases = [
+        "the golden ratio appears in nature ",
+        "fibonacci numbers grow exponentially ",
+        "symmetry underlies all of physics ",
+        "the icosahedron has twenty faces ",
+        "phi equals one plus one over phi ",
+        "geometry is the language of space ",
+        "five fold symmetry cannot tile a plane ",
+        "the dodecahedron has twelve faces ",
+    ]
+    text = ""
+    a, b = 1, 1
+    for _ in range(200):
+        phrase = base_phrases[a % len(base_phrases)]
+        text += phrase * (b % 3 + 1)
+        a, b = b, a + b
+    return text
+
+
+def _load_shakespeare():
+    """Download and return Tiny Shakespeare text."""
+    _ensure_data_dir()
+    cfg = DATASETS['shakespeare']
+    filepath = os.path.join(DATA_DIR, cfg['filename'])
+
+    if not os.path.exists(filepath):
+        if not _download_url(cfg['url'], filepath):
+            print("Shakespeare download failed, falling back to synthetic data.")
+            return None
+
+    with open(filepath, 'r', encoding='utf-8') as f:
+        text = f.read()
+    print(f"Loaded Shakespeare: {len(text):,} chars")
+    return text
+
+
+def _load_tinystories():
+    """Load TinyStories from HuggingFace datasets or cached files."""
+    _ensure_data_dir()
+    cfg = DATASETS['tinystories']
+    train_path = os.path.join(DATA_DIR, cfg['filename'])
+    val_path = os.path.join(DATA_DIR, cfg['val_filename'])
+
+    # Check cache first
+    if os.path.exists(train_path) and os.path.exists(val_path):
+        with open(train_path, 'r', encoding='utf-8') as f:
+            train_text = f.read()
+        with open(val_path, 'r', encoding='utf-8') as f:
+            val_text = f.read()
+        print(f"Loaded TinyStories from cache: train={len(train_text):,} chars, val={len(val_text):,} chars")
+        return train_text, val_text
+
+    # Try HuggingFace datasets library
+    try:
+        from datasets import load_dataset as hf_load_dataset
+        print("Loading TinyStories from HuggingFace (this may take a while)...")
+        ds = hf_load_dataset(cfg['path'])
+
+        # Extract text — TinyStories has a 'text' field
+        # Limit to first 5M chars for manageability on CPU
+        MAX_CHARS = 5_000_000
+        train_text = ""
+        for item in ds[cfg['split']]:
+            train_text += item['text'] + "\n"
+            if len(train_text) >= MAX_CHARS:
+                train_text = train_text[:MAX_CHARS]
+                break
+
+        val_text = ""
+        for item in ds[cfg['val_split']]:
+            val_text += item['text'] + "\n"
+            if len(val_text) >= MAX_CHARS // 10:
+                val_text = val_text[:MAX_CHARS // 10]
+                break
+
+        # Cache to disk
+        with open(train_path, 'w', encoding='utf-8') as f:
+            f.write(train_text)
+        with open(val_path, 'w', encoding='utf-8') as f:
+            f.write(val_text)
+
+        print(f"TinyStories loaded and cached: train={len(train_text):,} chars, val={len(val_text):,} chars")
+        return train_text, val_text
+
+    except ImportError:
+        print("HuggingFace 'datasets' library not installed.")
+        print("Install with: pip install datasets")
+        print("Falling back to synthetic data.")
+        return None
+    except Exception as e:
+        print(f"Failed to load TinyStories: {e}")
+        print("Falling back to synthetic data.")
+        return None
+
+
+def prepare_char_dataset(text, val_text=None):
+    """Prepare character-level dataset from text.
+
+    Returns:
+        (train_data, val_data, vocab_size, stoi, itos)
+    """
+    if val_text is not None:
+        # Pre-split data: build vocab from both
+        all_text = text + val_text
+    else:
+        all_text = text
+
+    chars = sorted(list(set(all_text)))
+    vocab_size = len(chars)
+    stoi = {ch: i for i, ch in enumerate(chars)}
+    itos = {i: ch for ch, i in stoi.items()}
+
+    if val_text is not None:
+        train_data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
+        val_data = torch.tensor([stoi[c] for c in val_text], dtype=torch.long)
+    else:
+        data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
+        n = int(0.9 * len(data))
+        train_data = data[:n]
+        val_data = data[n:]
+
+    return train_data, val_data, vocab_size, stoi, itos
+
+
+def load_dataset(name='shakespeare'):
+    """Load a dataset by name. Returns raw text (or tuple for pre-split datasets).
+
+    For use with train_cpu.py's load_text_data() replacement.
+
+    Args:
+        name: 'synthetic', 'shakespeare', or 'tinystories'
+
+    Returns:
+        text (str) for single-text datasets, or
+        (train_text, val_text) for pre-split datasets, or
+        None on failure (caller should fall back to synthetic)
+    """
+    if name == 'synthetic':
+        return _generate_synthetic_text()
+    elif name == 'shakespeare':
+        return _load_shakespeare()
+    elif name == 'tinystories':
+        return _load_tinystories()
+    else:
+        print(f"Unknown dataset: {name}. Available: {list(DATASETS.keys())}")
+        return None
+
+
+def load_and_prepare(name='shakespeare'):
+    """Full pipeline: download, tokenize, return ready-to-train tensors.
+
+    Returns:
+        (train_data, val_data, vocab_size, stoi, itos)
+    """
+    result = load_dataset(name)
+
+    if result is None:
+        # Fall back to synthetic
+        print("Using synthetic fallback data.")
+        text = _generate_synthetic_text()
+        return prepare_char_dataset(text)
+
+    if isinstance(result, tuple):
+        # Pre-split dataset (e.g., TinyStories)
+        train_text, val_text = result
+        return prepare_char_dataset(train_text, val_text)
+    else:
+        # Single text, will be split 90/10
+        return prepare_char_dataset(result)
+
+
+def list_datasets():
+    """Print available datasets."""
+    print("Available datasets:")
+    for name, cfg in DATASETS.items():
+        cached = ""
+        if cfg['source'] == 'url':
+            path = os.path.join(DATA_DIR, cfg.get('filename', ''))
+            if os.path.exists(path):
+                cached = f" [cached: {os.path.getsize(path):,} bytes]"
+        elif cfg['source'] == 'huggingface':
+            path = os.path.join(DATA_DIR, cfg.get('filename', ''))
+            if os.path.exists(path):
+                cached = f" [cached: {os.path.getsize(path):,} bytes]"
+        print(f"  {name:15s} — {cfg['description']}{cached}")
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser(description='Prepare datasets for H4 experiments')
+    parser.add_argument('dataset', nargs='?', default='shakespeare',
+                        choices=list(DATASETS.keys()),
+                        help='Dataset to prepare (default: shakespeare)')
+    parser.add_argument('--list', action='store_true', help='List available datasets')
+    args = parser.parse_args()
+
+    if args.list:
+        list_datasets()
+        sys.exit(0)
+
+    train_data, val_data, vocab_size, stoi, itos = load_and_prepare(args.dataset)
+    print(f"\nDataset: {args.dataset}")
+    print(f"Vocab size: {vocab_size}")
+    print(f"Train tokens: {len(train_data):,}")
+    print(f"Val tokens: {len(val_data):,}")
+    print(f"Sample chars: {''.join(itos[i] for i in train_data[:80].tolist())}")