lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Apply ruff format and document lint/type tooling in CLAUDE.md

Browse Source 
First repo-wide ruff format pass, plus a note in CLAUDE.md to run
ruff and ty periodically.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-06-18 17:40:27 +02:00
parent 53fd2e4405
commit 8cebc4809d
20 changed files with 668 additions and 307 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/steps_to_parquet.py
+3 -1
View File
@@ -96,7 +96,9 @@ def main() -> None:
help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).",
)
parser.add_argument(
"--tree", default="Steps", help="Tree name inside the ROOT file (default: Steps)"
"--tree",
default="Steps",
help="Tree name inside the ROOT file (default: Steps)",
)
parser.add_argument(
"--compression",
scripts/train.py
+74 -30
View File
@@ -10,7 +10,9 @@ from giant.pipeline import run_train_job
def main() -> None:
parser = argparse.ArgumentParser(description="Train GIANT surrogate model")
parser.add_argument("--config", default=None, help="Path to TOML config file")
parser.add_argument("--data", required=True, help="Path to parquet file or directory")
parser.add_argument(
"--data", required=True, help="Path to parquet file or directory"
)
parser.add_argument("--mode", choices=["flow", "ddpm"])
parser.add_argument("--epochs", type=int)
parser.add_argument("--batch-size", type=int)
@@ -18,49 +20,91 @@ def main() -> None:
parser.add_argument("--hidden-dim", type=int)
parser.add_argument("--n-blocks", type=int)
parser.add_argument("--emb-dim", type=int)
parser.add_argument("--dropout", type=float, help="Dropout probability in ResBlocks (default: 0.1)")
parser.add_argument(
"--dropout", type=float, help="Dropout probability in ResBlocks (default: 0.1)"
)
parser.add_argument("--val-fraction", type=float)
parser.add_argument("--seed", type=int, help="Random seed for reproducibility")
parser.add_argument("--validate-every", type=int,
help="Run marginal+KL validation every N epochs (0 disables)")
parser.add_argument("--shuffle-buffer", type=int, default=65536,
help="Rows held in RAM for shuffling per worker (default: 65536)")
parser.add_argument("--out", default=None, help="Checkpoint output directory (default: auto from hyperparams)")
parser.add_argument("--device", default=None, help="cpu | cuda | mps (default: auto)")
parser.add_argument(
"--validate-every",
type=int,
help="Run marginal+KL validation every N epochs (0 disables)",
)
parser.add_argument(
"--shuffle-buffer",
type=int,
default=65536,
help="Rows held in RAM for shuffling per worker (default: 65536)",
)
parser.add_argument(
"--out",
default=None,
help="Checkpoint output directory (default: auto from hyperparams)",
)
parser.add_argument(
"--device", default=None, help="cpu | cuda | mps (default: auto)"
)
parser.add_argument("--num-workers", type=int)
parser.add_argument("--resume", default=None, help="Checkpoint .pt to resume training from")
parser.add_argument(
"--resume", default=None, help="Checkpoint .pt to resume training from"
)
args = parser.parse_args()
cli_train = {k: v for k, v in {
"mode": args.mode, "epochs": args.epochs, "batch_size": args.batch_size,
"lr": args.lr, "val_fraction": args.val_fraction, "num_workers": args.num_workers,
"seed": args.seed, "validate_every": args.validate_every,
}.items() if v is not None}
cli_model = {k: v for k, v in {
"hidden_dim": args.hidden_dim, "n_blocks": args.n_blocks, "emb_dim": args.emb_dim,
"dropout": args.dropout,
}.items() if v is not None}
cli_train = {
k: v
for k, v in {
"mode": args.mode,
"epochs": args.epochs,
"batch_size": args.batch_size,
"lr": args.lr,
"val_fraction": args.val_fraction,
"num_workers": args.num_workers,
"seed": args.seed,
"validate_every": args.validate_every,
}.items()
if v is not None
}
cli_model = {
k: v
for k, v in {
"hidden_dim": args.hidden_dim,
"n_blocks": args.n_blocks,
"emb_dim": args.emb_dim,
"dropout": args.dropout,
}.items()
if v is not None
}
config_path = Path(args.config) if args.config else None
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, cli_train, cli_model)
cfg = gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG, config_path, cli_train, cli_model
)
t, m = cfg["train"], cfg["model"]
device = torch.device(args.device) if args.device else gconfig.auto_device()
out_dir = Path(args.out or (
f"checkpoints/{t['mode']}"
f"_h{m['hidden_dim']}"
f"_b{m['n_blocks']}"
f"_e{m['emb_dim']}"
f"_lr{t['lr']}"
f"_bs{t['batch_size']}"
))
out_dir = Path(
args.out
or (
f"checkpoints/{t['mode']}"
f"_h{m['hidden_dim']}"
f"_b{m['n_blocks']}"
f"_e{m['emb_dim']}"
f"_lr{t['lr']}"
f"_bs{t['batch_size']}"
)
)
print(f"device: {device}")
print(f"out_dir: {out_dir}")
run_train_job(
data=Path(args.data), cfg=cfg, out_dir=out_dir, device=device,
shuffle_buffer=args.shuffle_buffer, num_workers=t["num_workers"],
resume=Path(args.resume) if args.resume else None, echo=print,
data=Path(args.data),
cfg=cfg,
out_dir=out_dir,
device=device,
shuffle_buffer=args.shuffle_buffer,
num_workers=t["num_workers"],
resume=Path(args.resume) if args.resume else None,
echo=print,
)