added patient_splits_folder to evaluation (#83)

cumc-dbmi · Jan 15, 2025 · e62bfca · e62bfca
1 parent 5bcbd61
commit e62bfca
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 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,

diff --git a/src/cehrbert/evaluations/evaluation_parse_args.py b/src/cehrbert/evaluations/evaluation_parse_args.py
@@ -43,8 +43,8 @@ def create_evaluation_args():
         required=baseline_model_required,
     )
     main_parser.add_argument(
-        "--test_person_ids_path",
-        dest="test_person_ids_path",
+        "--patient_splits_folder",
+        dest="patient_splits_folder",
         action="store",
         help="The test person_ids data",
         required=False,