social-data-attribution-w2/src/unlearning/eval_harness.py

# pyright: reportPrivateImportUsage=false, reportCallIssue=false, reportOptionalSubscript=false, reportArgumentType=false
"""
Evaluation harness for unlearned models using the project OLMES pipeline.

Reuses:
  data_attribution.recipes.olmes_evaluation - command builder and task list
  data_attribution.recipes.olmes_predictions_export - per-prediction stitching

Flow
----
1. Optionally merge LoRA adapter into base model weights
2. Run OLMES eval using the same command the team baseline uses
3. Stitch per-prediction rows into a JSONL next to output_json
4. Aggregate accuracy per retained suite
5. Compute gamma against retained baselines
6. Compute Wikitext-2 word perplexity via lm-evaluation-harness (hf backend)
7. Write summary JSON to --output_json

Usage
-----
# Baseline (no adapter):
python -m unlearning.eval_harness \\
    --model_id allenai/OLMo-3-1025-7B \\
    --output_json runs/unlearn/baselines.json

# After unlearning (with LoRA adapter):
python -m unlearning.eval_harness \\
    --model_id allenai/OLMo-3-1025-7B \\
    --adapter_dir runs/unlearn/entertainment/adapter \\
    --output_json runs/unlearn/entertainment/eval_results.json \\
    --topic_bin entertainment
"""

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__)

# OLMo-3-1025-7B baseline values from project OLMES eval runs.
# Accuracy metrics are higher-is-better.
# gamma = (score - baseline) / |baseline|.
BASELINES: dict[str, float] = {
    "mmlu_social_science": 0.750846,
    "mmlu_stem": 0.597871,
    "socialiqa": 0.802900,
}

SUITE_TO_KEY: dict[str, str] = {
    "mmlu_social_sciences": "mmlu_social_science",
    "mmlu_social_science": "mmlu_social_science",
    "mmlu_stem": "mmlu_stem",
    "social_iqa": "socialiqa",
    "socialiqa": "socialiqa",
}

# Wikitext-2 word PPL baseline. Lower is better.
WIKITEXT_BASELINE_PPL: float = 9.1559


# ---------------------------------------------------------------------------
# Adapter merge
# ---------------------------------------------------------------------------


def merge_adapter(model_id: str, adapter_dir: str, output_dir: str) -> None:
    """Merge LoRA adapter weights into base model and save as a full model."""
    logger.info("Merging LoRA adapter %s into %s ...", adapter_dir, output_dir)
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    base = AutoModelForCausalLM.from_pretrained(
        model_id, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="cpu"
    )
    model = PeftModel.from_pretrained(base, adapter_dir)
    merged = model.merge_and_unload()
    merged.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    logger.info("Merged model saved to %s", output_dir)
    del merged, model, base
    import gc

    gc.collect()
    torch.cuda.empty_cache()


# ---------------------------------------------------------------------------
# OLMES evaluation (delegates to PR #26 recipe)
# ---------------------------------------------------------------------------


def run_olmes(
    model_path: str,
    olmes_run_dir: Path,
    gpus: int | None = 1,
    batch_size: int | None = 16,
    model_max_length: int = 8192,
) -> None:
    """
    Run OLMES evaluation using the same command builder as the team baseline.
    Writes prediction/request/metric files under olmes_run_dir/.
    """
    from data_attribution.recipes.olmes.evaluation import _build_command

    args = argparse.Namespace(
        model_id=model_path,
        model_max_length=model_max_length,
        batch_size=str(batch_size) if batch_size is not None else None,
        gpus=gpus,
        dry_run=False,
        instruct_model=False,
    )
    olmes_run_dir.mkdir(parents=True, exist_ok=True)
    cmd = _build_command(args, olmes_run_dir)
    logger.info("Running OLMES: %s", " ".join(cmd))

    import subprocess

    subprocess.run(cmd, check=True)
    logger.info("OLMES output written to %s", olmes_run_dir)


