File size: 5,019 Bytes
b6f7263 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | #!/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()
|