liqi / code /pmps /test_math_multi.py

Upload liqi files

0ab2566 verified 1 day ago

8.93 kB

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

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