| """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}`") |
| |
| |
| |
| |
| |
| 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: |
| |
| |
| 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 |
|
|