# ---------------------------------------------------------------------------
# Per-prediction export (delegates to PR #26 recipe)
# ---------------------------------------------------------------------------


def export_predictions(
    olmes_runs_root: Path,
    eval_model_path: str,
    predictions_jsonl: Path,
) -> None:
    """
    Stitch per-prediction rows from OLMES output into a single JSONL.
    Uses olmes_predictions_export._export_model_rows() directly.
    """
    from data_attribution.recipes.olmes.predictions_export import _export_model_rows

    predictions_jsonl.parent.mkdir(parents=True, exist_ok=True)
    counts = _export_model_rows(
        olmes_runs_root, eval_model_path, predictions_jsonl, instruct_mode=False
    )
    for suite, n in sorted(counts.items()):
        logger.info("  %s: %d prediction rows", suite, n)
    logger.info("Per-prediction JSONL written to %s", predictions_jsonl)


# ---------------------------------------------------------------------------
# Aggregate scores + gamma
# ---------------------------------------------------------------------------


def _suite_key(raw_suite: object) -> str | None:
    return SUITE_TO_KEY.get(str(raw_suite))


def compute_suite_scores(predictions_jsonl: Path) -> dict[str, dict]:
    """
    Read the stitched per-prediction JSONL and compute accuracy per suite.

    Aggregation:
      - MMLU suites: macro-average (mean of per-task accuracies) to match
        the standard MMLU evaluation protocol used in the baseline.
      - All other retained suites: micro-average.

    Returns {suite_key: {"correct": int, "total": int, "accuracy": float}}.
    """
    from collections import defaultdict

    # per-task tallies for macro-avg suites
    task_correct: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    task_total: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    # flat tallies for micro-avg suites
    correct: dict[str, int] = defaultdict(int)
    total: dict[str, int] = defaultdict(int)

    with predictions_jsonl.open(encoding="utf-8") as fh:
        for line in fh:
            if not line.strip():
                continue
            row = json.loads(line)
            suite = _suite_key(row.get("task_suite", ""))
            if not suite:
                continue
            task = row.get("task_name", suite)
            hit = int(bool(row.get("is_correct")))
            if "mmlu" in suite.lower():
                task_correct[suite][task] += hit
                task_total[suite][task] += 1
            else:
                correct[suite] += hit
                total[suite] += 1

    results: dict[str, dict] = {}

    # macro-average for MMLU
    for suite in task_total:
        per_task_acc = [
            task_correct[suite][t] / task_total[suite][t]
            for t in task_total[suite]
            if task_total[suite][t] > 0
        ]
        n_total = sum(task_total[suite].values())
        n_correct = sum(task_correct[suite].values())
        results[suite] = {
            "correct": n_correct,
            "total": n_total,
            "accuracy": (
                sum(per_task_acc) / len(per_task_acc) if per_task_acc else float("nan")
            ),
        }

    # micro-average for everything else
    for suite in total:
        results[suite] = {
            "correct": correct[suite],
            "total": total[suite],
            "accuracy": correct[suite] / total[suite] if total[suite] else float("nan"),
        }

    return results


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


def _add_wikitext_metric(
    metrics: dict[str, dict],
    eval_model_path: str,
    *,
    limit: int | None,
    fast_eval: bool = False,
) -> None:
    wikitext_ppl = run_wikitext_ppl(
        eval_model_path,
        limit=limit,
    )
    wikitext_baseline = WIKITEXT_BASELINE_PPL
    wikitext_gamma = compute_gamma(wikitext_ppl, wikitext_baseline)
    metrics["wikitext"] = {
        "word_perplexity": wikitext_ppl,
        "baseline": wikitext_baseline,
        "gamma": wikitext_gamma,
        "lower_is_better": True,
    }
    if fast_eval:
        metrics["wikitext"]["fast_eval"] = True
    logger.info(
        "  %-22s  ppl=%.4f  baseline=%.4f  gamma=%s",
        "wikitext",
        wikitext_ppl,
        wikitext_baseline,
        f"{wikitext_gamma:+.4f}" if wikitext_gamma is not None else "N/A",
    )


