File size: 3,917 Bytes
f4fc57b | 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 | """Tokenize text pools into uint16 token shards (one .bin per pool shard).
Documents are shuffled within each pool shard (seeded) and packed as
<bos> doc <eos> <bos> doc <eos> ... with no padding; the training loader
reads contiguous (seq_len+1)-token windows. Vocab is 65536 so uint16 is
exact.
"""
import json
import multiprocessing as mp
from pathlib import Path
import numpy as np
from ..paths import pool_dir, shard_dir, tokenizer_dir, ensure
ENCODE_BATCH = 512
def _iter_pool_docs(path: Path):
import zstandard
with open(path, "rb") as raw:
reader = zstandard.ZstdDecompressor().stream_reader(raw)
import io
for line in io.TextIOWrapper(reader, encoding="utf-8"):
if line.strip():
yield json.loads(line)["text"]
def _pack_one(args) -> tuple[str, int]:
pool_shard, out_path, tok_dir, seed = args
from ..tok.wrapper import Tok
import random
tok = Tok(tok_dir)
docs = list(_iter_pool_docs(Path(pool_shard)))
random.Random(seed).shuffle(docs)
n = 0
with open(out_path, "wb") as out:
for i in range(0, len(docs), ENCODE_BATCH):
batch = tok.encode_batch(docs[i:i + ENCODE_BATCH], bos=True, eos=True)
flat = np.concatenate([np.asarray(ids, dtype=np.uint16) for ids in batch])
flat.tofile(out)
n += len(flat)
return str(out_path), n
def pack(lang: str, tok_name: str, workers: int = 8, seed: int = 1234,
max_tokens: float | None = None) -> dict:
src = pool_dir(lang)
all_shards = sorted(src.glob("pool_*.jsonl.zst"))
if not all_shards:
raise FileNotFoundError(f"no pool shards in {src} - run `xscript pool --lang {lang}`")
# While `xscript pool` is still running, its writer may have the highest-
# indexed shard open (only closed on the next roll/checkpoint) -- hold it
# back so we never read a truncated file mid-download. Safe to include
# once the pool has reached its byte budget; picked up on the next
# incremental `pack` call otherwise.
pool_done = False
stats_path = src / "stats.json"
if stats_path.exists():
pst = json.loads(stats_path.read_text())
pool_done = pst["text_bytes"] >= pst["budget_bytes"] * 0.99
shards = all_shards if pool_done else all_shards[:-1]
out = ensure(shard_dir(lang, tok_name))
tok_dir = str(tokenizer_dir(tok_name))
index_path = out / "index.json"
index = json.loads(index_path.read_text()) if index_path.exists() else {}
import zlib
jobs = []
for s in shards:
dst = out / (s.name.replace(".jsonl.zst", ".bin"))
if dst.name not in index:
jobs.append((str(s), str(dst), tok_dir, seed ^ zlib.crc32(s.name.encode())))
total = sum(index.values())
if jobs:
with mp.Pool(workers) as pool:
for dst, n in pool.imap_unordered(_pack_one, jobs):
if n == 0:
# Empty pool checkpoint shards contain no documents. Do
# not expose a zero-byte file to np.memmap in PackedStream.
Path(dst).unlink(missing_ok=True)
continue
index[Path(dst).name] = n
total += n
index_path.write_text(json.dumps(index, indent=2, sort_keys=True))
print(f"[pack] {lang}/{tok_name}: {Path(dst).name} = {n/1e6:.1f}M tokens "
f"(total {total/1e9:.2f}B)")
if max_tokens and total >= max_tokens:
print("[pack] reached max_tokens; stopping")
pool.terminate()
break
meta = {"lang": lang, "tokenizer": tok_name, "total_tokens": total,
"n_shards": len(index)}
(out / "meta.json").write_text(json.dumps(meta, indent=2))
print(f"[pack] {lang}/{tok_name}: {total/1e9:.3f}B tokens in {len(index)} shards")
return meta
|