rzen
c202bada6a
- 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)
97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
"""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()
|