From 58b761c9fcf25807b5835fd456ccf80ea64627ee Mon Sep 17 00:00:00 2001 From: Chao Pang Date: Mon, 9 Sep 2024 09:27:29 -0400 Subject: [PATCH] added the backward compatibility if the finetuning data is constructed from the OMOP data, where the age_at_index is labelled as age --- .../data_generators/hf_data_generator/hf_dataset_mapping.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 2a413bb..576817e 100644 --- a/src/cehrbert/data_generators/hf_data_generator/hf_dataset_mapping.py +++ b/src/cehrbert/data_generators/hf_data_generator/hf_dataset_mapping.py @@ -453,6 +453,6 @@ class HFFineTuningMapping(DatasetMapping): def transform(self, record: Dict[str, Any]) -> Dict[str, Any]: return { - "age_at_index": record["age_at_index"], + "age_at_index": record["age"] if "age" in record else record["age_at_index"], "classifier_label": record["label"], }