Upload code/scripts/run_ablation_eval.py with huggingface_hub

b6f7263 verified 5 months ago

5.02 kB

#!/usr/bin/env python3
"""Ablation evaluation — loads models from JSON file instead of hardcoded dict.
Usage:
    python scripts/run_ablation_eval.py --models-json ablation_eval_models.json \
        --num-gpus 8 --output-dir ablation_eval
"""

import argparse
import json
import logging
import os
import sys
import time

# Reuse everything from the comprehensive eval script
sys.path.insert(0, os.path.dirname(__file__))
from run_comprehensive_eval import (
    select_prompts,
    get_ref_audios,
    build_task_list,
    run_generation_phase,
    score_all_audio,
    build_html_report,
    VAP_DIR,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


def main():
    parser = argparse.ArgumentParser(description="Run ablation evaluation with models from JSON")
    parser.add_argument("--models-json", required=True, help="Path to JSON file with models dict")
    parser.add_argument("--output-dir", default=os.path.join(VAP_DIR, "ablation_eval"))
    parser.add_argument("--refs-dir", default="/home/deployer/laion/test-refs")
    parser.add_argument("--num-gpus", type=int, default=8)
    parser.add_argument("--seeds", type=int, nargs="+", default=[42])
    parser.add_argument("--classifiers-dir", default=os.path.join(VAP_DIR, "classifiers"))
    parser.add_argument("--skip-generation", action="store_true")
    parser.add_argument("--skip-scoring", action="store_true")
    args = parser.parse_args()

    # Load models from JSON
    with open(args.models_json) as f:
        models = json.load(f)
    logging.info(f"Loaded {len(models)} models from {args.models_json}")
    for key, minfo in models.items():
        lora_str = os.path.basename(minfo["lora"]) if minfo.get("lora") else "None"
        logging.info(f"  {key}: {minfo['name']} (lora={lora_str})")

    # Verify all LoRA files exist
    for key, minfo in models.items():
        if minfo.get("lora") and not os.path.exists(minfo["lora"]):
            logging.error(f"LoRA file not found for {key}: {minfo['lora']}")
            sys.exit(1)

    args.output_dir = os.path.abspath(args.output_dir)
    os.makedirs(os.path.join(args.output_dir, "wavs"), exist_ok=True)

    # 1. Select prompts
    prompts = select_prompts()
    logging.info(f"Selected {len(prompts)} prompts")

    with open(os.path.join(args.output_dir, "prompts.json"), "w") as f:
        json.dump(prompts, f, indent=2)

    # 2. Get reference audios
    refs = get_ref_audios(args.refs_dir)
    logging.info(f"Found {len(refs)} reference audios")

    # 3. Build task list
    tasks = build_task_list(models, prompts, refs, args.seeds, args.output_dir)
    logging.info(f"Total tasks: {len(tasks)} ({len(models)} models × {len(prompts)} prompts × ({len(refs)} refs + 1 uncond) × {len(args.seeds)} seeds)")

    # 4. Generate (skip tasks where WAV already exists)
    if not args.skip_generation:
        tasks_to_gen = [t for t in tasks if not os.path.exists(t["output"])]
        tasks_existing = [t for t in tasks if os.path.exists(t["output"])]
        logging.info(f"Skipping {len(tasks_existing)} existing WAVs, generating {len(tasks_to_gen)} new samples")
        if tasks_to_gen:
            gen_results = run_generation_phase(tasks_to_gen, args.num_gpus)
        else:
            gen_results = []
        existing_results = [{**t, "success": True, "elapsed": 0, "gpu": -1} for t in tasks_existing]
        results = existing_results + gen_results
        with open(os.path.join(args.output_dir, "gen_results.json"), "w") as f:
            json.dump(results, f, indent=2, default=str)
    else:
        with open(os.path.join(args.output_dir, "gen_results.json")) as f:
            results = json.load(f)

    # 5. Score
    if not args.skip_scoring:
        scored = score_all_audio(results, args.classifiers_dir)
        with open(os.path.join(args.output_dir, "scored_results.json"), "w") as f:
            json.dump(scored, f, indent=2, default=str)
    else:
        with open(os.path.join(args.output_dir, "scored_results.json")) as f:
            scored = json.load(f)

    # 6. Build HTML
    model_means = build_html_report(
        scored, models, prompts, refs,
        os.path.join(args.output_dir, "eval_report.html")
    )

    # Print summary
    print("\n" + "="*70)
    print("ABLATION EVALUATION SUMMARY")
    print("="*70)
    for mk, minfo in models.items():
        means = model_means.get(mk, {})
        print(f"\n{minfo['name']}:")
        for metric in ["naturalness_small", "naturalness_large", "centroid", "quality", "speaker_sim"]:
            val = means.get(metric)
            print(f"  {metric:20s}: {val:.4f}" if val else f"  {metric:20s}: N/A")
    print("="*70)

    # Save summary JSON
    summary = {mk: model_means.get(mk, {}) for mk in models}
    with open(os.path.join(args.output_dir, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2)
    logging.info(f"Summary saved to {os.path.join(args.output_dir, 'summary.json')}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Ablation evaluation — loads models from JSON file instead of hardcoded dict.

	Usage:
	python scripts/run_ablation_eval.py --models-json ablation_eval_models.json \
	--num-gpus 8 --output-dir ablation_eval
	"""

	import argparse
	import json
	import logging
	import os
	import sys
	import time

	# Reuse everything from the comprehensive eval script
	sys.path.insert(0, os.path.dirname(__file__))
	from run_comprehensive_eval import (
	select_prompts,
	get_ref_audios,
	build_task_list,
	run_generation_phase,
	score_all_audio,
	build_html_report,
	VAP_DIR,
	)

	logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


	def main():
	parser = argparse.ArgumentParser(description="Run ablation evaluation with models from JSON")
	parser.add_argument("--models-json", required=True, help="Path to JSON file with models dict")
	parser.add_argument("--output-dir", default=os.path.join(VAP_DIR, "ablation_eval"))
	parser.add_argument("--refs-dir", default="/home/deployer/laion/test-refs")
	parser.add_argument("--num-gpus", type=int, default=8)
	parser.add_argument("--seeds", type=int, nargs="+", default=[42])
	parser.add_argument("--classifiers-dir", default=os.path.join(VAP_DIR, "classifiers"))
	parser.add_argument("--skip-generation", action="store_true")
	parser.add_argument("--skip-scoring", action="store_true")
	args = parser.parse_args()

	# Load models from JSON
	with open(args.models_json) as f:
	models = json.load(f)
	logging.info(f"Loaded {len(models)} models from {args.models_json}")
	for key, minfo in models.items():
	lora_str = os.path.basename(minfo["lora"]) if minfo.get("lora") else "None"
	logging.info(f" {key}: {minfo['name']} (lora={lora_str})")

	# Verify all LoRA files exist
	for key, minfo in models.items():
	if minfo.get("lora") and not os.path.exists(minfo["lora"]):
	logging.error(f"LoRA file not found for {key}: {minfo['lora']}")
	sys.exit(1)

	args.output_dir = os.path.abspath(args.output_dir)
	os.makedirs(os.path.join(args.output_dir, "wavs"), exist_ok=True)

	# 1. Select prompts
	prompts = select_prompts()
	logging.info(f"Selected {len(prompts)} prompts")

	with open(os.path.join(args.output_dir, "prompts.json"), "w") as f:
	json.dump(prompts, f, indent=2)

	# 2. Get reference audios
	refs = get_ref_audios(args.refs_dir)
	logging.info(f"Found {len(refs)} reference audios")

	# 3. Build task list
	tasks = build_task_list(models, prompts, refs, args.seeds, args.output_dir)
	logging.info(f"Total tasks: {len(tasks)} ({len(models)} models × {len(prompts)} prompts × ({len(refs)} refs + 1 uncond) × {len(args.seeds)} seeds)")

	# 4. Generate (skip tasks where WAV already exists)
	if not args.skip_generation:
	tasks_to_gen = [t for t in tasks if not os.path.exists(t["output"])]
	tasks_existing = [t for t in tasks if os.path.exists(t["output"])]
	logging.info(f"Skipping {len(tasks_existing)} existing WAVs, generating {len(tasks_to_gen)} new samples")
	if tasks_to_gen:
	gen_results = run_generation_phase(tasks_to_gen, args.num_gpus)
	else:
	gen_results = []
	existing_results = [{**t, "success": True, "elapsed": 0, "gpu": -1} for t in tasks_existing]
	results = existing_results + gen_results
	with open(os.path.join(args.output_dir, "gen_results.json"), "w") as f:
	json.dump(results, f, indent=2, default=str)
	else:
	with open(os.path.join(args.output_dir, "gen_results.json")) as f:
	results = json.load(f)

	# 5. Score
	if not args.skip_scoring:
	scored = score_all_audio(results, args.classifiers_dir)
	with open(os.path.join(args.output_dir, "scored_results.json"), "w") as f:
	json.dump(scored, f, indent=2, default=str)
	else:
	with open(os.path.join(args.output_dir, "scored_results.json")) as f:
	scored = json.load(f)

	# 6. Build HTML
	model_means = build_html_report(
	scored, models, prompts, refs,
	os.path.join(args.output_dir, "eval_report.html")
	)

	# Print summary
	print("\n" + "="*70)
	print("ABLATION EVALUATION SUMMARY")
	print("="*70)
	for mk, minfo in models.items():
	means = model_means.get(mk, {})
	print(f"\n{minfo['name']}:")
	for metric in ["naturalness_small", "naturalness_large", "centroid", "quality", "speaker_sim"]:
	val = means.get(metric)
	print(f" {metric:20s}: {val:.4f}" if val else f" {metric:20s}: N/A")
	print("="*70)

	# Save summary JSON
	summary = {mk: model_means.get(mk, {}) for mk in models}
	with open(os.path.join(args.output_dir, "summary.json"), "w") as f:
	json.dump(summary, f, indent=2)
	logging.info(f"Summary saved to {os.path.join(args.output_dir, 'summary.json')}")


	if __name__ == "__main__":
	main()