diff --git a/src/custom_models/cli.py b/src/custom_models/cli.py index bc55e02..8d846df 100644 --- a/src/custom_models/cli.py +++ b/src/custom_models/cli.py @@ -260,7 +260,7 @@ def _cmd_build(args: argparse.Namespace) -> int: Args: args: Parsed arguments (arch, budget, stride, d-embed, - budget-tol, n-heads, n-classes, seed). + budget-tol, n-heads, n-classes, seed, out). Returns: Exit code 0 on success. @@ -292,7 +292,14 @@ def _cmd_build(args: argparse.Namespace) -> int: "n_classes": args.n_classes, "count_params": count_params(probe), } - print(json.dumps(payload, indent=2, sort_keys=True)) + text = json.dumps(payload, indent=2, sort_keys=True) + if args.out: + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(text + "\n", encoding="utf-8") + print(f"wrote {out}") + else: + print(text) return 0 @@ -480,6 +487,7 @@ def build_parser() -> argparse.ArgumentParser: _add_model_args(p_build, required=True) p_build.add_argument("--n-classes", type=int, default=11) p_build.add_argument("--seed", type=int, default=0) + p_build.add_argument("--out", default=None, help="Write the JSON payload to PATH instead of stdout") p_build.set_defaults(func=_cmd_build) p_store = sub.add_parser("store", help="build a signal store (pod5+labels or synthetic)")