Upload source/files/needle_gen.py with huggingface_hub

source/files/needle_gen.py

@@ -0,0 +1,287 @@
+#!/usr/bin/env python3
+"""
+needle_gen.py -- synthesize retrieval-probe token streams for the existing
+CP_EVAL_PPL / CP_EVAL_TOKENS_BIN path. No C++ changes are needed: the eval
+already windows a raw int32 token file and already reports PER-POSITION
+mean_nll, so planting an answer span at a known offset lets us read retrieval
+NLL straight out of the CSV.
+
+Stdlib only (array + random) -- deliberately no numpy, because the local pyenv
+python3 lacks it while /usr/bin/python3 has it, and this must run identically on
+the local box and the server regardless of which `python3` resolves first.
+
+FIXED-SLOT LAYOUT (the construction that makes the differencing valid)
+---------------------------------------------------------------------
+Every variant of a trial has byte-identical structure except for WHERE the
+needle sits (or whether it is present). The probe prefix and the answer span are
+always at the SAME absolute offsets, so we always difference NLLs measured at
+the same point on a position-varying curve:
+
+  trial (T tokens):
+    [ filler .......... N0..N(k-1) ......... filler ][ probe prefix ][ ANSWER ][pad]
+                        ^needle, k tokens                 k/2           k/2     >=2
+                        slot_start = round(p_frac * max_slot_start)
+
+    p_frac = 0.0 -> needle as far from the probe as possible
+    p_frac = 1.0 -> needle immediately adjacent to the probe (the RETRIEVAL FLOOR)
+
+VARIANTS (7 per model/T/cache)
+  p0, p0.25, p0.5, p0.75, p1   needle planted at that slot; p1 is the floor
+  absent                       no needle planted; probe prefix and answer tokens
+                               still present (the no-information ceiling)
+  null                         needle planted at p_frac=0 but the probe prefix is
+                               an UNRELATED k/2-gram, so retrieval is impossible
+                               -> supplies the significance threshold (RAW NLL)
+
+CRITICAL OFF-BY-ONE
+  The eval scores input[t] = toks[base+t] against target[t] = toks[base+t+1] and
+  accumulates the NLL at the position of the INPUT. So the NLL of the answer token
+  living at trial offset `a` appears in the CSV at row `a-1`. The manifest records
+  BOTH; needle_score.py must read `csv_positions`, never `answer_offsets`.
+
+Usage
+  python3 needle_gen.py --T 32768 --variant p0.5 --block 0 --out-dir needles
+  python3 needle_gen.py --T 32768 --all --blocks 5 --out-dir needles
+"""
+import argparse, array, json, os, random, sys
+
+VOCAB = 50304          # must match cfg.vocab_size; the C++ side bound-checks too
+PAD_TAIL = 4           # >=2: windows stride by exactly T, so the target at T-1
+                       # belongs to the NEXT trial. Keep the answer clear of it.
+P_FRACS = [0.0, 0.25, 0.5, 0.75, 1.0]
+MIX = 1_000_003        # odd multiplier for deterministic seed derivation
+
+
+# Active p_frac grid. Defaults to P_FRACS; --p-fracs overrides it for the S>R
+# ablation, whose predicted dead gap [R,S) falls BETWEEN the default points and
+# would otherwise be unobservable. Overriding changes only which slots are probed,
+# never the layout, so answer-span offsets stay invariant.
+_ACTIVE_P_FRACS = list(P_FRACS)
+
+
+def set_p_fracs(vals):
+    global _ACTIVE_P_FRACS
+    _ACTIVE_P_FRACS = list(vals)
+
+
+def variant_names():
+    return [f"p{p:g}" for p in _ACTIVE_P_FRACS] + ["absent", "null"]
+
+
+def parse_variant(name):
+    """-> (p_frac or None, is_null). absent -> (None, False); null -> (0.0, True)."""
+    if name == "absent":
+        return None, False
+    if name == "null":
+        return 0.0, True
+    if name.startswith("p"):
+        p = float(name[1:])
+        if not 0.0 <= p <= 1.0:
+            raise SystemExit(f"p_frac must be in [0,1] (got {p})")
+        return p, False
+    raise SystemExit(f"unknown variant '{name}' (expected one of {variant_names()})")
+
+
+def check_platform():
+    if sys.byteorder != "little":
+        raise SystemExit("CP_EVAL_TOKENS_BIN expects int32 LITTLE-endian; this host is big-endian")
+    if array.array("i").itemsize != 4 or array.array("H").itemsize != 2:
+        raise SystemExit("unexpected array itemsizes ('i' must be 4B, 'H' must be 2B)")
+
+
+def load_filler(path, need, rng):
+    """Real tokens, so the LM is not off-distribution (fabricated filler would
+    inflate NLL everywhere and mask the effect). Source is uint16; the eval file
+    must be int32 LE, so widen explicitly and bound-check."""
+    ntok = os.path.getsize(path) // 2
+    if ntok < need + 1:
+        raise SystemExit(f"{path} has {ntok} tokens, need >= {need + 1}")
+    # One contiguous run per trial, so filler keeps local textual statistics
+    # instead of being a shuffled bag of tokens.
+    start = rng.randrange(0, ntok - need)
+    buf = array.array("H")
+    with open(path, "rb") as f:
+        f.seek(start * 2)
+        buf.fromfile(f, need)
+    out = array.array("i", buf)                       # C-level uint16 -> int32
+    lo, hi = min(out), max(out)
+    if lo < 0 or hi >= VOCAB:
+        raise SystemExit(f"{path} has ids outside [0,{VOCAB}): min={lo} max={hi}")
+    return out
+
+
+def build_layout(T, k, probe_pos=None):
+    """Offsets are a pure function of (T, k, probe_pos) -- identical for every variant.
+
+    probe_pos is None (default): probe+answer sit at the END of the trial, so the
+    needle-to-probe DISTANCE is T - slot. Distance and the needle's ABSOLUTE
+    position are then perfectly anti-correlated -- one variable, not two -- which
+    makes "far from the question" and "early in the document" indistinguishable.
+
+    probe_pos = P: probe+answer are planted at absolute offset P and everything
+    after them is ordinary filler. Pinning the needle's slot while varying P (or
+    vice versa) breaks that confound. Legal because the eval reports NLL at EVERY
+    position, so the scored span does not have to be at the end -- and the
+    manifest's csv_positions carries the offsets, so needle_score.py follows
+    automatically with no change.
+    """
+    if k % 2:
+        raise SystemExit(f"--needle-len must be even (got {k})")
+    if probe_pos is None:
+        answer_end = T - PAD_TAIL
+        answer_start = answer_end - k // 2
+        probe_start = answer_start - k // 2           # probe prefix: [probe_start, answer_start)
+    else:
+        probe_start = probe_pos
+        answer_start = probe_start + k // 2
+        answer_end = probe_start + k
+        if answer_end > T - PAD_TAIL:
+            raise SystemExit(f"--probe-pos {probe_pos} puts the answer span end at "
+                             f"{answer_end}, past T-PAD_TAIL={T - PAD_TAIL}; the last "
+                             "position's target belongs to the next trial")
+    max_slot_start = probe_start - k                   # needle must fit before the probe
+    if max_slot_start < 0:
+        raise SystemExit(f"probe at {probe_start} leaves no room for a needle of "
+                         f"{k} tokens before it (T={T})")
+    return dict(
+        T=T, needle_len=k, pad_tail=PAD_TAIL, probe_pos=probe_pos,
+        probe_start=probe_start, answer_start=answer_start, answer_end=answer_end,
+        max_slot_start=max_slot_start,
+        answer_offsets=[answer_start, answer_end],
+        # NLL is reported at the INPUT position, so the answer token at offset a is
+        # scored on CSV row a-1. See "CRITICAL OFF-BY-ONE" above.
+        csv_positions=[answer_start - 1, answer_end - 1],
+    )
+
+
+def make_trial(lay, p_frac, is_null, filler, rng, reverse=False):
+    """One trial of exactly T int32 tokens. `filler` is consumed (modified)."""
+    k = lay["needle_len"]
+    half = k // 2
+    trial = filler                                    # already a fresh copy per trial
+    needle = array.array("i", [rng.randrange(VOCAB) for _ in range(k)])
+
+    # The answer span is the needle's second half, at fixed offsets -- unless
+    # `reverse`, which swaps the halves: probe = needle's SECOND half, answer = its
+    # FIRST, so the model must retrieve what PRECEDES the probe. Forward retrieval
+    # (answer following the keyword) is an induction-head copy and "an easily
+    # learned pattern" (arXiv:2404.16811 s3.1), so it is exactly what training on
+    # forward data would Goodhart; backward retrieval cannot be satisfied that way,
+    # which makes it the transfer test. Layout OFFSETS are untouched, so
+    # csv_positions stays valid and needle_score.py needs no change.
+    trial[lay["answer_start"]:lay["answer_end"]] = needle[:half] if reverse else needle[half:]
+
+    # Probe prefix = the needle's other half, EXCEPT for `null`, where an unrelated
+    # k/2-gram makes retrieval impossible while the answer TARGETS stay identical.
+    if is_null:
+        prefix = array.array("i", [rng.randrange(VOCAB) for _ in range(half)])
+    else:
+        prefix = needle[half:] if reverse else needle[:half]
+    trial[lay["probe_start"]:lay["answer_start"]] = prefix
+
+    # Plant the needle in exactly one slot (or none, for `absent`).
+    slot_start = None
+    if p_frac is not None:
+        slot_start = int(round(p_frac * lay["max_slot_start"]))
+        trial[slot_start:slot_start + k] = needle
+    return trial, slot_start
+
+
+def gen(args, variant, block):
+    p_frac, is_null = parse_variant(variant)
+    lay = build_layout(args.T, args.needle_len, args.probe_pos)
+    # Seed depends on (seed, T, block, trial) and NOT on the variant, so every
+    # variant of a block shares identical filler and identical needle tokens --
+    # the whole point of the fixed-slot construction.
+    base_seed = (args.seed * MIX + args.T) * MIX + block
+
+    stream = array.array("i")
+    slots = []
+    for i in range(args.trials):
+        rng = random.Random(base_seed * MIX + i)
+        filler = load_filler(args.filler, args.T, rng)
+        trial, slot = make_trial(lay, p_frac, is_null, filler, rng, args.reverse)
+        if len(trial) != args.T:
+            raise SystemExit(f"trial length {len(trial)} != T={args.T} -- bug")
+        stream.extend(trial)
+        slots.append(slot)
+
+    # +1 token: window w reads toks[base+t] and toks[base+t+1] for t in [0,T), so
+    # the final window needs index N*T to exist.
+    stream.append(stream[-1])
+    if len(stream) != args.trials * args.T + 1:
+        raise SystemExit("stream length wrong -- bug")
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    tag = f"T{args.T}_{variant}_b{block}"
+    bin_path = os.path.join(args.out_dir, f"needle_{tag}.bin")
+    with open(bin_path, "wb") as f:
+        stream.tofile(f)
+
+    man = dict(bin=bin_path, variant=variant, block=block, p_frac=p_frac, is_null=is_null,
+               trials=args.trials, seed=args.seed, vocab=VOCAB, reverse=args.reverse,
+           slot_starts=slots, **lay)
+    man_path = os.path.join(args.out_dir, f"needle_{tag}.json")
+    with open(man_path, "w") as f:
+        json.dump(man, f, indent=2)
+    return bin_path, man_path, man
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--T", type=int, required=True, help="tokens per trial (= CP_T)")
+    ap.add_argument("--trials", type=int, default=None,
+                    help="trials per block (= CP_EVAL_WINDOWS). Default is the plan's "
+                         "pinned N: 8 at T<=131072, 4 above.")
+    ap.add_argument("--blocks", type=int, default=5,
+                    help="blocks = SEPARATE files, each needing its own eval invocation; "
+                         "the only source of spread (CP_EVAL_PPL emits no variance)")
+    ap.add_argument("--block", type=int, default=None, help="generate just this block index")
+    ap.add_argument("--variant", default=None, help=f"one of {variant_names()}")
+    ap.add_argument("--all", action="store_true", help="every variant x every block")
+    ap.add_argument("--needle-len", type=int, default=8, help="k; probe=k/2, answer=k/2")
+    ap.add_argument("--filler", default="Data_Loader/Data/edufineweb_val_00000.bin")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--probe-pos", type=int, default=None,
+                    help="absolute offset for the probe+answer span. Default: the end "
+                         "of the trial. Setting it lets DISTANCE vary independently of "
+                         "the needle's absolute position -- the confound the default "
+                         "layout cannot separate. Isolate runs by --out-dir: filenames "
+                         "are unchanged, so needle_score.py needs no flag.")
+    ap.add_argument("--out-dir", default="needles")
+    ap.add_argument("--reverse", action="store_true",
+                    help="BACKWARD retrieval (transfer test): probe = needle's 2nd half, "
+                         "answer = its 1st. Use a SEPARATE --out-dir -- filenames do not "
+                         "encode it, and absent/null must be regenerated in reverse too.")
+    ap.add_argument("--p-fracs", default=None,
+                    help="comma-separated p_frac list overriding the default "
+                         f"{P_FRACS}. Needed by the STRING S>R ablation, whose predicted "
+                         "dead gap lies between the default points.")
+    args = ap.parse_args()
+
+    if args.p_fracs:
+        vals = [float(x) for x in args.p_fracs.split(",") if x.strip()]
+        if not vals or any(not 0.0 <= v <= 1.0 for v in vals):
+            raise SystemExit(f"--p-fracs must be values in [0,1] (got {args.p_fracs})")
+        set_p_fracs(vals)
+
+    check_platform()
+    if args.trials is None:                            # pinned N from the plan (A1)
+        args.trials = 8 if args.T <= 131072 else 4
+
+    if args.all:
+        jobs = [(v, b) for v in variant_names() for b in range(args.blocks)]
+    else:
+        if args.variant is None:
+            raise SystemExit("give --variant NAME (and optionally --block), or --all")
+        jobs = [(args.variant, args.block if args.block is not None else 0)]
+
+    for v, b in jobs:
+        bp, _, man = gen(args, v, b)
+        slot = man["slot_starts"][0]
+        print(f"{v:8s} block={b}  {bp}  answer_offsets={man['answer_offsets']} "
+              f"csv_positions={man['csv_positions']}  slot={slot}")
+
+
+if __name__ == "__main__":
+    main()