check whether patient_id exists in the meds_reader database because s…

…ubject_splits can contain more patients than the dataset
cumc-dbmi · Oct 25, 2024 · 8d453cb · 8d453cb
1 parent ec975cc
commit 8d453cb
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/cehrbert/data_generators/hf_data_generator/meds_utils.py b/src/cehrbert/data_generators/hf_data_generator/meds_utils.py
@@ -221,8 +221,9 @@ def _meds_to_cehrbert_generator(
     with meds_reader.SubjectDatabase(path_to_db) as patient_database:
         for shard in shards:
             for patient_id, prediction_time, label in shard:
-                patient = patient_database[patient_id]
-                yield convert_one_patient(patient, conversion, prediction_time, label)
+                if patient_id in patient_database:
+                    patient = patient_database[patient_id]
+                    yield convert_one_patient(patient, conversion, prediction_time, label)
 
 
 def _create_cehrbert_data_from_meds(