File size: 5,364 Bytes
d656a30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Dataset loading and text cleanup (D-008, D-009, D-018).
Every step that reads TinyStories (tokenizer training, tokenization) goes
through `load_clean`, so they all see exactly the same text.
"""
from pathlib import Path
import numpy as np
import torch
from datasets import Dataset, load_dataset
DATASET = "roneneldan/TinyStories"
DATA_DIR = Path(__file__).resolve().parents[2] / "data" # git-ignored
TOKEN_DTYPE = np.uint16 # vocab 8192 < 65536 (D-009)
# D-018: TinyStories has UTF-8 punctuation that was decoded as Windows-1252
# ("mojibake"), e.g. “ -> "“". Map it back to plain ASCII.
# Order matters: the 3-character sequences must run before the bare "â€".
_REPLACEMENTS: list[tuple[str, str]] = [
("“", '"'), # “ left double quote
("â€\x9d", '"'), # ” right double quote (when the 3rd byte survived)
("’", "'"), # ’ apostrophe / right single quote
("‘", "'"), # ‘ left single quote
("–", "-"), # – en dash
("—", "-"), # — em dash
("…", "..."), # … ellipsis
("â€", '"'), # ” right double quote whose 3rd byte (0x9D) was lost
("Â\xa0", " "), # no-break space
("Â\xad", ""), # soft hyphen
# Real (correctly encoded) typographic punctuation -> ASCII.
("“", '"'), ("”", '"'), ("‘", "'"), ("’", "'"),
("–", "-"), ("—", "-"), ("…", "..."), ("\xa0", " "), ("\xad", ""),
]
def clean_text(text: str) -> str | None:
"""Return cleaned ASCII text, or None if the story should be dropped.
Dropped: empty stories, and stories that still contain non-ASCII
characters after the replacements (rare leftovers such as 'é').
"""
if not text.isascii():
for bad, good in _REPLACEMENTS:
if bad in text:
text = text.replace(bad, good)
text = text.strip()
if not text or not text.isascii():
return None
return text
def _clean_batch(batch: dict) -> dict:
return {"text": [clean_text(t) or "" for t in batch["text"]]}
def _keep_batch(batch: dict) -> list[bool]:
return [t != "" for t in batch["text"]]
def load_clean(split: str, num_proc: int = 4) -> Dataset:
"""Load a TinyStories split ('train' or 'validation') with D-018 cleanup applied.
`datasets` caches the result, so only the first call does the work.
"""
ds = load_dataset(DATASET, split=split)
# With split=... (and no streaming) this is always a Dataset. The check narrows
# the declared Union return type for type checkers, and fails loudly otherwise.
assert isinstance(ds, Dataset), f"expected Dataset, got {type(ds).__name__}"
ds = ds.map(_clean_batch, batched=True, num_proc=num_proc, desc=f"clean {split}")
return ds.filter(_keep_batch, batched=True, num_proc=num_proc, desc=f"drop {split}")
def bin_path(split: str) -> Path:
"""'train' -> data/train.bin, 'validation' -> data/val.bin"""
return DATA_DIR / f"{'val' if split == 'validation' else split}.bin"
def write_tokens(split: str, tok, eos_id: int, batch_size: int = 10_000) -> tuple[int, int]:
"""Encode a cleaned split into one flat uint16 file: story EOS story EOS ...
Streams batch by batch, so RAM use stays small. Writes to a .tmp file and
renames at the end, so a crash never leaves a half-written .bin behind.
Returns (n_stories, n_tokens).
"""
ds = load_clean(split)
out = bin_path(split)
out.parent.mkdir(exist_ok=True)
tmp = out.with_suffix(".tmp")
n_tokens = 0
with open(tmp, "wb") as f:
for i in range(0, len(ds), batch_size):
encs = tok.encode_batch(ds[i: i + batch_size]["text"]) # Rust, multi-threaded
ids = np.fromiter(
(t for e in encs for t in (*e.ids, eos_id)), dtype=np.int64
)
assert ids.max() < np.iinfo(TOKEN_DTYPE).max
ids.astype(TOKEN_DTYPE).tofile(f)
n_tokens += len(ids)
tmp.replace(out)
return len(ds), n_tokens
def load_tokens(split: str) -> np.memmap:
"""Memory-map a .bin file: no RAM cost, the OS pages it in on demand."""
return np.memmap(bin_path(split), dtype=TOKEN_DTYPE, mode="r")
def get_batch(
tokens: np.ndarray,
batch_size: int,
seq_len: int,
rng: np.random.Generator,
device: str | torch.device = "cpu",
) -> tuple[torch.Tensor, torch.Tensor]:
"""Random training windows: x = tokens[s : s+T], y = tokens[s+1 : s+T+1].
Windows start anywhere and cross story boundaries; EOS tells the model where a
story ends. Sampling with replacement from a seeded `rng` makes runs reproducible,
and saving `rng.bit_generator.state` in a checkpoint makes resume exact.
On CUDA the batch goes through pinned memory so the copy can overlap with compute.
"""
if len(tokens) < seq_len + 1:
raise ValueError(f"need at least {seq_len + 1} tokens, got {len(tokens)}")
starts = rng.integers(0, len(tokens) - seq_len, size=batch_size) # last valid start: len-T-1
idx = starts[:, None] + np.arange(seq_len + 1) # (B, T+1)
chunk = torch.from_numpy(tokens[idx].astype(np.int64)) # gather, then widen for embedding
if torch.device(device).type == "cuda":
chunk = chunk.pin_memory().to(device, non_blocking=True)
return chunk[:, :-1], chunk[:, 1:]
|