added patient_splits_folder to evaluation

cumc-dbmi · Jan 15, 2025 · 8e9161b · 8e9161b
1 parent 5bcbd61
commit 8e9161b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/cehrbert/evaluations/evaluation.py b/src/cehrbert/evaluations/evaluation.py
@@ -318,8 +318,9 @@ def evaluate_baseline_models(args):
     # Load the training data
     dataset = pd.read_parquet(args.data_path)
     test_person_ids = None
-    if args.test_person_ids_path:
-        test_person_ids = pd.read_parquet(args.test_person_ids_path)
+    if args.patient_splits_folder:
+        patient_splits = pd.read_parquet(args.patient_splits_folder)
+        test_person_ids = patient_splits[patient_splits.split == "test"]
 
     LogisticRegressionModelEvaluator(
         dataset=dataset,