scripts/run_gpqa_d_runtime_latency.py

import argparse, json, os, re, sys, time, random
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):
    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
    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 count_tokens(tokenizer, text):
    return len(tokenizer(text, add_special_tokens=False)["input_ids"])


def think_tokens(tokenizer, cot):
    seg = cot.split("</think>")[0] if "</think>" in cot else cot
    return count_tokens(tokenizer, seg)


def cuda_sync():
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def cuda_reset_peak():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()


def cuda_mem_stats():
    if not torch.cuda.is_available():
        return {
            "max_mem_allocated_gb": None,
            "max_mem_reserved_gb": None,
            "current_mem_allocated_gb": None,
            "current_mem_reserved_gb": None,
        }
    return {
        "max_mem_allocated_gb": torch.cuda.max_memory_allocated() / (1024 ** 3),
        "max_mem_reserved_gb": torch.cuda.max_memory_reserved() / (1024 ** 3),
        "current_mem_allocated_gb": torch.cuda.memory_allocated() / (1024 ** 3),
        "current_mem_reserved_gb": torch.cuda.memory_reserved() / (1024 ** 3),
    }


def choose_subset_indices(n_total, limit, subset_seed):
    rng = random.Random(subset_seed)
    idxs = sorted(rng.sample(range(n_total), min(limit, n_total)))
    return idxs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dimension", default="monitoring")
    ap.add_argument("--seed", type=int, default=64)
    ap.add_argument("--subset-seed", type=int, default=64)
    ap.add_argument("--limit", type=int, default=50)
    ap.add_argument("--alphas", type=float, nargs="+", default=[1.0, 0.7, 0.3])
    ap.add_argument("--sel-suffix", default="_allmonoV2")
    ap.add_argument("--gen-max-tokens", type=int, default=None)
    ap.add_argument("--data-path", default=None)
    ap.add_argument("--out-suffix", default="_n50_s64")
    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",
    )

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

    log = setup_logger(
        "gpqa_runtime_latency",
        os.path.join(LOG_DIR, f"run_gpqa_d_runtime_latency{args.out_suffix}.log"),
    )

    log.info("=" * 80)
    log.info("GPQA-D wall-clock latency / throughput benchmark")
    log.info(f"seed={args.seed}")
    log.info(f"subset_seed={args.subset_seed}")
    log.info(f"limit={args.limit}")
    log.info(f"alphas={args.alphas}")
    log.info(f"gen_max={gen_max}")
    log.info(f"temperature={temperature}")
    log.info(f"top_p={top_p}")
    log.info(f"data_path={data_path}")
    log.info("=" * 80)

    items = read_jsonl(data_path)
    subset_indices = choose_subset_indices(len(items), args.limit, args.subset_seed)
    gt = {i: items[i]["answer"].strip().upper() for i in subset_indices}

    log.info(f"subset_indices={subset_indices}")

    need_steering = any(abs(float(a) - 1.0) > 1e-9 for a in args.alphas)

    directions = {}
    selected_layers = []
    sel_path = None

    if need_steering:
        if not os.path.exists(p.DIRECTIONS):
            raise FileNotFoundError(f"Missing directions file: {p.DIRECTIONS}")

        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):
            raise FileNotFoundError(f"Missing selected layer file: {sel_path}")

        selected_json = read_json(sel_path)
        selected_layers = [int(L) for L in selected_json["selected_layers"] if int(L) in directions_all]
        directions = {L: directions_all[L] for L in selected_layers}

        log.info(f"selected_layer_file={sel_path}")
        log.info(f"selected_layers({len(selected_layers)})={selected_layers}")

    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, encoding="utf-8"):
            if not line.strip():
                continue
            try:
                seen.add(json.loads(line)["_key"])
            except Exception:
                pass
        log.info(f"[resume] found {len(seen)} completed records")

    todo = []
    for pi in subset_indices:
        for alpha in args.alphas:
            key = f"P{pi}_A{float(alpha):.2f}"
            if key not in seen:
                todo.append((pi, float(alpha), key))

    log.info(f"records to compute: {len(todo)} / {len(subset_indices) * 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, alpha, key in tqdm(todo, desc="gpqa_runtime_latency", dynamic_ncols=True, mininterval=10):
        problem = items[pi]["problem"]
        prompt = build_chat_prompt(tokenizer, problem, enable_thinking=True, system="")
        prompt_tokens = count_tokens(tokenizer, prompt)
        gen_seed = args.seed * 1000 + pi

        cuda_reset_peak()
        cuda_sync()
        t0 = time.perf_counter()

        if abs(alpha - 1.0) < 1e-9:
            method = "baseline_no_hook"
            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:
            method = "crest_projection_removal"
            alpha_per_layer = {L: alpha for L in selected_layers}
            cot = generate_with_alpha(
                model,
                tokenizer,
                prompt,
                directions,
                alpha_per_layer,
                device,
                max_new_tokens=gen_max,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                seed=gen_seed,
            )

        cuda_sync()
        elapsed = time.perf_counter() - t0
        mem = cuda_mem_stats()

        pred = extract_boxed_letter(cot)
        det = detector.detect(cot)
        rep = repetition_score(cot)
        gen_tokens = count_tokens(tokenizer, cot)
        ttok = think_tokens(tokenizer, cot)

        rec = {
            "_key": key,
            "problem_idx": pi,
            "dataset": "gpqa_d",
            "subset_seed": args.subset_seed,
            "limit": args.limit,
            "seed": args.seed,
            "method": method,
            "alpha": alpha,
            "temperature": temperature,
            "top_p": top_p,
            "gen_max": gen_max,
            "prompt_tokens": prompt_tokens,
            "generated_tokens": gen_tokens,
            "think_tokens": ttok,
            "wall_clock_s": elapsed,
            "tokens_per_sec": gen_tokens / elapsed if elapsed > 0 else None,
            "think_tokens_per_sec": ttok / elapsed if elapsed > 0 else None,
            "max_mem_allocated_gb": mem["max_mem_allocated_gb"],
            "max_mem_reserved_gb": mem["max_mem_reserved_gb"],
            "current_mem_allocated_gb": mem["current_mem_allocated_gb"],
            "current_mem_reserved_gb": mem["current_mem_reserved_gb"],
            "n_layers": len(selected_layers) if method != "baseline_no_hook" else 0,
            "selected_layer_file": sel_path,
            "layers": selected_layers if method != "baseline_no_hook" else [],
            "pred": pred,
            "gt": gt.get(pi),
            "correct": (pred == gt.get(pi)) if pred and gt.get(pi) else False,
            "has_boxed": pred is not None,
            "has_think_end": "</think>" in cot,
            "mon_total": det["total"],
            "repetition_score": rep,
            "collapse": rep > 0.5,
            "near_32768_think": ttok >= 0.95 * gen_max,
            "problem": problem,
            "cot": cot,
        }

        fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
        fh.flush()

        log.info(
            f"{key}: method={method} pred={pred} gt={gt.get(pi)} "
            f"{'OK' if rec['correct'] else 'x'} "
            f"wall={elapsed:.1f}s gen_tok={gen_tokens} think_tok={ttok} "
            f"tok/s={rec['tokens_per_sec']:.2f} mem_alloc={mem['max_mem_allocated_gb']}GB "
            f"mon={det['total']} rep={rep:.2f}"
        )

    if fh:
        fh.close()

    records = []
    if os.path.exists(out_path):
        for line in open(out_path, encoding="utf-8"):
            if not line.strip():
                continue
            try:
                records.append(json.loads(line))
            except Exception:
                pass

    avg = lambda xs: sum(xs) / len(xs) if xs else 0.0

    summary = {
        "dataset": "GPQA-Diamond",
        "benchmark": "wall_clock_latency_throughput",
        "seed": args.seed,
        "subset_seed": args.subset_seed,
        "limit": args.limit,
        "alphas": args.alphas,
        "temperature": temperature,
        "top_p": top_p,
        "gen_max": gen_max,
        "selected_layer_file": sel_path,
        "selected_layers": selected_layers,
        "n_selected_layers": len(selected_layers),
        "subset_indices": subset_indices,
        "out_jsonl": out_path,
        "per_alpha": {},
    }

    log.info("\n=== SUMMARY: GPQA-D runtime latency / throughput ===")
    log.info(
        f"{'alpha':>6} {'n':>4} {'acc':>8} {'wall_s':>10} {'gen_tok':>10} "
        f"{'think_tok':>10} {'tok/s':>8} {'mem_alloc':>10} {'mem_resv':>10} {'collapse':>9}"
    )

    for alpha in args.alphas:
        rs = [r for r in records if abs(float(r["alpha"]) - float(alpha)) < 1e-9]
        if not rs:
            continue

        n = len(rs)
        n_correct = sum(bool(r["correct"]) for r in rs)
        collapse_rate = sum(bool(r["collapse"]) for r in rs) / n if n else 0.0

        item = {
            "n": n,
            "accuracy": n_correct / n if n else 0.0,
            "n_correct": n_correct,
            "mean_wall_clock_s": avg([r["wall_clock_s"] for r in rs]),
            "median_wall_clock_s": sorted([r["wall_clock_s"] for r in rs])[n // 2] if n else 0.0,
            "total_wall_clock_s": sum(r["wall_clock_s"] for r in rs),
            "mean_generated_tokens": avg([r["generated_tokens"] for r in rs]),
            "mean_think_tokens": avg([r["think_tokens"] for r in rs]),
            "mean_tokens_per_sec": avg([r["tokens_per_sec"] for r in rs if r["tokens_per_sec"] is not None]),
            "mean_think_tokens_per_sec": avg([r["think_tokens_per_sec"] for r in rs if r["think_tokens_per_sec"] is not None]),
            "mean_max_mem_allocated_gb": avg([r["max_mem_allocated_gb"] for r in rs if r["max_mem_allocated_gb"] is not None]),
            "mean_max_mem_reserved_gb": avg([r["max_mem_reserved_gb"] for r in rs if r["max_mem_reserved_gb"] is not None]),
            "mean_mon": avg([r["mon_total"] for r in rs]),
            "collapse_rate": collapse_rate,
            "near_32768_think_rate": sum(bool(r["near_32768_think"]) for r in rs) / n if n else 0.0,
        }

        summary["per_alpha"][str(alpha)] = item

        log.info(
            f"{alpha:>6.2f} {n:>4} {item['accuracy']:>7.1%} "
            f"{item['mean_wall_clock_s']:>10.1f} "
            f"{item['mean_generated_tokens']:>10.0f} "
            f"{item['mean_think_tokens']:>10.0f} "
            f"{item['mean_tokens_per_sec']:>8.2f} "
            f"{item['mean_max_mem_allocated_gb']:>10.2f} "
            f"{item['mean_max_mem_reserved_gb']:>10.2f} "
            f"{collapse_rate*100:>8.1f}%"
        )

    # Add relative speedup against alpha=1.
    base = summary["per_alpha"].get("1.0") or summary["per_alpha"].get("1")
    if base:
        base_wall = base["mean_wall_clock_s"]
        base_tok = base["mean_think_tokens"]
        for k, item in summary["per_alpha"].items():
            item["wall_clock_reduction_vs_alpha1"] = (
                (base_wall - item["mean_wall_clock_s"]) / base_wall if base_wall else None
            )
            item["think_token_reduction_vs_alpha1"] = (
                (base_tok - item["mean_think_tokens"]) / base_tok if base_tok else None
            )
            item["speedup_vs_alpha1"] = (
                base_wall / item["mean_wall_clock_s"] if item["mean_wall_clock_s"] else None
            )

    write_json(summary, sum_path)
    log.info(f"Saved {out_path}")
    log.info(f"Saved {sum_path}")
    log.info("Done.")


if __name__ == "__main__":
    main()