scripts/run_gpqa_d.py

"""
run_gpqa_d.py — v12 GPQA-Diamond inference (single seed).

Mirrors run_crest_aime25 structure: takes 03b_v2 _allmonoV2 selected layers,
sweeps a UNIFORM global alpha across them, judges by exact letter match.
"""
import argparse, json, os, re, sys, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from tqdm import tqdm
from configs import get_config
from configs.paths import LOG_DIR, dim_paths, ensure_dirs
from src.detectors import BehaviorDetector
from src.interventions import generate_plain, generate_with_alpha
from src.utils import (
    build_chat_prompt, get_device, load_model_and_tokenizer,
    read_json, read_jsonl, setup_logger, write_json,
)


def extract_boxed_letter(text):
    """Find LAST \\boxed{X} where X is one of A/B/C/D."""
    if not text:
        return None
    matches, idx = [], 0
    while True:
        i = text.find("\\boxed", idx)
        if i < 0: break
        j = text.find("{", i)
        if j < 0: break
        depth, end = 0, -1
        for k in range(j, len(text)):
            if text[k] == "{": depth += 1
            elif text[k] == "}":
                depth -= 1
                if depth == 0: end = k; break
        if end > j:
            matches.append(text[j+1:end].strip()); idx = end + 1
        else: break
    if not matches:
        return None
    # take last; first valid-looking letter wins
    last = matches[-1].strip().upper()
    m = re.match(r"\(?\s*([ABCD])", last)
    return m.group(1) if m else None


def repetition_score(text, tail_chars=400, ngram=30):
    tail = text[-tail_chars:] if len(text) > tail_chars else text
    if len(tail) < ngram * 2: return 0.0
    seen, repeated, total = {}, 0, 0
    for i in range(len(tail) - ngram):
        chunk = tail[i:i+ngram]
        total += 1
        if chunk in seen: repeated += 1
        else: seen[chunk] = 1
    return repeated / total if total else 0.0


