scripts/run_crest_math500.py

"""
run_crest_math500.py — same setup as run_crest_aime25.py but over MATH-500 full,
with sympy-based grading (MATH answers include fractions, surds, etc.).

Same CREST system prompt, same _allmonoV2 layer set (14 layers), same
uniform global alpha sweep. Single seed. Ground-truth grading via:
  exact -> int -> float -> sympy.simplify(p - g)==0 -> normalized string.
"""
import argparse, json, os, re, sys, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from configs import get_config
from configs.paths import MATH500_FULL_PATH, 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)

CREST_SYSTEM = ("Answer the following questions. You should think step-by-step "
                "and put your final answer within \\boxed{}.")


def last_boxed(text):
    if not text: return None
    i = text.rfind("\\boxed"); j = text.find("{", i) if i >= 0 else -1
    if j < 0: return None
    depth, k = 0, j
    while k < len(text):
        if text[k] == "{": depth += 1
        elif text[k] == "}":
            depth -= 1
            if depth == 0: return text[j + 1:k].strip()
        k += 1
    return text[j + 1:].strip()


def _norm(s):
    if s is None: return ""
    t = s.strip()
    for x in ["\\left", "\\right", "\\!", "\\,", "\\;", "$", " "]:
        t = t.replace(x, "")
    return t.lower()


def _as_int(s):
    if s is None: return None
    t = re.sub(r"[^\d\-]", "", str(s))
    try: return int(t)
    except (ValueError, TypeError): return None


def _as_float(s):
    if s is None: return None
    try: return float(str(s).replace(",", "").replace("$", ""))
    except (ValueError, TypeError): return None


def _latex_to_sympy_src(s):
    t = s
    t = t.replace("\\dfrac", "\\frac")
    t = re.sub(r"\\frac\{([^{}]+)\}\{([^{}]+)\}", r"((\1)/(\2))", t)
    t = re.sub(r"\\sqrt\{([^{}]+)\}", r"sqrt(\1)", t)
    t = re.sub(r"\\sqrt\s*(\d+)", r"sqrt(\1)", t)
    t = t.replace("\\cdot", "*").replace("\\times", "*")
    t = t.replace("^", "**")
    t = re.sub(r"\\pi\b", "pi", t)
    t = re.sub(r"\\(left|right|!|,|;|:)", "", t)
    t = re.sub(r"\\[a-zA-Z]+", "", t)
    t = t.replace("{", "(").replace("}", ")").replace("$", "")
    return t


def _sympy_eq(a, b):
    try:
        from sympy import sympify, simplify
    except ImportError:
        return None
    try:
        pa = sympify(_latex_to_sympy_src(a))
        pb = sympify(_latex_to_sympy_src(b))
        return bool(simplify(pa - pb) == 0)
    except Exception:
        return None


def is_correct(pred, gold):
    if pred is None or gold is None or not str(gold).strip():
        return False
    p, g = str(pred).strip(), str(gold).strip()
    if p == g: return True
    pi, gi = _as_int(p), _as_int(g)
    if pi is not None and gi is not None and "/" not in p and "/" not in g:
        return pi == gi
    pf, gf = _as_float(p), _as_float(g)
    if pf is not None and gf is not None and abs(pf) < 1e9 and abs(gf) < 1e9:
        if abs(pf - gf) < 1e-6: return True
    sym = _sympy_eq(p, g)
    if sym is not None: return sym
    np_, ng_ = _norm(p), _norm(g)
    return np_ == ng_ and np_ != ""


def rep_score(text, tail=400, ng=30):
    t = text[-tail:] if len(text) > tail else text
    if len(t) < ng*2: return 0.0
    seen, r, tot = {}, 0, 0
    for i in range(len(t)-ng):
        c = t[i:i+ng]; tot += 1
        if c in seen: r += 1
        else: seen[c] = 1
    return r/tot if tot 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 load_math500():
    out = []
    for it in read_jsonl(MATH500_FULL_PATH):
        prob = it.get("problem") or it.get("question") or ""
        ans = str(it.get("answer", ""))
        if prob: out.append((prob, ans))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dimension", default="monitoring")
    ap.add_argument("--alphas", type=float, nargs="+", default=[0.0, 0.3, 0.7, 1.0])
    ap.add_argument("--sel-suffix", default="_allmonoV2")
    ap.add_argument("--out-suffix", default="_allmonoV2")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--gen-max-tokens", type=int, default=None)
    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 getattr(cfg, "TOPN_GEN_MAX_TOKENS", None)
               or cfg.GEN_MAX_NEW_TOKENS)
    temperature = getattr(cfg, "DEFAULT_TEMPERATURE", 0.6)
    top_p = getattr(cfg, "DEFAULT_TOP_P", 0.95)

    log = setup_logger("run_crest_math500",
                       os.path.join(LOG_DIR, f"run_crest_math500{args.out_suffix}.log"))
    log.info("=" * 72)
    log.info(f"CREST-prompt MATH-500 FULL  alphas={args.alphas}  seed={args.seed}")
    log.info(f"  gen_max={gen_max}  temp={temperature}  top_p={top_p}")
    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())}")

    problems = load_math500()
    log.info(f"  MATH-500 problems loaded: {len(problems)}")

    out_path = os.path.join(p.RESULTS_DIR, f"crest_math500{args.out_suffix}.jsonl")
    sum_path = os.path.join(p.RESULTS_DIR, f"crest_math500{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, ans, a, f"P{pi}_A{a:.2f}")
            for pi, (prob, ans) 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
    n_done_in_session = 0
    for pi, prob, ans, a, key in todo:
        prompt = build_chat_prompt(tokenizer, prob, enable_thinking=True,
                                   system=CREST_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)
        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 = last_boxed(cot); correct = is_correct(pred, ans)
        det = detector.detect(cot); rep = rep_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": ans, "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()
        n_done_in_session += 1
        log.info(f"  [{n_done_in_session}/{len(todo)}] {key}: pred={pred!r} gt={ans!r} "
                 f"{'OK' if correct else 'x'} ttok={ttok} 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(f"\n=== SUMMARY (CREST prompt, MATH-500 FULL, sympy GT, seed {args.seed}) ===")
    log.info(f"{'alpha':>6} {'n':>4} {'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:>4} {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()),
                "prompt_system": CREST_SYSTEM,
                "benchmark": "MATH-500 full (500 problems)",
                "summary": summary}, sum_path)
    log.info(f"\nSaved {out_path}\n      {sum_path}\nDone.")


if __name__ == "__main__":
    main()