chore: import upstream snapshot with attribution

scripts/prepare_pretrain_data.py

@@ -0,0 +1,139 @@
+"""
+Download + tokenize Pile-uncopyrighted shards into a single flat-token HDF5 for
+pretraining. Faster than the original ``data_preprocess.py`` (which resizes the HDF5 per
+document): here we stream-decompress, batch-tokenize with tiktoken, and write tokens to
+the HDF5 in large chunks.
+
+Writes to /ephemeral by default (the 1.5TB disk).
+
+Examples:
+    # dev split from the Pile validation file
+    PYTHONPATH=. python scripts/prepare_pretrain_data.py --split val \
+        --out /ephemeral/data/pile_dev.h5
+    # one training shard
+    PYTHONPATH=. python scripts/prepare_pretrain_data.py --split train --num_shards 1 \
+        --out /ephemeral/data/pile_train.h5
+"""
+
+from __future__ import annotations
+
+import argparse
+import io
+import json
+import os
+
+import h5py
+import numpy as np
+import requests
+import tiktoken
+import zstandard as zstd
+from tqdm import tqdm
+
+BASE_URL = "https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main"
+EOT_ID = 50256
+WRITE_CHUNK = 8_000_000  # flush tokens to HDF5 in ~8M-token chunks
+ENC_BATCH = 1024          # documents per tiktoken batch-encode
+
+
+def shard_urls(split: str, num_shards: int) -> list[str]:
+    if split == "val":
+        return [f"{BASE_URL}/val.jsonl.zst"]
+    return [f"{BASE_URL}/train/{i:02d}.jsonl.zst" for i in range(num_shards)]
+
+
+def download(url: str, dest: str) -> str:
+    if os.path.exists(dest):
+        print(f"  cached: {dest}")
+        return dest
+    os.makedirs(os.path.dirname(dest), exist_ok=True)
+    print(f"  downloading {url}")
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        total = int(r.headers.get("content-length", 0))
+        with open(dest + ".part", "wb") as f:
+            for chunk in tqdm(r.iter_content(1 << 20), total=total >> 20, unit="MB", desc="dl"):
+                f.write(chunk)
+    os.replace(dest + ".part", dest)
+    return dest
+
+
+def iter_texts(zst_path: str):
+    dctx = zstd.ZstdDecompressor()
+    with open(zst_path, "rb") as fh:
+        reader = dctx.stream_reader(fh)
+        for line in io.TextIOWrapper(reader, encoding="utf-8"):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                txt = json.loads(line).get("text")
+            except json.JSONDecodeError:
+                continue
+            if txt:
+                yield txt
+
+
+def tokenize_to_h5(zst_paths: list[str], out_path: str, max_tokens: int | None) -> int:
+    enc = tiktoken.get_encoding("r50k_base")
+    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+    total = 0
+    buf: list[int] = []
+    with h5py.File(out_path, "w") as f:
+        dset = f.create_dataset("tokens", (0,), maxshape=(None,), dtype="i4", chunks=(WRITE_CHUNK,))
+
+        def flush():
+            nonlocal total, buf
+            if not buf:
+                return
+            arr = np.asarray(buf, dtype=np.int32)
+            dset.resize(total + arr.size, axis=0)
+            dset[total: total + arr.size] = arr
+            total += arr.size
+            buf = []
+
+        for zp in zst_paths:
+            print(f"  tokenizing {zp}")
+            docs: list[str] = []
+            pbar = tqdm(iter_texts(zp), unit="doc", desc="tok")
+            for txt in pbar:
+                docs.append(txt)
+                if len(docs) >= ENC_BATCH:
+                    for ids in enc.encode_ordinary_batch(docs):
+                        buf.extend(ids)
+                        buf.append(EOT_ID)
+                    docs = []
+                    if len(buf) >= WRITE_CHUNK:
+                        flush()
+                        pbar.set_postfix(tokens=f"{total/1e6:.1f}M")
+                    if max_tokens and total >= max_tokens:
+                        break
+            if docs:
+                for ids in enc.encode_ordinary_batch(docs):
+                    buf.extend(ids)
+                    buf.append(EOT_ID)
+            flush()
+            if max_tokens and total >= max_tokens:
+                break
+    print(f"  wrote {total:,} tokens -> {out_path}")
+    return total
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--split", choices=["train", "val"], required=True)
+    p.add_argument("--num_shards", type=int, default=1, help="train shards to use")
+    p.add_argument("--raw_dir", default="/ephemeral/data/pile_raw")
+    p.add_argument("--out", required=True)
+    p.add_argument("--max_tokens", type=int, default=None, help="stop after this many tokens")
+    args = p.parse_args()
+
+    urls = shard_urls(args.split, args.num_shards)
+    local = []
+    for u in urls:
+        name = "val.jsonl.zst" if args.split == "val" else os.path.basename(u)
+        local.append(download(u, os.path.join(args.raw_dir, name)))
+    tokenize_to_h5(local, args.out, args.max_tokens)
+
+
+if __name__ == "__main__":
+    main()