laionbox-ablation-checkpoints / code /scripts /run_ablation_overnight.py

Upload code/scripts/run_ablation_overnight.py with huggingface_hub

8018728 verified 5 months ago

15 kB

#!/usr/bin/env python3
"""Overnight ablation study runner.
Runs 4 ablation training experiments sequentially, then evaluates all models
in a single comprehensive HTML report.
Ablations (all from LaionBox v0.1-wip, 2 epochs, same hyperparameters):
  A: naturalness + quality MLP
  B: naturalness + centroid
  C: naturalness + quality MLP + speaker similarity
  D: naturalness + speaker similarity
After training, generates audio for the best-flow and best-naturalness
checkpoint from each ablation, scores all models (including baselines),
and produces a combined HTML comparison.
"""

import json
import os
import signal
import subprocess
import sys
import time

PYTHON = "/home/deployer/miniconda3/envs/ml-general/bin/python"
VAP_DIR = "/home/deployer/laion/Voice-Acting-Pipeline"
TRAIN_SCRIPT = os.path.join(VAP_DIR, "scripts", "dramabox_finetune_train_multi_aux.py")
EVAL_SCRIPT = os.path.join(VAP_DIR, "scripts", "run_ablation_eval.py")

ABLATIONS = [
    {
        "name": "A_nat_quality",
        "config": "configs/ablation_nat_quality.yaml",
        "output_dir": "finetune_output/ablation_nat_quality",
        "desc_prefix": "Nat+Quality",
        "losses": "naturalness + quality_mlp",
    },
    {
        "name": "B_nat_centroid",
        "config": "configs/ablation_nat_centroid.yaml",
        "output_dir": "finetune_output/ablation_nat_centroid",
        "desc_prefix": "Nat+Centroid",
        "losses": "naturalness + centroid",
    },
    {
        "name": "C_nat_quality_speaker",
        "config": "configs/ablation_nat_quality_speaker.yaml",
        "output_dir": "finetune_output/ablation_nat_quality_speaker",
        "desc_prefix": "Nat+Quality+Speaker",
        "losses": "naturalness + quality_mlp + speaker_sim",
    },
    {
        "name": "D_nat_speaker",
        "config": "configs/ablation_nat_speaker.yaml",
        "output_dir": "finetune_output/ablation_nat_speaker",
        "desc_prefix": "Nat+Speaker",
        "losses": "naturalness + speaker_sim",
    },
]


def log(msg):
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)


