Blame for code/pmps/test_math_multi.py

"""
Multi-folder MATH evaluation script for ablation studies.

Evaluates checkpoints across multiple experiment directories.
Each experiment runs in a separate subprocess to ensure clean CUDA state.

Usage:
    # Run all experiments:
    python test_math_multi.py --config_file=./args/gsm_codi_pmps_math_test_multi.yaml

    # Run a single experiment (used internally by subprocess):
    python test_math_multi.py --config_file=./args/gsm_codi_pmps_math_test_multi.yaml --single_exp=0
"""

import os
import glob
import re
import sys
import json
import argparse
import subprocess
import yaml
import torch
import transformers

from test_math import evaluation, write_json
from src.model import ModelArguments, DataArguments, TrainingArguments


DEFAULT_LAST_K = 10


def discover_checkpoints(ckpt_base_dir, last_k=DEFAULT_LAST_K):
    """Auto-discover checkpoint-* subdirectories, return the last `last_k` sorted by step number."""
    pattern = os.path.join(ckpt_base_dir, "checkpoint-*")
    ckpt_ids = []
    for d in glob.glob(pattern):
        if os.path.isdir(d):
            match = re.search(r'checkpoint-(\d+)$', d)
            if match:
                ckpt_ids.append(int(match.group(1)))
    ckpt_ids.sort()
    if last_k is not None and last_k > 0:
        ckpt_ids = ckpt_ids[-last_k:]
    return [str(x) for x in ckpt_ids]


def expand_grid(grid_config):
    """Expand a parameter grid into a list of experiment dicts."""
    base_dir = grid_config['base_dir']
    ckpt_subdir = grid_config.get('ckpt_subdir', 'gsm8k_gpt_pmps')
    output_subdir = grid_config.get('output_subdir', 'output_math')
    folder_template = grid_config.get(
        'folder_template',
        'k{k}_t0.1_mt0.1_fac0.5_l{l}_lr3e3_btz256_ep300'
    )
    latent_iters_list = grid_config['inf_latent_iterations']
    multipliers = grid_config['pmps_nmb_prototypes_multipliers']

    experiments = []
    for l in latent_iters_list:
        for m in multipliers:
            k = l * m
            folder = folder_template.format(k=k, l=l)
            experiments.append({
                'name': f"k{k}_l{l}",
                'ckpt_dir': os.path.join(base_dir, folder, ckpt_subdir, ''),
                'output_dir': os.path.join(base_dir, folder, output_subdir),
                'pmps_nmb_prototypes': k,
                'inf_latent_iterations': l,
            })
    return experiments


def load_experiments(config_file):
    """Load config and return (shared_config, experiments)."""
    with open(config_file, 'r') as f:
        config = yaml.safe_load(f)

    grid_config = config.pop('grid', None)
    experiments = config.pop('experiments', [])

    if grid_config:
        grid_experiments = expand_grid(grid_config)
        experiments = grid_experiments + experiments

    return config, experiments


def build_cmd_args(config_dict):
    """Convert a flat config dict into command-line arg list for HfArgumentParser."""
    cmd_args = []
    for key, value in config_dict.items():
        if value is None:
            continue
        if isinstance(value, bool):
            cmd_args.extend([f"--{key}", "True" if value else "False"])
        else:
            cmd_args.extend([f"--{key}", str(value)])
    return cmd_args


