File size: 1,703 Bytes
b2f3bf4 | 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 | #!/usr/bin/env python3
"""Train S1/S2/S3 (DESIGN §6.4) on one GPU.
python3 scripts/train.py --config configs/train.yaml --stage s2 --seed 42 --out checkpoints/s2_seed42
python3 scripts/train.py --config configs/train.yaml --stage s2 --subset-frac 0.1 --out checkpoints/s2_scan_a --set lambda_rps=1.0
"""
from __future__ import annotations
import argparse
import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
from jev_judge.train_loop import TrainConfig, Trainer # noqa: E402
def parse_set(items: list[str]) -> dict:
out = {}
for it in items or []:
k, v = it.split("=", 1)
try:
import json
out[k] = json.loads(v)
except Exception:
out[k] = v
return out
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--config", default="configs/train.yaml")
ap.add_argument("--stage", choices=["s1", "s2", "s3"])
ap.add_argument("--seed", type=int)
ap.add_argument("--out", dest="out_dir")
ap.add_argument("--model", dest="model_path")
ap.add_argument("--subset-frac", dest="subset_frac", type=float)
ap.add_argument("--epochs", type=int)
ap.add_argument("--max-steps", dest="max_steps", type=int)
ap.add_argument("--init-from", dest="init_from")
ap.add_argument("--set", nargs="*", help="extra overrides key=value (JSON values)")
args = ap.parse_args()
overrides = {k: v for k, v in vars(args).items() if k not in ("config", "set")}
overrides.update(parse_set(args.set))
cfg = TrainConfig.from_yaml(args.config, overrides)
Trainer(cfg).fit()
if __name__ == "__main__":
main()
|