Spaces:
No application file
No application file
File size: 4,051 Bytes
96d8696 | 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 | #!/usr/bin/env python
"""Stage 06 — fusion ablation.
Trains and evaluates all three fusion strategies (concat, gated,
cross_attention) at matched hyperparameters, using the SAME cached
embeddings from stage 02 (no re-encoding, no new GPU-heavy work). Produces
a comparison report so cross-attention's benefit over simpler fusion can
be quantified rather than just asserted.
Each variant gets its own checkpoint_dir (checkpoints/ablation_<type>) and
its own generated config (configs/ablation/<type>.yaml) so runs never
overwrite each other and remain individually reproducible.
Usage:
python scripts/06_fusion_ablation.py --config configs/base.yaml
"""
import argparse
import copy
import json
import subprocess
import sys
from pathlib import Path
import yaml
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.utils.config import load_config
from src.utils.exceptions import PricePredictorError
from src.utils.logging import get_logger
logger = get_logger(__name__)
FUSION_TYPES = ["concat", "gated", "cross_attention"]
def _run_subprocess(cmd: list) -> None:
result = subprocess.run(cmd)
if result.returncode != 0:
raise PricePredictorError(f"Command failed (exit {result.returncode}): {' '.join(cmd)}")
def run(base_config_path: str, output_report: str) -> dict:
base_config = load_config(base_config_path) # validates the base config once up front
repo_root = Path(__file__).resolve().parents[1]
ablation_configs_dir = repo_root / "configs" / "ablation"
ablation_configs_dir.mkdir(parents=True, exist_ok=True)
results = {}
for fusion_type in FUSION_TYPES:
logger.info("=== Ablation run: fusion.type=%s ===", fusion_type)
config = copy.deepcopy(base_config)
config["fusion"]["type"] = fusion_type
config["checkpoint_dir"] = f"checkpoints/ablation_{fusion_type}"
config_path = ablation_configs_dir / f"{fusion_type}.yaml"
with config_path.open("w") as f:
yaml.safe_dump(config, f)
_run_subprocess([sys.executable, "scripts/03_train.py", "--config", str(config_path)])
_run_subprocess([sys.executable, "scripts/04_evaluate.py", "--config", str(config_path)])
eval_report_path = repo_root / "reports" / f"ablation_{fusion_type}" / "eval_report.json"
if not eval_report_path.exists():
raise PricePredictorError(
f"Expected eval report not found at {eval_report_path} — "
"stage 04 may have failed silently for this fusion type."
)
with eval_report_path.open() as f:
metrics = json.load(f)
results[fusion_type] = metrics
logger.info("Result for %s: SMAPE=%.4f MAE=%.4f", fusion_type, metrics["smape"], metrics["mae"])
output_path = Path(output_report)
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w") as f:
json.dump(results, f, indent=2)
best_fusion = min(results, key=lambda k: results[k]["smape"])
logger.info("=== Fusion ablation comparison (lower SMAPE is better) ===")
for fusion_type, metrics in results.items():
marker = " <-- best" if fusion_type == best_fusion else ""
logger.info("%-16s SMAPE=%.4f MAE=%.4f%s", fusion_type, metrics["smape"], metrics["mae"], marker)
logger.info("Wrote comparison report to %s", output_path)
return results
def main() -> None:
parser = argparse.ArgumentParser(description="Stage 06: train+evaluate all fusion strategies for comparison")
parser.add_argument("--config", default="configs/base.yaml")
parser.add_argument("--output", default="reports/fusion_ablation_comparison.json")
args = parser.parse_args()
try:
run(args.config, args.output)
except PricePredictorError as e:
logger.error("Fusion ablation failed: %s", e)
sys.exit(1)
except Exception as e:
logger.exception("Unexpected error during fusion ablation: %s", e)
sys.exit(1)
if __name__ == "__main__":
main() |