niah1.py

#!/usr/bin/env python3
"""niah1.py — simple Needle-in-a-Haystack benchmark.

Plants a random N-digit passkey inside a haystack of English filler sentences
and asks the model to recall it. Sweeps over context lengths and depths and
prints a pass-rate grid at the end.

Typical use:

    # single-batch run
    python niah1.py --batch-size 1

    # multi-batch run (4 independent samples per forward)
    python niah1.py --batch-size 4

    # different digit count / seed / ctx range
    python niah1.py --digits 8 --seed 123 --start-ctx 2048 --end-ctx 8192
"""
from __future__ import annotations

import argparse
import random
import time
from typing import List, Tuple

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


# ---------------------------------------------------------------------------- #
#  English filler pool — intentionally neutral and information-free so the
#  needle is the only piece of "signal" in the prompt.
# ---------------------------------------------------------------------------- #
FILLER_POOL: List[str] = [
    "The grass is green and the sky is blue.",
    "The quick brown fox jumps over the lazy dog.",
    "Autumn leaves drift slowly from the old oak tree.",
    "Coffee tastes better when the rain taps the window.",
    "She walked along the river in complete silence.",
    "Empty shelves stared back from the dusty attic.",
    "Morning fog rolled over the sleeping harbor.",
    "Children laughed while kites danced on the wind.",
    "The librarian adjusted her glasses and sighed softly.",
    "Old maps told stories of forgotten continents.",
    "Warm bread filled the kitchen with a comforting smell.",
    "Streetlights flickered as the storm approached.",
    "Clouds gathered above the distant mountains.",
    "The train whistled as it disappeared into the tunnel.",
    "A small boat drifted across the quiet pond.",
    "Silver stars scattered across the midnight sky.",
    "The old clock ticked steadily in the empty hall.",
    "Waves crashed against the rocky shoreline.",
    "A cat sat patiently beside the window.",
    "Bookshelves reached all the way to the ceiling.",
    "Raindrops pattered softly against the tin roof.",
    "The kettle whistled from the warm kitchen.",
    "Footsteps echoed down the long marble corridor.",
    "Wildflowers grew between cracks in the sidewalk.",
    "The moon cast a silver path across the lake.",
    "Snow piled gently on the garden fence.",
    "Birds sang cheerfully in the morning sunshine.",
    "She turned the page and began a new chapter.",
    "The bakery opened early every single morning.",
    "A small dog chased leaves across the yard.",
    "The candle burned low as midnight approached.",
    "Thunder rolled across the open prairie.",
    "He closed the notebook and leaned back in his chair.",
    "Fireflies blinked above the summer meadow.",
    "The teacher wrote the date on the blackboard.",
    "Ships sailed past the lighthouse at dusk.",
    "Leaves rustled gently in the evening breeze.",
    "A warm cup of tea calmed her racing thoughts.",
]


def build_haystack_tokens(tokenizer, target_tokens: int, rng: random.Random) -> List[int]:
    """Return a list of token ids (no specials) of length >= target_tokens, then truncate."""
    tokens: List[int] = []
    while len(tokens) < target_tokens:
        sentence = rng.choice(FILLER_POOL) + " "
        tokens.extend(tokenizer.encode(sentence, add_special_tokens=False))
    return tokens[:target_tokens]


def needle_sentence(passkey: int) -> str:
    """The "needle" — one sentence containing the passkey to be retrieved."""
    return (
        f"\n\nIMPORTANT: The magic passkey is {passkey}. "
        f"Please remember this exact number.\n\n"
    )


QUESTION = (
    "\n\nQuestion: What is the magic passkey mentioned above? "
    "Answer with the number only."
)


def build_prompt(
    tokenizer,
    ctx_tokens: int,
    depth: float,
    passkey: int,
    overhead: int,
    rng: random.Random,
) -> str:
    """Construct one NIAH prompt whose total token length is ~ctx_tokens.

    Pieces (in order): [haystack prefix] + [needle] + [haystack suffix] + [question]
    Wrapped in the model's chat template.
    """
    needle = needle_sentence(passkey)
    needle_ids = tokenizer.encode(needle, add_special_tokens=False)
    question_ids = tokenizer.encode(QUESTION, add_special_tokens=False)

    # Budget haystack tokens so the final prompt lands near ctx_tokens.
    noise_budget = max(32, ctx_tokens - overhead - len(needle_ids) - len(question_ids))
    haystack = build_haystack_tokens(tokenizer, noise_budget, rng)

    # Insert needle at the requested relative depth.
    depth = max(0.0, min(1.0, depth))
    pos = int(round(len(haystack) * depth))
    body_ids = haystack[:pos] + needle_ids + haystack[pos:]
    body_text = tokenizer.decode(body_ids) + QUESTION

    return tokenizer.apply_chat_template(
        [{"role": "user", "content": body_text}],
        tokenize=False,
        add_generation_prompt=True,
    )


