|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import argparse
|
| import array
|
| import hashlib
|
| import json
|
| import os
|
| import random
|
| import struct
|
| import sys
|
| import time
|
|
|
|
|
|
|
|
|
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__))
|
| for _p in (_HERE, os.path.dirname(_HERE)):
|
| if _p not in sys.path:
|
| sys.path.insert(0, _p)
|
|
|
| MAGIC_FLAGS = 0
|
| TOKENS_PER_SHARD = 8_000_000
|
| MERGE_BATCH_DOCS = 32
|
| VAL_STRIDE = 50
|
|
|
| VAL_TOKENS_PER_SHARD = 2_000_000
|
| ROW_TOKEN_CAP = 16_384
|
| SHUFFLE_SEED = 20260919
|
| HASHES_FILE = "seen_hashes.bin"
|
| KEYS_FILE = "seen_keys.bin"
|
| STATE_NAME = "build_state.json"
|
|
|
|
|
|
|
| SOURCES = [
|
| dict(repo="HuggingFaceFW/fineweb-edu", config="sample-10BT", col="text", tokens=340_000_000,
|
| min_score=4, why="quality-scored general English"),
|
| dict(repo="HuggingFaceFW/finepdfs-edu", config="eng_Latn", col="text", tokens=180_000_000,
|
| en_filter=True, why="long-form educational PDFs; needs own en filter"),
|
| dict(repo="HuggingFaceTB/cosmopedia", config="stanford", col="text", tokens=60_000_000,
|
| why="textbook expository"),
|
| dict(repo="HuggingFaceTB/cosmopedia", config="openstax", col="text", tokens=30_000_000,
|
| why="textbook science"),
|
| dict(repo="HuggingFaceTB/cosmopedia", config="khanacademy", col="text", tokens=10_000_000,
|
| why="worked steps"),
|
| dict(repo="HuggingFaceTB/cosmopedia", config="auto_math_text", col="text", tokens=30_000_000,
|
| why="synthetic math exposition"),
|
| dict(repo="HuggingFaceTB/cosmopedia", config="wikihow", col="text", tokens=30_000_000,
|
| why="procedural / physical-world steps"),
|
| dict(repo="HuggingFaceFW/finewiki", config="en", col="rewritten", tokens=80_000_000,
|
| why="rewritten encyclopedic prose"),
|
| dict(repo="omarkamali/wikipedia-monthly", config="20260101.en", col="text", tokens=60_000_000,
|
| en_filter=True, why="current encyclopedic breadth"),
|
| dict(repo="HuggingFaceTB/finemath", config="finemath-4plus", col="text", tokens=160_000_000,
|
| why="math; upstream-decontaminated against GSM8K/MMLU/ARC"),
|
| dict(repo="open-web-math/open-web-math", config="default", col="text", tokens=70_000_000,
|
| en_filter=True, why="forum/webbook math prose"),
|
| dict(repo="HuggingFaceFW/finephrase", config="tutorial", col="text", tokens=60_000_000,
|
| en_filter=True, why="stepwise tutorial register"),
|
| dict(repo="HuggingFaceFW/finephrase", config="faq", col="text", tokens=30_000_000,
|
| en_filter=True, why="FAQ register"),
|
| dict(repo="HuggingFaceFW/finephrase", config="table", col="text", tokens=30_000_000,
|
| en_filter=True, why="tabular->prose"),
|
| dict(repo="SimpleStories/SimpleStories", config="default", col="story", tokens=40_000_000,
|
| why="long-range simple narrative coherence"),
|
| dict(repo="common-pile/arxiv_abstracts", config="default", col="text", tokens=20_000_000,
|
| en_filter=True, why="CC0 scientific abstracts"),
|
|
|
|
|
|
|
| ]
|
|
|
| STOP = (" the ", " and ", " of ", " to ", " that ", " is ", " in ", " for ", " with ", " as ")
|
| KEY_COLS = ("url", "id", "file_path", "doc_id", "generation_id", "repo_name")
|
|
|
|
|
| def src_key(src):
|
| return f"{src['repo'].split('/')[1]}__{src['config']}"
|
|
|
|
|
| def normalize(s):
|
|
|
|
|
| return " ".join(s.split())
|
|
|
|
|
| def looks_english(s):
|
| low = s[:600].lower()
|
| return sum(w in low for w in STOP) >= 3
|
|
|
|
|
| def h64(s):
|
| """64-bit digest of a string. digest()[:8] is 8 BYTES; hexdigest()[:16] is 16 *characters*
|
| and broke the first version (E-015)."""
|
| return struct.unpack("<Q", hashlib.sha256(s.encode("utf-8")).digest()[:8])[0]
|
|
|
|
|
| def key_of(row, cols):
|
| for c in KEY_COLS:
|
| if c in cols:
|
| v = row.get(c)
|
| if isinstance(v, str) and v:
|
| return v
|
| return None
|
|
|
|
|
|
|
| def write_shard(path, vocab, docs):
|
| """docs: iterable of array('H')/list of token ids. Returns (n_docs, n_tokens, sha256, bytes)."""
|
| offs = array.array("I", [0])
|
| ids = array.array("H")
|
| for d in docs:
|
| ids.extend(d)
|
| offs.append(len(ids))
|
| tmp = path + ".tmp"
|
| with open(tmp, "wb") as f:
|
| f.write(struct.pack("<IIII", vocab, len(offs) - 1, len(ids), MAGIC_FLAGS))
|
| offs.tofile(f)
|
| ids.tofile(f)
|
| os.replace(tmp, path)
|
| digest = hashlib.sha256(open(path, "rb").read()).hexdigest()
|
| return len(offs) - 1, len(ids), digest, os.path.getsize(path)
|
|
|
|
|
| def read_shard_header(path):
|
| with open(path, "rb") as f:
|
| vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", f.read(16))
|
| return vocab, n_docs, n_tokens, flags
|
|
|
|
|
| def iter_shard_docs(path):
|
| """Yield documents as array('H'). Offsets are read in full (4 B/doc) which is fine at 8M tok/shard."""
|
| with open(path, "rb") as f:
|
| vocab, n_docs, n_tokens, _flags = struct.unpack("<IIII", f.read(16))
|
| offs = array.array("I")
|
| offs.fromfile(f, n_docs + 1)
|
| ids = array.array("H")
|
| ids.fromfile(f, n_tokens)
|
| for i in range(n_docs):
|
| yield ids[offs[i]:offs[i + 1]]
|
|
|
|
|
|
|
| class State:
|
| def __init__(self, root):
|
| self.root = root
|
| self.stage_dir = os.path.join(root, "stage")
|
| self.mix_dir = os.path.join(root, "mix")
|
| os.makedirs(self.stage_dir, exist_ok=True)
|
| os.makedirs(self.mix_dir, exist_ok=True)
|
| self.path = os.path.join(root, STATE_NAME)
|
| self.hashes_path = os.path.join(root, HASHES_FILE)
|
| self.keys_path = os.path.join(root, KEYS_FILE)
|
| self.d = {"staged": {}, "merged_shards": [], "version": 3}
|
| self.seen_exact = set()
|
| self.seen_key = set()
|
| if os.path.exists(self.path):
|
| self.d = json.load(open(self.path))
|
| self.seen_exact = self._load_u64(self.hashes_path)
|
| self.seen_key = self._load_u64(self.keys_path)
|
| print(f"resuming: {len(self.d['staged'])} sources staged, "
|
| f"{len(self.d['merged_shards'])} mixed shards, "
|
| f"{len(self.seen_exact):,} hashes", flush=True)
|
|
|
| @staticmethod
|
| def _load_u64(p):
|
| if not os.path.exists(p):
|
| return set()
|
| a = array.array("Q")
|
| with open(p, "rb") as f:
|
| a.fromfile(f, os.path.getsize(p) // 8)
|
| return set(a)
|
|
|
| def commit(self, new_hashes, new_keys):
|
| """Append the dedup sets FIRST, then the state file: if the job dies between the two, the worst
|
| case is a hash set slightly ahead of the state (a document re-emitted is a duplicate, which the
|
| next run's dedup catches; a hash set behind the state is a permanent silent duplicate)."""
|
| if new_hashes:
|
| with open(self.hashes_path, "ab") as f:
|
| array.array("Q", new_hashes).tofile(f)
|
| if new_keys:
|
| with open(self.keys_path, "ab") as f:
|
| array.array("Q", new_keys).tofile(f)
|
| tmp = self.path + ".tmp"
|
| with open(tmp, "w") as f:
|
| json.dump(self.d, f, indent=1)
|
| os.replace(tmp, self.path)
|
|
|
| def staged_tokens(self, key):
|
| return sum(s["tokens"] for s in self.d["staged"].get(key, {}).get("shards", []))
|
|
|
| def staged_tokens_total(self):
|
| return sum(v.get("tokens", 0) for v in self.d["staged"].values())
|
|
|
| def merged_tokens(self):
|
| return sum(s["tokens"] for s in self.d["merged_shards"])
|
|
|
|
|
|
|
| def stage(args, st, tokenizer):
|
| remote = HS.remote_stage_files(args.hub_repo) if (HS and args.hub_repo
|
| and not args.merge_only) else set()
|
| if args.hub_repo:
|
| print(f"hub: {len(remote)} staged files already on {args.hub_repo}", flush=True)
|
| for src in SOURCES:
|
| key = src_key(src)
|
| rec = st.d["staged"].get(key, {})
|
| if rec.get("complete"):
|
| print(f"stage: skip {key} (complete)", flush=True)
|
| continue
|
| if HS and args.hub_repo:
|
|
|
|
|
| have = [p for p in remote if p.startswith(f"stage/{key}/") and p.endswith(".bin")]
|
| if have and f"stage/{key}/record.json" in remote:
|
| try:
|
| HS.download_missing(st.root, [key], args.hub_repo, remote=remote)
|
| rec_json = json.load(open(os.path.join(st.stage_dir, key, "record.json")))
|
| st.d["staged"][key] = {"complete": True, "shards": rec_json["shards"],
|
| "tokens": rec_json["tokens"],
|
| "docs_seen": rec_json["docs_seen"],
|
| "dropped": rec_json["dropped"],
|
| "target": rec_json["tokens"], "restored_from_hub": True}
|
| st.commit([], [])
|
| print(f"stage: restored {key} from Hub ({rec_json['tokens']:,} tok), skipping",
|
| flush=True)
|
| continue
|
| except Exception as e:
|
| print(f"!! restore failed for {key}: {type(e).__name__}: {str(e)[:140]}", flush=True)
|
| want = min(src["tokens"], args.cap_tokens_per_source or src["tokens"])
|
| if args.budget_tokens and st.staged_tokens_total() >= args.budget_tokens:
|
| print("stage: overall budget reached", flush=True)
|
| break
|
| import datasets
|
|
|
| try:
|
| ds = datasets.load_dataset(src["repo"], src["config"], split="train", streaming=True)
|
| except Exception as e:
|
| print(f"!! stage: {key} could not load: {type(e).__name__}: {str(e)[:200]}", flush=True)
|
| rec["error"] = f"{type(e).__name__}: {str(e)[:200]}"
|
| st.d["staged"][key] = rec
|
| st.commit([], [])
|
| continue
|
|
|
| got = docs_seen = 0
|
| drops = {"short_or_low_score": 0, "non_english": 0, "exact_dup": 0, "key_dup": 0,
|
| "too_few_tokens": 0, "no_column": 0}
|
| cols = None
|
| buf, new_h, new_k = [], [], []
|
| sd = os.path.join(st.stage_dir, key)
|
| os.makedirs(sd, exist_ok=True)
|
| shards = rec.get("shards", [])
|
| t0 = time.monotonic()
|
| for row in ds:
|
| if got >= want:
|
| break
|
| if args.sec_per_source and time.monotonic() - t0 > args.sec_per_source:
|
| break
|
| if cols is None:
|
| cols = set(row.keys())
|
| if src["col"] not in cols:
|
| print(f"!! stage: {key} has no column {src['col']!r}; seen {sorted(cols)[:10]}",
|
| flush=True)
|
| drops["no_column"] = -1
|
| break
|
| raw = row.get(src["col"])
|
| if not isinstance(raw, str):
|
| continue
|
| docs_seen += 1
|
| s = normalize(raw)
|
| if len(s) < args.min_chars:
|
| drops["short_or_low_score"] += 1
|
| continue
|
| if src.get("min_score") is not None:
|
| sc = row.get("int_score", row.get("score"))
|
| try:
|
| if sc is not None and float(sc) < src["min_score"]:
|
| drops["short_or_low_score"] += 1
|
| continue
|
| except (TypeError, ValueError):
|
| pass
|
| if src.get("en_filter") and not looks_english(s):
|
| drops["non_english"] += 1
|
| continue
|
| h = h64(s)
|
| if h in st.seen_exact:
|
| drops["exact_dup"] += 1
|
| continue
|
| u = key_of(row, cols)
|
| ku = h64(u) if u else None
|
| if ku is not None and ku in st.seen_key:
|
| drops["key_dup"] += 1
|
| continue
|
| ids = tokenizer.encode(s[:args.max_chars], add_special_tokens=False).ids[:ROW_TOKEN_CAP]
|
| if len(ids) < 16:
|
| drops["too_few_tokens"] += 1
|
| continue
|
| st.seen_exact.add(h)
|
| new_h.append(h)
|
| if ku is not None:
|
| st.seen_key.add(ku)
|
| new_k.append(ku)
|
| buf.append(array.array("H", ids))
|
| got += len(ids)
|
| if sum(len(b) for b in buf) >= TOKENS_PER_SHARD:
|
| idx = len(shards)
|
| p = os.path.join(sd, f"shard-{idx:04d}.bin")
|
| nd, nt, dg, bs = write_shard(p, tokenizer.get_vocab_size(), buf)
|
| shards.append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
|
| "sha256": dg, "bytes": bs})
|
| st.seen_exact.update(new_h)
|
| st.d["staged"][key] = {"complete": False, "shards": shards, "tokens": got,
|
| "docs_seen": docs_seen, "dropped": drops}
|
| st.commit(new_h, new_k)
|
| new_h, new_k = [], []
|
| buf = []
|
| print(f" staged {key} shard {idx:04d}: {nt:,} tok [{time.monotonic() - t0:.0f}s]",
|
| flush=True)
|
|
|
| if buf:
|
| idx = len(shards)
|
| p = os.path.join(sd, f"shard-{idx:04d}.bin")
|
| nd, nt, dg, bs = write_shard(p, tokenizer.get_vocab_size(), buf)
|
| shards.append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
|
| "sha256": dg, "bytes": bs})
|
| print(f" staged {key} shard {idx:04d}: {nt:,} tok (tail)", flush=True)
|
|
|
| complete = got >= want
|
| st.d["staged"][key] = {"complete": bool(complete), "shards": shards, "tokens": got,
|
| "docs_seen": docs_seen, "dropped": dict(drops),
|
| "target": want, "seconds": round(time.monotonic() - t0, 1)}
|
| st.commit(new_h, new_k)
|
|
|
|
|
| if complete and HS and args.hub_repo:
|
| with open(os.path.join(st.stage_dir, key, "record.json"), "w") as f:
|
| json.dump({"key": key, "src": {k: v for k, v in src.items()},
|
| "shards": shards, "tokens": got, "docs_seen": docs_seen,
|
| "dropped": dict(drops)}, f, indent=1)
|
| try:
|
| n = HS.upload_source(st.root, key, args.hub_repo)
|
| print(f" published {key} -> {args.hub_repo}/stage/{key} ({n} files)", flush=True)
|
| except Exception as e:
|
| print(f"!! publish FAILED for {key}: {type(e).__name__}: {str(e)[:160]} -- "
|
| f"this source will have to be re-staged if the session dies", flush=True)
|
| el = time.monotonic() - t0
|
| print(f"stage: {key} {got:,} tok of {want:,} from {docs_seen:,} docs; "
|
| f"dup={drops['exact_dup']} keydup={drops['key_dup']} lang={drops['non_english']} "
|
| f"short={drops['short_or_low_score']} [{got / max(1e-9, el):,.0f} tok/s]", flush=True)
|
|
|
|
|
|
|
| def merge(args, st, tokenizer):
|
| """Interleave staged documents proportionally to each source's staged token total.
|
|
|
| --exclude-dir applies per-source uint8 masks written by audit_contamination --write-filter: a 1 byte
|
| means "this document's 13-grams appear in a benchmark train/validation/dev split, do not emit it"
|
| (§3.3). Masks are indexed by document ordinal within the source, which is stable because both the
|
| audit and the merge replay the source's own record.json shard order.
|
| """
|
| vocab = tokenizer.get_vocab_size()
|
| masks, dropped, filter_digest = {}, {}, None
|
| if args.exclude_dir:
|
| fj = os.path.join(args.exclude_dir, "filter.json")
|
| if not os.path.exists(fj):
|
| print(f"FATAL: --exclude-dir {args.exclude_dir} contains no filter.json", flush=True)
|
| sys.exit(5)
|
| filter_digest = hashlib.sha256(open(fj, "rb").read()).hexdigest()
|
| fj_obj = json.load(open(fj))
|
| if not fj_obj.get("scanned_sources") and not fj_obj.get("sources"):
|
| print("FATAL: filter.json names no sources at all -- nothing was scanned, so merging with it "
|
| "would publish an unfiltered mix under a filter claim", flush=True)
|
| sys.exit(5)
|
| if not fj_obj.get("scanned_sources"):
|
|
|
|
|
| print(f"WARNING: filter.json predates the scanned_sources field; trusting its "
|
| f"{len(fj_obj.get('sources', {}))} per-source entries because every merged source must "
|
| "still carry a mask or the merge refuses", flush=True)
|
| dropped = fj_obj.get("sources", {})
|
| for key in sorted(os.listdir(args.exclude_dir)):
|
| if key.endswith(".u8"):
|
| with open(os.path.join(args.exclude_dir, key), "rb") as fh:
|
| masks[key[:-3]] = fh.read()
|
| print(f"exclude: {len(masks)} mask(s), "
|
| f"{sum(len(m) for m in masks.values()):,} docs covered, "
|
| f"{sum(int(v.get('dropped_docs', 0)) for v in dropped.values()):,} flagged", flush=True)
|
| streams = []
|
| for src in SOURCES:
|
| key = src_key(src)
|
| rec = st.d["staged"].get(key)
|
| if not rec or not rec.get("shards"):
|
| continue
|
| staged_docs = sum(int(s.get("docs", 0)) for s in rec["shards"])
|
| m = masks.get(key)
|
| if m is not None and len(m) != staged_docs:
|
|
|
| print(f"FATAL: mask for {key} covers {len(m)} docs but {staged_docs} are staged", flush=True)
|
| sys.exit(5)
|
| weight = rec.get("tokens") or sum(s["tokens"] for s in rec["shards"])
|
| weight = weight - int((dropped.get(key) or {}).get("dropped_tokens", 0))
|
| if weight <= 0:
|
| print(f" {key}: fully excluded, dropped from the merge", flush=True)
|
| continue
|
| iters = [iter_shard_docs(os.path.join(st.root, s["file"])) for s in rec["shards"]]
|
|
|
| def gen(its=iters, m=m):
|
| if m is None:
|
| for it in its:
|
| for d in it:
|
| yield d
|
| return
|
| i = 0
|
| for it in its:
|
| for d in it:
|
| if not m[i]:
|
| yield d
|
| i += 1
|
|
|
| streams.append({"key": key, "weight": float(weight), "credit": 0.0, "it": gen(iters),
|
| "pending": None, "emitted_tokens": 0, "exhausted": False})
|
| if not streams:
|
|
|
|
|
|
|
| print("FATAL merge: nothing staged. Refusing to emit an empty mix.", flush=True)
|
| sys.exit(4)
|
| if args.exclude_dir:
|
|
|
|
|
| unscanned = [s["key"] for s in streams if s["key"] not in masks]
|
| if unscanned:
|
| print(f"FATAL merge: --exclude-dir has no mask for {unscanned}. Refusing to merge sources "
|
| "the contamination audit did not read.", flush=True)
|
| sys.exit(5)
|
|
|
|
|
|
|
|
|
|
|
|
|
| digest = hashlib.sha256(json.dumps({
|
| "streams": [[s["key"], round(s["weight"]) ] for s in streams], "seed": SHUFFLE_SEED,
|
| "tokens_per_shard": TOKENS_PER_SHARD, "val_stride": VAL_STRIDE,
|
| "filter": filter_digest}, sort_keys=True).encode()).hexdigest()
|
| prior = st.d.get("merge_inputs")
|
| if prior is not None and prior != digest and st.d["merged_shards"]:
|
| if not (args.rebuild_merged or args.exclude_dir):
|
| print("FATAL merge: the mix root already holds shards merged from DIFFERENT inputs "
|
| f"(recorded {prior[:12]}, now {digest[:12]}). Continuing would be a silent no-op. "
|
| "Pass --rebuild-merged to discard them deliberately.", flush=True)
|
| sys.exit(6)
|
| n_old = len(st.d["merged_shards"]) + len(st.d.get("val_shards", []))
|
| for grp in ("merged_shards", "val_shards"):
|
| for rec in st.d.get(grp, []):
|
| p = os.path.join(st.root, rec["file"])
|
| if os.path.exists(p):
|
| os.remove(p)
|
| st.d["merged_shards"] = []
|
| st.d["val_shards"] = []
|
| print(f"merge: inputs changed ({prior[:12]} -> {digest[:12]}); discarded {n_old} old shard(s) "
|
| "and rebuilding from the staged sources", flush=True)
|
| st.d["merge_inputs"] = digest
|
| total_w = sum(s["weight"] for s in streams)
|
| print(f"merge: {len(streams)} sources, "
|
| f"{sum(s['weight'] for s in streams) / 1e6:,.0f}M staged tokens", flush=True)
|
|
|
| def next_doc(s):
|
| if s["pending"] is None and not s["exhausted"]:
|
| try:
|
| s["pending"] = next(s["it"])
|
| except StopIteration:
|
| s["exhausted"] = True
|
| return s["pending"]
|
|
|
| cap = args.merge_budget_tokens or int(total_w)
|
| buf, buf_src, buf_tokens = [], [], 0
|
| val_buf, val_tokens, docs_seq = [], 0, 0
|
| out_dir = os.path.join(st.root, "mix")
|
| os.makedirs(out_dir, exist_ok=True)
|
| t0 = time.monotonic()
|
|
|
| def flush_val():
|
| nonlocal val_buf, val_tokens
|
| if not val_buf:
|
| return
|
| idx = len(st.d.get("val_shards", []))
|
| p = os.path.join(st.root, "val", f"shard-{idx:04d}.bin")
|
| os.makedirs(os.path.dirname(p), exist_ok=True)
|
| tmp = p + ".tmp"
|
| offs = array.array("I", [0])
|
| ids = array.array("H")
|
| for d in val_buf:
|
| ids.extend(d)
|
| offs.append(len(ids))
|
| with open(tmp, "wb") as f:
|
| f.write(struct.pack("<IIII", vocab, len(val_buf), len(ids), MAGIC_FLAGS))
|
| offs.tofile(f)
|
| ids.tofile(f)
|
| os.replace(tmp, p)
|
| st.d.setdefault("val_shards", []).append({
|
| "file": os.path.relpath(p, st.root), "docs": len(val_buf), "tokens": len(ids),
|
| "sha256": hashlib.sha256(open(p, "rb").read()).hexdigest(), "bytes": os.path.getsize(p)})
|
| st.commit([], [])
|
| print(f" val shard {idx:04d}: {len(ids):,} tok / {len(val_buf):,} docs", flush=True)
|
| val_buf, val_tokens = [], 0
|
|
|
| def flush():
|
| nonlocal buf, buf_src, buf_tokens
|
| if not buf:
|
| return
|
| random.Random(SHUFFLE_SEED + len(st.d["merged_shards"])).shuffle(buf)
|
| docs = [d for d, _ in buf]
|
| p = os.path.join(out_dir, f"shard-{len(st.d['merged_shards']):04d}.bin")
|
| nd, nt, dg, bs = write_shard(p, vocab, docs)
|
| st.d["merged_shards"].append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
|
| "sha256": dg, "bytes": bs,
|
| "sources": sorted(set(s for _, s in buf))})
|
| st.commit([], [])
|
| print(f" mix shard {len(st.d['merged_shards']) - 1:04d}: {nt:,} tok / {nd:,} docs "
|
| f"from {len(st.d['merged_shards'][-1]['sources'])} sources "
|
| f"[{time.monotonic() - t0:.0f}s]", flush=True)
|
| buf, buf_src, buf_tokens = [], [], 0
|
|
|
| while st.merged_tokens() + buf_tokens < cap:
|
| live = [s for s in streams if next_doc(s) is not None]
|
| if not live:
|
| break
|
|
|
|
|
| emitted = sum(s["emitted_tokens"] for s in streams) or 1
|
| pick = max(live, key=lambda s: s["weight"] / total_w - s["emitted_tokens"] / emitted)
|
| took = 0
|
| while took < MERGE_BATCH_DOCS:
|
| d = next_doc(pick)
|
| if d is None:
|
| break
|
| docs_seq += 1
|
| if docs_seq % VAL_STRIDE == 0:
|
|
|
| val_buf.append(d)
|
| val_tokens += len(d)
|
| if val_tokens >= VAL_TOKENS_PER_SHARD:
|
| flush_val()
|
| pick["emitted_tokens"] += len(d)
|
| pick["pending"] = None
|
| took += 1
|
| continue
|
| buf.append((d, pick["key"]))
|
| buf_tokens += len(d)
|
| pick["emitted_tokens"] += len(d)
|
| pick["pending"] = None
|
| took += 1
|
| if buf_tokens >= TOKENS_PER_SHARD:
|
| flush()
|
| break
|
| if all(s["pending"] is None and s["exhausted"] for s in streams):
|
| break
|
| flush()
|
| flush_val()
|
| print(f"merge: {len(st.d['val_shards']) if 'val_shards' in st.d else 0} val shards, "
|
| f"{sum(s['tokens'] for s in st.d.get('val_shards', [])):,} held-out tokens", flush=True)
|
| print(f"merge: {len(st.d['merged_shards'])} shards, "
|
| f"{sum(s['tokens'] for s in st.d['merged_shards']):,} tokens", flush=True)
|
|
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--root", default="/kaggle/working/mixroot")
|
| ap.add_argument("--budget-tokens", type=int, default=1_310_000_000, help="stage cap")
|
| ap.add_argument("--merge-budget-tokens", type=int, default=0, help="0 = take everything staged")
|
| ap.add_argument("--cap-tokens-per-source", type=int, default=0, help="smoke test knob")
|
| ap.add_argument("--min-chars", type=int, default=300)
|
| ap.add_argument("--max-chars", type=int, default=100_000)
|
| ap.add_argument("--sec-per-source", type=int, default=0)
|
| ap.add_argument("--only", default="", help="comma list of config names to stage")
|
| ap.add_argument("--hub-repo", default="Cion-lab/ounce100m-mix-stage",
|
| help="staged sources are published here so a session kill cannot lose them; "
|
| "empty string disables sync")
|
| ap.add_argument("--stage-only", action="store_true")
|
| ap.add_argument("--merge-only", action="store_true")
|
| ap.add_argument("--exclude-dir", default="",
|
| help="directory of per-source <key>.u8 contamination masks + filter.json written by "
|
| "audit_contamination --write-filter; documents flagged there are not merged (§3.3)")
|
| ap.add_argument("--rebuild-merged", action="store_true",
|
| help="discard shards already recorded in build_state.json before merging. Implied by "
|
| "--exclude-dir, because a filter changes every shard's contents")
|
| args = ap.parse_args()
|
|
|
| if args.only:
|
| keep = set(args.only.split(","))
|
| SOURCES[:] = [s for s in SOURCES if s["config"] in keep or src_key(s) in keep]
|
| print(f"--only: {len(SOURCES)} source(s)", flush=True)
|
|
|
| from tokenizers import Tokenizer
|
| from huggingface_hub import hf_hub_download
|
| tok_path = hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json")
|
| tk = Tokenizer.from_file(tok_path)
|
| tk.no_truncation()
|
| assert tk.get_vocab_size() <= 65536, "uint16 shard layout breaks above 65,536"
|
| print(f"tokenizer vocab={tk.get_vocab_size()} "
|
| f"sha256={hashlib.sha256(open(tok_path, 'rb').read()).hexdigest()[:12]}", flush=True)
|
|
|
| global HS
|
| HS = None
|
|
|
|
|
|
|
| if args.hub_repo:
|
| try:
|
| import hubsync as HS_mod
|
| HS = HS_mod
|
| HS.ensure_repo(args.hub_repo)
|
| except Exception as e:
|
| print(f"FATAL hub sync could not initialise: {type(e).__name__}: {str(e)[:200]}",
|
| flush=True)
|
| print("Refusing to run a multi-hour build that is not interruption-safe. Resolve the Hub "
|
| "write path, or pass --hub-repo '' only if you accept losing all staged work when the "
|
| "session ends.", flush=True)
|
| sys.exit(3)
|
|
|
| st = State(args.root)
|
| if args.hub_repo and not st.d["staged"]:
|
|
|
|
|
|
|
| try:
|
| remote = HS.remote_stage_files(args.hub_repo)
|
| keys = sorted({p.split("/")[1] for p in remote if p.startswith("stage/") and "/" in p})
|
| pulled = HS.download_missing(args.root, keys, args.hub_repo, remote=remote)
|
| for k in keys:
|
| rp = os.path.join(args.root, "stage", k, "record.json")
|
| if not os.path.exists(rp):
|
| continue
|
| j = json.load(open(rp))
|
| st.d["staged"][k] = {"complete": True, "shards": j["shards"], "tokens": j["tokens"],
|
| "docs_seen": j["docs_seen"], "dropped": j["dropped"],
|
| "target": j["tokens"], "restored_from_hub": True}
|
| st.commit([], [])
|
| print(f"hub: restored {len(st.d['staged'])} staged source(s) "
|
| f"({st.staged_tokens_total():,} tokens)", flush=True)
|
| except Exception as e:
|
| print(f"FATAL hub restore failed: {type(e).__name__}: {str(e)[:200]}", flush=True)
|
| sys.exit(3)
|
|
|
| t0 = time.monotonic()
|
| if not args.merge_only:
|
| stage(args, st, tk)
|
| if not args.stage_only:
|
| merge(args, st, tk)
|
|
|
| manifest = {
|
| "tokenizer": {"id": "HuggingFaceTB/SmolLM2-135M:tokenizer.json",
|
| "vocab_size": tk.get_vocab_size(),
|
| "sha256": hashlib.sha256(open(tok_path, "rb").read()).hexdigest()},
|
| "layout": {"dtype": "uint16", "header": "<IIII (vocab, n_docs, n_tokens, flags)",
|
| "offsets": "uint32 × (n_docs+1) cumulative",
|
| "doc_slice": "ids[offsets[i]:offsets[i+1]]",
|
| "tokens_per_shard_target": TOKENS_PER_SHARD, "row_token_cap": ROW_TOKEN_CAP,
|
| "merge_batch_docs": MERGE_BATCH_DOCS, "shuffle_seed_base": SHUFFLE_SEED,
|
| "resume_cursor": "(shard_index, token_offset) into the concatenated id stream"},
|
| "stage": {k: {kk: vv for kk, vv in v.items() if kk != "shards"}
|
| for k, v in st.d["staged"].items()},
|
| "staged_tokens": st.staged_tokens_total(),
|
| "shards": st.d["merged_shards"],
|
| "n_shards": len(st.d["merged_shards"]),
|
| "val_shards": st.d.get("val_shards", []),
|
| "total_tokens": sum(s["tokens"] for s in st.d["merged_shards"]),
|
| "val_tokens": sum(s["tokens"] for s in st.d.get("val_shards", [])),
|
| "val_holdout": {"stride": VAL_STRIDE, "note": "every 50th merged document, sampled across "
|
| "sources rather than whole documents, excluded from training shards"},
|
| "distinct_sources_per_shard_min": min((len(s["sources"]) for s in st.d["merged_shards"]),
|
| default=0),
|
| "seconds": round(time.monotonic() - t0, 1),
|
| }
|
| if args.exclude_dir:
|
|
|
|
|
| try:
|
| fj = json.load(open(os.path.join(args.exclude_dir, "filter.json")))
|
| manifest["contamination_filter"] = {
|
| "criterion": fj.get("criterion"), "k": fj.get("k"),
|
| "docs_dropped": sum(int(v.get("dropped_docs", 0)) for v in fj["sources"].values()),
|
| "tokens_dropped": sum(int(v.get("dropped_tokens", 0) or 0)
|
| for v in fj["sources"].values()),
|
| "per_source": {k2: {"docs": v2.get("dropped_docs"),
|
| "tokens": v2.get("dropped_tokens")}
|
| for k2, v2 in fj["sources"].items()}}
|
| except Exception as e:
|
| print(f"FATAL: --exclude-dir set but filter.json unreadable: {e}", flush=True)
|
| sys.exit(5)
|
| with open(os.path.join(args.root, "manifest.json"), "w") as f:
|
| json.dump(manifest, f, indent=1)
|
|
|
|
|
| print("BUILD_JSON_BEGIN")
|
| print(json.dumps({k: manifest[k] for k in
|
| ("staged_tokens", "total_tokens", "val_tokens", "n_shards",
|
| "distinct_sources_per_shard_min", "seconds")}, indent=1))
|
| print("BUILD_JSON_END")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
