Blame for scripts/peer/peer_gen.py

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