Text Generation
PyTorch
English
research
semsimula
conservative-language-model
scalar-potential
lagrangian-mechanics
energy-based-model
physics-informed
codeparrot
mechanism-ablation
ladder-study
pre-registered
live-gradients
generation-3
first-vs-second-order
first-order
code
matched-token-budget
fock-space
parflm
anisotropic-gaussian-potential
cfc
baoab
fully-conservative
riemannian-geometry
riemannian-geodesics
jacobi-metric
non-transformer
attention-free
constant-memory-inference
Eval Results (legacy)
| """ | |
| Data loading utilities for conservative-by-construction language models. | |
| Provides two corpora: | |
| 1. tiny_shakespeare (~1.1 MB, plain text, downloads on first call) | |
| 2. tiny_stories (~few hundred MB of parquet, downloaded via HF Hub) | |
| Both are tokenised with the GPT-2 BPE tokenizer (matching the §1 E-init | |
| experiments). Tokens are cached on disk as uint16 numpy arrays for fast | |
| reload. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import urllib.request | |
| from pathlib import Path | |
| from typing import Tuple | |
| import numpy as np | |
| DATA_DIR = Path(__file__).parent / "data" | |
| DATA_DIR.mkdir(exist_ok=True) | |
| _TINY_SHAKESPEARE_URL = ( | |
| "https://raw.githubusercontent.com/karpathy/char-rnn/master/" | |
| "data/tinyshakespeare/input.txt" | |
| ) | |
| def _ensure_shakespeare_text() -> str: | |
| txt_path = DATA_DIR / "tinyshakespeare.txt" | |
| if not txt_path.exists(): | |
| print(f"[data_module] Downloading tiny_shakespeare to {txt_path} ...") | |
| urllib.request.urlretrieve(_TINY_SHAKESPEARE_URL, txt_path) | |
| return txt_path.read_text(encoding="utf-8") | |
| def _get_gpt2_tokenizer(): | |
| """Return cached GPT-2 tokenizer (loaded once per process).""" | |
| if not hasattr(_get_gpt2_tokenizer, "_tok"): | |
| from transformers import AutoTokenizer | |
| _get_gpt2_tokenizer._tok = AutoTokenizer.from_pretrained("gpt2") | |
| return _get_gpt2_tokenizer._tok | |
| def _gpt2_tokenize(text: str) -> np.ndarray: | |
| """GPT-2 BPE tokenise `text`; return uint16 numpy array of token ids.""" | |
| tok = _get_gpt2_tokenizer() | |
| ids = tok.encode(text) | |
| ids = np.asarray(ids, dtype=np.uint16) | |
| return ids | |
| def _gpt2_tokenize_stories(stories: list[str], | |
| batch_size: int = 2000) -> np.ndarray: | |
| """Tokenize a list of stories in batches to limit peak RAM. | |
| Instead of joining all stories into one giant string (which can | |
| exceed available RAM on Colab), we process small batches and | |
| concatenate the resulting token arrays. | |
| """ | |
| tok = _get_gpt2_tokenizer() | |
| chunks: list[np.ndarray] = [] | |
| for start in range(0, len(stories), batch_size): | |
| batch = stories[start:start + batch_size] | |
| text = "\n\n".join(batch) | |
| ids = tok.encode(text) | |
| chunks.append(np.asarray(ids, dtype=np.uint16)) | |
| del text, ids, batch | |
| return np.concatenate(chunks) | |
| def load_tiny_shakespeare(val_frac: float = 0.05) -> Tuple[np.ndarray, np.ndarray]: | |
| """Return (train_ids, val_ids) as uint16 numpy arrays.""" | |
| cache = DATA_DIR / "tinyshakespeare_gpt2.npz" | |
| if cache.exists(): | |
| z = np.load(cache) | |
| return z["train"], z["val"] | |
| text = _ensure_shakespeare_text() | |
| ids = _gpt2_tokenize(text) | |
| n_val = int(len(ids) * val_frac) | |
| val_ids = ids[-n_val:] | |
| train_ids = ids[:-n_val] | |
| np.savez(cache, train=train_ids, val=val_ids) | |
| print(f"[data_module] Cached tokens: train={len(train_ids):,} val={len(val_ids):,}") | |
| return train_ids, val_ids | |
| def _download_hf_parquet(repo_id: str, filename: str, local_name: str) -> Path: | |
| from huggingface_hub import hf_hub_download | |
| dest = DATA_DIR / local_name | |
| if dest.exists(): | |
| return dest | |
| src = hf_hub_download( | |
| repo_id=repo_id, filename=filename, | |
| repo_type="dataset", local_dir=str(DATA_DIR), | |
| ) | |
| src_path = Path(src) | |
| if src_path != dest: | |
| src_path.rename(dest) | |
| return dest | |
| def _resolve_tinystories_shard(prefix: str) -> str: | |
| """Return the actual repo path for a TinyStories shard. | |
| The HF dataset stores parquet shards under names like | |
| data/train-00001-of-00004-<hash>.parquet | |
| where the hash suffix can change across uploads. This helper looks | |
| up the live file list and returns the unique shard whose name starts | |
| with `prefix` (e.g. ``data/train-00001-of-00004``). | |
| """ | |
| from huggingface_hub import HfApi | |
| files = HfApi().list_repo_files("roneneldan/TinyStories", repo_type="dataset") | |
| matches = [f for f in files if f.startswith(prefix) and f.endswith(".parquet")] | |
| if len(matches) != 1: | |
| raise FileNotFoundError( | |
| f"TinyStories shard prefix {prefix!r} resolved to " | |
| f"{len(matches)} files (expected exactly 1): {matches!r}" | |
| ) | |
| return matches[0] | |
| def load_tiny_stories( | |
| n_train_files: int = 1, | |
| val_frac: float = 0.01, | |
| max_train_tokens: int | None = None, | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """Return (train_ids, val_ids) tokenised with GPT-2 BPE. | |
| Uses `roneneldan/TinyStories` parquet files directly via HF Hub, | |
| avoiding the `datasets` dependency. `n_train_files` selects how | |
| many training parquet shards to ingest (each ~300 MB uncompressed). | |
| """ | |
| import pyarrow.parquet as pq | |
| cache = DATA_DIR / ( | |
| f"tinystories_gpt2_{n_train_files}files" | |
| + (f"_{max_train_tokens}toks" if max_train_tokens else "") | |
| + ".npz" | |
| ) | |
| if cache.exists(): | |
| print(f"[data_module] Loading cached tokens from {cache}") | |
| z = np.load(cache) | |
| print(f"[data_module] train={len(z['train']):,} val={len(z['val']):,}") | |
| return z["train"], z["val"] | |
| # Fetch train parquet shard(s) and a separate validation shard. | |
| # Tokenize shard-by-shard to keep peak memory manageable (avoids | |
| # joining ~1 GB of text into a single string for multi-shard loads). | |
| train_id_chunks: list[np.ndarray] = [] | |
| n_stories = 0 | |
| for i in range(n_train_files): | |
| fname = _resolve_tinystories_shard( | |
| f"data/train-{i:05d}-of-00004" | |
| ) | |
| p = _download_hf_parquet("roneneldan/TinyStories", fname, | |
| f"tinystories_train_{i:05d}.parquet") | |
| print(f"[data_module] Reading {p}") | |
| table = pq.read_table(p, columns=["text"]) | |
| shard_texts = table["text"].to_pylist() | |
| n_stories += len(shard_texts) | |
| del table | |
| print(f"[data_module] tokenising shard {i} " | |
| f"({len(shard_texts):,} stories) in batches ...") | |
| chunk = _gpt2_tokenize_stories(shard_texts) | |
| del shard_texts | |
| train_id_chunks.append(chunk) | |
| # Early exit if we already have enough tokens. | |
| if (max_train_tokens is not None | |
| and sum(len(c) for c in train_id_chunks) >= max_train_tokens): | |
| break | |
| train_ids = np.concatenate(train_id_chunks) | |
| del train_id_chunks | |
| val_fname = _resolve_tinystories_shard("data/validation-00000-of-00001") | |
| p = _download_hf_parquet( | |
| "roneneldan/TinyStories", | |
| val_fname, | |
| "tinystories_val.parquet", | |
| ) | |
| print(f"[data_module] Reading {p}") | |
| val_texts = pq.read_table(p, columns=["text"])["text"].to_pylist() | |
| print(f"[data_module] Tokenising {len(val_texts):,} val stories " | |
| f"with GPT-2 BPE in batches ...") | |
| val_ids = _gpt2_tokenize_stories(val_texts) | |
| del val_texts | |
| if max_train_tokens is not None and len(train_ids) > max_train_tokens: | |
| train_ids = train_ids[:max_train_tokens] | |
| np.savez(cache, train=train_ids, val=val_ids) | |
| print(f"[data_module] Cached tokens: train={len(train_ids):,} " | |
| f"val={len(val_ids):,} -> {cache}") | |
| return train_ids, val_ids | |
| def get_batch(ids: np.ndarray, batch_size: int, block_size: int, | |
| rng: np.random.Generator) -> Tuple[np.ndarray, np.ndarray]: | |
| """Uniform random (x, y) pairs of shape (B, block_size) each. | |
| y[b, t] = ids[start[b] + t + 1] is the next-token target for | |
| x[b, t] = ids[start[b] + t]. | |
| """ | |
| n = len(ids) - block_size - 1 | |
| starts = rng.integers(0, n, size=batch_size) | |
| x = np.stack([ids[s:s + block_size] for s in starts]) | |
| y = np.stack([ids[s + 1:s + 1 + block_size] for s in starts]) | |
| return x.astype(np.int64), y.astype(np.int64) | |
| if __name__ == "__main__": | |
| import argparse | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--corpus", choices=["shakespeare", "tinystories"], | |
| default="shakespeare") | |
| args = ap.parse_args() | |
| if args.corpus == "shakespeare": | |
| tr, va = load_tiny_shakespeare() | |
| else: | |
| tr, va = load_tiny_stories(max_train_tokens=5_000_000) | |
| print(f"train tokens: {len(tr):,} val tokens: {len(va):,}") | |
| print(f"train first 32 ids: {tr[:32]}") | |