| """Merge several DocumentDataset packs into one corpus. |
| |
| Concatenates the token streams and shifts each pack's document offsets, so no |
| re-tokenisation is needed. All packs must share a tokenizer and seq_len. |
| |
| Training orderings (epoch_0..) are reshuffled across the union of every pack's |
| documents; the holdout (epoch_9) is the union of the packs' holdouts, so the |
| validation mix matches the training mix. |
| |
| python tools/merge_document_packs.py \ |
| --inputs data/dolma3-mix-document-1024 data/fineweb-edu-tokenized-document-1024 \ |
| --out data/mix-dolma3-fineweb-1024 --epochs 9 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import shutil |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def load_pack(path: Path, holdout_epoch: int) -> dict: |
| meta = json.loads((path / "metadata.json").read_text()) |
| train_dir = path / "epoch_0" |
| hold_dir = path / f"epoch_{holdout_epoch}" |
| pack = { |
| "path": path, |
| "meta": meta, |
| "tokens": np.load(path / "tokens.npy", mmap_mode="r"), |
| "train_start": np.load(train_dir / "doc_start.npy"), |
| "train_len": np.load(train_dir / "doc_len.npy"), |
| } |
| if hold_dir.is_dir(): |
| pack["hold_start"] = np.load(hold_dir / "doc_start.npy") |
| pack["hold_len"] = np.load(hold_dir / "doc_len.npy") |
| else: |
| pack["hold_start"] = np.zeros(0, dtype=np.int64) |
| pack["hold_len"] = np.zeros(0, dtype=np.int32) |
| return pack |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--inputs", nargs="+", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--epochs", type=int, default=9) |
| ap.add_argument("--holdout-epoch", type=int, default=9) |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
|
|
| packs = [load_pack(Path(p), args.holdout_epoch) for p in args.inputs] |
| seq_lens = {p["meta"]["max_seq_len"] for p in packs} |
| vocabs = {p["meta"]["tokenizer_info"]["vocab_size"] for p in packs} |
| tokenizers = {p["meta"]["tokenizer_info"].get("name_or_path") for p in packs} |
| if len(seq_lens) != 1 or len(vocabs) != 1 or len(tokenizers) != 1: |
| raise SystemExit(f"packs disagree: seq_len={seq_lens} vocab={vocabs} tokenizer={tokenizers}") |
|
|
| out = Path(args.out) |
| out.mkdir(parents=True, exist_ok=True) |
| total_tokens = sum(int(p["tokens"].size) for p in packs) |
| print(f"[merge] {len(packs)} packs, {total_tokens/1e9:.3f}B tokens total", flush=True) |
|
|
| merged = np.lib.format.open_memmap( |
| out / "tokens.npy", mode="w+", dtype=np.uint16, shape=(total_tokens,) |
| ) |
| offsets = [] |
| cursor = 0 |
| step = 1 << 26 |
| for pack in packs: |
| offsets.append(cursor) |
| src = pack["tokens"] |
| for i in range(0, src.size, step): |
| chunk = src[i : i + step] |
| merged[cursor : cursor + chunk.size] = chunk |
| cursor += int(chunk.size) |
| print(f"[merge] copied {pack['path'].name}: {src.size/1e9:.3f}B tokens", flush=True) |
| merged.flush() |
| del merged |
|
|
| train_start = np.concatenate( |
| [p["train_start"].astype(np.int64) + off for p, off in zip(packs, offsets)] |
| ) |
| train_len = np.concatenate([p["train_len"] for p in packs]) |
| hold_start = np.concatenate( |
| [p["hold_start"].astype(np.int64) + off for p, off in zip(packs, offsets)] |
| ) |
| hold_len = np.concatenate([p["hold_len"] for p in packs]) |
|
|
| rng = np.random.default_rng(args.seed) |
| for epoch in range(args.holdout_epoch): |
| epoch_dir = out / f"epoch_{epoch}" |
| epoch_dir.mkdir(exist_ok=True) |
| order = rng.permutation(train_start.size) |
| np.save(epoch_dir / "doc_start.npy", train_start[order]) |
| np.save(epoch_dir / "doc_len.npy", train_len[order]) |
|
|
| hold_dir = out / f"epoch_{args.holdout_epoch}" |
| hold_dir.mkdir(exist_ok=True) |
| order = rng.permutation(hold_start.size) |
| np.save(hold_dir / "doc_start.npy", hold_start[order]) |
| np.save(hold_dir / "doc_len.npy", hold_len[order]) |
|
|
| meta = dict(packs[0]["meta"]) |
| meta["total_length"] = int(total_tokens) |
| meta["sources"] = [ |
| { |
| "path": str(p["path"]), |
| "tokens": int(p["tokens"].size), |
| "train_docs": int(p["train_start"].size), |
| "holdout_docs": int(p["hold_start"].size), |
| "offset": off, |
| } |
| for p, off in zip(packs, offsets) |
| ] |
| (out / "metadata.json").write_text(json.dumps(meta, indent=2)) |
|
|
| print( |
| f"[merge] done: {train_start.size} train docs " |
| f"({int(train_len.sum())/1e9:.3f}B tokens), " |
| f"{hold_start.size} holdout docs ({int(hold_len.sum())/1e6:.1f}M tokens) -> {out}", |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
