scripts/train.py

"""CLI entry point: train a single model variant."""
from __future__ import annotations

import argparse
import os

from dotenv import load_dotenv

load_dotenv()
os.environ.setdefault("HF_TOKEN", os.environ.get("HUGGINGFACE_API_KEY", ""))
os.environ.setdefault("WANDB_API_KEY", os.environ.get("WANDB_API_KEY", ""))

from physiojepa.trainer import TrainConfig, load_yaml_config, train


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--run_name", type=str, default=None)
    ap.add_argument("--model", type=str, default=None, choices=["A", "B", "C", "F"])
    ap.add_argument("--epochs", type=int, default=None)
    ap.add_argument("--batch_size", type=int, default=None)
    ap.add_argument("--index_path", type=str, default=None)
    ap.add_argument("--shard_roots_json", type=str, default=None,
                    help="JSON file listing shard roots")
    ap.add_argument("--wandb_mode", type=str, default=None)
    ap.add_argument("--num_workers", type=int, default=None)
    ap.add_argument("--output_dir", type=str, default=None)
    ap.add_argument("--subset_frac", type=float, default=None)
    ap.add_argument("--log_every", type=int, default=None)
    ap.add_argument("--ema_start", type=float, default=None)
    ap.add_argument("--ema_end", type=float, default=None)
    ap.add_argument("--ema_warmup_frac", type=float, default=None)
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--pred_depth", type=int, default=None)
    ap.add_argument("--query_mode", type=str, default=None, choices=["learned", "sinusoidal"])
    ap.add_argument("--mask_ratio", type=float, default=None)
    ap.add_argument("--fast_cache_dir", type=str, default=None)
    args = ap.parse_args()

    cfg = load_yaml_config(args.config)
    overrides = {k: v for k, v in vars(args).items() if v is not None and k not in ("config",)}
    if "shard_roots_json" in overrides:
        import json
        cfg.shard_roots = json.loads(open(overrides.pop("shard_roots_json")).read())
    for k, v in overrides.items():
        setattr(cfg, k, v)
    print(f"[train] resolved config: model={cfg.model} run={cfg.run_name} "
          f"epochs={cfg.epochs} bs={cfg.batch_size} shards={len(cfg.shard_roots)}")
    res = train(cfg)
    print(f"[train] done: {res}")


if __name__ == "__main__":
    main()