Upload scripts/longeval_gen.py with huggingface_hub

scripts/longeval_gen.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""
+longeval_gen.py -- build GENERATIVE long-context benchmark prompts for bluscriptCP.
+
+Unlike needle_gen.py (which plants a random-token key/value pair and is scored by
+teacher-forced NLL over an answer span that is ALREADY IN THE INPUT), everything
+here is scored by what the model actually GENERATES. That is what RULER, NIAH and
+VAL Probing measure, and it is the only form comparable to published numbers.
+
+WHY THE PROMPTS LOOK LIKE THIS
+  The target is a ~1B BASE language model with no instruction tuning, so a bare
+  "Question: ...\\nAnswer:" often fails for reasons that have nothing to do with
+  retrieval. Every task therefore ends with an answer PREFIX that makes the
+  correct answer the natural continuation, e.g.
+
+      ... haystack ...
+      Question: What is the special magic number for aardvark?
+      Answer: The special magic number for aardvark is
+
+  The model only has to emit the digits. Scoring looks for the expected string
+  anywhere in the generated tokens, so trailing chatter does not cost a point.
+
+TASKS
+  niah_single      RULER NIAH-single      one needle, one key, retrieve its value
+  niah_multikey    RULER NIAH-multikey    K needles, retrieve ONE named key
+  niah_multivalue  RULER NIAH-multivalue  one key with V values, retrieve all
+  vt               RULER VT               chain of variable assignments, resolve
+  val_forward      VAL forward retrieval  key stated before the query (default)
+  val_backward     VAL backward retrieval value stated first, retrieve its KEY
+
+OUTPUT (per task, into --out-dir)
+  <task>_prompts.bin   int32 token stream, N prompts x --prompt-len tokens, end to end
+                       (exactly what CP_GEN_TOKENS_BIN expects)
+  <task>_meta.json     per-prompt expected answers + placement metadata
+
+The prompt length is EXACT: the haystack is trimmed or extended so every prompt is
+--prompt-len tokens, because the C++ reader slices the stream at a fixed stride.
+"""
+
+import argparse
+import json
+import os
+import random
+import struct
+import sys
+
+try:
+    import tiktoken
+except ImportError:
+    sys.exit("longeval_gen.py needs tiktoken:  pip install tiktoken")
+
+ENC = tiktoken.get_encoding("gpt2")
+
+
+# ---------------------------------------------------------------------------
+# Haystack: real text from the flux shards (uint16 GPT-2 token stream), so the
+# distractor context has the same distribution the model was trained on. Falls
+# back to a synthetic filler sentence when no shard is given, which keeps the
+# generator usable on a box with no data.
+# ---------------------------------------------------------------------------
+FALLBACK = ENC.encode(
+    "The grass swayed in the long afternoon light while the river ran on "
+    "unhurried past the old stone bridge and the fields beyond. "
+)
+
+
+def load_haystack(path, need):
+    """Return >= `need` GPT-2 token ids of filler.
+
+    SLICE BEFORE CONVERTING. The flux val shard is 200 MB = 100M uint16 tokens;
+    `toks.astype(np.int64).tolist()` over the whole array costs ~800 MB of numpy
+    plus a ~2.8 GB Python list. Callers now ask for prompt_len*(n+2) so each
+    prompt can start at its own offset, which makes the slice essential.
+    """
+    if not path:
+        reps = need // len(FALLBACK) + 2
+        return (FALLBACK * reps)[:need + len(FALLBACK)]
+    import numpy as np
+    toks = np.fromfile(path, dtype=np.uint16)
+    if toks.size < need:
+        reps = int(need // toks.size) + 2
+        toks = np.tile(toks, reps)
+    return toks[:need].astype(np.int64).tolist()
+
+
+# vt's chain needs room AFTER its anchor for the remaining links, so the anchor
+# cannot sit at depth 1.0. Both t_vt and the metadata writer call this, so the
+# recorded depth can never drift from the depth actually built.
+VT_ANCHOR_MAX = 0.8
+
+
+def effective_depth(task, depth):
+    return min(depth, VT_ANCHOR_MAX) if task == "vt" else depth
+
+
+# ---------------------------------------------------------------------------
+# Needle vocabulary. Words are common English so they tokenize compactly and the
+# model has seen them; numbers are 7 digits so a lucky guess is ~1e-7.
+# ---------------------------------------------------------------------------
+WORDS = [
+    "aardvark", "buffalo", "cactus", "dolphin", "emerald", "falcon", "glacier",
+    "harbor", "iceberg", "jaguar", "kestrel", "lantern", "meadow", "nutmeg",
+    "orchid", "penguin", "quartz", "raven", "sapphire", "tundra", "umbrella",
+    "violet", "walnut", "xylophone", "yarrow", "zephyr", "anchor", "bramble",
+    "cinder", "driftwood", "ember", "fjord", "granite", "hollow", "ivory",
+]
+
+
+def rnd_number(rng):
+    return str(rng.randint(1000000, 9999999))
+
+
+# ---------------------------------------------------------------------------
+# Task builders. Each returns (needle_sentences, question, answer_prefix, expected)
+#   needle_sentences : list of (sentence_text, depth_fraction) to plant
+#   expected         : list of strings; ALL must appear in the generation to score 1
+# ---------------------------------------------------------------------------
+def t_niah_single(rng, depth):
+    k, v = rng.choice(WORDS), rnd_number(rng)
+    s = f" The special magic number for {k} is {v}. "
+    q = f"\nQuestion: What is the special magic number for {k}?\n"
+    a = f"Answer: The special magic number for {k} is"
+    return [(s, depth)], q, a, [v]
+
+
+def t_niah_multikey(rng, depth, n_keys=4):
+    keys = rng.sample(WORDS, n_keys)
+    vals = [rnd_number(rng) for _ in keys]
+    # Distractors spread across the whole context; the TARGET sits at `depth`.
+    sents = []
+    tgt = rng.randrange(n_keys)
+    for i, (k, v) in enumerate(zip(keys, vals)):
+        d = depth if i == tgt else rng.random()
+        sents.append((f" The special magic number for {k} is {v}. ", d))
+    k, v = keys[tgt], vals[tgt]
+    q = f"\nQuestion: What is the special magic number for {k}?\n"
+    a = f"Answer: The special magic number for {k} is"
+    return sents, q, a, [v]
+
+
+def t_niah_multivalue(rng, depth, n_vals=3):
+    k = rng.choice(WORDS)
+    vals = [rnd_number(rng) for _ in range(n_vals)]
+    # All values share one key, scattered; depth anchors the FIRST one.
+    sents = [(f" One of the special magic numbers for {k} is {v}. ",
+              depth if i == 0 else rng.random())
+             for i, v in enumerate(vals)]
+    q = f"\nQuestion: What are all the special magic numbers for {k}?\n"
+    a = f"Answer: The special magic numbers for {k} are"
+    return sents, q, a, vals          # every value must appear
+
+
+def t_vt(rng, depth, hops=4):
+    """RULER variable tracking: X1 = <num>, X2 = X1, X3 = X2 ... resolve the chain.
+
+    Depths must INCREASE along the chain. build_one sorts planted sentences by
+    depth, so random depths would scatter the chain in random document order --
+    that is backward-reference resolution, strictly harder than RULER's ordered
+    VT and not comparable to published numbers. Monotone depths put document
+    order = chain order (Python's sort is stable, so ties keep chain order).
+    """
+    names = rng.sample([w.upper()[:3] + str(rng.randint(10, 99)) for w in WORDS], hops)
+    val = rnd_number(rng)
+    anchor = effective_depth("vt", depth)
+    sents = [(f" VAR {names[0]} = {val}. ", anchor)]
+    for i in range(1, hops):
+        sents.append((f" VAR {names[i]} = {names[i-1]}. ",
+                      anchor + (1.0 - anchor) * i / hops))
+    q = f"\nQuestion: What is the value of VAR {names[-1]}?\n"
+    a = f"Answer: VAR {names[-1]} ="
+    return sents, q, a, [val]
+
+
+def t_val_forward(rng, depth):
+    return t_niah_single(rng, depth)
+
+
+def t_val_backward(rng, depth):
+    """VAL backward retrieval: the VALUE is the query, the KEY is the answer."""
+    k, v = rng.choice(WORDS), rnd_number(rng)
+    s = f" The special magic number for {k} is {v}. "
+    q = f"\nQuestion: Which item has the special magic number {v}?\n"
+    a = f"Answer: The item with the special magic number {v} is"
+    return [(s, depth)], q, a, [k]
+
+
+TASKS = {
+    "niah_single":     t_niah_single,
+    "niah_multikey":   t_niah_multikey,
+    "niah_multivalue": t_niah_multivalue,
+    "vt":              t_vt,
+    "val_forward":     t_val_forward,
+    "val_backward":    t_val_backward,
+}
+
+
+def build_one(task, rng, prompt_len, hay, depth, hay_off=0):
+    """One prompt of EXACTLY prompt_len tokens.
+
+    hay_off shifts this prompt's filler window. Without it every prompt starts
+    the haystack at 0 and differs only in the needle, making the N prompts
+    replicates rather than independent draws -- which inflates apparent
+    precision because the effective sample size is far below N.
+    """
+    sents, q, a, expected = TASKS[task](rng, depth)
+    q_ids, a_ids = ENC.encode(q), ENC.encode(a)
+    tail = len(q_ids) + len(a_ids)
+    body_budget = prompt_len - tail
+    if body_budget <= 0:
+        raise SystemExit(f"--prompt-len {prompt_len} too small for task {task}")
+
+    # Plant each needle at its depth fraction of the BODY, then fill with haystack.
+    planted = sorted(((max(0.0, min(1.0, d)), ENC.encode(s)) for s, d in sents),
+                     key=lambda x: x[0])
+    total_needle = sum(len(ids) for _, ids in planted)
+    fill = body_budget - total_needle
+    if fill <= 0:
+        raise SystemExit(f"needles do not fit in --prompt-len {prompt_len}")
+
+    H = lambda lo, n: hay[hay_off + lo: hay_off + lo + max(0, n)]
+
+    body, used, hay_pos = [], 0, 0
+    for frac, ids in planted:
+        want = int(frac * fill)
+        take = max(0, want - used)
+        body.extend(H(hay_pos, take)); hay_pos += take; used += take
+        body.extend(ids)
+    body.extend(H(hay_pos, fill - used))
+
+    out = body + q_ids + a_ids
+    # Exact length: the C++ reader slices at a fixed stride, so drift is fatal.
+    if len(out) < prompt_len:
+        out = body + H(hay_pos, prompt_len - len(out)) + q_ids + a_ids
+    out = out[:prompt_len - tail] + q_ids + a_ids
+    assert len(out) == prompt_len, (len(out), prompt_len)
+    return out, expected
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--task", required=True, choices=sorted(TASKS) + ["all"])
+    ap.add_argument("--prompt-len", type=int, required=True,
+                    help="exact tokens per prompt; must leave room for CP_GEN_MAX_NEW under CP_T")
+    ap.add_argument("--n", type=int, default=50, help="prompts per task")
+    ap.add_argument("--depths", default="0.1,0.3,0.5,0.7,0.9",
+                    help="needle depth fractions, cycled across prompts")
+    ap.add_argument("--haystack", default="", help="flux .bin shard for filler text")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--out-dir", required=True)
+    args = ap.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    depths = [float(d) for d in args.depths.split(",")]
+    tasks = sorted(TASKS) if args.task == "all" else [args.task]
+
+    for task in tasks:
+        rng = random.Random(args.seed)          # per-task, so tasks are independent
+        # (n + 2) windows so every prompt gets its own filler offset
+        hay = load_haystack(args.haystack, args.prompt_len * (args.n + 2))
+        bin_p = os.path.join(args.out_dir, f"{task}_prompts.bin")
+        meta = []
+        with open(bin_p, "wb") as f:
+            for i in range(args.n):
+                want = depths[i % len(depths)]
+                eff = effective_depth(task, want)
+                ids, expected = build_one(task, rng, args.prompt_len, hay, want,
+                                          hay_off=i * args.prompt_len)
+                f.write(struct.pack(f"<{len(ids)}i", *ids))
+                # Record the EFFECTIVE depth: longeval_score.py groups on this
+                # field, and depth-pooling is the primary unit of analysis, so a
+                # clamped vt row must not be counted in the depth-1.0 cell.
+                meta.append({"i": i, "depth": eff, "depth_requested": want,
+                             "expected": expected})
+        with open(os.path.join(args.out_dir, f"{task}_meta.json"), "w") as f:
+            json.dump({"task": task, "prompt_len": args.prompt_len, "n": args.n,
+                       "seed": args.seed, "items": meta}, f, indent=1)
+        print(f"  {task:<16} {args.n} x {args.prompt_len} tok -> {bin_p}")
+
+
+if __name__ == "__main__":
+    main()