File size: 3,208 Bytes
867d0f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """
pass@1 vs pass@k on the synthetic held-out set, broken down by step count.
This is the go/no-go diagnostic for RL with verifiable rewards. RLVR (GRPO/RLOO) reweights samples
the model ALREADY produces: if a problem is never solved in k tries, every sample in the group gets
reward 0, the advantage is 0, and there is no gradient. So the headroom RL can capture is bounded
by (pass@k - pass@1), and only on problems where pass@k > 0.
"""
import json, os, sys, collections
import torch, pyarrow.parquet as pq
from transformers import AutoModelForCausalLM, AutoTokenizer
sys.path.insert(0, ".")
from eval_reasoning import numbers, parse
MODEL = sys.argv[1] if len(sys.argv) > 1 else "./Nawah-Reasoning-v5"
PER_BUCKET = int(os.environ.get("PER_BUCKET", 60))
K = int(os.environ.get("K", 8))
TEMP = float(os.environ.get("TEMP", 1.0))
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).cuda().eval()
rows = pq.read_table("out_merged/arabic_math_reasoning_synth.parquet").to_pylist()[:2000]
buckets = collections.defaultdict(list)
for r in rows:
buckets[r["axis_steps"]].append(r)
sample = [r for b in buckets.values() for r in b[:PER_BUCKET]]
print(f"[*] {len(sample)} problems x k={K} @ T={TEMP} -> {len(sample)*K} generations", flush=True)
def ref_number(r):
ns = numbers(r["answer"])
return ns[-1] if ns else None
stats = collections.defaultdict(lambda: {"n": 0, "p1": 0, "pk": 0})
BATCH = 16
for start in range(0, len(sample), BATCH):
chunk = sample[start:start + BATCH]
prompts = [tok.apply_chat_template([{"role": "user", "content": r["instruction"]}],
tokenize=False, add_generation_prompt=True) for r in chunk]
enc = tok(prompts, return_tensors="pt", padding=True, padding_side="left").to("cuda")
torch.manual_seed(1234 + start)
out = model.generate(**enc, max_new_tokens=320, do_sample=True, temperature=TEMP,
top_p=0.95, num_return_sequences=K)
gen = tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
for i, r in enumerate(chunk):
ref = ref_number(r)
hits = []
for j in range(K):
_, ans, _ = parse(gen[i * K + j])
ns = numbers(ans or "")
hits.append(bool(ns) and ref is not None and ns[-1] == ref)
s = stats[r["axis_steps"]]
s["n"] += 1
s["p1"] += hits[0]
s["pk"] += any(hits)
print(f" {start + len(chunk)}/{len(sample)}", flush=True)
print("\n| steps | n | pass@1 | pass@%d | headroom |" % K)
print("|---|---:|---:|---:|---:|")
tot = {"n": 0, "p1": 0, "pk": 0}
for k, s in sorted(stats.items(), key=lambda kv: kv[1]["n"], reverse=True):
for f in tot: tot[f] += s[f]
print(f"| {k} | {s['n']} | {100*s['p1']/s['n']:.1f}% | {100*s['pk']/s['n']:.1f}% | "
f"{100*(s['pk']-s['p1'])/s['n']:+.1f} |")
print(f"| **all** | {tot['n']} | {100*tot['p1']/tot['n']:.1f}% | {100*tot['pk']/tot['n']:.1f}% | "
f"{100*(tot['pk']-tot['p1'])/tot['n']:+.1f} |")
print(f"\nnever-solved (pass@{K}=0): {100*(tot['n']-tot['pk'])/tot['n']:.1f}% of problems "
f"-> zero RL gradient on these")
|