Upload source/files/in2_gen.py with huggingface_hub

2ff33ec verified 2 months ago

15.5 kB

#!/usr/bin/env python3
"""
in2_gen.py -- synthesise IN2-style TRAINING shards for bluscriptCP.
IN2 (arXiv:2404.16811) teaches long-range retrieval by making the answer
impossible to produce without a fact planted at a random distant position. We
cannot reproduce its "loss on the answer only" masking (see the plan: TI has no
maskable CE, `narrow` has no autograd node, and sentinel target ids are unsafe
because the fwd/bwd kernels do not bounds-check). So we remove the NEED to mask:
  * our filler is CONTIGUOUS REAL TEXT, not IN2's shuffled-segment bag, so
    unmasked LM loss on it is ordinary harmless LM training; and
  * dilution is fixed by construction -- K triples per example raise the answer
    tokens from ~0.1% of positions to several percent.
Output is what DataLoaderLite already eats: a HEADERLESS little-endian uint16
token stream in a file whose name contains "train"/"val", under a directory
passed as CP_DATA_ROOT. No C++ training/loss/loader change is required.
LAYOUT (one example = exactly T tokens)
    [ filler, with needles and decoys planted throughout ]
       ^needle_i (key_i + val_i)        ^decoy (key + val, NO probe)
    ... and at each site p_i:  [ key_i ][ val_i ]
                                probe    answer
  * SITES p_i are drawn from an ARM-INDEPENDENT rng stream, so arms B and C have
    bit-identical answer-position histograms and differ ONLY in needle distance.
    Without this we would difference NLLs measured at different points on a
    position-varying curve -- the error needle_gen.py's fixed-slot design exists
    to prevent.
  * Sites are INTERLEAVED over the sequence, never concatenated at the tail:
    tail-packing would confine every supervised position to the last K*SPAN
    offsets and train a position-specific circuit indistinguishable from
    retrieval.
  * PAIRING is shuffled and DECOY needles (no probe) are planted, so "the i-th
    probe pairs with the i-th needle" is not a winning heuristic. That shortcut
    is order/counting, not content retrieval, and would earn full reward at any
    distance.
  * KEY ALPHABET is capped (default 512, not the full 50304). Uniform draws from
    the vocab cost ln(50304)=10.83 nats IRREDUCIBLY -- pure noise that inflates
    the pre-clip gradient norm and, against the GLOBAL clip_grad_norm_(1.0),
    could silently leave the arms not LR-matched. ln(512)=6.24 nats gives the
    same discriminability (512^4 = 6.9e10 keys) with 42% less noise.
ARMS
  a  plain text, no injection            -- drift control
  b  needle immediately before its probe -- the objective change at distance ~0
  c  needle uniform in [0, p_i - SPAN)   -- DISTANCE, the actual IN2 claim
Stdlib only, matching needle_gen.py (the local pyenv python3 lacks numpy).
Usage
  python3 in2_gen.py --arm c --T 16384 --examples 64 --split train --out-dir /tmp/in2c
  python3 in2_gen.py --arm b --T 16384 --examples 64 --split train --out-dir /tmp/in2b
"""
import argparse, array, json, os, random, sys

VOCAB = 50304          # must match cfg.vocab_size; the C++ side bound-checks too
PAD_TAIL = 4           # >=2: the loader's y-shift makes the last position of each
                       # example target the NEXT example's first token. Keep the
                       # answer clear of it.
MIX = 1_000_003        # odd multiplier for deterministic seed derivation
MAX_SHARD_TOKENS = 100_000_000   # DataLoaderLite is constructed with this cap;
                                 # anything past it in a shard is SILENTLY invisible


def check_platform():
    if sys.byteorder != "little":
        raise SystemExit("training shards must be uint16 LITTLE-endian; this host is big-endian")
    if array.array("H").itemsize != 2:
        raise SystemExit("unexpected array itemsize ('H' must be 2B)")


