Formatting

Author: ae-foster

README.md

@@ -64,7 +64,7 @@ We recommend using distinct output directories for every training run.
 Regarding other optional arguments, run `python scripts/transferable.py -h` for more information.
 
 **Note**: Checkpoints are stored using `pickle`. This means that opening a checkpoint from an untrusted source is a major security risk. Only ever open checkpoint files from trusted sources. To prevent reading untrusted pickle files, checkpoint reading is disabled by default and can be
-re-enabled by settings the environment variables `ORBFORMER_PICKLE_LOADING=1`.
+re-enabled by settings the environment variable `ORBFORMER_PICKLE_LOADING=1`.
 
 ### Preparing new structure data for fine-tuning
 

src/oneqmc/entrypoint.py

@@ -19,7 +19,9 @@
 
 def load_chkpt_file(chkpt: str, discard_sampler_state: bool) -> Tuple[TrainState, int]:
     if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
-        raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
+        raise PermissionError(
+            "Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow"
+        )
     with open(chkpt, "rb") as chkpt_file:
         init_step, (smpl_state, param_state, opt_state) = pickle.load(chkpt_file)
         if discard_sampler_state:
@@ -31,7 +33,9 @@ def load_chkpt_file(chkpt: str, discard_sampler_state: bool) -> Tuple[TrainState
 
 def load_density_chkpt_file(chkpt: str, discard_sampler_state: bool) -> Tuple[Tuple, int]:
     if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
-        raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
+        raise PermissionError(
+            "Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow"
+        )
     with open(chkpt, "rb") as chkpt_file:
         init_step, (param_state, opt_state) = pickle.load(chkpt_file)
     return (param_state, opt_state), init_step

src/oneqmc/log.py

@@ -134,7 +134,9 @@ def last(self):
         step_fast, step_slow = -1, -1  # account for the case where a queue is not initialized
         while self.fast_chkpts:
             if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
-                raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
+                raise PermissionError(
+                    "Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow"
+                )
             with self.fast_chkpts.pop(-1).path.open("rb") as f:
                 step_fast, last_chkpt_fast = pickle.load(f)
             if not jax.tree.reduce(
@@ -145,7 +147,9 @@ def last(self):
                 break
         while self.slow_chkpts:
             if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
-                raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
+                raise PermissionError(
+                    "Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow"
+                )
             with self.slow_chkpts.pop(-1).path.open("rb") as f:
                 step_slow, last_chkpt_slow = pickle.load(f)
             if not jax.tree.reduce(

tests/integration_tests/test_scripts.py

@@ -52,7 +52,7 @@ def runner(extra_args):
             ],
             cwd=project_root,
             capture_output=True,
-            env={"ORBFORMER_PICKLE_LOADING": "1"}
+            env={"ORBFORMER_PICKLE_LOADING": "1"},
         )
         if result.returncode != 0:
             raise OneQMCProcessError(result.stderr.decode())
@@ -91,7 +91,7 @@ def runner(extra_args):
             ],
             cwd=project_root,
             capture_output=True,
-            env={"ORBFORMER_PICKLE_LOADING": "1"}
+            env={"ORBFORMER_PICKLE_LOADING": "1"},
         )
         if result.returncode != 0:
             raise OneQMCProcessError(result.stderr.decode())