Upload scripts/mix_buckets.py with huggingface_hub

scripts/mix_buckets.py

@@ -79,6 +79,11 @@ def main():
     ap.add_argument("--window", type=int, default=65536, help="CP_T")
     ap.add_argument("--grad-accum", type=int, default=8)
     ap.add_argument("--split", default="train")
+    ap.add_argument("--skip-tokens", type=float, default=0,
+                    help="discard this many tokens of the interleaved stream before "
+                         "writing. CP_INIT_FROM restarts the loader at token 0 "
+                         "(skip_batches is bypassed, bluscriptCP.cpp:2047), so a "
+                         "continuation must skip in the DATA instead.")
     a = ap.parse_args()
 
     srcs = []
@@ -112,6 +117,8 @@ def main():
 
     got = {n: 0 for n, _ in srcs}
     written = 0
+    skip = int(a.skip_tokens)
+    emitted = 0
     shard_idx = 0
     buf = []
     buf_n = 0
@@ -147,12 +154,17 @@ def main():
             exhausted.add(pick)
             print(f"    {pick} exhausted at {got[pick]/1e6:.1f} M tokens")
             continue
-        buf.append(doc)
-        buf_n += doc.size
         got[pick] += doc.size
         written += doc.size
-        if buf_n >= shard_tokens:
-            flush()
+        # The picker is driven by RUNNING SHARES and never references the total,
+        # so the document sequence is identical for any --total-tokens. That is
+        # what makes a skip equivalent to resuming the same stream.
+        if written > skip:
+            if emitted == 0 and written - doc.size < skip:
+                doc = doc[skip - (written - doc.size):]   # partial first document
+            buf.append(doc); buf_n += doc.size; emitted += doc.size
+            if buf_n >= shard_tokens:
+                flush()
     flush()
 
     # --- emulate DataLoaderLite so STEPS is right ---------------------------
@@ -173,7 +185,8 @@ def main():
     for n, w in srcs:
         print(f"    {n:<22} target {w:5.1%}   actual {got[n]/written:5.1%}  "
               f"({got[n]/1e6:.1f} M tokens)")
-    print(f"    total {written/1e6:.1f} M tokens in {len(paths)} shards")
+    print(f"    stream {written/1e6:.1f} M tokens; skipped {skip/1e6:.1f} M; "
+          f"WROTE {emitted/1e6:.1f} M in {len(paths)} shards")
     print(f"\n  USABLE STEPS at window {BT:,} x grad_accum {a.grad_accum}: {steps}")
     print(f"    (naive total//global_batch would say "
           f"{written // (BT * a.grad_accum)} -- use {steps}, it accounts for the")