code/repro_all.py

"""
JustGRPO reproduction script — all three claims in one GPU run.

Claim 1: JustGRPO achieves 89.1% on GSM8K (gen_length=256, steps=256, block_length=32).
Claim 2: Retains parallel decoding (fully-parallel block_length=256 still works).
Claim 3: Arbitrary order (parallel) collapses solution coverage (Pass@k) vs AR order.

Single-GPU (no torchrun). Uses the official JustGRPO utils (generate, grader, parser).
"""
import argparse
import json
import os
import sys
import time

import torch
from datasets import load_dataset
from transformers import AutoModel, AutoTokenizer
from tqdm import tqdm

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from utils.generate import generate
from utils.grader import math_equal
from utils.parser import extract_answer


def extract_answer_gsm8k(answer: str):
    """Extract the final answer from GSM8K format (after ####)."""
    return answer.split('####')[-1].strip()


def load_model_tokenizer(ckpt_path, base_path, device):
    print(f"[load] tokenizer from {base_path}", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(base_path, trust_remote_code=True)
    print(f"[load] model from {ckpt_path} (bf16)", flush=True)
    model = AutoModel.from_pretrained(
        ckpt_path, trust_remote_code=True, torch_dtype=torch.bfloat16
    )
    model.eval().requires_grad_(False).to(device)
    print(f"[load] done. #params={sum(p.numel() for p in model.parameters())/1e9:.2f}B", flush=True)
    return model, tokenizer


def build_prompt_ids(tokenizer, question, device):
    msgs = [[{"role": "user", "content": question}]]
    prompt_text = tokenizer.apply_chat_template(
        msgs, add_generation_prompt=True, tokenize=False
    )
    prompt_ids = tokenizer(prompt_text, return_tensors="pt")["input_ids"].to(device)
    return prompt_ids


def grade_gsm8k(gt_answer, response):
    """gt_answer: raw GSM8K 'answer' field (with ####). response: model text."""
    return bool(math_equal(extract_answer_gsm8k(gt_answer), extract_answer(response)))



# ---------------- Claim 1 ----------------
def claim1(model, tokenizer, device, args, out_dir):
    """GSM8K accuracy with paper config: block_length=32, gen_length=256, steps=256, temp=0."""
    print("\n" + "=" * 70, flush=True)
    print("CLAIM 1: GSM8K accuracy (block_length=32, gen=256, steps=256, temp=0)", flush=True)
    print("=" * 70, flush=True)
    ds = load_dataset("gsm8k", "main", split="test")
    # Seeded representative subset (no clustering bias) of size N.
    n = min(args.c1_n, len(ds))
    rng = torch.Generator().manual_seed(args.seed)
    idxs = torch.randperm(len(ds), generator=rng)[:n].tolist()
    print(f"[c1] evaluating {n} / {len(ds)} examples", flush=True)

    correct, total = 0, 0
    t0 = time.time()
    rows = []
    for k, i in enumerate(tqdm(idxs, desc="c1")):
        ex = ds[int(i)]
        q, a = ex["question"], ex["answer"]
        pid = build_prompt_ids(tokenizer, q, device)
        gen = generate(model=model, prompt=pid, steps=args.gen_length,
                       gen_length=args.gen_length, block_length=32, temperature=0.0)
        resp = tokenizer.batch_decode(gen[:, pid.shape[1]:], skip_special_tokens=True)[0]
        ok = grade_gsm8k(a, resp)
        correct += int(ok); total += 1
        rows.append({"idx": int(i), "correct": int(ok),
                     "gt": extract_answer_gsm8k(a), "pred": extract_answer(resp)[:120]})
        if (k + 1) % 50 == 0:
            acc = correct / total
            el = time.time() - t0
            eta = el / (k + 1) * (n - k - 1)
            print(f"[c1] {k+1}/{n} acc={acc*100:.2f}% elapsed={el:.0f}s eta={eta:.0f}s", flush=True)
    acc = correct / total
    el = time.time() - t0
    res = {"claim": 1, "task": "gsm8k", "ckpt": args.ckpt_path,
           "block_length": 32, "gen_length": args.gen_length, "steps": args.gen_length,
           "temperature": 0.0, "n": n, "correct": correct, "total": total,
           "accuracy": acc, "paper_reported": 0.891, "elapsed_sec": el,
           "sample_results": rows[:20]}
    with open(os.path.join(out_dir, "claim1_gsm8k.json"), "w") as f:
        json.dump(res, f, indent=2)
    print(f"[c1] FINAL acc={acc*100:.2f}% ({correct}/{total})  paper=89.10%  elapsed={el:.0f}s", flush=True)
    return res


# ---------------- Claim 2 ----------------
def claim2(model, tokenizer, device, args, out_dir):
    """Parallel decoding retained: compare block_length=256 (fully parallel) vs
    block_length=1 (pure AR) vs block_length=32 (semi-AR paper default) on same subset."""
    print("\n" + "=" * 70, flush=True)
    print("CLAIM 2: Parallel decoding retained (block_length sweep, temp=0)", flush=True)
    print("=" * 70, flush=True)
    ds = load_dataset("gsm8k", "main", split="test")
    n = min(args.c2_n, len(ds))
    rng = torch.Generator().manual_seed(args.seed)
    idxs = torch.randperm(len(ds), generator=rng)[:n].tolist()
    print(f"[c2] evaluating {n} examples under 3 block_length settings", flush=True)

    results = {}
    for bl in [32, 256, 1]:  # semi-AR (paper default), fully-parallel, pure-AR
        label = {32: "semi_AR_blk32", 256: "fully_parallel_blk256", 1: "pure_AR_blk1"}[bl]
        correct, total = 0, 0
        t0 = time.time()
        for k, i in enumerate(tqdm(idxs, desc=f"c2-{label}")):
            ex = ds[int(i)]
            q, a = ex["question"], ex["answer"]
            pid = build_prompt_ids(tokenizer, q, device)
            gen = generate(model=model, prompt=pid, steps=args.gen_length,
                           gen_length=args.gen_length, block_length=bl, temperature=0.0)
            resp = tokenizer.batch_decode(gen[:, pid.shape[1]:], skip_special_tokens=True)[0]
            ok = grade_gsm8k(a, resp)
            correct += int(ok); total += 1
        acc = correct / total
        el = time.time() - t0
        results[label] = {"block_length": bl, "n": n, "correct": correct,
                          "total": total, "accuracy": acc, "elapsed_sec": el}
        print(f"[c2] {label}: acc={acc*100:.2f}% ({correct}/{total}) elapsed={el:.0f}s", flush=True)
    res = {"claim": 2, "task": "gsm8k", "ckpt": args.ckpt_path,
           "gen_length": args.gen_length, "steps": args.gen_length,
           "n": n, "modes": results,
           "interpretation": (
               "If fully_parallel_blk256 accuracy is close to semi_AR_blk32 and pure_AR_blk1, "
               "the model trained with AR order RETAINS parallel decoding ability at inference.")}
    with open(os.path.join(out_dir, "claim2_parallel.json"), "w") as f:
        json.dump(res, f, indent=2)
    return res


# ---------------- Claim 3 ----------------
def pass_at_k(n, c, k):
    """Unbiased pass@k (Chen et al. 2021). n=samples, c=correct, k=requested."""
    import math
    if n - c < k:
        return 1.0
    return 1.0 - math.prod((n - c - j) / (n - j) for j in range(k))


def claim3(model, tokenizer, device, args, out_dir):
    """Solution coverage collapse: Pass@k of AR order (block_length=1) vs
    arbitrary order (block_length=256, low_confidence) at temperature>0."""
    print("\n" + "=" * 70, flush=True)
    print(f"CLAIM 3: Solution coverage — Pass@{args.c3_k} AR vs arbitrary (temp={args.c3_temp})", flush=True)
    print("=" * 70, flush=True)
    ds = load_dataset("gsm8k", "main", split="test")
    n_prob = min(args.c3_problems, len(ds))
    rng = torch.Generator().manual_seed(args.seed)
    idxs = torch.randperm(len(ds), generator=rng)[:n_prob].tolist()
    K = args.c3_k
    print(f"[c3] {n_prob} problems x {K} samples x 2 modes (AR blk1, arbitrary blk256)", flush=True)

    modes = {"AR_order_blk1": 1, "arbitrary_order_blk256": 256}
    summary = {}
    for mode, bl in modes.items():
        per_problem = []  # list of (n_correct, n_total)
        t0 = time.time()
        for pi, i in enumerate(tqdm(idxs, desc=f"c3-{mode}")):
            ex = ds[int(i)]
            q, a = ex["question"], ex["answer"]
            pid = build_prompt_ids(tokenizer, q, device)
            # generate K samples in one batch (identical prompt -> supported by generate)
            gen = generate(model=model, prompt=pid.repeat(K, 1), steps=args.gen_length,
                           gen_length=args.gen_length, block_length=bl, temperature=args.c3_temp)
            resps = tokenizer.batch_decode(gen[:, pid.shape[1]:], skip_special_tokens=True)
            n_correct = sum(grade_gsm8k(a, r) for r in resps)
            per_problem.append((n_correct, K))
        # aggregate pass@k
        pk = [pass_at_k(K, c, K) for (c, _) in per_problem]  # pass@K = any correct in K
        passk = sum(pk) / len(pk)
        # also greedy-style acc (sample-1 correctness rate)
        acc1 = sum(c / K for c, _ in per_problem) / len(per_problem)
        # exact pass@1 over all samples
        all_correct = sum(c for c, _ in per_problem)
        all_total = sum(t for _, t in per_problem)
        sample_acc = all_correct / all_total
        el = time.time() - t0
        summary[mode] = {"block_length": bl, "n_problems": n_prob, "k": K,
                         "pass_at_k": passk, "sample_accuracy": sample_acc,
                         "mean_per_problem_acc": acc1,
                         "n_correct_samples": all_correct, "n_total_samples": all_total,
                         "elapsed_sec": el,
                         "per_problem": [{"n_correct": c, "n_total": t} for c, t in per_problem]}
        print(f"[c3] {mode}: pass@{K}={passk*100:.2f}% sample_acc={sample_acc*100:.2f}% elapsed={el:.0f}s", flush=True)

    ar_pk = summary["AR_order_blk1"]["pass_at_k"]
    arb_pk = summary["arbitrary_order_blk256"]["pass_at_k"]
    collapse = arb_pk < ar_pk
    res = {"claim": 3, "task": "gsm8k", "ckpt": args.ckpt_path,
           "gen_length": args.gen_length, "steps": args.gen_length,
           "n_problems": n_prob, "k": K, "temperature": args.c3_temp,
           "modes": summary,
           "AR_pass_at_k": ar_pk, "arbitrary_pass_at_k": arb_pk,
           "arbitrary_collapses_coverage": bool(collapse),
           "relative_drop_pct": (ar_pk - arb_pk) / (ar_pk + 1e-9) * 100 if ar_pk > 0 else 0.0,
           "interpretation": (
               "If arbitrary_order pass@k < AR_order pass@k, arbitrary-order generation "
               "collapses solution coverage, supporting the paper's core claim.")}
    with open(os.path.join(out_dir, "claim3_coverage.json"), "w") as f:
        json.dump(res, f, indent=2)
    print(f"[c3] AR pass@{K}={ar_pk*100:.2f}%  arbitrary pass@{K}={arb_pk*100:.2f}%  "
          f"collapse={collapse}  drop={res['relative_drop_pct']:.1f}%", flush=True)
    return res


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt_path", default="/ckpt")
    p.add_argument("--base_path", default="/base")
    p.add_argument("--output", default="/outputs")
    p.add_argument("--gen_length", type=int, default=256)
    p.add_argument("--seed", type=int, default=113)
    p.add_argument("--c1_n", type=int, default=500)
    p.add_argument("--c2_n", type=int, default=100)
    p.add_argument("--c3_problems", type=int, default=30)
    p.add_argument("--c3_k", type=int, default=4)
    p.add_argument("--c3_temp", type=float, default=0.7)
    p.add_argument("--claims", default="1,2,3")
    args = p.parse_args()

    os.makedirs(args.output, exist_ok=True)
    device = torch.device("cuda:0")
    torch.manual_seed(args.seed)
    model, tokenizer = load_model_tokenizer(args.ckpt_path, args.base_path, device)

    claims = [int(c) for c in args.claims.split(",")]
    all_res = {}
    if 1 in claims:
        all_res["c1"] = claim1(model, tokenizer, device, args, args.output)
    if 2 in claims:
        all_res["c2"] = claim2(model, tokenizer, device, args, args.output)
    if 3 in claims:
        all_res["c3"] = claim3(model, tokenizer, device, args, args.output)

    with open(os.path.join(args.output, "summary.json"), "w") as f:
        json.dump(all_res, f, indent=2)
    print("\n" + "=" * 70, flush=True)
    print("ALL CLAIMS COMPLETE — summary:", flush=True)
    print(json.dumps({k: {kk: vv for kk, vv in v.items() if kk in
                          ("accuracy", "pass_at_k", "modes", "AR_pass_at_k",
                           "arbitrary_pass_at_k", "arbitrary_collapses_coverage",
                           "relative_drop_pct", "paper_reported", "n", "correct", "total")}
                      for k, v in all_res.items()}, indent=2), flush=True)
    print("=" * 70, flush=True)


if __name__ == "__main__":
    main()