added num_of_concepts and num_of_visits to the saved pretraining data…

…set so we can filter examples by the num of tokens
cumc-dbmi · Oct 11, 2024 · 3375240 · 3375240
1 parent 8c40a27
commit 3375240
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/cehrbert/data_generators/hf_data_generator/hf_dataset.py b/src/cehrbert/data_generators/hf_data_generator/hf_dataset.py
@@ -21,6 +21,8 @@
     "concept_values",
     "concept_value_masks",
     "mlm_skip_values",
+    "num_of_concepts",
+    "num_of_visits",
 ]
 
 TRANSFORMER_COLUMNS = ["input_ids", "labels"]