Upload scripts/peer/peer_gen.py with huggingface_hub

scripts/peer/peer_gen.py

@@ -0,0 +1,213 @@
+#!/usr/bin/env python3
+"""
+peer_gen.py -- build NIAH prompts for a PEER model, in that model's own tokenizer.
+
+Why this exists rather than reusing longeval_gen.py: that script emits GPT-2 token
+ids, and the peer models do not share GPT-2's vocabulary (Qwen3 151,936,
+Spark-X2.5 131,072). Feeding GPT-2 ids to them is meaningless, and decoding to
+text then letting each model re-tokenize would make "16,384 tokens" mean a
+different amount of text per model, which is exactly the comparison we are trying
+to make fair.
+
+THE INVARIANT: a prompt is EXACTLY --prompt-len tokens OF THE TARGET TOKENIZER,
+and the haystack TEXT is byte-identical across models. So every model sees the
+same story, cut to the same number of its own tokens.
+
+Prompts are emitted as token IDS, not text. A decode->re-encode round trip is not
+length-preserving for BPE, so building in token space and handing ids straight to
+the model is the only way to guarantee the exact length.
+
+Tasks and needle vocabulary are copied from longeval_gen.py deliberately, so a
+peer number is directly comparable to ours. Same WORDS list, same 7-digit values,
+same answer-prefix prompting, same depth semantics (0.0 = window start, farthest
+from the question; 1.0 = adjacent to it).
+
+  usage:
+    peer_gen.py --tokenizer Qwen/Qwen3-4B-Base --haystack hay.txt \
+                --task niah_single --prompt-len 16384 --n 100 \
+                --depths 0.0,0.25,0.5,0.75,1.0 --out-dir out/
+"""
+
+import argparse
+import json
+import os
+import random
+import sys
+
+# --- copied verbatim from longeval_gen.py so peer numbers are comparable -------
+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",
+]
+VT_ANCHOR_MAX = 0.8
+
+
+def rnd_number(rng):
+    return str(rng.randint(1000000, 9999999))
+
+
+def effective_depth(task, depth):
+    return min(depth, VT_ANCHOR_MAX) if task == "vt" else depth
+
+
+def t_niah_single(rng, depth):
+    k, v = rng.choice(WORDS), rnd_number(rng)
+    return ([(f" The special magic number for {k} is {v}. ", depth)],
+            f"\nQuestion: What is the special magic number for {k}?\n",
+            f"Answer: The special magic number for {k} is", [v])
+
+
+def t_niah_multikey(rng, depth, n_keys=4):
+    keys = rng.sample(WORDS, n_keys)
+    vals = [rnd_number(rng) for _ in keys]
+    tgt = rng.randrange(n_keys)
+    sents = [(f" The special magic number for {k} is {v}. ",
+              depth if i == tgt else rng.random())
+             for i, (k, v) in enumerate(zip(keys, vals))]
+    k, v = keys[tgt], vals[tgt]
+    return (sents, f"\nQuestion: What is the special magic number for {k}?\n",
+            f"Answer: The special magic number for {k} is", [v])
+
+
+def t_niah_multivalue(rng, depth, n_vals=3):
+    k = rng.choice(WORDS)
+    vals = [rnd_number(rng) for _ in range(n_vals)]
+    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)]
+    return (sents, f"\nQuestion: What are all the special magic numbers for {k}?\n",
+            f"Answer: The special magic numbers for {k} are", vals)
+
+
+def t_vt(rng, depth, hops=4):
+    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))
+    return (sents, f"\nQuestion: What is the value of VAR {names[-1]}?\n",
+            f"Answer: VAR {names[-1]} =", [val])
+
+
+def t_val_backward(rng, depth):
+    k, v = rng.choice(WORDS), rnd_number(rng)
+    return ([(f" The special magic number for {k} is {v}. ", depth)],
+            f"\nQuestion: Which item has the special magic number {v}?\n",
+            f"Answer: The item with the special magic number {v} is", [k])
+
+
+TASKS = {"niah_single": t_niah_single, "niah_multikey": t_niah_multikey,
+         "niah_multivalue": t_niah_multivalue, "vt": t_vt,
+         "val_forward": t_niah_single, "val_backward": t_val_backward}
+
+
+def build_one(task, rng, enc, prompt_len, hay_ids, depth, hay_off):
+    """One prompt of EXACTLY prompt_len tokens in `enc`'s vocabulary."""
+    sents, q, a, expected = TASKS[task](rng, depth)
+    q_ids, a_ids = enc(q), enc(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}")
+
+    planted = sorted(((max(0.0, min(1.0, d)), enc(s)) for s, d in sents),
+                     key=lambda x: x[0])
+    fill = body_budget - sum(len(i) for _, i in planted)
+    if fill <= 0:
+        raise SystemExit(f"needles do not fit in --prompt-len {prompt_len}")
+
+    H = lambda lo, n: hay_ids[hay_off + lo: hay_off + lo + max(0, n)]
+    body, used, pos = [], 0, 0
+    for frac, ids in planted:
+        take = max(0, int(frac * fill) - used)
+        body.extend(H(pos, take)); pos += take; used += take
+        body.extend(ids)
+    body.extend(H(pos, fill - used))
+
+    out = body + q_ids + a_ids
+    if len(out) < prompt_len:                      # haystack ran short
+        out = body + H(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("--tokenizer", required=True,
+                    help="HF repo id, or the literal 'gpt2' for tiktoken parity")
+    ap.add_argument("--haystack", required=True, help="UTF-8 text file of filler")
+    ap.add_argument("--task", required=True, choices=sorted(TASKS) + ["all"])
+    ap.add_argument("--prompt-len", type=int, required=True)
+    ap.add_argument("--n", type=int, default=100)
+    ap.add_argument("--depths", default="0.0,0.25,0.5,0.75,1.0")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--out-dir", required=True)
+    args = ap.parse_args()
+
+    if args.tokenizer == "gpt2":
+        import tiktoken
+        _e = tiktoken.get_encoding("gpt2")
+        # disallowed_special=() matches the corpus's own encode_ordinary: any
+        # "<|endoftext|>"-looking text is ordinary text, never a special id.
+        # Without this, tiktoken raises on the first such string in the haystack.
+        enc = lambda s: _e.encode(s, disallowed_special=())
+        name = "tiktoken-gpt2"
+    else:
+        from transformers import AutoTokenizer
+        _t = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
+        enc = lambda s: _t.encode(s, add_special_tokens=False)
+        name = args.tokenizer
+
+    text = open(args.haystack, encoding="utf-8", errors="ignore").read()
+    hay_ids = enc(text)
+    print(f"  tokenizer {name}: haystack {len(text):,} chars -> {len(hay_ids):,} tokens")
+
+    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]
+
+    need = args.prompt_len * (args.n + 2)
+    if len(hay_ids) < need:
+        # tile rather than fail: a short haystack repeats, which is weaker
+        # distraction and INFLATES accuracy, so it is reported loudly.
+        reps = need // len(hay_ids) + 2
+        print(f"  WARN: haystack has {len(hay_ids):,} tokens, need {need:,} -- "
+              f"tiling {reps}x. Repetitive filler inflates accuracy; record this.")
+        hay_ids = hay_ids * reps
+
+    for task in tasks:
+        rng = random.Random(args.seed)     # per task, so tasks are independent
+        meta, rows = [], []
+        for i in range(args.n):
+            want = depths[i % len(depths)]
+            eff = effective_depth(task, want)
+            ids, expected = build_one(task, rng, enc, args.prompt_len, hay_ids,
+                                      want, hay_off=i * args.prompt_len)
+            rows.append({"i": i, "ids": ids})
+            meta.append({"i": i, "depth": eff, "depth_requested": want,
+                         "expected": expected})
+        # A special id inside the body would act as a document reset and silently
+        # shorten the effective context -- the exact thing NIAH measures.
+        if args.tokenizer != "gpt2":
+            bad = sorted(set(_t.all_special_ids or []))
+            if bad:
+                hit = {t for r in rows for t in r["ids"] if t in set(bad)}
+                if hit:
+                    print(f"  WARN: {len(hit)} special id(s) {sorted(hit)} appear "
+                          f"inside {task} prompts -- these truncate attention")
+        with open(os.path.join(args.out_dir, f"{task}_prompts.jsonl"), "w") as f:
+            for r in rows:
+                f.write(json.dumps(r) + "\n")
+        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, "tokenizer": name, "items": meta}, f)
+        print(f"  {task:<16} {args.n} x {args.prompt_len} tok -> {args.out_dir}")
+
+
+if __name__ == "__main__":
+    main()