Added --out optional flag for build subcommand to make usage in snakemake cleaner

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)")