added patient_id to the pretraining data for debugging

cumc-dbmi · Sep 12, 2024 · 34b16d8 · 34b16d8
1 parent 9ecaca5
commit 34b16d8
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 1 deletion.
diff --git a/src/cehrbert/data_generators/hf_data_generator/hf_dataset.py b/src/cehrbert/data_generators/hf_data_generator/hf_dataset.py
@@ -12,6 +12,7 @@
 from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments
 
 CEHRBERT_COLUMNS = [
+    "person_id",
     "concept_ids",
     "ages",
     "dates",

diff --git a/src/cehrbert/data_generators/hf_data_generator/hf_dataset_mapping.py b/src/cehrbert/data_generators/hf_data_generator/hf_dataset_mapping.py
@@ -115,7 +115,7 @@ def __init__(self, data_args: DataTrainingArguments, is_pretraining: bool = True
 
     def remove_columns(self):
         if self._is_pretraining:
-            return ["visits", "patient_id", "birth_datetime", "index_date"]
+            return ["visits", "birth_datetime", "index_date"]
         else:
             return [
                 "visits",

diff --git a/src/cehrbert/utils/stat_utils.py b/src/cehrbert/utils/stat_utils.py