Upload prepare_data.py with huggingface_hub

prepare_data.py

@@ -0,0 +1,137 @@
+"""
+Prepare İvme's pretraining data: stream the dense mix, tokenize, pack to disk.
+
+Output is a flat uint16 memmap (vocab 16384 < 65536, so uint16 is exact). We
+write documents in ASCENDING quality order so a sequential read during training
+acts as a curriculum — the model sees noisier web text first and the densest
+material (textbooks, then Wikipedia) last. Research shows this ordering plus a
+moderate LR decay beats random shuffling for free.
+
+The mix mirrors what made Archaea-74M punch so far above its weight, pushed a
+little denser (more math, stricter web filter):
+
+    FineWeb-HQ (score-gated web)   45%   ~710M tokens   [first / lowest density]
+    Python stack (filtered)        10%   ~160M tokens
+    FineMath-4+                    15%   ~235M tokens
+    Cosmopedia (stanford+wikihow)  25%   ~395M tokens
+    Wikipedia EN                    5%    ~80M tokens   [last / highest density]
+                                   ----  -----------
+                                   100%  ~1.57B tokens  (Chinchilla-optimal)
+
+Usage:
+    python prepare_data.py                 # full ~1.57B token build
+    python prepare_data.py --smoke         # tiny build to test the pipeline
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+import numpy as np
+
+from huggingface_hub import login
+
+login(token="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")
+
+TOKENIZER_PATH = "ivme_tokenizer.json"
+OUT_DIR = "data"
+DTYPE = np.uint16
+
+# (source_key, target_tokens) in ASCENDING quality order — written in this order.
+TOKEN_BUDGET = [
+    ("fineweb_hq", 710_000_000),
+    ("python",     160_000_000),
+    ("finemath",   235_000_000),
+    ("cosmopedia", 395_000_000),
+    ("wikipedia",   80_000_000),
+]
+SMOKE_BUDGET = [(k, 200_000) for k, _ in TOKEN_BUDGET]
+
+VAL_TOKENS = 2_000_000  # held out from the tail of each source proportionally
+
+
+def make_stream(source_key):
+    """Return (iterable_of_text, text_field) for a source."""
+    from datasets import load_dataset
+
+    if source_key == "fineweb_hq":
+        ds = load_dataset("epfml/FineWeb-HQ", split="train", streaming=True)
+        return ds, "text"
+    if source_key == "cosmopedia":
+        # Two dense subsets concatenated.
+        a = load_dataset("HuggingFaceTB/cosmopedia", "stanford", split="train", streaming=True)
+        b = load_dataset("HuggingFaceTB/cosmopedia", "wikihow", split="train", streaming=True)
+        from itertools import chain
+        return chain(a, b), "text"
+    if source_key == "finemath":
+        ds = load_dataset("HuggingFaceTB/finemath", "finemath-4plus", split="train", streaming=True)
+        return ds, "text"
+    if source_key == "python":
+        ds = load_dataset("bigcode/python-stack-v1-functions-filtered", split="train", streaming=True)
+        return ds, "content"
+    if source_key == "wikipedia":
+        ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
+        return ds, "text"
+    raise ValueError(source_key)
+
+
+def build(budget):
+    from tokenizers import Tokenizer
+
+    os.makedirs(OUT_DIR, exist_ok=True)
+    tok = Tokenizer.from_file(TOKENIZER_PATH)
+    eos_id = tok.token_to_id("<|eos|>")
+
+    train_path = os.path.join(OUT_DIR, "train.bin")
+    val_path = os.path.join(OUT_DIR, "val.bin")
+
+    total_target = sum(n for _, n in budget)
+    print(f"[data] target ~{total_target/1e6:.0f}M tokens across {len(budget)} sources")
+
+    train_f = open(train_path, "wb")
+    val_buf = []  # small, held in memory
+
+    written_train = 0
+    for source_key, target in budget:
+        stream, field = make_stream(source_key)
+        src_written = 0
+        # Reserve a slice of each source's tail for validation.
+        val_target = int(VAL_TOKENS * (target / total_target))
+        print(f"[data] {source_key}: target {target/1e6:.0f}M (val {val_target/1e6:.2f}M)")
+
+        for row in stream:
+            text = row.get(field)
+            if not text:
+                continue
+            ids = tok.encode(text).ids
+            ids.append(eos_id)  # document boundary
+            arr = np.array(ids, dtype=DTYPE)
+
+            if len(val_buf) * 0 + src_written >= target:
+                break
+            # Send the first val_target tokens of this source to val, rest to train.
+            if src_written < val_target:
+                val_buf.append(arr)
+            else:
+                arr.tofile(train_f)
+                written_train += len(arr)
+            src_written += len(arr)
+
+            if src_written % 5_000_000 < len(arr):
+                print(f"  [{source_key}] {src_written/1e6:.1f}M / {target/1e6:.0f}M")
+
+    train_f.close()
+    val_arr = np.concatenate(val_buf) if val_buf else np.array([], dtype=DTYPE)
+    val_arr.tofile(val_path)
+
+    print(f"[data] train.bin : {written_train:,} tokens -> {train_path}")
+    print(f"[data] val.bin   : {len(val_arr):,} tokens -> {val_path}")
+    print(f"[data] curriculum order preserved (sequential read = ascending quality)")
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--smoke", action="store_true", help="tiny build to test the pipeline")
+    args = ap.parse_args()
+    build(SMOKE_BUDGET if args.smoke else TOKEN_BUDGET)