def think_tokens(tok, cot):
    seg = cot.split("</think>")[0] if "</think>" in cot else cot
    return len(tok(seg, add_special_tokens=False)["input_ids"])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dimension", default="monitoring")
    ap.add_argument("--alphas", type=float, nargs="+",
                    default=[1.0, 0.7, 0.3, 0.0])
    ap.add_argument("--sel-suffix", default="_allmonoV2")
    ap.add_argument("--out-suffix", default="_gpqa_d_s64")
    ap.add_argument("--seed", type=int, default=64)
    ap.add_argument("--gen-max-tokens", type=int, default=None)
    ap.add_argument("--data-path", default=None,
                    help="Override (default data/gpqa_d.jsonl)")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    ensure_dirs(args.dimension)
    cfg = get_config(args.dimension)
    p = dim_paths(args.dimension)
    gen_max = args.gen_max_tokens or cfg.GEN_MAX_NEW_TOKENS
    temperature = getattr(cfg, "DEFAULT_TEMPERATURE", 0.6)
    top_p = getattr(cfg, "DEFAULT_TOP_P", 0.95)

    data_path = args.data_path or os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        "data", "gpqa_d.jsonl")

    log = setup_logger("run_gpqa_d",
                       os.path.join(LOG_DIR, f"run_gpqa_d{args.out_suffix}.log"))
    log.info("=" * 72)
    log.info(f"GPQA-Diamond run on v12 (30B)")
    log.info(f"  alphas={args.alphas}  seed={args.seed}")
    log.info(f"  gen_max={gen_max}  temp={temperature}  top_p={top_p}")
    log.info(f"  data_path = {data_path}")
    log.info("=" * 72)

    if not os.path.exists(p.DIRECTIONS):
        log.error(f"missing {p.DIRECTIONS}"); sys.exit(1)
    dblob = torch.load(p.DIRECTIONS, map_location="cpu", weights_only=False)
    directions_all = {int(L): v for L, v in dblob["directions"].items()}

    base, ext = os.path.splitext(p.SELECTED_LAYERS)
    sel_path = f"{base}{args.sel_suffix}{ext}"
    if not os.path.exists(sel_path):
        log.error(f"missing {sel_path}. Run 03b_v2_allmono.py first."); sys.exit(1)
    sel = read_json(sel_path)
    selected = [int(L) for L in sel["selected_layers"]]
    directions = {L: directions_all[L] for L in selected if L in directions_all}
    log.info(f"  selected layers ({len(directions)}): {sorted(directions.keys())}")
    if not directions:
        log.error("No directions for selected layers."); sys.exit(2)

    items = read_jsonl(data_path)
    problems = [it["problem"] for it in items]
    gt = {i: it["answer"].strip().upper() for i, it in enumerate(items)}
    log.info(f"  problems: {len(problems)}")
    if not problems:
        log.error(f"No problems loaded from {data_path}"); sys.exit(3)

    out_path = os.path.join(p.RESULTS_DIR, f"run_gpqa_d{args.out_suffix}.jsonl")
    sum_path = os.path.join(p.RESULTS_DIR, f"run_gpqa_d{args.out_suffix}_summary.json")

    if args.force and os.path.exists(out_path): os.remove(out_path)
    seen = set()
    if os.path.exists(out_path):
        for line in open(out_path):
            line = line.strip()
            if line:
                try: seen.add(json.loads(line)["_key"])
                except Exception: pass
        log.info(f"  [resume] {len(seen)} records cached")

    todo = [(pi, prob, a, f"P{pi}_A{a:.2f}")
            for pi, prob in enumerate(problems)
            for a in args.alphas
            if f"P{pi}_A{a:.2f}" not in seen]
    log.info(f"  records to compute: {len(todo)} / {len(problems)*len(args.alphas)}")

    detector = BehaviorDetector(cfg)
    device = get_device()
    model = tokenizer = None
    if todo:
        log.info("Loading model...")
        model, tokenizer = load_model_and_tokenizer(device=device)

    fh = open(out_path, "a", encoding="utf-8") if todo else None
    for pi, prob, a, key in tqdm(todo, desc="gpqa_d", dynamic_ncols=True, mininterval=10):
        # build_chat_prompt 在 v12 里支持 enable_thinking 但 system 我们留空
        # (gpqa_d.jsonl 的 problem 字段里已经写了 instruction)
        prompt = build_chat_prompt(tokenizer, prob, enable_thinking=True, system="")
        gen_seed = args.seed * 1000 + pi
        t0 = time.time()
        if a >= 1.0 - 1e-6:
            cot = generate_plain(model, tokenizer, prompt, device,
                                 max_new_tokens=gen_max, do_sample=True,
                                 temperature=temperature, top_p=top_p, seed=gen_seed)
            eff = {int(L): 1.0 for L in directions}
        else:
            eff = {int(L): float(a) for L in directions}
            cot = generate_with_alpha(model, tokenizer, prompt, directions, eff,
                                      device, max_new_tokens=gen_max, do_sample=True,
                                      temperature=temperature, top_p=top_p, seed=gen_seed)
        elapsed = time.time() - t0

        pred = extract_boxed_letter(cot)
        gtv = gt.get(pi)
        correct = (pred == gtv) if (pred and gtv) else False
        det = detector.detect(cot)
        rep = repetition_score(cot)
        ttok = think_tokens(tokenizer, cot)

        rec = {
            "_key": key, "problem_idx": pi, "alpha": a, "seed": args.seed,
            "problem": prob, "cot": cot,
            "pred": pred, "gt": gtv, "correct": correct,
            "has_boxed": pred is not None,
            "think_tokens": ttok, "n_chars": len(cot),
            "mon_total": det["total"], "repetition_score": rep,
            "collapse": rep > 0.5, "elapsed_s": elapsed,
        }
        if fh:
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n"); fh.flush()
        log.info(f"  {key}: pred={pred} gt={gtv} {'OK' if correct else 'x'}  "
                 f"think_tok={ttok} mon={det['total']} rep={rep:.2f} t={elapsed:.0f}s")
    if fh: fh.close()

    recs = []
    for line in open(out_path):
        line = line.strip()
        if line:
            try: recs.append(json.loads(line))
            except Exception: pass

    avg = lambda xs: sum(xs)/len(xs) if xs else 0.0
    summary = {}
    log.info("\n=== SUMMARY (GPQA-D letter grading, seed %d) ===" % args.seed)
    log.info(f"{'alpha':>6} {'n':>3} {'acc':>8} {'correct':>8} {'noBox':>6} "
             f"{'think_tok':>10} {'mon':>6} {'collapse':>9}")
    for a in sorted(args.alphas, reverse=True):
        rs = [r for r in recs if abs(r["alpha"]-a) < 1e-6]
        if not rs: continue
        n = len(rs)
        acc = sum(r["correct"] for r in rs) / n
        summary[f"{a:.2f}"] = {
            "n": n, "accuracy": acc,
            "n_correct": sum(r["correct"] for r in rs),
            "n_no_boxed": n - sum(r["has_boxed"] for r in rs),
            "mean_think_tokens": avg([r["think_tokens"] for r in rs]),
            "mean_chars": avg([r["n_chars"] for r in rs]),
            "mean_mon": avg([r["mon_total"] for r in rs]),
            "collapse_rate": sum(r["collapse"] for r in rs) / n,
        }
        log.info(f"{a:>6.2f} {n:>3} {acc:>7.1%} "
                 f"{sum(r['correct'] for r in rs):>8} "
                 f"{n - sum(r['has_boxed'] for r in rs):>6} "
                 f"{avg([r['think_tokens'] for r in rs]):>10.0f} "
                 f"{avg([r['mon_total'] for r in rs]):>6.1f} "
                 f"{sum(r['collapse'] for r in rs)/n*100:>8.1f}%")

    write_json({"seed": args.seed, "alphas": args.alphas,
                "selected_layers": sorted(directions.keys()),
                "dataset": "GPQA-Diamond", "n_problems": len(problems),
                "summary": summary}, sum_path)
    log.info(f"\nSaved {out_path}\n      {sum_path}\nDone.")


if __name__ == "__main__":
    main()