Use yaml.safe_load in training scripts. Fixes #215.

openvax · Aug 1, 2023 · 57538e3 · 57538e3
1 parent f91d40c
commit 57538e3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/mhcflurry/train_allele_specific_models_command.py b/mhcflurry/train_allele_specific_models_command.py
@@ -142,7 +142,7 @@ def run(argv=sys.argv[1:]):
 
     configure_logging(verbose=args.verbosity > 1)
 
-    hyperparameters_lst = yaml.load(open(args.hyperparameters))
+    hyperparameters_lst = yaml.safe_load(open(args.hyperparameters))
     assert isinstance(hyperparameters_lst, list), hyperparameters_lst
     print("Loaded hyperparameters list: %s" % str(hyperparameters_lst))
 

diff --git a/mhcflurry/train_pan_allele_models_command.py b/mhcflurry/train_pan_allele_models_command.py
@@ -303,7 +303,7 @@ def initialize_training(args):
             parser.error("Missing required arg: %s" % arg)
 
     print("Initializing training.")
-    hyperparameters_lst = yaml.load(open(args.hyperparameters))
+    hyperparameters_lst = yaml.safe_load(open(args.hyperparameters))
     assert isinstance(hyperparameters_lst, list)
     print("Loaded hyperparameters list:")
     pprint.pprint(hyperparameters_lst)