""" 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}%)")