load different formats for parse_dataset_eval

mir-group · Jan 24, 2023 · 2c98613 · 2c98613
1 parent a2c0dcc
commit 2c98613
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/orchard/data.py b/orchard/data.py
@@ -56,7 +56,11 @@ def get_accdb_formulas(dataset_eval_name):
     return formulas
 
 def parse_dataset_eval(subdb_names, dataset_fname):
-    formulas = get_accdb_formulas(dataset_fname)
+    if dataset_fname.endswith('.yaml'):
+        with open(dataset_fname, 'r') as f:
+            formulas = yaml.load(f, Loader=yaml.Loader)
+    else:
+        formulas = get_accdb_formulas(dataset_fname)
     cats = {}
     sumabs = {}
     counts = {}