def load_filler(path, need, rng):
    """One CONTIGUOUS run of real tokens per example, so filler keeps local
    textual statistics. Source shards are already uint16, so unlike
    needle_gen.py (which must widen to int32 for CP_EVAL_TOKENS_BIN) we keep the
    native width -- training shards are uint16."""
    ntok = os.path.getsize(path) // 2
    if ntok < need + 1:
        raise SystemExit(f"{path} has {ntok} tokens, need >= {need + 1}")
    start = rng.randrange(0, ntok - need)
    buf = array.array("H")
    with open(path, "rb") as f:
        f.seek(start * 2)
        buf.fromfile(f, need)
    lo, hi = min(buf), max(buf)
    if lo < 0 or hi >= VOCAB:
        raise SystemExit(f"{path} has ids outside [0,{VOCAB}): min={lo} max={hi}")
    return buf


def draw_sites(T, n, span, rng, lo_margin):
    """n site starts in [lo_margin, T-PAD_TAIL-span], separated by >= 2*span.
    Drawn from an ARM-INDEPENDENT rng so every arm sees the same sites. Returned
    sorted; the pairing to needles is shuffled separately.
    Separation is 2*span, not span, and the reason is load-bearing: arm B plants
    its needle immediately before the probe, at p-span. At only span separation,
    two adjacent sites can sit exactly span apart and B's needle then lands on
    top of the earlier site, so B silently drops that triple -- which would make
    B's answer_start histogram DIFFER from C's and destroy the matched-position
    discipline the B-vs-C contrast depends on. The constraint is applied for all
    arms so sites stay arm-independent.
    """
    hi = T - PAD_TAIL - span
    if hi <= lo_margin:
        raise SystemExit(f"T={T} too small for span={span}")
    sep = 2 * span
    chosen = []
    guard = 0
    while len(chosen) < n:
        guard += 1
        if guard > 400 * n:
            raise SystemExit(f"could not place {n} sites in T={T} at separation {sep}; "
                             f"lower --K (need roughly K*{sep} << {hi - lo_margin})")
        p = rng.randrange(lo_margin, hi + 1)
        if all(abs(p - q) >= sep for q in chosen):
            chosen.append(p)
    return sorted(chosen)


def place_free(occupied, T, span, rng, hi_bound):
    """A start in [0, hi_bound-span] disjoint from `occupied`; None if it fails."""
    for _ in range(400):
        if hi_bound - span < 0:
            return None
        s = rng.randrange(0, hi_bound - span + 1)
        if all(s + span <= a or s >= a + span for a, _ in occupied):
            return s
    return None