def run_batch(
    model,
    tokenizer,
    prompts: List[str],
    passkeys: List[int],
    max_new_tokens: int,
) -> Tuple[List[Tuple[bool, str]], float, int]:
    """Run one batched forward+generate. Returns [(pass, pred_text)], time, prompt_len."""
    enc = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
    prompt_len = enc.input_ids.shape[1]
    t0 = time.time()
    with torch.no_grad():
        out = model.generate(
            **enc,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.pad_token_id,
        )
    dt = time.time() - t0

    results: List[Tuple[bool, str]] = []
    for b, pk in enumerate(passkeys):
        gen_ids = out[b, prompt_len:]
        gen_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
        ok = str(pk) in gen_text
        results.append((ok, gen_text))
    return results, dt, prompt_len


def main() -> None:
    p = argparse.ArgumentParser(description="Simple NIAH benchmark")
    p.add_argument("--model-path", default="/home/llm/gemma-4-26B-A4B-Text")
    p.add_argument("--start-ctx", type=int, default=64)
    p.add_argument("--end-ctx", type=int, default=16384)
    p.add_argument("--step", type=int, default=1024)
    p.add_argument(
        "--depths", nargs="+", type=float, default=[0.0, 0.5, 1.0],
        help="Relative depths to insert the needle at (0=start, 1=end).",
    )
    p.add_argument("--trials", type=int, default=1,
                   help="Batches per (ctx, depth). Total samples = trials * batch_size.")
    p.add_argument("--batch-size", type=int, default=1,
                   help="Samples per forward pass. 1 = single-batch, >1 = multi-batch.")
    p.add_argument("--digits", type=int, default=6, help="Passkey digit count.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--max-new-tokens", type=int, default=64)
    args = p.parse_args()

    print(f"[niah1] loading tokenizer & model from {args.model_path}")
    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    tokenizer.padding_side = "left"
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = 0

    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()

    rng = random.Random(args.seed)

    # Chat-template overhead (tokens injected around the user content).
    overhead = len(
        tokenizer.apply_chat_template(
            [{"role": "user", "content": "x"}],
            tokenize=True,
            add_generation_prompt=True,
        )
    )
    print(f"[niah1] chat-template overhead: {overhead} tokens")
    print(f"[niah1] seed={args.seed} batch_size={args.batch_size} trials={args.trials} "
          f"digits={args.digits}")

    low = 10 ** (args.digits - 1) if args.digits > 1 else 0
    high = 10 ** args.digits

    ctxs = list(range(args.start_ctx, args.end_ctx + 1, args.step))
    results: dict[Tuple[int, float], dict[str, int]] = {
        (c, d): {"pass": 0, "total": 0} for c in ctxs for d in args.depths
    }

    for ctx_len in ctxs:
        for depth in args.depths:
            for tr in range(args.trials):
                prompts: List[str] = []
                passkeys: List[int] = []
                for _ in range(args.batch_size):
                    pk = rng.randrange(low, high)
                    prompts.append(
                        build_prompt(tokenizer, ctx_len, depth, pk, overhead, rng)
                    )
                    passkeys.append(pk)

                bres, dt, plen = run_batch(
                    model, tokenizer, prompts, passkeys, args.max_new_tokens
                )

                for b, (ok, gen) in enumerate(bres):
                    results[(ctx_len, depth)]["total"] += 1
                    results[(ctx_len, depth)]["pass"] += int(ok)
                    short = gen.strip().replace("\n", " ")[:80]
                    tag = "PASS" if ok else "FAIL"
                    print(
                        f"  ctx={ctx_len:<5d} depth={depth:.2f} trial={tr} b={b} "
                        f"plen={plen:<5d} passkey={passkeys[b]} -> {tag}  | pred={short!r}"
                    )
                print(f"  [batch took {dt:.2f}s]")

    # ---- summary grid ------------------------------------------------------ #
    print("\n=== NIAH pass-rate grid ===")
    header = "ctx_len |" + "|".join(f" d={d:4.2f} " for d in args.depths) + "|"
    print(header)
    print("-" * len(header))
    for c in ctxs:
        row = f"{c:>7d} |"
        for d in args.depths:
            r = results[(c, d)]
            rate = r["pass"] / max(r["total"], 1)
            row += f" {r['pass']:2d}/{r['total']:<2d} "
        row += "|"
        print(row)

    print()
    total_pass = sum(r["pass"] for r in results.values())
    total_runs = sum(r["total"] for r in results.values())
    print(f"overall: {total_pass}/{total_runs} = {total_pass / max(total_runs, 1):.1%}")


if __name__ == "__main__":
    main()