SlayerLab/gollem-v5-expand

1#!/usr/bin/env python32"""merge_expanded.py — GoLLeM-v5 expanded-corpus merge (crown-run data, 16M@expanded).3 4Concat (canonical uint16-LE, EOS=12285):5  corpus train.bin (5,396,605,407 tok)6  + fineweb_edu_clean.bin MINUS 42,846 overlap-docs (byte-range-skip) = 2,756,248,149 tok7  + openstax_clean.bin x OS_REPEAT (35,185,551 tok x4 = 140,742,204 tok)8  = expanded-train.bin ~8,293,595,760 tok (16.59 GB)9 10Dataloader = random-offset (train_gpt_ref.py get_batch: torch.randint over whole train.bin)11  -> concat-order IRRELEVANT, NO doc-shuffle needed. Verified 2026-09-22.12 13Overlap byte-ranges from fineweb_overlap_docs.jsonl (bin_byte_start/end, token-aligned;14embedded-EOS doc fineweb-edu:1662841 already merged into one range).15NIE dotyka val.bin (held-out, unchanged). NIE GPU.16"""17import argparse, json, os, time18BPT = 2  # bytes/token (uint16)19 20def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)21 22def load_overlap_ranges(jsonl):23    ranges = []24    with open(jsonl, encoding="utf-8") as f:25        for line in f:26            line = line.strip()27            if not line:28                continue29            d = json.loads(line)30            if "_meta" in d or "bin_byte_start" not in d:31                continue32            ranges.append((int(d["bin_byte_start"]), int(d["bin_byte_end"])))33    ranges.sort()34    merged = []35    for s, e in ranges:36        if merged and s <= merged[-1][1]:37            merged[-1] = (merged[-1][0], max(merged[-1][1], e))38        else:39            merged.append((s, e))40    return merged41 42def copy_whole(src, out, buf=64 * 1024 * 1024):43    n = 044    with open(src, "rb") as f:45        while True:46            b = f.read(buf)47            if not b:48                break49            out.write(b); n += len(b)50    return n51 52def copy_skip(src, out, skip_ranges, buf=64 * 1024 * 1024):53    written = 0; pos = 054    size = os.path.getsize(src)55    with open(src, "rb") as f:56        for s, e in skip_ranges:57            while pos < s:58                b = f.read(min(buf, s - pos))59                if not b:60                    break61                out.write(b); written += len(b); pos += len(b)62            f.seek(e); pos = e63        while pos < size:64            b = f.read(buf)65            if not b:66                break67            out.write(b); written += len(b); pos += len(b)68    return written69 70def main():71    ap = argparse.ArgumentParser()72    ap.add_argument("--corpus", required=True)73    ap.add_argument("--fineweb", required=True)74    ap.add_argument("--openstax", required=True)75    ap.add_argument("--overlap", required=True)76    ap.add_argument("--out", required=True)77    ap.add_argument("--os-repeat", type=int, default=4)78    a = ap.parse_args()79 80    skip = load_overlap_ranges(a.overlap)81    skip_bytes = sum(e - s for s, e in skip)82    log(f"overlap ranges={len(skip)} skip_bytes={skip_bytes:,} (={skip_bytes // BPT:,} tok)")83    assert skip_bytes // BPT == 57241804, f"overlap tokens {skip_bytes // BPT} != 57,241,804 expected"84 85    t0 = time.time()86    with open(a.out, "wb", buffering=64 * 1024 * 1024) as out:87        n_corpus = copy_whole(a.corpus, out)88        log(f"corpus {n_corpus:,}B ({n_corpus // BPT:,} tok) {time.time() - t0:.0f}s")89        n_fw = copy_skip(a.fineweb, out, skip)90        log(f"fineweb-dedup {n_fw:,}B ({n_fw // BPT:,} tok) {time.time() - t0:.0f}s")91        n_os = 092        for r in range(a.os_repeat):93            n = copy_whole(a.openstax, out); n_os += n94            log(f"openstax {r + 1}/{a.os_repeat} {n:,}B {time.time() - t0:.0f}s")95    total = n_corpus + n_fw + n_os96 97    exp_fw = os.path.getsize(a.fineweb) - skip_bytes98    assert n_fw == exp_fw, f"fineweb mismatch {n_fw} != {exp_fw}"99    assert n_corpus // BPT == 5396605407, f"corpus tok {n_corpus // BPT} != 5,396,605,407"100    assert n_fw // BPT == 2756248149, f"fineweb-dedup tok {n_fw // BPT} != 2,756,248,149"101    log(f"VALIDATION OK: corpus={n_corpus // BPT:,} fineweb_dedup={n_fw // BPT:,} "102        f"openstax_x{a.os_repeat}={n_os // BPT:,} TOTAL={total // BPT:,} tok ({total:,}B)")103    log(f"DONE -> {a.out}")104 105if __name__ == "__main__":106    main()107