experiment-2107/scripts/run_ablation_20k.py

"""Run the baseline + CTM ablation sweep on the full /data/processed (~20k sample) dataset.

Eight configs x {fine, coarse} label spaces, matching experiment_04's design:
    baseline, CTM-10 (fixed), CTM-20 (fixed), CTM-ACT (<=30, adaptive)

Each run trains to completion (or early stop) and appends its metrics dict to
/workspace/runs_20k_logs/results.json, so a crash partway through the sweep doesn't
lose earlier results. Per-run stdout goes to its own log file.
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

sys.path.insert(0, "/workspace")

import model as base_mod
import model_ctm as ctm_mod

DATA_DIR = Path("/data/processed")
LOG_DIR = Path("/workspace/runs_20k_logs")
RESULTS_PATH = LOG_DIR / "results.json"

CONFIGS = [
    dict(name="baseline", coarse=False, kind="baseline"),
    dict(name="baseline_coarse", coarse=True, kind="baseline"),
    dict(name="ctm10", coarse=False, kind="ctm", iterations=10, memory=20, use_act=False),
    dict(name="ctm10_coarse", coarse=True, kind="ctm", iterations=10, memory=20, use_act=False),
    dict(name="ctm20", coarse=False, kind="ctm", iterations=20, memory=20, use_act=False),
    dict(name="ctm20_coarse", coarse=True, kind="ctm", iterations=20, memory=20, use_act=False),
    dict(name="ctm_act", coarse=False, kind="ctm", iterations=30, memory=50, use_act=True),
    dict(name="ctm_act_coarse", coarse=True, kind="ctm", iterations=30, memory=50, use_act=True),
]


class Tee:
    def __init__(self, path: Path):
        self.file = open(path, "w")
        self.stdout = sys.stdout

    def write(self, data):
        self.stdout.write(data)
        self.file.write(data)
        self.file.flush()

    def flush(self):
        self.stdout.flush()
        self.file.flush()


def load_results() -> dict:
    if RESULTS_PATH.exists():
        return json.loads(RESULTS_PATH.read_text())
    return {}


def save_results(results: dict) -> None:
    RESULTS_PATH.write_text(json.dumps(results, indent=2))


def main() -> int:
    LOG_DIR.mkdir(exist_ok=True)
    results = load_results()

    for cfg in CONFIGS:
        name = cfg["name"]
        if name in results:
            print(f"=== SKIP {name} (already in results.json) ===", flush=True)
            continue

        log_path = LOG_DIR / f"{name}.log"
        real_stdout = sys.stdout
        sys.stdout = Tee(log_path)
        t0 = time.perf_counter()
        try:
            print(f"=== START {name} | cfg={cfg} ===", flush=True)
            if cfg["kind"] == "baseline":
                metrics = base_mod.train(DATA_DIR, coarse=cfg["coarse"])
            else:
                metrics = ctm_mod.train(
                    DATA_DIR,
                    coarse=cfg["coarse"],
                    max_iterations=cfg["iterations"],
                    memory_length=cfg["memory"],
                    use_act=cfg["use_act"],
                    ponder_lambda=(0.01 if cfg["use_act"] else 0.0),
                    tag=name,
                )
            elapsed = time.perf_counter() - t0
            metrics["wall_seconds"] = elapsed
            print(f"=== DONE {name} in {elapsed:.1f}s ===", flush=True)
        finally:
            sys.stdout.file.close()
            sys.stdout = real_stdout

        results[name] = metrics
        save_results(results)
        print(f"[driver] {name}: {metrics}")

    print("[driver] ALL CONFIGS DONE")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())