Upload scripts/peer/peer_run.py with huggingface_hub

scripts/peer/peer_run.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""
+peer_run.py -- generate completions for peer-model NIAH prompts.
+
+Backends, tried in this order unless --backend forces one:
+  vllm   fast, batched. Cannot load Spark-X2.5 (model_type spark2_5 is
+         custom_code and has no vLLM backend), so it is skipped automatically.
+  hf     transformers with trust_remote_code. Works for everything, slower.
+
+Prompts are fed as TOKEN IDS, never text: peer_gen.py guarantees an exact token
+count and a decode->re-encode round trip would silently break it.
+
+Greedy decoding, matching our own harness (CP_GEN_GREEDY=1). Generation stops at
+--max-new tokens; no early stop, because the peer models have no shared
+end-of-answer convention and truncation is harmless -- scoring is substring
+match, so trailing chatter costs nothing.
+
+  usage:
+    peer_run.py --model Qwen/Qwen3-4B-Base --prompts out/niah_single_prompts.jsonl \
+                --out gen/qwen3-4b_niah_single.jsonl --max-new 24
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+
+
+def load_prompts(p):
+    rows = []
+    with open(p) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def run_vllm(model, rows, max_new, gpu_frac, max_len):
+    from vllm import LLM, SamplingParams
+    llm = LLM(model=model, dtype="bfloat16", trust_remote_code=True,
+              gpu_memory_utilization=gpu_frac, max_model_len=max_len,
+              enforce_eager=False)
+    sp = SamplingParams(temperature=0.0, max_tokens=max_new)
+    outs = llm.generate(prompt_token_ids=[r["ids"] for r in rows],
+                        sampling_params=sp)
+    return [o.outputs[0].text for o in outs]
+
+
+def run_hf(model, rows, max_new, batch):
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    m = AutoModelForCausalLM.from_pretrained(
+        model, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True)
+    m.eval()
+    texts = []
+    # batch=1 by default: these are 16k-token prompts and padding a ragged batch
+    # would change each sequence's positions, which is exactly what NIAH measures.
+    for s in range(0, len(rows), batch):
+        chunk = rows[s:s + batch]
+        ids = torch.tensor([r["ids"] for r in chunk], device="cuda")
+        with torch.no_grad():
+            out = m.generate(input_ids=ids, max_new_tokens=max_new,
+                             do_sample=False, num_beams=1,
+                             pad_token_id=getattr(tok, "pad_token_id", None)
+                             or getattr(tok, "eos_token_id", 0))
+        for j in range(len(chunk)):
+            texts.append(tok.decode(out[j][ids.shape[1]:], skip_special_tokens=True))
+        done = min(s + batch, len(rows))
+        print(f"    {done}/{len(rows)}", flush=True)
+    return texts
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--prompts", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--max-new", type=int, default=24)
+    ap.add_argument("--backend", default="auto", choices=["auto", "vllm", "hf"])
+    ap.add_argument("--gpu-frac", type=float, default=0.90)
+    ap.add_argument("--batch", type=int, default=1, help="hf backend only")
+    args = ap.parse_args()
+
+    rows = load_prompts(args.prompts)
+    if not rows:
+        sys.exit(f"no prompts in {args.prompts}")
+    plen = len(rows[0]["ids"])
+    max_len = plen + args.max_new + 8
+    # Never ask vLLM for more than the checkpoint's trained window: Qwen3-Base
+    # ships max_position_embeddings=32768, so a 32,744-token prompt + 24 new + 8
+    # slack would exceed it and vLLM refuses to start. Clamp, and say so.
+    try:
+        from transformers import AutoConfig
+        mp = getattr(AutoConfig.from_pretrained(args.model, trust_remote_code=True),
+                     "max_position_embeddings", None)
+        if mp and max_len > mp:
+            print(f"  clamping max_model_len {max_len} -> {mp} (model's native window)")
+            max_len = mp
+            if plen + args.max_new > mp:
+                sys.exit(f"  FATAL: prompt {plen} + max_new {args.max_new} exceeds "
+                         f"the model's window {mp}; lower --prompt-len")
+    except Exception as e:
+        print(f"  (could not read max_position_embeddings: {type(e).__name__})")
+    print(f"  {args.model}  {len(rows)} prompts x {plen} tok  max_new={args.max_new}")
+
+    backend = args.backend
+    if backend == "auto":
+        backend = "hf"
+        try:
+            import vllm  # noqa: F401
+            # spark2_5 is custom_code with no vLLM backend; do not waste a load
+            backend = "hf" if "spark" in args.model.lower() else "vllm"
+        except ImportError:
+            pass
+    print(f"  backend: {backend}")
+
+    t0 = time.time()
+    if backend == "vllm":
+        try:
+            texts = run_vllm(args.model, rows, args.max_new, args.gpu_frac, max_len)
+        except Exception as e:
+            print(f"  vLLM failed ({type(e).__name__}: {str(e)[:140]}) -- falling back to hf")
+            texts = run_hf(args.model, rows, args.max_new, args.batch)
+    else:
+        texts = run_hf(args.model, rows, args.max_new, args.batch)
+    dt = time.time() - t0
+
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    with open(args.out, "w") as f:
+        for r, t in zip(rows, texts):
+            f.write(json.dumps({"i": r["i"], "generated": t}) + "\n")
+    print(f"  wrote {args.out}  ({len(texts)} rows, {dt:.0f}s, "
+          f"{dt/max(1,len(texts)):.2f}s/prompt)")
+
+
+if __name__ == "__main__":
+    main()