File size: 2,673 Bytes
880f286 | 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 | """Tokenize a HuggingFace text dataset into verifiable uint16 shards.
Default: FineWeb-Edu sample-10BT (best small-scale benchmark movement, max
comparability to Pythia/SmolLM). Dataset-agnostic by design so DCLM /
Nemotron-CC become drop-in data-ablation rows:
python scripts/prepare_data.py --target-tokens 3_000_000_000 \
--dataset HuggingFaceFW/fineweb-edu --name sample-10BT --out-dir data/fwedu
# ablation rows (same flags, different source):
--dataset mlfoundations/dclm-baseline-1.0 --out-dir data/dclm
--dataset nvidia/Nemotron-CC --out-dir data/nemotron
Streams the source (no full download), so disk holds only the tokenized output.
"""
from __future__ import annotations
import sys
import argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
from matilda.data import ShardWriter, verify_manifest # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
ap.add_argument("--name", default="sample-10BT")
ap.add_argument("--split", default="train")
ap.add_argument("--text-key", default="text")
ap.add_argument("--tokenizer", default="gpt2")
ap.add_argument("--target-tokens", type=int, default=3_000_000_000)
ap.add_argument("--shard-tokens", type=int, default=100_000_000)
ap.add_argument("--out-dir", default="data/fwedu")
args = ap.parse_args()
import tiktoken
from datasets import load_dataset
enc = tiktoken.get_encoding(args.tokenizer)
eot = enc.eot_token
assert enc.n_vocab <= 65535, "vocab > uint16; use a smaller tokenizer"
ds = load_dataset(args.dataset, name=args.name, split=args.split,
streaming=True)
writer = ShardWriter(args.out_dir, shard_tokens=args.shard_tokens)
n_docs = 0
for doc in ds:
ids = enc.encode_ordinary(doc[args.text_key])
ids.append(eot) # document boundary
writer.add(ids)
n_docs += 1
if n_docs % 1000 == 0:
print(f"\rdocs={n_docs:,} tokens={writer.total_tokens:,}", end="")
if writer.total_tokens >= args.target_tokens:
break
manifest = writer.close(meta={
"dataset": args.dataset, "name": args.name, "split": args.split,
"tokenizer": args.tokenizer, "eot_token": eot, "n_docs": n_docs,
})
print(f"\nwrote {manifest['total_tokens']:,} tokens in "
f"{len(manifest['shards'])} shards -> {args.out_dir}")
verify_manifest(args.out_dir)
print("manifest verified (checksums + sizes OK)")
if __name__ == "__main__":
main()
|