set None unit to a default value N/A

Author: ChaoPang

src/cehrbert/data_generators/hf_data_generator/hf_dataset_mapping.py

@@ -17,7 +17,7 @@
 from meds.schema import birth_code, death_code
 from pandas import Series
 
-from cehrbert.med_extension.schema_extension import Event, Visit
+from cehrbert.med_extension.schema_extension import Event
 from cehrbert.models.hf_models.tokenization_hf_cehrbert import CehrBertTokenizer
 from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments
 
@@ -578,6 +578,13 @@ def transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
         input_ids = self._concept_tokenizer.encode(record["concept_ids"])
         record["input_ids"] = input_ids
         concept_value_masks = record["concept_value_masks"]
+        units = record["units"]
+        none_values = np.array([x is None for x in units])
+        if none_values.any():
+            units = record["units"].copy()
+            units[none_values] = NA
+            record["units"] = units
+
         # Backward compatibility
         if "concept_values" not in record:
             record["concept_values"] = record["number_as_values"]