Phase 2: token dataset + corpus encoder (code only, not yet run)
- src/dataset.py: memmap uint16 token loader + random-crop batch sampler (x, y shifted-by-one, reproducible via generator) - scripts/encode_corpus.py: stream stories -> encode -> uint16 .bin, one EOT token appended per story, flat memory. encode_file() factored out for testing. - tests/test_dataset.py: batch shapes/shift/bounds/reproducibility + encoder round-trip/separators/uint16 on synthetic data (5 tests)
This commit is contained in:
@@ -0,0 +1,96 @@
|
||||
"""Encode the raw corpus into flat uint16 token files.
|
||||
|
||||
Each story in the source is delimited by a line containing only
|
||||
``<|endoftext|>``. We stream the file one line at a time, encode each story
|
||||
with the trained tokenizer, append the end-of-text token id after it, and flush
|
||||
the ids to disk in batches — so memory stays flat regardless of corpus size.
|
||||
|
||||
Output: data/tokens/train.bin and data/tokens/val.bin.
|
||||
|
||||
Usage:
|
||||
python scripts/encode_corpus.py
|
||||
"""
|
||||
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
|
||||
from dataset import TOKEN_DTYPE # noqa: E402
|
||||
from tokenizer import BPETokenizer # noqa: E402
|
||||
|
||||
EOT = "<|endoftext|>"
|
||||
TOK_PATH = ROOT / "artifacts" / "tokenizer.json"
|
||||
RAW = ROOT / "data" / "raw"
|
||||
OUT = ROOT / "data" / "tokens"
|
||||
FLUSH_EVERY = 20_000 # stories between disk writes / progress updates
|
||||
|
||||
JOBS = [
|
||||
("TinyStoriesV2-GPT4-train.txt", "train.bin"),
|
||||
("TinyStoriesV2-GPT4-valid.txt", "val.bin"),
|
||||
]
|
||||
|
||||
|
||||
def encode_file(tok: BPETokenizer, src: Path, dst: Path, progress: bool = True) -> tuple[int, int]:
|
||||
"""Encode one corpus file to a uint16 .bin. Returns (n_stories, n_tokens)."""
|
||||
eot_id = tok.special_tokens[EOT]
|
||||
buf: list[int] = []
|
||||
story: list[str] = []
|
||||
n_stories = n_tokens = 0
|
||||
t0 = time.perf_counter()
|
||||
|
||||
def emit(text: str) -> None:
|
||||
nonlocal n_stories, n_tokens
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return # skip blank runs between separators
|
||||
ids = tok.encode(text)
|
||||
ids.append(eot_id)
|
||||
buf.extend(ids)
|
||||
n_stories += 1
|
||||
n_tokens += len(ids)
|
||||
|
||||
with open(src, "r", encoding="utf-8", errors="ignore") as fin, open(dst, "wb") as fout:
|
||||
for line in fin:
|
||||
if line.strip() == EOT:
|
||||
emit("".join(story))
|
||||
story = []
|
||||
if n_stories % FLUSH_EVERY == 0 and buf:
|
||||
np.asarray(buf, dtype=TOKEN_DTYPE).tofile(fout)
|
||||
buf.clear()
|
||||
if progress:
|
||||
rate = n_stories / (time.perf_counter() - t0)
|
||||
print(f"\r {dst.name}: {n_stories:,} stories, "
|
||||
f"{n_tokens:,} tokens ({rate:,.0f}/s)", end="", flush=True)
|
||||
else:
|
||||
story.append(line)
|
||||
emit("".join(story)) # trailing story with no final separator
|
||||
if buf:
|
||||
np.asarray(buf, dtype=TOKEN_DTYPE).tofile(fout)
|
||||
if progress:
|
||||
print()
|
||||
return n_stories, n_tokens
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if not TOK_PATH.exists():
|
||||
sys.exit(f"missing {TOK_PATH} — run scripts/train_tokenizer.py first")
|
||||
tok = BPETokenizer.load(TOK_PATH)
|
||||
OUT.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for src_name, dst_name in JOBS:
|
||||
src, dst = RAW / src_name, OUT / dst_name
|
||||
if not src.exists():
|
||||
sys.exit(f"missing {src} — run scripts/download_data.py first")
|
||||
print(f"encoding {src_name} -> {dst_name}")
|
||||
n_stories, n_tokens = encode_file(tok, src, dst)
|
||||
size_mb = dst.stat().st_size / 1e6
|
||||
print(f" done: {n_stories:,} stories, {n_tokens:,} tokens, {size_mb:,.0f} MB")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Memory-mapped token dataset and random-crop batch sampler.
|
||||
|
||||
Tokens live on disk as a flat array of uint16 (our vocab is < 65536, so two
|
||||
bytes each). We memory-map the file so the multi-hundred-MB token array is
|
||||
never loaded into RAM as Python objects — the OS pages in only the slices each
|
||||
batch actually touches.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Vocab < 65536, so every token id fits in an unsigned 16-bit integer.
|
||||
TOKEN_DTYPE = np.uint16
|
||||
|
||||
|
||||
def load_tokens(path: str | Path) -> np.memmap:
|
||||
"""Memory-map a .bin token file as a read-only 1-D array (never loaded whole)."""
|
||||
return np.memmap(path, dtype=TOKEN_DTYPE, mode="r")
|
||||
|
||||
|
||||
def get_batch(
|
||||
data: np.ndarray,
|
||||
batch_size: int,
|
||||
context_length: int,
|
||||
device: str,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Sample a batch of random crops from the token stream.
|
||||
|
||||
Returns (x, y), each shape (batch_size, context_length), dtype long.
|
||||
y is x shifted left by one: the target at position t is the token that
|
||||
should follow x[:t+1]. A crop of context_length + 1 tokens gives both.
|
||||
"""
|
||||
# Highest start index that still leaves room for context_length + 1 tokens.
|
||||
hi = len(data) - context_length - 1
|
||||
starts = torch.randint(0, hi, (batch_size,), generator=generator).tolist()
|
||||
x = torch.empty((batch_size, context_length), dtype=torch.long)
|
||||
y = torch.empty((batch_size, context_length), dtype=torch.long)
|
||||
for row, s in enumerate(starts):
|
||||
crop = data[s : s + context_length + 1].astype(np.int64)
|
||||
x[row] = torch.from_numpy(crop[:-1])
|
||||
y[row] = torch.from_numpy(crop[1:])
|
||||
return x.to(device), y.to(device)
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Dataset + corpus-encoding tests on synthetic data (no real corpus needed).
|
||||
|
||||
Run directly: `.venv/bin/python tests/test_dataset.py`
|
||||
"""
|
||||
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
sys.path.insert(0, str(ROOT / "scripts"))
|
||||
|
||||
from dataset import TOKEN_DTYPE, get_batch, load_tokens # noqa: E402
|
||||
from tokenizer import BPETokenizer # noqa: E402
|
||||
from encode_corpus import encode_file # noqa: E402
|
||||
|
||||
EOT = "<|endoftext|>"
|
||||
|
||||
|
||||
def test_get_batch_shapes_and_shift():
|
||||
# Tokens 0..999 in order, so the shift property is easy to verify.
|
||||
data = np.arange(1000, dtype=TOKEN_DTYPE)
|
||||
gen = torch.Generator().manual_seed(0)
|
||||
x, y = get_batch(data, batch_size=8, context_length=16, device="cpu", generator=gen)
|
||||
assert x.shape == (8, 16) and y.shape == (8, 16)
|
||||
assert x.dtype == torch.long
|
||||
# y is x shifted by one; with sequential data each y == x + 1.
|
||||
assert torch.equal(y, x + 1)
|
||||
|
||||
|
||||
def test_get_batch_is_reproducible():
|
||||
data = np.arange(1000, dtype=TOKEN_DTYPE)
|
||||
a = get_batch(data, 4, 8, "cpu", torch.Generator().manual_seed(42))
|
||||
b = get_batch(data, 4, 8, "cpu", torch.Generator().manual_seed(42))
|
||||
assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1])
|
||||
|
||||
|
||||
def test_get_batch_in_bounds():
|
||||
data = np.arange(50, dtype=TOKEN_DTYPE)
|
||||
x, y = get_batch(data, 32, 8, "cpu", torch.Generator().manual_seed(1))
|
||||
assert int(x.max()) < 50 and int(y.max()) < 50
|
||||
|
||||
|
||||
def _tiny_tokenizer() -> BPETokenizer:
|
||||
corpus = ("The cat sat. Tom ran fast. Lily played.\n" + EOT + "\n") * 40
|
||||
return BPETokenizer.train(corpus, vocab_size=350, special_tokens=[EOT])
|
||||
|
||||
|
||||
def test_encode_file_roundtrip_and_separators():
|
||||
tok = _tiny_tokenizer()
|
||||
eot_id = tok.special_tokens[EOT]
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
src = Path(d) / "mini.txt"
|
||||
dst = Path(d) / "mini.bin"
|
||||
# Two stories, separator on its own line, with stray blank lines.
|
||||
src.write_text(
|
||||
"The cat sat.\nTom ran fast.\n" + EOT + "\n\n"
|
||||
"Lily played.\n" + EOT + "\n"
|
||||
)
|
||||
n_stories, n_tokens = encode_file(tok, src, dst, progress=False)
|
||||
|
||||
assert n_stories == 2
|
||||
ids = load_tokens(dst).tolist()
|
||||
assert len(ids) == n_tokens
|
||||
# Exactly one EOT per story, and each story ends with one.
|
||||
assert ids.count(eot_id) == 2
|
||||
assert ids[-1] == eot_id
|
||||
|
||||
# Content between separators decodes back to the original stories.
|
||||
first = ids[: ids.index(eot_id)]
|
||||
assert tok.decode(first) == "The cat sat.\nTom ran fast."
|
||||
|
||||
|
||||
def test_encoded_ids_fit_uint16():
|
||||
tok = _tiny_tokenizer()
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
src = Path(d) / "m.txt"
|
||||
dst = Path(d) / "m.bin"
|
||||
src.write_text("The cat sat.\n" + EOT + "\n")
|
||||
encode_file(tok, src, dst, progress=False)
|
||||
arr = load_tokens(dst)
|
||||
assert arr.dtype == TOKEN_DTYPE # stored as uint16
|
||||
assert int(arr.max()) < 65536
|
||||
|
||||
|
||||
def main() -> None:
|
||||
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
||||
for t in tests:
|
||||
t()
|
||||
print(f"ok {t.__name__}")
|
||||
print(f"\n{len(tests)} passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user