cehrgpt 0.1.1__tar.gz → 0.1.2__tar.gz

{cehrgpt-0.1.1 → cehrgpt-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cehrgpt
-Version: 0.1.1
+Version: 0.1.2
 Summary: CEHR-GPT: Generating Electronic Health Records with Chronological Patient Timelines
 Author-email: Chao Pang <chaopang229@gmail.com>, Xinzhuo Jiang <xj2193@cumc.columbia.edu>, Krishna Kalluri <kk3326@cumc.columbia.edu>, Elise Minto <em3697@cumc.columbia.edu>, Jason Patterson <jp3477@cumc.columbia.edu>, Nishanth Parameshwar Pavinkurve <np2689@cumc.columbia.edu>, Karthik Natarajan <kn2174@cumc.columbia.edu>
 License: MIT License
@@ -105,6 +105,100 @@ sh scripts/omop_pipeline.sh \
   $OMOP_VOCAB_DIR
 ```
 
+# MEDS Support
+
+This section demonstrates how to pretrain CEHR-GPT using MIMIC-IV data in the MEDS (Medical Event Data Standard) format.
+
+## Prerequisites
+
+Set up the required environment variables before beginning:
+
+```bash
+export CEHR_GPT_MODEL_DIR=""    # Path to CEHR-GPT model directory
+export MEDS_DIR=""              # Path to MEDS data directory
+export MEDS_READER_DIR=""       # Path to MEDS reader output directory
+```
+
+## Step 1: Create MIMIC MEDS Data
+
+Transform your MIMIC files into MEDS format by following the instructions in the [MEDS_transforms](https://github.com/mmcdermott/MEDS_transforms/) repository.
+
+## Step 2: Create the MEDS Reader
+
+Convert the MEDS data for use with CEHR-GPT:
+
+```bash
+meds_reader_convert $MEDS_DIR $MEDS_READER_DIR --num_threads 10
+```
+
+## Step 3: Pretrain CEHR-GPT
+
+Run the pretraining process using the prepared MEDS data:
+
+```bash
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $CEHR_GPT_MODEL_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --do_train true --seed 42 \
+  --dataloader_num_workers 16 --dataloader_prefetch_factor 8 \
+  --hidden_size 768 --num_hidden_layers 14 --max_position_embeddings 8192 \
+  --evaluation_strategy epoch --save_strategy epoch \
+  --sample_packing --max_tokens_per_batch 16384 \
+  --warmup_steps 500 --weight_decay 0.01 \
+  --num_train_epochs 50 --learning_rate 0.0002 \
+  --use_early_stopping --early_stopping_threshold 0.001 \
+  --is_data_in_meds --inpatient_att_function_type day \
+  --att_function_type day --include_inpatient_hour_token \
+  --include_auxiliary_token --include_demographic_prompt \
+  --meds_to_cehrbert_conversion_type "MedsToBertMimic4"
+```
+
+## Step 4: Generate MEDS Trajectories
+
+### Environment Setup for Trajectory Generation
+
+Configure additional environment variables for trajectory generation with task labels (`subject_id`, `prediction_time`, `boolean_value` [optional]):
+
+```bash
+# MEDS_LABEL_COHORT_DIR must contain a set of parquet files
+export MEDS_LABEL_COHORT_DIR=""     # Path to cohort labels directory
+export MEDS_TRAJECTORY_DIR=""       # Path for trajectory output
+```
+
+### Generate Trajectories
+
+Create synthetic patient trajectories using the trained model:
+
+> **Important:** The total sequence length (`generation_input_length` + `generation_max_new_tokens`) cannot exceed the `max_position_embeddings` value (8192) defined during pretraining.
+
+```bash
+python -u -m cehrgpt.generation.cehrgpt_conditional_generation \
+  --cohort_folder $MEDS_LABEL_COHORT_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $MEDS_TRAJECTORY_DIR \
+  --per_device_eval_batch_size 16 \
+  --num_of_trajectories_per_sample 2 \
+  --generation_input_length 4096 \
+  --generation_max_new_tokens 4096 \
+  --is_data_in_meds \
+  --att_function_type day --inpatient_att_function_type day \
+  --meds_to_cehrbert_conversion_type MedsToBertMimic4 \
+  --include_auxiliary_token --include_demographic_prompt \
+  --include_inpatient_hour_token
+```
+
+### Parameters Explanation
+
+- `generation_input_length`: Controls the length of input context for generation
+- `generation_max_new_tokens`: Maximum number of new tokens to generate
+- `num_of_trajectories_per_sample`: Number of trajectories to generate per patient sample
+
 ## Citation
 ```
 @article{cehrgpt2024,

cehrgpt-0.1.2/README.md

@@ -0,0 +1,174 @@
+# CEHRGPT
+
+[![PyPI - Version](https://img.shields.io/pypi/v/cehrgpt)](https://pypi.org/project/cehrgpt/)
+![Python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)
+[![tests](https://github.com/knatarajan-lab/cehrgpt/actions/workflows/tests.yaml/badge.svg)](https://github.com/knatarajan-lab/cehrgpt/actions/workflows/tests.yaml)
+[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/knatarajan-lab/cehrgpt/blob/main/LICENSE)
+[![contributors](https://img.shields.io/github/contributors/knatarajan-lab/cehrgpt.svg)](https://github.com/knatarajan-lab/cehrgpt/graphs/contributors)
+
+## Description
+CEHRGPT is a synthetic data generation model developed to handle structured electronic health records (EHR) with enhanced privacy and reliability. It leverages state-of-the-art natural language processing techniques to create realistic, anonymized patient data that can be used for research and development without compromising patient privacy.
+
+## Features
+- **Synthetic Patient Data Generation**: Generates comprehensive patient profiles including demographics, medical history, treatment courses, and outcomes.
+- **Privacy-Preserving**: Implements techniques to ensure the generated data does not reveal identifiable information.
+- **Compatibility with OMOP**: Fully compatible with the OMOP common data model, allowing seamless integration with existing healthcare data systems.
+- **Extensible**: Designed to be adaptable to new datasets and different EHR systems.
+
+## Installation
+To install CEHRGPT, clone this repository and install the required dependencies.
+
+```bash
+git clone https://github.com/knatarajan-lab/cehrgpt.git
+cd cehrgpt
+pip install .
+```
+
+## Pretrain
+Pretrain cehrgpt using the Hugging Face trainer, the parameters can be found in the sample configuration yaml
+```bash
+mkdir test_results
+# This is NOT required when streaming is set to true
+mkdir test_dataset_prepared
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner sample_configs/cehrgpt_pretrain_sample_config.yaml
+```
+
+## Generate synthetic sequences
+Generate synthetic sequences using the trained model
+```bash
+export TRANSFORMERS_VERBOSITY=info
+export CUDA_VISIBLE_DEVICES="0"
+python -u -m cehrgpt.generation.generate_batch_hf_gpt_sequence \
+  --model_folder test_results \
+  --tokenizer_folder test_results \
+  --output_folder test_results \
+  --num_of_patients 128 \
+  --batch_size 32 \
+  --buffer_size 128 \
+  --context_window 1024 \
+  --sampling_strategy TopPStrategy \
+  --top_p 1.0 --temperature 1.0 --repetition_penalty 1.0 \
+  --epsilon_cutoff 0.00 \
+  --demographic_data_path sample_data/pretrain
+```
+
+## Convert synthetic sequences to OMOP
+```bash
+# omop converter requires the OHDSI vocabulary
+export OMOP_VOCAB_DIR = ""
+# the omop derived tables need to be built using pyspark
+export SPARK_WORKER_INSTANCES="1"
+export SPARK_WORKER_CORES="8"
+export SPARK_EXECUTOR_CORES="2"
+export SPARK_DRIVER_MEMORY="2g"
+export SPARK_EXECUTOR_MEMORY="2g"
+
+# Convert the sequences, create the omop derived tables
+sh scripts/omop_pipeline.sh \
+  test_results/top_p10000/generated_sequences/ \
+  test_results/top_p10000/restored_omop/ \
+  $OMOP_VOCAB_DIR
+```
+
+# MEDS Support
+
+This section demonstrates how to pretrain CEHR-GPT using MIMIC-IV data in the MEDS (Medical Event Data Standard) format.
+
+## Prerequisites
+
+Set up the required environment variables before beginning:
+
+```bash
+export CEHR_GPT_MODEL_DIR=""    # Path to CEHR-GPT model directory
+export MEDS_DIR=""              # Path to MEDS data directory
+export MEDS_READER_DIR=""       # Path to MEDS reader output directory
+```
+
+## Step 1: Create MIMIC MEDS Data
+
+Transform your MIMIC files into MEDS format by following the instructions in the [MEDS_transforms](https://github.com/mmcdermott/MEDS_transforms/) repository.
+
+## Step 2: Create the MEDS Reader
+
+Convert the MEDS data for use with CEHR-GPT:
+
+```bash
+meds_reader_convert $MEDS_DIR $MEDS_READER_DIR --num_threads 10
+```
+
+## Step 3: Pretrain CEHR-GPT
+
+Run the pretraining process using the prepared MEDS data:
+
+```bash
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $CEHR_GPT_MODEL_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --do_train true --seed 42 \
+  --dataloader_num_workers 16 --dataloader_prefetch_factor 8 \
+  --hidden_size 768 --num_hidden_layers 14 --max_position_embeddings 8192 \
+  --evaluation_strategy epoch --save_strategy epoch \
+  --sample_packing --max_tokens_per_batch 16384 \
+  --warmup_steps 500 --weight_decay 0.01 \
+  --num_train_epochs 50 --learning_rate 0.0002 \
+  --use_early_stopping --early_stopping_threshold 0.001 \
+  --is_data_in_meds --inpatient_att_function_type day \
+  --att_function_type day --include_inpatient_hour_token \
+  --include_auxiliary_token --include_demographic_prompt \
+  --meds_to_cehrbert_conversion_type "MedsToBertMimic4"
+```
+
+## Step 4: Generate MEDS Trajectories
+
+### Environment Setup for Trajectory Generation
+
+Configure additional environment variables for trajectory generation with task labels (`subject_id`, `prediction_time`, `boolean_value` [optional]):
+
+```bash
+# MEDS_LABEL_COHORT_DIR must contain a set of parquet files
+export MEDS_LABEL_COHORT_DIR=""     # Path to cohort labels directory
+export MEDS_TRAJECTORY_DIR=""       # Path for trajectory output
+```
+
+### Generate Trajectories
+
+Create synthetic patient trajectories using the trained model:
+
+> **Important:** The total sequence length (`generation_input_length` + `generation_max_new_tokens`) cannot exceed the `max_position_embeddings` value (8192) defined during pretraining.
+
+```bash
+python -u -m cehrgpt.generation.cehrgpt_conditional_generation \
+  --cohort_folder $MEDS_LABEL_COHORT_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $MEDS_TRAJECTORY_DIR \
+  --per_device_eval_batch_size 16 \
+  --num_of_trajectories_per_sample 2 \
+  --generation_input_length 4096 \
+  --generation_max_new_tokens 4096 \
+  --is_data_in_meds \
+  --att_function_type day --inpatient_att_function_type day \
+  --meds_to_cehrbert_conversion_type MedsToBertMimic4 \
+  --include_auxiliary_token --include_demographic_prompt \
+  --include_inpatient_hour_token
+```
+
+### Parameters Explanation
+
+- `generation_input_length`: Controls the length of input context for generation
+- `generation_max_new_tokens`: Maximum number of new tokens to generate
+- `num_of_trajectories_per_sample`: Number of trajectories to generate per patient sample
+
+## Citation
+```
+@article{cehrgpt2024,
+  title={CEHRGPT: Synthetic Data Generation for Electronic Health Records},
+  author={Natarajan, K and others},
+  journal={arXiv preprint arXiv:2402.04400},
+  year={2024}
+}

{cehrgpt-0.1.1 → cehrgpt-0.1.2}/src/cehrgpt/data/hf_cehrgpt_dataset_collator.py

@@ -162,6 +162,22 @@ class CehrGptDataCollator:
             f"batch['input_ids']: {batch['input_ids']} "
         )
 
+        if "epoch_times" in examples[0]:
+            batch_epoch_times = [
+                self._try_reverse_tensor(
+                    self._convert_to_tensor(example["epoch_times"])
+                )
+                for example in examples
+            ]
+            # Pad sequences to the max length in the batch
+            batch["epoch_times"] = self._try_reverse_tensor(
+                pad_sequence(
+                    batch_epoch_times,
+                    batch_first=True,
+                    padding_value=0,
+                ).to(torch.float32)
+            )
+
         if "position_ids" in examples[0]:
             batch_position_ids = [
                 self._try_reverse_tensor(
@@ -663,7 +679,9 @@ class CehrGptDataCollator:
 
         # Subtract one for the [END] token when sample_packing is not enabled
         new_max_length = (
-            max_length_allowed if sample_packing else max_length_allowed - 1
+            max_length_allowed - 1
+            if not sample_packing and self.pretraining
+            else max_length_allowed
         )
 
         if self.include_ttv_prediction:
@@ -685,13 +703,20 @@ class CehrGptDataCollator:
 
         # Return the record directly if the actual sequence length is less than the max sequence
         if seq_length <= new_max_length:
-            if not sample_packing:
+            if not sample_packing and self.pretraining:
                 record["input_ids"] = torch.concat(
                     [
                         self._convert_to_tensor(record["input_ids"]),
                         self._convert_to_tensor([eos_token]),
                     ]
                 )
+                if "epoch_times" in record:
+                    record["epoch_times"] = torch.concat(
+                        [
+                            self._convert_to_tensor(record["epoch_times"]),
+                            self._convert_to_tensor([record["epoch_times"][-1]]),
+                        ]
+                    )
                 if self.include_values:
                     record["value_indicators"] = torch.concat(
                         [
@@ -727,6 +752,10 @@ class CehrGptDataCollator:
                     record["input_ids"] = self._convert_to_tensor(
                         record["input_ids"][start_index : end_index + 1]
                     )
+                    if "epoch_times" in record:
+                        record["epoch_times"] = self._convert_to_tensor(
+                            record["epoch_times"][start_index : end_index + 1]
+                        )
                     if self.include_values:
                         record["value_indicators"] = self._convert_to_tensor(
                             record["value_indicators"][start_index : end_index + 1]
@@ -760,6 +789,11 @@ class CehrGptDataCollator:
             if sample_packing and "position_ids" in record:
                 record["position_ids"] = record["position_ids"][0:end_index]
 
+            if "epoch_times" in record:
+                record["epoch_times"] = self._convert_to_tensor(
+                    record["epoch_times"][0:end_index]
+                )
+
             if self.include_values:
                 record["value_indicators"] = self._convert_to_tensor(
                     record["value_indicators"][0:end_index]
@@ -792,6 +826,17 @@ class CehrGptDataCollator:
                                 ),
                             ]
                         )
+                        if "epoch_times" in record:
+                            record["epoch_times"] = torch.concat(
+                                [
+                                    torch.zeros(
+                                        [record["epoch_times"][0]], dtype=torch.float32
+                                    ),
+                                    self._convert_to_tensor(
+                                        record["epoch_times"][token_index:seq_length]
+                                    ),
+                                ]
+                            )
                         if self.include_values:
                             record["value_indicators"] = torch.concat(
                                 [
@@ -830,7 +875,7 @@ class CehrGptDataCollator:
                             )
                         break
             else:
-                start_index = seq_length - new_max_length
+                start_index = max(seq_length - new_max_length, 0)
                 end_index = seq_length
                 for i in range(start_index, end_index):
                     current_token = record["input_ids"][i]
@@ -842,6 +887,11 @@ class CehrGptDataCollator:
                             ]
                         if sample_packing and "position_ids" in record:
                             record["position_ids"] = record["position_ids"][i:end_index]
+
+                        if "epoch_times" in record:
+                            record["epoch_times"] = self._convert_to_tensor(
+                                record["epoch_times"][i:end_index]
+                            )
                         if self.include_values:
                             record["value_indicators"] = record["value_indicators"][
                                 i:end_index
@@ -863,6 +913,10 @@ class CehrGptDataCollator:
                     ]
                 if sample_packing and "position_ids" in record:
                     record["position_ids"] = record["position_ids"][-new_max_length:]
+                if "epoch_times" in record:
+                    record["epoch_times"] = self._convert_to_tensor(
+                        record["epoch_times"][-new_max_length:]
+                    )
                 if self.include_values:
                     record["value_indicators"] = record["value_indicators"][
                         -new_max_length:
@@ -873,36 +927,6 @@ class CehrGptDataCollator:
                         -new_max_length:
                     ]
 
-            if not sample_packing:
-                # Finally we add the end token to the end of the sequence
-                record["input_ids"] = torch.concat(
-                    [
-                        self._convert_to_tensor(record["input_ids"]),
-                        self._convert_to_tensor([eos_token]),
-                    ]
-                )
-                if self.include_values:
-                    record["value_indicators"] = torch.concat(
-                        [
-                            self._convert_to_tensor(record["value_indicators"]),
-                            self._convert_to_tensor([False]),
-                        ]
-                    ).to(torch.bool)
-                    record["values"] = torch.concat(
-                        [
-                            self._convert_to_tensor(record["values"]),
-                            self._convert_to_tensor(
-                                [self.tokenizer.pad_value_token_id]
-                            ),
-                        ]
-                    )
-                if self.include_ttv_prediction:
-                    record["time_to_visits"] = torch.concat(
-                        [
-                            record["time_to_visits"],
-                            self._convert_to_tensor([-100.0]),
-                        ]
-                    )
             return record
 
 

{cehrgpt-0.1.1 → cehrgpt-0.1.2}/src/cehrgpt/data/hf_cehrgpt_dataset_mapping.py

@@ -21,7 +21,6 @@ from cehrbert_data.const.artificial_tokens import (
     DISCHARGE_UNKNOWN_TOKEN,
     GENDER_UNKNOWN_TOKEN,
     RACE_UNKNOWN_TOKEN,
-    VISIT_UNKNOWN_TOKEN,
 )
 from cehrbert_data.const.common import NA
 from cehrbert_data.decorators.patient_event_decorator_base import get_att_function
@@ -47,10 +46,16 @@ CEHRGPT_COLUMNS = [
 ]
 
 
-def convert_date_to_posix_time(index_date: datetime.date) -> float:
-    return datetime.datetime.combine(
-        index_date, datetime.datetime.min.time()
-    ).timestamp()
+def convert_date_to_posix_time(index_date: Union[datetime.date, int, float]) -> float:
+    if isinstance(index_date, datetime.date):
+        return (
+            datetime.datetime.combine(index_date, datetime.datetime.min.time())
+            .replace(tzinfo=datetime.timezone.utc)
+            .timestamp()
+        )
+    elif isinstance(index_date, datetime.datetime):
+        return index_date.replace(tzinfo=datetime.timezone.utc).timestamp()
+    return index_date
 
 
 class DatasetMappingDecorator(DatasetMapping):
@@ -128,7 +133,9 @@ class MedToCehrGPTDatasetMapping(DatasetMappingDecorator):
         cehrgpt_record["concept_as_values"].append(concept_as_value)
         cehrgpt_record["units"].append(unit)
         cehrgpt_record["is_numeric_types"].append(is_numeric_type)
-        cehrgpt_record["epoch_times"].append(time.timestamp())
+        cehrgpt_record["epoch_times"].append(
+            time.replace(tzinfo=datetime.timezone.utc).timestamp()
+        )
 
     def transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
         cehrgpt_record = {
@@ -360,7 +367,9 @@ class MedToCehrGPTDatasetMapping(DatasetMappingDecorator):
         cehrgpt_record["num_of_visits"] = len(visits)
 
         if record.get("index_date", None) is not None:
-            cehrgpt_record["index_date"] = record["index_date"]
+            cehrgpt_record["index_date"] = (
+                record["index_date"].replace(tzinfo=datetime.timezone.utc).timestamp()
+            )
         if record.get("label", None) is not None:
             cehrgpt_record["label"] = record["label"]
         if record.get("age_at_index", None) is not None:
@@ -529,9 +538,13 @@ class ExtractTokenizedSequenceDataMapping:
         prediction_start_end_times = [
             (
                 self._calculate_prediction_start_time(
-                    prediction_time_label_map["index_date"].timestamp()
+                    prediction_time_label_map["index_date"]
+                    .replace(tzinfo=datetime.timezone.utc)
+                    .timestamp()
                 ),
-                prediction_time_label_map["index_date"].timestamp(),
+                prediction_time_label_map["index_date"]
+                .replace(tzinfo=datetime.timezone.utc)
+                .timestamp(),
                 prediction_time_label_map["label"],
             )
             for prediction_time_label_map in prediction_times