Merge pull request #795 from WillBaldwin0/refactor-data

configargparse for the preprocessing parser
ACEsuit · Jan 22, 2025 · 251bb43 · 251bb43
2 parents 8579c26 + 8f0c6f3
commit 251bb43
Showing 1 changed file with 25 additions and 3 deletions.
diff --git a/mace/tools/arg_parser.py b/mace/tools/arg_parser.py
@@ -713,9 +713,24 @@ def build_default_arg_parser() -> argparse.ArgumentParser:
 
 
 def build_preprocess_arg_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
+    try:
+        import configargparse
+
+        parser = configargparse.ArgumentParser(
+            config_file_parser_class=configargparse.YAMLConfigFileParser,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        )
+        parser.add(
+            "--config",
+            type=str,
+            is_config_file=True,
+            help="config file to agregate options",
+        )
+    except ImportError:
+        parser = argparse.ArgumentParser(
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        )
+
     parser.add_argument(
         "--train_file",
         help="Training set h5 file",
@@ -859,6 +874,13 @@ def build_preprocess_arg_parser() -> argparse.ArgumentParser:
         type=str,
         default="head",
     )
+    parser.add_argument(
+        "--heads",
+        help="Dict of heads: containing individual files and E0s",
+        type=str,
+        default=None,
+        required=False,
+    )
     return parser