| """Single, thin command-line entry point for GNN4Colliders.""" |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import os |
| import sys |
| from importlib import resources |
|
|
| from hydra import compose, initialize_config_dir |
| from omegaconf import OmegaConf |
|
|
| from gnn4colliders.config.application import ( |
| export_onnx, |
| predict_or_evaluate, |
| prepare, |
| train, |
| ) |
|
|
| _COMMANDS = {"prepare", "train", "evaluate", "predict", "export"} |
|
|
|
|
| def _is_main_process() -> bool: |
| return int(os.environ.get("RANK", "0")) == 0 |
|
|
|
|
| def _help(command: str | None = None) -> None: |
| if command: |
| print(f"Usage: gnn4colliders {command} [key=value ...]") |
| print("Compose Hydra configuration with semantic overrides.") |
| else: |
| print( |
| "Usage: gnn4colliders <prepare|train|evaluate|predict|export> " |
| "[key=value ...]" |
| ) |
| print("Use gnn4colliders <command> --help for command help.") |
|
|
|
|
| def _config(overrides: list[str]): |
| config_resource = resources.files("gnn4colliders.configs") |
| with resources.as_file(config_resource) as config_dir: |
| with initialize_config_dir(version_base=None, config_dir=str(config_dir)): |
| return compose(config_name="config", overrides=overrides) |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| args = list(sys.argv[1:] if argv is None else argv) |
| if not args or args[0] in {"-h", "--help"}: |
| _help() |
| return 0 |
| command = args.pop(0) |
| if command not in _COMMANDS: |
| print( |
| f"unknown command {command!r}; choose from {', '.join(sorted(_COMMANDS))}", |
| file=sys.stderr, |
| ) |
| return 2 |
| if "--help" in args or "-h" in args: |
| _help(command) |
| return 0 |
| try: |
| config = _config(args) |
| logging.basicConfig( |
| level=getattr(logging, str(config.logging.level).upper(), logging.INFO) |
| ) |
| if command == "prepare": |
| print(f"prepared cache: {prepare(config)}") |
| elif command == "train": |
| if _is_main_process(): |
| print(f"saved checkpoint: {train(config)}") |
| elif command == "export": |
| if _is_main_process(): |
| print(f"wrote ONNX model: {export_onnx(config)}") |
| elif command == "evaluate": |
| result = predict_or_evaluate(config, evaluate=True) |
| if _is_main_process(): |
| print(OmegaConf.to_yaml(result.metrics)) |
| else: |
| result = predict_or_evaluate(config) |
| if _is_main_process(): |
| print( |
| f"wrote predictions: {config.inference.output} " |
| f"({len(result.sample_ids)} events)" |
| ) |
| return 0 |
| except (ValueError, FileNotFoundError, KeyError, ImportError) as error: |
| print(f"gnn4colliders: {error}", file=sys.stderr) |
| return 2 |
|
|
|
|
| __all__ = ["main"] |
|
|