# ---------------------------------------------------------------------------
# Wikitext PPL (lm-evaluation-harness, hf backend)
# ---------------------------------------------------------------------------


def run_wikitext_ppl(
    model_path: str, batch_size: int = 4, limit: int | None = None
) -> float:
    """
    Compute Wikitext-2 word perplexity using lm-evaluation-harness (hf backend).
    Returns word_perplexity (lower is better).
    Loads the model independently from OLMES (which uses vLLM via subprocess).
    """
    import lm_eval
    import os

    os.environ["HF_DATASETS_TRUST_REMOTE_CODE"] = "1"

    logger.info("Running Wikitext-2 word PPL on %s ...", model_path)
    results = lm_eval.simple_evaluate(
        model="hf",
        model_args=f"pretrained={model_path},dtype=bfloat16,trust_remote_code=True,max_length=2048",
        tasks=["wikitext"],
        batch_size=1,
        log_samples=False,
        limit=limit,
    )
    wikitext_results = results["results"]["wikitext"]
    # lm_eval stores metrics as "metric_name,filter_name"
    ppl = wikitext_results.get("word_perplexity,none") or wikitext_results.get(
        "word_perplexity"
    )
    if ppl is None:
        raise ValueError(
            f"word_perplexity not found in wikitext results: {wikitext_results}"
        )
    logger.info("Wikitext word PPL: %.4f", ppl)
    return float(ppl)


# ---------------------------------------------------------------------------
# Fast eval (lm-eval hf backend, no OLMES/vLLM, supports --limit)
# ---------------------------------------------------------------------------

# Map from BASELINES keys to lm-eval task names
_FAST_EVAL_TASKS: dict[str, str] = {
    "mmlu_social_science": "mmlu_social_sciences",
    "mmlu_stem": "mmlu_stem",
    "socialiqa": "social_iqa",
}


