Blame for source/files/in2_gen.py

#!/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()