eval_tombench_unlearn.py

"""Evaluate a Comma unlearned LoRA adapter on ToMBench and compute gamma.

The core-4 unlearning harness (``unlearning.eval_harness``) runs a fixed OLMES
task list and only aggregates suites in ``SUITE_TO_KEY``; ToMBench is neither in
that task list nor wired for gamma. This standalone wrapper closes that gap
without touching the proven core-4 path:

1. merge the LoRA adapter into a full model (reuses ``eval_harness.merge_adapter``);
2. run OLMES on the single task ``tombench:mc::socialtda`` via the team-standard
   command builder (``recipes.olmes.evaluation._build_command`` with a --task
   override), so the task config (num_shots 0, split test, acc_raw, max_length
   8192) matches the base ToMBench eval exactly;
3. read ``primary_score`` (== ``acc_raw``) from the produced metrics JSON;
4. gamma = (acc - baseline) / |baseline| against the base Comma-2T ToMBench score.

Usage:
    python scripts/unlearning/eval_tombench_unlearn.py \
        --model-id common-pile/comma-v0.1-2t \
        --adapter-dir .../adapter \
        --output-json .../cell_tombench_eval.json \
        --baseline 0.5132867 --topic-bin social_life --seed 42
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import logging
import os
import shutil
import sys
import tempfile
from pathlib import Path

logger = logging.getLogger(__name__)

DEFAULT_TASK = "tombench:mc::socialtda"


def read_primary_score(metrics_path: Path) -> float:
    """Read the primary metric (acc_raw) from an OLMES task metrics JSON."""
    data = json.loads(Path(metrics_path).read_text(encoding="utf-8"))
    metrics = data.get("metrics", {})
    for key in ("primary_score", "acc_raw", "acc_per_char"):
        if key in metrics and metrics[key] is not None:
            return float(metrics[key])
    raise ValueError(
        f"no primary score in {metrics_path}: metrics keys {list(metrics)}"
    )


def find_metrics_file(run_dir: Path, task_substr: str = "tombench") -> Path:
    """Locate the single task metrics JSON matching task_substr under run_dir."""
    matches = sorted(
        p for p in Path(run_dir).rglob("*-metrics.json") if task_substr in p.name
    )
    if not matches:
        raise FileNotFoundError(f"no '*{task_substr}*-metrics.json' under {run_dir}")
    if len(matches) > 1:
        logger.warning("multiple metrics files, using first: %s", matches)
    return matches[0]


def compute_gamma(score: float, baseline: float) -> float | None:
    """Relative score change against a baseline (matches eval_harness)."""
    if baseline == 0:
        return None
    return (score - baseline) / abs(baseline)


def _olmes_args(model_path: str, task: str, batch_size: int, model_max_length: int):
    return argparse.Namespace(
        model_id=model_path,
        model_type="vllm",
        model_max_length=model_max_length,
        batch_size=str(batch_size),
        gpus=1,
        limit=None,
        random_subsample_seed=None,
        dry_run=False,
        task=[[task]],
        suite=None,
        instruct_model=False,
        instruct_base=False,
    )


def run_tombench_olmes(
    model_path: str,
    run_dir: Path,
    task: str,
    batch_size: int,
    model_max_length: int,
) -> None:
    import subprocess

    from data_attribution.recipes.olmes.evaluation import _build_command

    run_dir.mkdir(parents=True, exist_ok=True)
    args = _olmes_args(model_path, task, batch_size, model_max_length)
    cmd = _build_command(args, run_dir)
    logger.info("Running OLMES ToMBench: %s", " ".join(cmd))
    subprocess.run(cmd, check=True)


def evaluate(args: argparse.Namespace) -> dict:
    output_json = Path(args.output_json)
    output_json.parent.mkdir(parents=True, exist_ok=True)
    run_dir = (
        output_json.parent
        / "olmes_runs"
        / (output_json.stem + "_" + dt.datetime.now(dt.UTC).strftime("%Y%m%d_%H%M%S"))
    )

    merged_dir: str | None = None
    try:
        if args.adapter_dir:
            from unlearning.eval_harness import merge_adapter

            merged_dir = tempfile.mkdtemp(prefix="merged_tom_")
            merge_adapter(args.model_id, args.adapter_dir, merged_dir)
            eval_model_path = merged_dir
        else:
            eval_model_path = args.model_id

        run_tombench_olmes(
            eval_model_path,
            run_dir,
            args.task,
            args.batch_size,
            args.model_max_length,
        )
        metrics_path = find_metrics_file(run_dir, "tombench")
        acc = read_primary_score(metrics_path)
        gamma = compute_gamma(acc, args.baseline)
        result = {
            "topic_bin": args.topic_bin,
            "seed": args.seed,
            "model_id": args.model_id,
            "adapter_dir": args.adapter_dir,
            "task": args.task,
            "tombench_acc": acc,
            "baseline": args.baseline,
            "gamma": gamma,
            "metrics_path": str(metrics_path),
        }
        output_json.write_text(json.dumps(result, indent=2), encoding="utf-8")
        logger.info(
            "  tombench  acc=%.4f  baseline=%.4f  gamma=%s -> %s",
            acc,
            args.baseline,
            f"{gamma:+.4f}" if gamma is not None else "N/A",
            output_json,
        )
        return result
    finally:
        if merged_dir and os.path.exists(merged_dir):
            shutil.rmtree(merged_dir, ignore_errors=True)


def _parse_args(argv=None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--model-id", default="common-pile/comma-v0.1-2t")
    ap.add_argument("--adapter-dir", default=None, help="omit for base baseline")
    ap.add_argument("--output-json", required=True)
    ap.add_argument("--baseline", type=float, required=True)
    ap.add_argument("--task", default=DEFAULT_TASK)
    ap.add_argument("--topic-bin", default=None)
    ap.add_argument("--seed", default=None)
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--model-max-length", type=int, default=8192)
    return ap.parse_args(argv)


def main(argv=None) -> int:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s: %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    evaluate(_parse_args(argv))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())