def run_single_experiment(shared_config, exp):
    """Run evaluation for a single experiment. Called in a clean subprocess."""
    ckpt_base_dir = exp['ckpt_dir']
    output_dir = exp['output_dir']
    ckpt_list = exp.get('ckpt_list', None)
    last_k = exp.get('last_k', DEFAULT_LAST_K)

    if not ckpt_list:
        ckpt_list = discover_checkpoints(ckpt_base_dir, last_k=last_k)
        if not ckpt_list:
            print(f"[WARNING] No checkpoints found in {ckpt_base_dir}, skipping.")
            return None
    else:
        ckpt_list = [str(x) for x in ckpt_list]

    exp_name = exp.get('name', os.path.basename(os.path.dirname(ckpt_base_dir.rstrip('/'))))

    # Merge shared config with experiment-specific overrides
    exp_config = dict(shared_config)
    exp_config['ckpt_dir'] = ckpt_base_dir
    exp_config['output_dir'] = output_dir
    if 'pmps_nmb_prototypes' in exp:
        exp_config['pmps_nmb_prototypes'] = exp['pmps_nmb_prototypes']
    if 'inf_latent_iterations' in exp:
        exp_config['inf_latent_iterations'] = exp['inf_latent_iterations']

    cmd_args = build_cmd_args(exp_config)
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args)

    print(f"\n{'#'*70}")
    print(f"# Experiment: {exp_name}")
    print(f"# ckpt_dir:  {ckpt_base_dir}")
    print(f"# output_dir: {output_dir}")
    print(f"# pmps_nmb_prototypes: {training_args.pmps_nmb_prototypes}")
    print(f"# inf_latent_iterations: {training_args.inf_latent_iterations}")
    print(f"# checkpoints ({len(ckpt_list)}): {ckpt_list}")
    print(f"{'#'*70}")

    all_ckpt_results = {}
    base_dir = model_args.ckpt_dir

    for ckpt_id in ckpt_list:
        print(f"\n{'='*60}")
        print(f"[{exp_name}] Evaluating checkpoint: {ckpt_id}")
        print(f"{'='*60}")

        accu_list = []
        for _ in range(training_args.inf_num_iterations):
            model_args.ckpt_dir = os.path.join(base_dir, f"checkpoint-{ckpt_id}")
            accu = evaluation(model_args, data_args, training_args)
            accu_list.append(accu)
            torch.cuda.empty_cache()

        avg_accu = sum(accu_list) / len(accu_list)
        all_ckpt_results[ckpt_id] = avg_accu
        print(f"[{exp_name}] Checkpoint {ckpt_id} | Average accuracy: {avg_accu:.2f}%")

    os.makedirs(output_dir, exist_ok=True)
    write_json(all_ckpt_results, f"{output_dir}/all_ckpt_results.json")

    print(f"\n{'='*60}")
    print(f"Summary for {exp_name}:")
    for ckpt_id, avg_accu in all_ckpt_results.items():
        print(f"  checkpoint-{ckpt_id}: {avg_accu:.2f}%")

    if all_ckpt_results:
        best_ckpt = max(all_ckpt_results, key=all_ckpt_results.get)
        best_accu = all_ckpt_results[best_ckpt]
        print(f"  Best: checkpoint-{best_ckpt} ({best_accu:.2f}%)")
        return {
            "name": exp_name,
            "ckpt_dir": ckpt_base_dir,
            "output_dir": output_dir,
            "best_checkpoint": best_ckpt,
            "best_accuracy": best_accu,
            "all_results": all_ckpt_results,
        }
    return None


if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(description="Multi-folder MATH evaluation")
    arg_parser.add_argument("--config_file", type=str, required=True)
    arg_parser.add_argument("--single_exp", type=int, default=None,
                            help="Run only the i-th experiment (used by subprocess)")
    args = arg_parser.parse_args()

    shared_config, experiments = load_experiments(args.config_file)

    if not experiments:
        raise ValueError("Config must contain a 'grid' or 'experiments' section.")

    # ---- Single-experiment mode (subprocess) ----
    if args.single_exp is not None:
        idx = args.single_exp
        if idx < 0 or idx >= len(experiments):
            raise ValueError(f"--single_exp={idx} out of range [0, {len(experiments)})")
        result = run_single_experiment(shared_config, experiments[idx])
        if result:
            # Write result to a temp file for the master process to collect
            result_file = f"/tmp/_math_eval_result_{idx}.json"
            with open(result_file, 'w') as f:
                json.dump(result, f)
        sys.exit(0)

    # ---- Master mode: spawn one subprocess per experiment ----
    print(f"Total experiments: {len(experiments)}")
    print("Each experiment runs in a separate subprocess for CUDA isolation.\n")

    global_summary = {}

    for i, exp in enumerate(experiments):
        exp_name = exp.get('name', f"exp_{i}")
        print(f"\n{'*'*70}")
        print(f"* Launching experiment [{i+1}/{len(experiments)}]: {exp_name}")
        print(f"{'*'*70}")

        result_file = f"/tmp/_math_eval_result_{i}.json"
        if os.path.exists(result_file):
            os.remove(result_file)

        ret = subprocess.run(
            [sys.executable, __file__, '--config_file', args.config_file, '--single_exp', str(i)],
        )

        if ret.returncode != 0:
            print(f"[ERROR] Experiment {exp_name} failed (exit code {ret.returncode}), continuing...")
            continue

        if os.path.exists(result_file):
            with open(result_file, 'r') as f:
                result = json.load(f)
            global_summary[exp_name] = result
            os.remove(result_file)

    # ---- Global summary ----
    print(f"\n{'#'*70}")
    print(f"# GLOBAL SUMMARY ({len(global_summary)}/{len(experiments)} succeeded)")
    print(f"{'#'*70}")
    for exp_name, info in global_summary.items():
        print(f"  {exp_name}: best=checkpoint-{info['best_checkpoint']} ({info['best_accuracy']:.2f}%)")