File size: 1,877 Bytes
1c0d385 | 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 | """Streaming encode of a large raw-text corpus to a uint16 .bin file.
Memory-safe: encodes line-by-line, flushes chunks of ~16M tokens.
Usage:
.venv/bin/python data/encode_full.py --raw data/TinyStoriesV2-GPT4-train.txt \
--out data/train_full.bin --tok data/tokenizer.json
"""
import argparse, time
from pathlib import Path
import numpy as np
from data.tokenizer import load_tokenizer
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--raw", default="data/TinyStoriesV2-GPT4-train.txt")
ap.add_argument("--out", default="data/train_full.bin")
ap.add_argument("--tok", default="data/tokenizer.json")
ap.add_argument("--chunk", type=int, default=16_000_000)
ap.add_argument("--max-tokens", type=int, default=600_000_000)
args = ap.parse_args()
tok = load_tokenizer(args.tok)
eot = tok.token_to_id("<|endoftext|>")
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
total, buf = 0, []
t0 = time.time()
with open(args.raw, "rb") as f, open(out, "wb") as g:
for raw in f:
line = raw.decode("utf-8", errors="replace").strip()
if not line:
continue
ids = tok.encode(line).ids
buf.extend(ids)
buf.append(eot)
total += len(ids) + 1
if len(buf) >= args.chunk or total >= args.max_tokens:
np.asarray(buf, dtype=np.uint16).tofile(g)
buf.clear()
print(f"encoded {total:,} tokens in {time.time()-t0:.0f}s "
f"({total/(time.time()-t0):,.0f} tok/s)", flush=True)
if total >= args.max_tokens:
break
if buf:
np.asarray(buf, dtype=np.uint16).tofile(g)
print(f"done: {total:,} tokens -> {out} ({out.stat().st_size/1e9:.2f} GB)", flush=True)
if __name__ == "__main__":
main()
|