def find_best_checkpoints(output_dir):
    """Parse metrics.jsonl to find best flow loss and best naturalness checkpoints."""
    metrics_path = os.path.join(VAP_DIR, output_dir, "metrics.jsonl")
    if not os.path.exists(metrics_path):
        log(f"  WARNING: {metrics_path} not found")
        return None, None

    best_flow_loss = float("inf")
    best_flow_step = None
    best_nat = float("-inf")
    best_nat_step = None

    with open(metrics_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                m = json.loads(line)
            except json.JSONDecodeError:
                continue
            step = m.get("step", 0)
            flow = m.get("flow_loss", m.get("loss", float("inf")))
            nat = m.get("naturalness_reward", float("-inf"))

            if flow < best_flow_loss:
                best_flow_loss = flow
                best_flow_step = step
            if nat is not None and nat > best_nat:
                best_nat = nat
                best_nat_step = step

    log(f"  Best flow: step {best_flow_step} (loss={best_flow_loss:.4f})")
    log(f"  Best nat:  step {best_nat_step} (nat={best_nat:.4f})")

    # Find closest saved checkpoint (save_every=10)
    def find_ckpt(step):
        if step is None:
            return None, None
        # Try exact step first, then nearby
        for offset in [0, -1, 1, -2, 2, -3, 3, -4, 4, -5, 5]:
            candidate_step = (step + offset * 10) if offset != 0 else step
            # Try step checkpoint
            fname = f"lora_step_{candidate_step:05d}.safetensors"
            path = os.path.join(VAP_DIR, output_dir, fname)
            if os.path.exists(path):
                return path, candidate_step
            # Also try without zero padding
            fname = f"lora_step_{candidate_step:05d}.safetensors"
            path = os.path.join(VAP_DIR, output_dir, fname)
            if os.path.exists(path):
                return path, candidate_step
        # Try epoch checkpoints
        for epoch in [1, 2]:
            path = os.path.join(VAP_DIR, output_dir, f"lora_epoch{epoch}.safetensors")
            if os.path.exists(path):
                return path, f"epoch{epoch}"
        return None, None

    # Round step to nearest save_every=10
    def snap_step(step):
        if step is None:
            return None
        return round(step / 10) * 10

    flow_path, flow_step = find_ckpt(snap_step(best_flow_step))
    nat_path, nat_step = find_ckpt(snap_step(best_nat_step))

    if flow_path:
        log(f"  Flow checkpoint: {os.path.basename(flow_path)}")
    else:
        log(f"  WARNING: No flow checkpoint found near step {best_flow_step}")

    if nat_path:
        log(f"  Nat checkpoint:  {os.path.basename(nat_path)}")
    else:
        log(f"  WARNING: No nat checkpoint found near step {best_nat_step}")

    return (flow_path, flow_step, best_flow_loss), (nat_path, nat_step, best_nat)


def run_training(ablation):
    """Run a single ablation training."""
    name = ablation["name"]
    config = ablation["config"]
    log_file = f"/tmp/ablation_{name}.log"

    log(f"{'='*70}")
    log(f"STARTING ABLATION: {name} ({ablation['losses']})")
    log(f"Config: {config}")
    log(f"Log: {log_file}")
    log(f"{'='*70}")

    cmd = [
        "env", "-u", "LD_LIBRARY_PATH",
        "accelerate", "launch", "--num_processes=8",
        TRAIN_SCRIPT,
        "--config", config,
    ]

    env = os.environ.copy()
    env.pop("LD_LIBRARY_PATH", None)
    # Ensure accelerate from ml-general is used
    env["PATH"] = "/home/deployer/miniconda3/envs/ml-general/bin:" + env.get("PATH", "")

    with open(log_file, "w") as lf:
        proc = subprocess.Popen(
            cmd, stdout=lf, stderr=subprocess.STDOUT,
            cwd=VAP_DIR, env=env,
        )

    log(f"Training PID: {proc.pid}")

    # Wait for completion
    t0 = time.time()
    while True:
        ret = proc.poll()
        if ret is not None:
            break
        elapsed = time.time() - t0
        # Print progress every 5 minutes
        if int(elapsed) % 300 == 0 and elapsed > 10:
            # Read last line of status.json if it exists
            status_path = os.path.join(VAP_DIR, ablation["output_dir"], "status.json")
            if os.path.exists(status_path):
                try:
                    with open(status_path) as sf:
                        status = json.load(sf)
                    step = status.get("step", "?")
                    total = status.get("total_steps", "?")
                    loss = status.get("flow_loss", status.get("loss", "?"))
                    eta = status.get("eta_sec", 0)
                    eta_m = int(eta // 60) if isinstance(eta, (int, float)) else "?"
                    log(f"  Progress: step {step}/{total}, loss={loss}, ETA={eta_m}m")
                except Exception:
                    pass
        time.sleep(10)

    elapsed_min = (time.time() - t0) / 60
    if ret == 0:
        log(f"Training {name} COMPLETED in {elapsed_min:.1f} min (exit code 0)")
    else:
        log(f"Training {name} FAILED with exit code {ret} after {elapsed_min:.1f} min")
        log(f"  Check log: {log_file}")

    return ret == 0


def write_eval_models_json(all_models, path):
    """Write the models dict as JSON for the eval script."""
    with open(path, "w") as f:
        json.dump(all_models, f, indent=2)
    log(f"Wrote {len(all_models)} models to {path}")


def upload_to_hf(abl_name, output_dir, flow_info, nat_info):
    """Upload best checkpoints + metrics to HF."""
    try:
        from huggingface_hub import HfApi
        api = HfApi(token="HF_TOKEN_REDACTED")
        repo = "TTS-AGI/laionbox-ablation-checkpoints"

        uploads = []
        metrics_path = os.path.join(VAP_DIR, output_dir, "metrics.jsonl")
        if os.path.exists(metrics_path):
            uploads.append((metrics_path, f"{abl_name}/metrics.jsonl"))

        if flow_info and flow_info[0]:
            path, step, _ = flow_info
            uploads.append((path, f"{abl_name}/best_flow_step{step}.safetensors"))
        if nat_info and nat_info[0]:
            path, step, _ = nat_info
            uploads.append((path, f"{abl_name}/best_nat_step{step}.safetensors"))

        for local, remote in uploads:
            if os.path.exists(local):
                sz = os.path.getsize(local) / 1024 / 1024
                log(f"  HF upload: {remote} ({sz:.0f} MB)")
                api.upload_file(path_or_fileobj=local, path_in_repo=remote, repo_id=repo)
        log(f"  HF upload complete for {abl_name}")
    except Exception as e:
        log(f"  HF upload FAILED for {abl_name}: {e}")


def run_eval(models_json_path, output_dir):
    """Run the comprehensive evaluation."""
    log(f"{'='*70}")
    log("STARTING COMPREHENSIVE EVALUATION")
    log(f"{'='*70}")

    log_file = "/tmp/ablation_eval.log"

    cmd = [
        "env", "-u", "LD_LIBRARY_PATH",
        PYTHON, EVAL_SCRIPT,
        "--models-json", models_json_path,
        "--num-gpus", "8",
        "--output-dir", output_dir,
    ]

    env = os.environ.copy()
    env.pop("LD_LIBRARY_PATH", None)

    with open(log_file, "w") as lf:
        proc = subprocess.Popen(
            cmd, stdout=lf, stderr=subprocess.STDOUT,
            cwd=VAP_DIR, env=env,
        )

    log(f"Eval PID: {proc.pid}, log: {log_file}")

    proc.wait()
    elapsed_min = proc.returncode
    if proc.returncode == 0:
        log("Evaluation COMPLETED successfully")
    else:
        log(f"Evaluation FAILED with exit code {proc.returncode}")
        log(f"  Check log: {log_file}")

    return proc.returncode == 0


def main():
    log("="*70)
    log("ABLATION STUDY: Auxiliary Loss Comparison")
    log("4 training runs + comprehensive evaluation")
    log("="*70)
    log("")
    log("Ablations:")
    for i, abl in enumerate(ABLATIONS):
        log(f"  {abl['name']}: {abl['losses']}")
    log("")

    # Phase 1: Training
    results = {}
    for abl in ABLATIONS:
        # Check if this ablation already completed (metrics.jsonl has >= 20 entries)
        metrics_path = os.path.join(VAP_DIR, abl["output_dir"], "metrics.jsonl")
        if os.path.exists(metrics_path):
            with open(metrics_path) as f:
                n_lines = sum(1 for _ in f)
            if n_lines >= 20:
                log(f"SKIPPING {abl['name']}: already has {n_lines} metric entries")
                log(f"Finding best checkpoints for {abl['name']}...")
                flow_info, nat_info = find_best_checkpoints(abl["output_dir"])
                results[abl["name"]] = {
                    "ablation": abl,
                    "flow": flow_info,
                    "nat": nat_info,
                }
                log("")
                continue

        success = run_training(abl)
        if success:
            log(f"\nFinding best checkpoints for {abl['name']}...")
            flow_info, nat_info = find_best_checkpoints(abl["output_dir"])
            results[abl["name"]] = {
                "ablation": abl,
                "flow": flow_info,  # (path, step, loss)
                "nat": nat_info,    # (path, step, nat_score)
            }
            # Upload to HF
            log(f"Uploading {abl['name']} checkpoints to HF...")
            upload_to_hf(abl["name"], abl["output_dir"], flow_info, nat_info)
        else:
            log(f"\nWARNING: {abl['name']} failed, skipping checkpoint extraction")
            results[abl["name"]] = None
        log("")

    # Phase 2: Build models dict for evaluation
    log("\n" + "="*70)
    log("TRAINING PHASE COMPLETE — BUILDING EVALUATION")
    log("="*70)

    models = {
        "vanilla": {
            "name": "Vanilla DramaBox",
            "lora": None,
            "desc": "Base DramaBox model without any fine-tuning",
        },
        "laionbox_v01": {
            "name": "LaionBox v0.1-wip",
            "lora": "/home/deployer/.cache/huggingface/hub/models--laion--laionbox-v0.1-wip/snapshots/66176d2a653a013a7b71c1ccb7a7a4d4cf514b0d/lora_epoch5.safetensors",
            "desc": "Previous best LoRA (5-epoch diff reward, DramaBox+Emolia data)",
        },
        "nat_only_best_flow": {
            "name": "Nat-Only Best-Flow (s160)",
            "lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00160.safetensors"),
            "desc": "Naturalness-only (CLAP-7B), best flow=0.528 @s160",
        },
        "nat_only_best_nat": {
            "name": "Nat-Only Best-Nat (s190)",
            "lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00190.safetensors"),
            "desc": "Naturalness-only (CLAP-7B), best nat=0.111 @s190",
        },
    }

    for abl_name, res in results.items():
        if res is None:
            continue
        abl = res["ablation"]
        prefix = abl["desc_prefix"]
        losses_short = abl["losses"]

        if res["flow"] and res["flow"][0]:
            path, step, loss = res["flow"]
            key = f"{abl_name}_best_flow"
            models[key] = {
                "name": f"{prefix} Best-Flow (s{step})",
                "lora": path,
                "desc": f"{losses_short}, best flow={loss:.4f} @s{step}",
            }

        if res["nat"] and res["nat"][0]:
            path, step, nat_score = res["nat"]
            key = f"{abl_name}_best_nat"
            models[key] = {
                "name": f"{prefix} Best-Nat (s{step})",
                "lora": path,
                "desc": f"{losses_short}, best nat={nat_score:.4f} @s{step}",
            }

    log(f"\nTotal models for evaluation: {len(models)}")
    for key, minfo in models.items():
        log(f"  {key}: {minfo['name']}")

    # Write models JSON
    models_json = os.path.join(VAP_DIR, "ablation_eval_models.json")
    write_eval_models_json(models, models_json)

    # Phase 3: Evaluation
    eval_output = os.path.join(VAP_DIR, "ablation_eval")
    success = run_eval(models_json, eval_output)

    # Phase 4: Summary
    log("\n" + "="*70)
    log("ABLATION STUDY COMPLETE")
    log("="*70)

    if success:
        report_path = os.path.join(eval_output, "eval_report.html")
        log(f"HTML report: {report_path}")
        log("Serve with: python -m http.server 8780 --directory " + eval_output)
    else:
        log("Evaluation failed — check /tmp/ablation_eval.log")

    # Print training summary
    log("\nTraining Summary:")
    for abl_name, res in results.items():
        if res is None:
            log(f"  {abl_name}: FAILED")
            continue
        flow_str = f"flow={res['flow'][2]:.4f}@s{res['flow'][1]}" if res["flow"] and res["flow"][0] else "N/A"
        nat_str = f"nat={res['nat'][2]:.4f}@s{res['nat'][1]}" if res["nat"] and res["nat"][0] else "N/A"
        log(f"  {abl_name}: {flow_str}, {nat_str}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Overnight ablation study runner.

	Runs 4 ablation training experiments sequentially, then evaluates all models
	in a single comprehensive HTML report.

	Ablations (all from LaionBox v0.1-wip, 2 epochs, same hyperparameters):
	A: naturalness + quality MLP
	B: naturalness + centroid
	C: naturalness + quality MLP + speaker similarity
	D: naturalness + speaker similarity

	After training, generates audio for the best-flow and best-naturalness
	checkpoint from each ablation, scores all models (including baselines),
	and produces a combined HTML comparison.
	"""

	import json
	import os
	import signal
	import subprocess
	import sys
	import time

	PYTHON = "/home/deployer/miniconda3/envs/ml-general/bin/python"
	VAP_DIR = "/home/deployer/laion/Voice-Acting-Pipeline"
	TRAIN_SCRIPT = os.path.join(VAP_DIR, "scripts", "dramabox_finetune_train_multi_aux.py")
	EVAL_SCRIPT = os.path.join(VAP_DIR, "scripts", "run_ablation_eval.py")

	ABLATIONS = [
	{
	"name": "A_nat_quality",
	"config": "configs/ablation_nat_quality.yaml",
	"output_dir": "finetune_output/ablation_nat_quality",
	"desc_prefix": "Nat+Quality",
	"losses": "naturalness + quality_mlp",
	},
	{
	"name": "B_nat_centroid",
	"config": "configs/ablation_nat_centroid.yaml",
	"output_dir": "finetune_output/ablation_nat_centroid",
	"desc_prefix": "Nat+Centroid",
	"losses": "naturalness + centroid",
	},
	{
	"name": "C_nat_quality_speaker",
	"config": "configs/ablation_nat_quality_speaker.yaml",
	"output_dir": "finetune_output/ablation_nat_quality_speaker",
	"desc_prefix": "Nat+Quality+Speaker",
	"losses": "naturalness + quality_mlp + speaker_sim",
	},
	{
	"name": "D_nat_speaker",
	"config": "configs/ablation_nat_speaker.yaml",
	"output_dir": "finetune_output/ablation_nat_speaker",
	"desc_prefix": "Nat+Speaker",
	"losses": "naturalness + speaker_sim",
	},
	]


	def log(msg):
	ts = time.strftime("%Y-%m-%d %H:%M:%S")
	line = f"[{ts}] {msg}"
	print(line, flush=True)


	def find_best_checkpoints(output_dir):
	"""Parse metrics.jsonl to find best flow loss and best naturalness checkpoints."""
	metrics_path = os.path.join(VAP_DIR, output_dir, "metrics.jsonl")
	if not os.path.exists(metrics_path):
	log(f" WARNING: {metrics_path} not found")
	return None, None

	best_flow_loss = float("inf")
	best_flow_step = None
	best_nat = float("-inf")
	best_nat_step = None

	with open(metrics_path) as f:
	for line in f:
	line = line.strip()
	if not line:
	continue
	try:
	m = json.loads(line)
	except json.JSONDecodeError:
	continue
	step = m.get("step", 0)
	flow = m.get("flow_loss", m.get("loss", float("inf")))
	nat = m.get("naturalness_reward", float("-inf"))

	if flow < best_flow_loss:
	best_flow_loss = flow
	best_flow_step = step
	if nat is not None and nat > best_nat:
	best_nat = nat
	best_nat_step = step

	log(f" Best flow: step {best_flow_step} (loss={best_flow_loss:.4f})")
	log(f" Best nat: step {best_nat_step} (nat={best_nat:.4f})")

	# Find closest saved checkpoint (save_every=10)
	def find_ckpt(step):
	if step is None:
	return None, None
	# Try exact step first, then nearby
	for offset in [0, -1, 1, -2, 2, -3, 3, -4, 4, -5, 5]:
	candidate_step = (step + offset * 10) if offset != 0 else step
	# Try step checkpoint
	fname = f"lora_step_{candidate_step:05d}.safetensors"
	path = os.path.join(VAP_DIR, output_dir, fname)
	if os.path.exists(path):
	return path, candidate_step
	# Also try without zero padding
	fname = f"lora_step_{candidate_step:05d}.safetensors"
	path = os.path.join(VAP_DIR, output_dir, fname)
	if os.path.exists(path):
	return path, candidate_step
	# Try epoch checkpoints
	for epoch in [1, 2]:
	path = os.path.join(VAP_DIR, output_dir, f"lora_epoch{epoch}.safetensors")
	if os.path.exists(path):
	return path, f"epoch{epoch}"
	return None, None

	# Round step to nearest save_every=10
	def snap_step(step):
	if step is None:
	return None
	return round(step / 10) * 10

	flow_path, flow_step = find_ckpt(snap_step(best_flow_step))
	nat_path, nat_step = find_ckpt(snap_step(best_nat_step))

	if flow_path:
	log(f" Flow checkpoint: {os.path.basename(flow_path)}")
	else:
	log(f" WARNING: No flow checkpoint found near step {best_flow_step}")

	if nat_path:
	log(f" Nat checkpoint: {os.path.basename(nat_path)}")
	else:
	log(f" WARNING: No nat checkpoint found near step {best_nat_step}")

	return (flow_path, flow_step, best_flow_loss), (nat_path, nat_step, best_nat)


	def run_training(ablation):
	"""Run a single ablation training."""
	name = ablation["name"]
	config = ablation["config"]
	log_file = f"/tmp/ablation_{name}.log"

	log(f"{'='*70}")
	log(f"STARTING ABLATION: {name} ({ablation['losses']})")
	log(f"Config: {config}")
	log(f"Log: {log_file}")
	log(f"{'='*70}")

	cmd = [
	"env", "-u", "LD_LIBRARY_PATH",
	"accelerate", "launch", "--num_processes=8",
	TRAIN_SCRIPT,
	"--config", config,
	]

	env = os.environ.copy()
	env.pop("LD_LIBRARY_PATH", None)
	# Ensure accelerate from ml-general is used
	env["PATH"] = "/home/deployer/miniconda3/envs/ml-general/bin:" + env.get("PATH", "")

	with open(log_file, "w") as lf:
	proc = subprocess.Popen(
	cmd, stdout=lf, stderr=subprocess.STDOUT,
	cwd=VAP_DIR, env=env,
	)

	log(f"Training PID: {proc.pid}")

	# Wait for completion
	t0 = time.time()
	while True:
	ret = proc.poll()
	if ret is not None:
	break
	elapsed = time.time() - t0
	# Print progress every 5 minutes
	if int(elapsed) % 300 == 0 and elapsed > 10:
	# Read last line of status.json if it exists
	status_path = os.path.join(VAP_DIR, ablation["output_dir"], "status.json")
	if os.path.exists(status_path):
	try:
	with open(status_path) as sf:
	status = json.load(sf)
	step = status.get("step", "?")
	total = status.get("total_steps", "?")
	loss = status.get("flow_loss", status.get("loss", "?"))
	eta = status.get("eta_sec", 0)
	eta_m = int(eta // 60) if isinstance(eta, (int, float)) else "?"
	log(f" Progress: step {step}/{total}, loss={loss}, ETA={eta_m}m")
	except Exception:
	pass
	time.sleep(10)

	elapsed_min = (time.time() - t0) / 60
	if ret == 0:
	log(f"Training {name} COMPLETED in {elapsed_min:.1f} min (exit code 0)")
	else:
	log(f"Training {name} FAILED with exit code {ret} after {elapsed_min:.1f} min")
	log(f" Check log: {log_file}")

	return ret == 0


	def write_eval_models_json(all_models, path):
	"""Write the models dict as JSON for the eval script."""
	with open(path, "w") as f:
	json.dump(all_models, f, indent=2)
	log(f"Wrote {len(all_models)} models to {path}")


	def upload_to_hf(abl_name, output_dir, flow_info, nat_info):
	"""Upload best checkpoints + metrics to HF."""
	try:
	from huggingface_hub import HfApi
	api = HfApi(token="HF_TOKEN_REDACTED")
	repo = "TTS-AGI/laionbox-ablation-checkpoints"

	uploads = []
	metrics_path = os.path.join(VAP_DIR, output_dir, "metrics.jsonl")
	if os.path.exists(metrics_path):
	uploads.append((metrics_path, f"{abl_name}/metrics.jsonl"))

	if flow_info and flow_info[0]:
	path, step, _ = flow_info
	uploads.append((path, f"{abl_name}/best_flow_step{step}.safetensors"))
	if nat_info and nat_info[0]:
	path, step, _ = nat_info
	uploads.append((path, f"{abl_name}/best_nat_step{step}.safetensors"))

	for local, remote in uploads:
	if os.path.exists(local):
	sz = os.path.getsize(local) / 1024 / 1024
	log(f" HF upload: {remote} ({sz:.0f} MB)")
	api.upload_file(path_or_fileobj=local, path_in_repo=remote, repo_id=repo)
	log(f" HF upload complete for {abl_name}")
	except Exception as e:
	log(f" HF upload FAILED for {abl_name}: {e}")


	def run_eval(models_json_path, output_dir):
	"""Run the comprehensive evaluation."""
	log(f"{'='*70}")
	log("STARTING COMPREHENSIVE EVALUATION")
	log(f"{'='*70}")

	log_file = "/tmp/ablation_eval.log"

	cmd = [
	"env", "-u", "LD_LIBRARY_PATH",
	PYTHON, EVAL_SCRIPT,
	"--models-json", models_json_path,
	"--num-gpus", "8",
	"--output-dir", output_dir,
	]

	env = os.environ.copy()
	env.pop("LD_LIBRARY_PATH", None)

	with open(log_file, "w") as lf:
	proc = subprocess.Popen(
	cmd, stdout=lf, stderr=subprocess.STDOUT,
	cwd=VAP_DIR, env=env,
	)

	log(f"Eval PID: {proc.pid}, log: {log_file}")

	proc.wait()
	elapsed_min = proc.returncode
	if proc.returncode == 0:
	log("Evaluation COMPLETED successfully")
	else:
	log(f"Evaluation FAILED with exit code {proc.returncode}")
	log(f" Check log: {log_file}")

	return proc.returncode == 0


	def main():
	log("="*70)
	log("ABLATION STUDY: Auxiliary Loss Comparison")
	log("4 training runs + comprehensive evaluation")
	log("="*70)
	log("")
	log("Ablations:")
	for i, abl in enumerate(ABLATIONS):
	log(f" {abl['name']}: {abl['losses']}")
	log("")

	# Phase 1: Training
	results = {}
	for abl in ABLATIONS:
	# Check if this ablation already completed (metrics.jsonl has >= 20 entries)
	metrics_path = os.path.join(VAP_DIR, abl["output_dir"], "metrics.jsonl")
	if os.path.exists(metrics_path):
	with open(metrics_path) as f:
	n_lines = sum(1 for _ in f)
	if n_lines >= 20:
	log(f"SKIPPING {abl['name']}: already has {n_lines} metric entries")
	log(f"Finding best checkpoints for {abl['name']}...")
	flow_info, nat_info = find_best_checkpoints(abl["output_dir"])
	results[abl["name"]] = {
	"ablation": abl,
	"flow": flow_info,
	"nat": nat_info,
	}
	log("")
	continue

	success = run_training(abl)
	if success:
	log(f"\nFinding best checkpoints for {abl['name']}...")
	flow_info, nat_info = find_best_checkpoints(abl["output_dir"])
	results[abl["name"]] = {
	"ablation": abl,
	"flow": flow_info, # (path, step, loss)
	"nat": nat_info, # (path, step, nat_score)
	}
	# Upload to HF
	log(f"Uploading {abl['name']} checkpoints to HF...")
	upload_to_hf(abl["name"], abl["output_dir"], flow_info, nat_info)
	else:
	log(f"\nWARNING: {abl['name']} failed, skipping checkpoint extraction")
	results[abl["name"]] = None
	log("")

	# Phase 2: Build models dict for evaluation
	log("\n" + "="*70)
	log("TRAINING PHASE COMPLETE — BUILDING EVALUATION")
	log("="*70)

	models = {
	"vanilla": {
	"name": "Vanilla DramaBox",
	"lora": None,
	"desc": "Base DramaBox model without any fine-tuning",
	},
	"laionbox_v01": {
	"name": "LaionBox v0.1-wip",
	"lora": "/home/deployer/.cache/huggingface/hub/models--laion--laionbox-v0.1-wip/snapshots/66176d2a653a013a7b71c1ccb7a7a4d4cf514b0d/lora_epoch5.safetensors",
	"desc": "Previous best LoRA (5-epoch diff reward, DramaBox+Emolia data)",
	},
	"nat_only_best_flow": {
	"name": "Nat-Only Best-Flow (s160)",
	"lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00160.safetensors"),
	"desc": "Naturalness-only (CLAP-7B), best flow=0.528 @s160",
	},
	"nat_only_best_nat": {
	"name": "Nat-Only Best-Nat (s190)",
	"lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00190.safetensors"),
	"desc": "Naturalness-only (CLAP-7B), best nat=0.111 @s190",
	},
	}

	for abl_name, res in results.items():
	if res is None:
	continue
	abl = res["ablation"]
	prefix = abl["desc_prefix"]
	losses_short = abl["losses"]

	if res["flow"] and res["flow"][0]:
	path, step, loss = res["flow"]
	key = f"{abl_name}_best_flow"
	models[key] = {
	"name": f"{prefix} Best-Flow (s{step})",
	"lora": path,
	"desc": f"{losses_short}, best flow={loss:.4f} @s{step}",
	}

	if res["nat"] and res["nat"][0]:
	path, step, nat_score = res["nat"]
	key = f"{abl_name}_best_nat"
	models[key] = {
	"name": f"{prefix} Best-Nat (s{step})",
	"lora": path,
	"desc": f"{losses_short}, best nat={nat_score:.4f} @s{step}",
	}

	log(f"\nTotal models for evaluation: {len(models)}")
	for key, minfo in models.items():
	log(f" {key}: {minfo['name']}")

	# Write models JSON
	models_json = os.path.join(VAP_DIR, "ablation_eval_models.json")
	write_eval_models_json(models, models_json)

	# Phase 3: Evaluation
	eval_output = os.path.join(VAP_DIR, "ablation_eval")
	success = run_eval(models_json, eval_output)

	# Phase 4: Summary
	log("\n" + "="*70)
	log("ABLATION STUDY COMPLETE")
	log("="*70)

	if success:
	report_path = os.path.join(eval_output, "eval_report.html")
	log(f"HTML report: {report_path}")
	log("Serve with: python -m http.server 8780 --directory " + eval_output)
	else:
	log("Evaluation failed — check /tmp/ablation_eval.log")

	# Print training summary
	log("\nTraining Summary:")
	for abl_name, res in results.items():
	if res is None:
	log(f" {abl_name}: FAILED")
	continue
	flow_str = f"flow={res['flow'][2]:.4f}@s{res['flow'][1]}" if res["flow"] and res["flow"][0] else "N/A"
	nat_str = f"nat={res['nat'][2]:.4f}@s{res['nat'][1]}" if res["nat"] and res["nat"][0] else "N/A"
	log(f" {abl_name}: {flow_str}, {nat_str}")


	if __name__ == "__main__":
	main()