def run_fast_eval(
    model_path: str,
    num_samples: int = 200,
    batch_size: int = 4,
    tasks_filter: list | None = None,
) -> dict[str, float]:
    """
    Fast benchmark eval using lm-eval hf backend directly (no OLMES/vLLM subprocess).
    Samples up to num_samples examples per task.
    Returns {baseline_key: accuracy} mapping compatible with BASELINES dict.
    """
    import lm_eval
    import os

    os.environ["HF_DATASETS_TRUST_REMOTE_CODE"] = "1"

    task_map = {
        k: v
        for k, v in _FAST_EVAL_TASKS.items()
        if tasks_filter is None or k in tasks_filter
    }
    tasks = list(task_map.values())
    if not tasks:
        logger.warning("No supported fast-eval tasks selected.")
        return {}
    logger.info(
        "Running fast eval on %s (limit=%d per task) ...",
        model_path,
        num_samples,
    )
    results = lm_eval.simple_evaluate(
        model="hf",
        model_args=f"pretrained={model_path},dtype=bfloat16,trust_remote_code=True",
        tasks=tasks,
        limit=num_samples,
        batch_size=batch_size,
        log_samples=False,
    )

    scores: dict[str, float] = {}
    for baseline_key, lm_task in task_map.items():
        task_results = results["results"].get(lm_task, {})
        acc = (
            task_results.get("acc,none")
            or task_results.get("exact_match,none")
            or task_results.get("exact_match,strict-match")
            or task_results.get("exact_match,flexible-extract")
            or task_results.get("acc")
        )
        if acc is not None:
            scores[baseline_key] = float(acc)
            logger.info("  %-22s  acc=%.4f (n=%d)", baseline_key, acc, num_samples)
        else:
            logger.warning(
                "  %-22s  no accuracy found in: %s", baseline_key, task_results
            )
    return scores


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _parse_args(argv=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate model with OLMES (PR #26 pipeline) + Wikitext PPL"
    )
    parser.add_argument("--model_id", default="allenai/OLMo-3-1025-7B")
    parser.add_argument(
        "--adapter_dir", default=None, help="LoRA adapter dir; omit for baseline eval"
    )
    parser.add_argument("--output_json", required=True, help="Summary JSON output path")
    parser.add_argument(
        "--topic_bin", default=None, help="Topic bin that was unlearned (metadata only)"
    )
    parser.add_argument("--gpus", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--model_max_length", type=int, default=8192)
    parser.add_argument(
        "--skip_wikitext",
        action="store_true",
        help="Skip Wikitext PPL eval (faster, for debugging)",
    )
    parser.add_argument(
        "--wikitext_only",
        action="store_true",
        help="Run only Wikitext PPL (skip all lm-eval benchmarks). Avoids GPU OOM.",
    )
    parser.add_argument(
        "--fast_eval",
        action="store_true",
        help="Use lm-eval hf backend directly instead of OLMES/vLLM. "
        "Much faster; use with --fast_eval_samples.",
    )
    parser.add_argument(
        "--fast_eval_samples",
        type=int,
        default=200,
        help="Examples per task in fast eval mode (default: 200)",
    )
    parser.add_argument(
        "--fast_eval_tasks",
        nargs="+",
        default=None,
        help="Subset of tasks in fast eval (e.g. socialiqa mmlu_stem). Default: all.",
    )
    parser.add_argument(
        "--wikitext_samples",
        type=int,
        default=None,
        help="Limit wikitext PPL eval to N chunks (None = full set).",
    )
    # Kept for backward compatibility with existing wrappers; not used for gamma.
    parser.add_argument(
        "--baseline_json",
        default=None,
        help="(unused) kept for compatibility with existing wrapper commands",
    )
    return parser.parse_args(argv)


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

    output_json = Path(args.output_json)
    output_json.parent.mkdir(parents=True, exist_ok=True)

    olmes_runs_root = output_json.parent / "olmes_runs"
    stamp = dt.datetime.now(dt.UTC).strftime("%Y%m%d_%H%M%S")

    merged_dir: str | None = None
    try:
        # ------------------------------------------------------------------ #
        # 1. Merge adapter (if provided)                                       #
        # ------------------------------------------------------------------ #
        if args.adapter_dir is not None:
            merged_dir = tempfile.mkdtemp(prefix="merged_model_")
            merge_adapter(args.model_id, args.adapter_dir, merged_dir)
            eval_model_path = merged_dir
        else:
            eval_model_path = args.model_id

        metrics: dict[str, dict] = {}

        if args.fast_eval:
            # -------------------------------------------------------------- #
            # Fast path: lm-eval hf backend, subsampled (no OLMES/vLLM)      #
            # -------------------------------------------------------------- #
            if not getattr(args, "wikitext_only", False):
                fast_scores = run_fast_eval(
                    eval_model_path,
                    num_samples=args.fast_eval_samples,
                    batch_size=args.batch_size,
                    tasks_filter=getattr(args, "fast_eval_tasks", None),
                )
            else:
                fast_scores = {}
            for suite, acc in sorted(fast_scores.items()):
                baseline = BASELINES.get(suite)
                gamma = compute_gamma(acc, baseline) if baseline is not None else None
                metrics[suite] = {
                    "accuracy": acc,
                    "total": args.fast_eval_samples,
                    "baseline": baseline,
                    "gamma": gamma,
                    "lower_is_better": False,
                    "fast_eval": True,
                }
                logger.info(
                    "  %-22s  acc=%.4f  baseline=%s  gamma=%s",
                    suite,
                    acc,
                    f"{baseline:.4f}" if baseline is not None else "N/A",
                    f"{gamma:+.4f}" if gamma is not None else "N/A",
                )
            if not args.skip_wikitext:
                try:
                    _add_wikitext_metric(
                        metrics,
                        eval_model_path,
                        limit=getattr(args, "wikitext_samples", None),
                        fast_eval=True,
                    )
                except Exception as e:
                    logger.warning("Wikitext PPL failed: %s", e)
            output = {
                "topic_bin": args.topic_bin,
                "model_id": args.model_id,
                "adapter_dir": args.adapter_dir,
                "fast_eval": True,
                "fast_eval_samples": args.fast_eval_samples,
                "metrics": metrics,
            }
        else:
            # -------------------------------------------------------------- #
            # Full path: OLMES/vLLM + per-prediction export                   #
            # -------------------------------------------------------------- #
            safe_model = (
                eval_model_path.replace("/", "-").replace(os.sep, "-").strip("-")
            )
            olmes_run_dir = olmes_runs_root / f"{safe_model}_{stamp}"
            predictions_jsonl = output_json.parent / (
                output_json.stem + "_predictions.jsonl"
            )

            # ---------------------------------------------------------------- #
            # 2. Run OLMES eval                                                 #
            # ---------------------------------------------------------------- #
            if not args.wikitext_only:
                run_olmes(
                    model_path=eval_model_path,
                    olmes_run_dir=olmes_run_dir,
                    gpus=args.gpus,
                    batch_size=args.batch_size,
                    model_max_length=args.model_max_length,
                )

                # ---------------------------------------------------------------- #
                # 3. Stitch per-prediction rows (PR #26 schema)                    #
                # ---------------------------------------------------------------- #
                export_predictions(olmes_runs_root, eval_model_path, predictions_jsonl)

                # ---------------------------------------------------------------- #
                # 4. Aggregate accuracy per suite                                   #
                # ---------------------------------------------------------------- #
                suite_scores = compute_suite_scores(predictions_jsonl)

                # ---------------------------------------------------------------- #
                # 5. Compute gamma for accuracy metrics                             #
                # ---------------------------------------------------------------- #
                for suite, stats in sorted(suite_scores.items()):
                    baseline = BASELINES.get(suite)
                    if baseline is None:
                        logger.info(
                            "Skipping unsupported unlearning eval suite: %s", suite
                        )
                        continue
                    gamma = compute_gamma(stats["accuracy"], baseline)
                    metrics[suite] = {
                        "accuracy": stats["accuracy"],
                        "correct": stats["correct"],
                        "total": stats["total"],
                        "baseline": baseline,
                        "gamma": gamma,
                        "lower_is_better": False,
                    }
                    logger.info(
                        "  %-22s  acc=%.4f  baseline=%.4f  gamma=%s",
                        suite,
                        stats["accuracy"],
                        baseline,
                        f"{gamma:+.4f}" if gamma is not None else "N/A",
                    )

            # ---------------------------------------------------------------- #
            # 6. Wikitext PPL                                                   #
            # ---------------------------------------------------------------- #
            if not args.skip_wikitext:
                try:
                    _add_wikitext_metric(
                        metrics,
                        eval_model_path,
                        limit=getattr(args, "wikitext_samples", None),
                    )
                except Exception as exc:
                    logger.warning("Wikitext PPL eval failed (non-fatal): %s", exc)
                    metrics["wikitext"] = {"error": str(exc)}

            output = {
                "topic_bin": args.topic_bin,
                "model_id": args.model_id,
                "adapter_dir": args.adapter_dir,
                "wikitext_only": args.wikitext_only,
                "metrics": metrics,
            }
            if not args.wikitext_only:
                output["olmes_run_dir"] = str(olmes_run_dir)
                output["predictions_jsonl"] = str(predictions_jsonl)
        with output_json.open("w", encoding="utf-8") as fh:
            json.dump(output, fh, indent=2)
        logger.info("Summary written to %s", output_json)

    finally:
        if merged_dir and os.path.exists(merged_dir):
            logger.info("Cleaning up merged model at %s", merged_dir)
            shutil.rmtree(merged_dir, ignore_errors=True)

    return 0


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