def make_example(arm, T, K, key_len, val_len, alphabet, decoys, d_min,
                 filler, site_rng, arm_rng):
    """Returns (tokens uint16 array of length T, list of records).
    site_rng  -- ARM-INDEPENDENT: sites, keys, values, pairing.
    arm_rng   -- arm-specific: needle placement only.
    """
    span = key_len + val_len
    toks = filler                                    # fresh copy per example
    recs = []

    if arm == "a":                                   # plain text, no injection
        return toks, recs, 0

    sites = draw_sites(T, K, span, site_rng, lo_margin=span + d_min)

    # Distinct key/value pairs, then a shuffled site<->needle pairing so the
    # i-th probe does not correspond to the i-th needle by position.
    keys = [[site_rng.randrange(alphabet) for _ in range(key_len)] for _ in range(K)]
    vals = [[site_rng.randrange(alphabet) for _ in range(val_len)] for _ in range(K)]
    order = list(range(K))
    site_rng.shuffle(order)

    occupied = [(p, "site") for p in sites]

    for si, p in enumerate(sites):
        j = order[si]                                # which key/value this site queries
        key, val = keys[j], vals[j]
        toks[p:p + key_len] = array.array("H", key)              # probe
        toks[p + key_len:p + span] = array.array("H", val)       # answer

        if arm == "b":
            ns = p - span                            # immediately before the probe
            # Guaranteed free by the 2*span site separation; assert rather than
            # skip, because a silent skip would desynchronise B's histogram from C's.
            if ns < 0 or any(ns + span > a and ns < a + span
                             for a, _ in occupied if a != p):
                raise SystemExit(f"arm b: needle slot {ns} for site {p} is occupied -- "
                                 "site separation invariant violated")
        else:                                        # arm c: uniform, far
            hi = p - d_min
            ns = place_free(occupied, T, span, arm_rng, hi)
            if ns is None:
                raise SystemExit(f"arm c: could not place a needle before site {p} "
                                 f"(d_min={d_min}); lower --K or --d-min")
        toks[ns:ns + span] = array.array("H", key + val)
        occupied.append((ns, "needle"))

        answer_start = p + key_len
        recs.append(dict(needle_start=ns, probe_start=p, answer_start=answer_start,
                         distance=p - ns,
                         csv_first=answer_start - 1))   # see off-by-one note below

    n_decoy = 0
    # Decoy needles: no probe queries them, so needle count != probe count and
    # the i-th-of-K index carries no information.
    for _ in range(decoys):
        key = [site_rng.randrange(alphabet) for _ in range(key_len)]
        val = [site_rng.randrange(alphabet) for _ in range(val_len)]
        ds = place_free(occupied, T, span, arm_rng, T - PAD_TAIL)
        if ds is None:
            continue
        toks[ds:ds + span] = array.array("H", key + val)
        occupied.append((ds, "decoy"))
        n_decoy += 1

    return toks, recs, n_decoy


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=["a", "b", "c"])
    ap.add_argument("--T", type=int, required=True)
    ap.add_argument("--examples", type=int, required=True)
    ap.add_argument("--split", default="train", choices=["train", "val"])
    ap.add_argument("--K", type=int, default=64, help="triples per example")
    ap.add_argument("--key-len", type=int, default=4, help="probe/key tokens")
    ap.add_argument("--val-len", type=int, default=8, help="answer/value tokens")
    ap.add_argument("--alphabet", type=int, default=512, help="key/value alphabet size")
    ap.add_argument("--decoys", type=int, default=None, help="default K//4")
    ap.add_argument("--d-min", type=int, default=64, help="arm c: min needle->probe distance")
    ap.add_argument("--filler-shards", nargs="+", required=True)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--shards", type=int, default=1,
                    help="split the examples across this many files. REQUIRED above "
                         "100M tokens: DataLoaderLite is constructed with a "
                         "max_tokens_per_shard of 1e8 and anything past it in a single "
                         "shard is SILENTLY invisible. The example index stays GLOBAL "
                         "across shards, so the arm-independent RNG stream (sites, keys, "
                         "pairing, filler) is unchanged by how the split falls -- arms "
                         "remain comparable.")
    args = ap.parse_args()

    check_platform()
    if args.alphabet > VOCAB:
        raise SystemExit(f"--alphabet {args.alphabet} exceeds VOCAB {VOCAB}")
    decoys = args.decoys if args.decoys is not None else args.K // 4
    span = args.key_len + args.val_len
    for f in args.filler_shards:
        if not os.path.exists(f):
            raise SystemExit(f"filler shard not found: {f}")
        if args.split == "train" and "val" in os.path.basename(f):
            raise SystemExit(f"refusing to build TRAIN data from a val shard: {f}\n"
                             "  eval filler must stay disjoint from training data")
    os.makedirs(args.out_dir, exist_ok=True)

    total_tok = args.examples * args.T
    per_shard = (args.examples + args.shards - 1) // args.shards
    if per_shard * args.T > MAX_SHARD_TOKENS:
        raise SystemExit(f"{per_shard * args.T} tokens/shard exceeds the loader's "
                         f"{MAX_SHARD_TOKENS} cap; raise --shards to at least "
                         f"{(total_tok + MAX_SHARD_TOKENS - 1) // MAX_SHARD_TOKENS}")

    all_recs, hist, n_decoy_total = [], {}, 0
    shard_paths = []
    stream = array.array("H")
    shard_i = 0
    for i in range(args.examples):
        # Site/key/value/pairing rng is ARM-INDEPENDENT -- identical across arms.
        site_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 7)
        # Needle placement is arm-specific.
        arm_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX
                                + 1000 + ord(args.arm))
        # Filler draw is arm-independent too, so arms see the same documents.
        fill_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 99)
        shard = args.filler_shards[i % len(args.filler_shards)]
        filler = load_filler(shard, args.T, fill_rng)

        toks, recs, n_dec = make_example(args.arm, args.T, args.K, args.key_len, args.val_len,
                                  args.alphabet, decoys, args.d_min,
                                  filler, site_rng, arm_rng)
        if len(toks) != args.T:
            raise SystemExit(f"example {i} is {len(toks)} tokens, expected {args.T}")
        stream.extend(toks)
        for r in recs:
            hist[r["answer_start"]] = hist.get(r["answer_start"], 0) + 1
        all_recs.append(recs)
        n_decoy_total += n_dec

        # Flush a shard whenever it is full, or on the last example. The global `i`
        # above drives every RNG, so the split point cannot change the data.
        if len(stream) // args.T >= per_shard or i == args.examples - 1:
            sp = os.path.join(args.out_dir, f"in2{args.arm}_{args.split}_{shard_i:05d}.bin")
            with open(sp, "wb") as f:
                stream.tofile(f)
            shard_paths.append(sp)
            stream = array.array("H")
            shard_i += 1

    path = shard_paths[0]

    n_ans = sum(len(r) for r in all_recs) * args.val_len
    # Counted from ACTUAL placements, never derived from parameters: arm A plants
    # nothing, so a parameter-derived figure reported phantom decoy tokens for it.
    n_syn = sum(len(r) for r in all_recs) * (2 * span) + n_decoy_total * span
    man = dict(arm=args.arm, T=args.T, examples=args.examples, split=args.split,
               K=args.K, key_len=args.key_len, val_len=args.val_len, span=span,
               alphabet=args.alphabet, decoys=decoys, d_min=args.d_min,
               seed=args.seed, vocab=VOCAB, pad_tail=PAD_TAIL,
               filler_shards=args.filler_shards, bin=path, shards=shard_paths,
               total_tokens=total_tok, dtype="uint16",
               answer_tokens=n_ans, answer_frac=n_ans / total_tok,
               synthetic_tokens=n_syn, synthetic_frac=n_syn / total_tok,
               decoys_placed=n_decoy_total,
               answer_start_hist=dict(sorted(hist.items())),
               records=all_recs)
    with open(path.replace(".bin", ".json"), "w") as f:
        json.dump(man, f, indent=1)

    print(f"wrote {len(shard_paths)} shard(s) in {args.out_dir}  {total_tok} tokens total")
    print(f"  answer tokens {n_ans} ({n_ans/total_tok:.2%})  "
          f"synthetic {n_syn} ({n_syn/total_tok:.2%})")
    if all_recs and all_recs[0]:
        ds = [r["distance"] for rs in all_recs for r in rs]
        ds.sort()
        print(f"  distance: min {ds[0]}  median {ds[len(ds)//2]}  max {ds[-1]}  n={len(ds)}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	in2_gen.py -- synthesise IN2-style TRAINING shards for bluscriptCP.

	IN2 (arXiv:2404.16811) teaches long-range retrieval by making the answer
	impossible to produce without a fact planted at a random distant position. We
	cannot reproduce its "loss on the answer only" masking (see the plan: TI has no
	maskable CE, `narrow` has no autograd node, and sentinel target ids are unsafe
	because the fwd/bwd kernels do not bounds-check). So we remove the NEED to mask:

	* our filler is CONTIGUOUS REAL TEXT, not IN2's shuffled-segment bag, so
	unmasked LM loss on it is ordinary harmless LM training; and
	* dilution is fixed by construction -- K triples per example raise the answer
	tokens from ~0.1% of positions to several percent.

	Output is what DataLoaderLite already eats: a HEADERLESS little-endian uint16
	token stream in a file whose name contains "train"/"val", under a directory
	passed as CP_DATA_ROOT. No C++ training/loss/loader change is required.

	LAYOUT (one example = exactly T tokens)

	[ filler, with needles and decoys planted throughout ]
	^needle_i (key_i + val_i) ^decoy (key + val, NO probe)
	... and at each site p_i: [ key_i ][ val_i ]
	probe answer

	* SITES p_i are drawn from an ARM-INDEPENDENT rng stream, so arms B and C have
	bit-identical answer-position histograms and differ ONLY in needle distance.
	Without this we would difference NLLs measured at different points on a
	position-varying curve -- the error needle_gen.py's fixed-slot design exists
	to prevent.
	* Sites are INTERLEAVED over the sequence, never concatenated at the tail:
	tail-packing would confine every supervised position to the last K*SPAN
	offsets and train a position-specific circuit indistinguishable from
	retrieval.
	* PAIRING is shuffled and DECOY needles (no probe) are planted, so "the i-th
	probe pairs with the i-th needle" is not a winning heuristic. That shortcut
	is order/counting, not content retrieval, and would earn full reward at any
	distance.
	* KEY ALPHABET is capped (default 512, not the full 50304). Uniform draws from
	the vocab cost ln(50304)=10.83 nats IRREDUCIBLY -- pure noise that inflates
	the pre-clip gradient norm and, against the GLOBAL clip_grad_norm_(1.0),
	could silently leave the arms not LR-matched. ln(512)=6.24 nats gives the
	same discriminability (512^4 = 6.9e10 keys) with 42% less noise.

	ARMS
	a plain text, no injection -- drift control
	b needle immediately before its probe -- the objective change at distance ~0
	c needle uniform in [0, p_i - SPAN) -- DISTANCE, the actual IN2 claim

	Stdlib only, matching needle_gen.py (the local pyenv python3 lacks numpy).

	Usage
	python3 in2_gen.py --arm c --T 16384 --examples 64 --split train --out-dir /tmp/in2c
	python3 in2_gen.py --arm b --T 16384 --examples 64 --split train --out-dir /tmp/in2b
	"""
	import argparse, array, json, os, random, sys

	VOCAB = 50304 # must match cfg.vocab_size; the C++ side bound-checks too
	PAD_TAIL = 4 # >=2: the loader's y-shift makes the last position of each
	# example target the NEXT example's first token. Keep the
	# answer clear of it.
	MIX = 1_000_003 # odd multiplier for deterministic seed derivation
	MAX_SHARD_TOKENS = 100_000_000 # DataLoaderLite is constructed with this cap;
	# anything past it in a shard is SILENTLY invisible


	def check_platform():
	if sys.byteorder != "little":
	raise SystemExit("training shards must be uint16 LITTLE-endian; this host is big-endian")
	if array.array("H").itemsize != 2:
	raise SystemExit("unexpected array itemsize ('H' must be 2B)")


	def load_filler(path, need, rng):
	"""One CONTIGUOUS run of real tokens per example, so filler keeps local
	textual statistics. Source shards are already uint16, so unlike
	needle_gen.py (which must widen to int32 for CP_EVAL_TOKENS_BIN) we keep the
	native width -- training shards are uint16."""
	ntok = os.path.getsize(path) // 2
	if ntok < need + 1:
	raise SystemExit(f"{path} has {ntok} tokens, need >= {need + 1}")
	start = rng.randrange(0, ntok - need)
	buf = array.array("H")
	with open(path, "rb") as f:
	f.seek(start * 2)
	buf.fromfile(f, need)
	lo, hi = min(buf), max(buf)
	if lo < 0 or hi >= VOCAB:
	raise SystemExit(f"{path} has ids outside [0,{VOCAB}): min={lo} max={hi}")
	return buf


	def draw_sites(T, n, span, rng, lo_margin):
	"""n site starts in [lo_margin, T-PAD_TAIL-span], separated by >= 2*span.

	Drawn from an ARM-INDEPENDENT rng so every arm sees the same sites. Returned
	sorted; the pairing to needles is shuffled separately.

	Separation is 2*span, not span, and the reason is load-bearing: arm B plants
	its needle immediately before the probe, at p-span. At only span separation,
	two adjacent sites can sit exactly span apart and B's needle then lands on
	top of the earlier site, so B silently drops that triple -- which would make
	B's answer_start histogram DIFFER from C's and destroy the matched-position
	discipline the B-vs-C contrast depends on. The constraint is applied for all
	arms so sites stay arm-independent.
	"""
	hi = T - PAD_TAIL - span
	if hi <= lo_margin:
	raise SystemExit(f"T={T} too small for span={span}")
	sep = 2 * span
	chosen = []
	guard = 0
	while len(chosen) < n:
	guard += 1
	if guard > 400 * n:
	raise SystemExit(f"could not place {n} sites in T={T} at separation {sep}; "
	f"lower --K (need roughly K*{sep} << {hi - lo_margin})")
	p = rng.randrange(lo_margin, hi + 1)
	if all(abs(p - q) >= sep for q in chosen):
	chosen.append(p)
	return sorted(chosen)


	def place_free(occupied, T, span, rng, hi_bound):
	"""A start in [0, hi_bound-span] disjoint from `occupied`; None if it fails."""
	for _ in range(400):
	if hi_bound - span < 0:
	return None
	s = rng.randrange(0, hi_bound - span + 1)
	if all(s + span <= a or s >= a + span for a, _ in occupied):
	return s
	return None


	def make_example(arm, T, K, key_len, val_len, alphabet, decoys, d_min,
	filler, site_rng, arm_rng):
	"""Returns (tokens uint16 array of length T, list of records).

	site_rng -- ARM-INDEPENDENT: sites, keys, values, pairing.
	arm_rng -- arm-specific: needle placement only.
	"""
	span = key_len + val_len
	toks = filler # fresh copy per example
	recs = []

	if arm == "a": # plain text, no injection
	return toks, recs, 0

	sites = draw_sites(T, K, span, site_rng, lo_margin=span + d_min)

	# Distinct key/value pairs, then a shuffled site<->needle pairing so the
	# i-th probe does not correspond to the i-th needle by position.
	keys = [[site_rng.randrange(alphabet) for _ in range(key_len)] for _ in range(K)]
	vals = [[site_rng.randrange(alphabet) for _ in range(val_len)] for _ in range(K)]
	order = list(range(K))
	site_rng.shuffle(order)

	occupied = [(p, "site") for p in sites]

	for si, p in enumerate(sites):
	j = order[si] # which key/value this site queries
	key, val = keys[j], vals[j]
	toks[p:p + key_len] = array.array("H", key) # probe
	toks[p + key_len:p + span] = array.array("H", val) # answer

	if arm == "b":
	ns = p - span # immediately before the probe
	# Guaranteed free by the 2*span site separation; assert rather than
	# skip, because a silent skip would desynchronise B's histogram from C's.
	if ns < 0 or any(ns + span > a and ns < a + span
	for a, _ in occupied if a != p):
	raise SystemExit(f"arm b: needle slot {ns} for site {p} is occupied -- "
	"site separation invariant violated")
	else: # arm c: uniform, far
	hi = p - d_min
	ns = place_free(occupied, T, span, arm_rng, hi)
	if ns is None:
	raise SystemExit(f"arm c: could not place a needle before site {p} "
	f"(d_min={d_min}); lower --K or --d-min")
	toks[ns:ns + span] = array.array("H", key + val)
	occupied.append((ns, "needle"))

	answer_start = p + key_len
	recs.append(dict(needle_start=ns, probe_start=p, answer_start=answer_start,
	distance=p - ns,
	csv_first=answer_start - 1)) # see off-by-one note below

	n_decoy = 0
	# Decoy needles: no probe queries them, so needle count != probe count and
	# the i-th-of-K index carries no information.
	for _ in range(decoys):
	key = [site_rng.randrange(alphabet) for _ in range(key_len)]
	val = [site_rng.randrange(alphabet) for _ in range(val_len)]
	ds = place_free(occupied, T, span, arm_rng, T - PAD_TAIL)
	if ds is None:
	continue
	toks[ds:ds + span] = array.array("H", key + val)
	occupied.append((ds, "decoy"))
	n_decoy += 1

	return toks, recs, n_decoy


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--arm", required=True, choices=["a", "b", "c"])
	ap.add_argument("--T", type=int, required=True)
	ap.add_argument("--examples", type=int, required=True)
	ap.add_argument("--split", default="train", choices=["train", "val"])
	ap.add_argument("--K", type=int, default=64, help="triples per example")
	ap.add_argument("--key-len", type=int, default=4, help="probe/key tokens")
	ap.add_argument("--val-len", type=int, default=8, help="answer/value tokens")
	ap.add_argument("--alphabet", type=int, default=512, help="key/value alphabet size")
	ap.add_argument("--decoys", type=int, default=None, help="default K//4")
	ap.add_argument("--d-min", type=int, default=64, help="arm c: min needle->probe distance")
	ap.add_argument("--filler-shards", nargs="+", required=True)
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--out-dir", required=True)
	ap.add_argument("--shards", type=int, default=1,
	help="split the examples across this many files. REQUIRED above "
	"100M tokens: DataLoaderLite is constructed with a "
	"max_tokens_per_shard of 1e8 and anything past it in a single "
	"shard is SILENTLY invisible. The example index stays GLOBAL "
	"across shards, so the arm-independent RNG stream (sites, keys, "
	"pairing, filler) is unchanged by how the split falls -- arms "
	"remain comparable.")
	args = ap.parse_args()

	check_platform()
	if args.alphabet > VOCAB:
	raise SystemExit(f"--alphabet {args.alphabet} exceeds VOCAB {VOCAB}")
	decoys = args.decoys if args.decoys is not None else args.K // 4
	span = args.key_len + args.val_len
	for f in args.filler_shards:
	if not os.path.exists(f):
	raise SystemExit(f"filler shard not found: {f}")
	if args.split == "train" and "val" in os.path.basename(f):
	raise SystemExit(f"refusing to build TRAIN data from a val shard: {f}\n"
	" eval filler must stay disjoint from training data")
	os.makedirs(args.out_dir, exist_ok=True)

	total_tok = args.examples * args.T
	per_shard = (args.examples + args.shards - 1) // args.shards
	if per_shard * args.T > MAX_SHARD_TOKENS:
	raise SystemExit(f"{per_shard * args.T} tokens/shard exceeds the loader's "
	f"{MAX_SHARD_TOKENS} cap; raise --shards to at least "
	f"{(total_tok + MAX_SHARD_TOKENS - 1) // MAX_SHARD_TOKENS}")

	all_recs, hist, n_decoy_total = [], {}, 0
	shard_paths = []
	stream = array.array("H")
	shard_i = 0
	for i in range(args.examples):
	# Site/key/value/pairing rng is ARM-INDEPENDENT -- identical across arms.
	site_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 7)
	# Needle placement is arm-specific.
	arm_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX
	+ 1000 + ord(args.arm))
	# Filler draw is arm-independent too, so arms see the same documents.
	fill_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 99)
	shard = args.filler_shards[i % len(args.filler_shards)]
	filler = load_filler(shard, args.T, fill_rng)

	toks, recs, n_dec = make_example(args.arm, args.T, args.K, args.key_len, args.val_len,
	args.alphabet, decoys, args.d_min,
	filler, site_rng, arm_rng)
	if len(toks) != args.T:
	raise SystemExit(f"example {i} is {len(toks)} tokens, expected {args.T}")
	stream.extend(toks)
	for r in recs:
	hist[r["answer_start"]] = hist.get(r["answer_start"], 0) + 1
	all_recs.append(recs)
	n_decoy_total += n_dec

	# Flush a shard whenever it is full, or on the last example. The global `i`
	# above drives every RNG, so the split point cannot change the data.
	if len(stream) // args.T >= per_shard or i == args.examples - 1:
	sp = os.path.join(args.out_dir, f"in2{args.arm}_{args.split}_{shard_i:05d}.bin")
	with open(sp, "wb") as f:
	stream.tofile(f)
	shard_paths.append(sp)
	stream = array.array("H")
	shard_i += 1

	path = shard_paths[0]

	n_ans = sum(len(r) for r in all_recs) * args.val_len
	# Counted from ACTUAL placements, never derived from parameters: arm A plants
	# nothing, so a parameter-derived figure reported phantom decoy tokens for it.
	n_syn = sum(len(r) for r in all_recs) * (2 * span) + n_decoy_total * span
	man = dict(arm=args.arm, T=args.T, examples=args.examples, split=args.split,
	K=args.K, key_len=args.key_len, val_len=args.val_len, span=span,
	alphabet=args.alphabet, decoys=decoys, d_min=args.d_min,
	seed=args.seed, vocab=VOCAB, pad_tail=PAD_TAIL,
	filler_shards=args.filler_shards, bin=path, shards=shard_paths,
	total_tokens=total_tok, dtype="uint16",
	answer_tokens=n_ans, answer_frac=n_ans / total_tok,
	synthetic_tokens=n_syn, synthetic_frac=n_syn / total_tok,
	decoys_placed=n_decoy_total,
	answer_start_hist=dict(sorted(hist.items())),
	records=all_recs)
	with open(path.replace(".bin", ".json"), "w") as f:
	json.dump(man, f, indent=1)

	print(f"wrote {len(shard_paths)} shard(s) in {args.out_dir} {total_tok} tokens total")
	print(f" answer tokens {n_ans} ({n_ans/total_tok:.2%}) "
	f"synthetic {n_syn} ({n_syn/total_tok:.2%})")
	if all_recs and all_recs[0]:
	ds = [r["distance"] for rs in all_recs for r in rs]
	ds.sort()
	print(f" distance: min {ds[0]} median {ds[len(ds)//2]} max {ds[-1]} n={len(ds)}")


	if __name__ == "__main__":
	main()