"""Document-level order randomization for the fineweb_edu bins (user directive 2026-07-13:
防止顺序相似文本导致 eval OOD). Two things it fixes:
  1. val was the HEAD of shard 0 (a narrow, possibly topically-clustered slice) -> val becomes a
     seeded random document sample from the WHOLE corpus.
  2. train.bin document order was shard-sequential -> becomes a seeded global permutation
     (future-proofs sequential/streaming readers; note the current trainer's get_batch already
     samples uniform-random offsets, so batch order was never sequential).
Method: full corpus fits RAM (~21GB, box has ~450G free). concat(val, train) restores the exact
original stream (heals the doc split at the old 20M cut), split on <|eot|>(id 0), permute docs,
re-draw val (~20M tokens), rewrite via tmp + atomic replace. Marker: DONE_SHUFFLE.
"""
import os, time
import numpy as np
from pathlib import Path

D = Path('/home/yurenh2/ept/ep_run/data/fineweb_edu')
EOT, VAL_TOKENS, SEED = 0, 20_000_000, 1234
t0 = time.time()

val = np.fromfile(D / 'val.bin', dtype=np.uint16)
train = np.fromfile(D / 'train.bin', dtype=np.uint16)
full = np.concatenate([val, train]); del val, train
print(f"loaded {len(full)/1e9:.3f}B tokens {time.time()-t0:.0f}s", flush=True)

ends = np.flatnonzero(full == EOT) + 1          # each doc ends right after its <|eot|>
starts = np.concatenate([[0], ends[:-1]])
if ends[-1] != len(full):                        # trailing partial doc (no eot) -> keep as one doc
    starts = np.append(starts, ends[-1]); ends = np.append(ends, len(full))
ndoc = len(ends)
print(f"{ndoc/1e6:.2f}M docs, mean {len(full)/ndoc:.0f} tok/doc {time.time()-t0:.0f}s", flush=True)

rng = np.random.default_rng(SEED)
perm = rng.permutation(ndoc)
lens = (ends - starts)[perm]
cut = int(np.searchsorted(np.cumsum(lens), VAL_TOKENS)) + 1   # first `cut` permuted docs -> val

def write(idx, path):
    out = np.empty(int((ends[idx] - starts[idx]).sum()), dtype=np.uint16)
    o = 0
    for i in idx:
        n = ends[i] - starts[i]
        out[o:o + n] = full[starts[i]:ends[i]]; o += n
    tmp = path.with_suffix('.tmp')
    out.tofile(tmp); os.replace(tmp, path)
    return len(out)

nv = write(perm[:cut], D / 'val.bin')
nt = write(perm[cut:], D / 'train.bin')
print(f"DONE_SHUFFLE train {nt/1e9:.3f}B val {nv/1e6:.1f}M (seed {SEED}, doc-level) {time.time()-t0:.0f}s", flush=True)