cehrgpt 0.0.1__tar.gz → 0.1.0__tar.gz

cehrgpt-0.1.0/.gitignore

@@ -0,0 +1,28 @@
+.DS_Store
+.idea/
+.vscode/
+venv*
+dist/*
+
+*png
+*json
+
+*ipynb_checkpoints/
+*h5
+*logs
+*nohup.out
+*ipynb
+
+*__pycache__/
+.eggs/
+*.dat
+.metastore_db/
+
+build/
+
+*.out
+*.egg-info/
+
+test_data
+test_dataset_prepared
+test*results

{cehrgpt-0.0.1/src/cehrgpt.egg-info → cehrgpt-0.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: cehrgpt
-Version: 0.0.1
+Version: 0.1.0
 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
@@ -12,12 +12,14 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: cehrbert==1.3.3
+Requires-Dist: cehrbert==1.4.1
+Requires-Dist: cehrbert_data==0.0.7
 Requires-Dist: openai==1.54.3
 Requires-Dist: optuna==4.0.0
-Requires-Dist: transformers==4.40.0
-Requires-Dist: tokenizers==0.19
-Requires-Dist: trl==0.11.4
+Requires-Dist: transformers==4.44.0
+Requires-Dist: tokenizers==0.19.0
+Requires-Dist: peft==0.10.0
+Requires-Dist: lightgbm
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
@@ -28,6 +30,7 @@ Requires-Dist: hypothesis; extra == "dev"
 Requires-Dist: black; extra == "dev"
 Provides-Extra: flash-attn
 Requires-Dist: flash_attn; extra == "flash-attn"
+Dynamic: license-file
 
 # CEHRGPT
 
@@ -50,11 +53,57 @@ CEHRGPT is a synthetic data generation model developed to handle structured elec
 To install CEHRGPT, clone this repository and install the required dependencies.
 
 ```bash
-git clone https://github.com/knatarajan-lab/cehrgpt-public.git
-cd cehrgpt-public
+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
+```
+
 ## Citation
 ```
 @article{cehrgpt2024,
@@ -63,4 +112,3 @@ pip install .
   journal={arXiv preprint arXiv:2402.04400},
   year={2024}
 }
-```

{cehrgpt-0.0.1 → cehrgpt-0.1.0}/README.md

@@ -19,11 +19,57 @@ CEHRGPT is a synthetic data generation model developed to handle structured elec
 To install CEHRGPT, clone this repository and install the required dependencies.
 
 ```bash
-git clone https://github.com/knatarajan-lab/cehrgpt-public.git
-cd cehrgpt-public
+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
+```
+
 ## Citation
 ```
 @article{cehrgpt2024,
@@ -31,5 +77,4 @@ pip install .
   author={Natarajan, K and others},
   journal={arXiv preprint arXiv:2402.04400},
   year={2024}
-}
-```
+}

{cehrgpt-0.0.1 → cehrgpt-0.1.0}/pyproject.toml

@@ -28,12 +28,14 @@ classifiers = [
 ]
 
 dependencies = [
-    "cehrbert==1.3.3",
+    "cehrbert==1.4.1",
+    "cehrbert_data==0.0.7",
     "openai==1.54.3",
     "optuna==4.0.0",
-    "transformers==4.40.0",
-    "tokenizers==0.19",
-    "trl==0.11.4",
+    "transformers==4.44.0",
+    "tokenizers==0.19.0",
+    "peft==0.10.0",
+    "lightgbm",
 ]
 
 [tool.setuptools_scm]

cehrgpt-0.1.0/sample_configs/cehrgpt_pretrain_sample_config.yaml

@@ -0,0 +1,51 @@
+model_name_or_path: "test_results"
+tokenizer_name_or_path: "test_results"
+
+data_folder: "sample_data/pretrain"
+dataset_prepared_path: "test_dataset_prepared"
+validation_split_percentage: 0.05
+validation_split_num: 10
+preprocessing_num_workers: 4
+preprocessing_batch_size: 1000
+streaming: true
+
+#Tokenizer
+vocab_size: 50000
+min_frequency: 0
+
+do_train: true
+overwrite_output_dir: false
+resume_from_checkpoint: # path to the checkpoint folder
+seed: 42
+
+num_hidden_layers: 6
+hidden_size: 768
+n_head: 12
+max_position_embeddings: 1024
+
+# torch dataloader configs
+dataloader_num_workers: 4
+dataloader_prefetch_factor: 2
+
+output_dir: "test_results"
+save_strategy: "steps"
+evaluation_strategy: "no"
+learning_rate: 0.00005
+per_device_train_batch_size: 4
+per_device_eval_batch_size: 4
+gradient_accumulation_steps: 1
+num_train_epochs: 1
+# When streaming is set to True, max_steps needs to be provided
+max_steps: 1000
+save_steps: 500
+
+warmup_steps: 100
+weight_decay: 0.01
+logging_dir: "./logs"
+logging_steps: 100
+save_total_limit: 5
+load_best_model_at_end: false
+metric_for_best_model: "eval_loss"
+greater_is_better: false
+
+report_to: "none"

{cehrgpt-0.0.1 → cehrgpt-0.1.0}/scripts/level_three_evaluation.sh

@@ -29,7 +29,8 @@ python -u -m cehrbert_data.prediction_cohorts.cad_cabg_cohort \
     -dl 1985-01-01 -du 2023-12-31 \
     -l 18 -u 100 -ow 360 -ps 0 -pw 360 -f \
     --att_type cehr_bert \
-    --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+    --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+    --is_remove_index_prediction_starts
 
 # Run Predictions on CAD CABG
 echo "Run predictions on cad_cabg"
@@ -56,9 +57,10 @@ python -u -m cehrbert_data.prediction_cohorts.hf_readmission \
   -c hf_readmission_bow \
   -i "$OMOP_FOLDER" \
   -o "$OMOP_FOLDER/cohorts/hf_readmission" \
-  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 360 -ps 0 -pw 30 -f \
+  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 360 -ps 1 -pw 30 -f \
   --att_type cehr_bert \
-  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+  --is_remove_index_prediction_starts
 
 # Run predictions on HF Readmission
 echo "Run predictions on hf_readmission"
@@ -85,9 +87,10 @@ python -u -m cehrbert_data.prediction_cohorts.copd_readmission \
   -c copd_readmission_bow \
   -i "$OMOP_FOLDER" \
   -o "$OMOP_FOLDER/cohorts/copd_readmission" \
-  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 720 -ps 0 -pw 360 -f \
+  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 360 -ps 1 -pw 30 -f \
   --att_type cehr_bert \
-  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+  --is_remove_index_prediction_starts
 
 # Run predictions on COPD Readmission
 echo "Run predictions on copd_readmission"
@@ -145,7 +148,8 @@ python -u -m cehrbert_data.prediction_cohorts.afib_ischemic_stroke \
   -o "$OMOP_FOLDER/cohorts/afib_ischemic_stroke" \
   -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 720 -ps 0 -pw 360 -f \
   --att_type cehr_bert \
-  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+  --is_remove_index_prediction_starts
 
 # Run predictions on AFIB Ischemic Stroke
 echo "Run predictions on afib_ischemic_stroke"

cehrgpt-0.1.0/scripts/omop_pipeline.sh

@@ -0,0 +1,55 @@
+#!/bin/bash
+
+# Exporting input arguments as environment variables
+export PATIENT_SEQUENCE_FOLDER="$1"
+export OMOP_FOLDER="$2"
+export SOURCE_OMOP_FOLDER="$3"
+export PATIENT_SPLITS_FOLDER="$SOURCE_OMOP_FOLDER/patient_splits"
+
+# Echoing the values of the environment variables
+echo "PATIENT_SEQUENCE_FOLDER=$PATIENT_SEQUENCE_FOLDER"
+echo "OMOP_FOLDER=$OMOP_FOLDER"
+echo "SOURCE_OMOP_FOLDER=$SOURCE_OMOP_FOLDER"
+
+# Ensure OMOP_FOLDER exists
+if [ ! -d "$OMOP_FOLDER" ]; then
+    echo "Creating $OMOP_FOLDER"
+    mkdir -p "$OMOP_FOLDER"
+fi
+
+# Removing existing OMOP tables
+rm -rf $OMOP_FOLDER/{person,visit_occurrence,condition_occurrence,procedure_occurrence,drug_exposure,death,measurement,observation_period,condition_era}
+
+# Removing existing OMOP concept tables
+rm -rf $OMOP_FOLDER/{concept,concept_ancestor,concept_relationship}
+
+# Copying OMOP concept tables if they don't already exist
+for table in concept concept_relationship concept_ancestor; do
+    if [ ! -d "$OMOP_FOLDER/$table" ]; then
+        echo "Creating $OMOP_FOLDER/$table"
+        cp -r "$SOURCE_OMOP_FOLDER/$table" "$OMOP_FOLDER/$table"
+    fi
+done
+
+# Reconstructing the OMOP instance from patient sequences
+echo "Reconstructing the OMOP instance from patient sequences in $OMOP_FOLDER"
+python -m cehrgpt.generation.omop_converter_batch \
+  --patient_sequence_path "$PATIENT_SEQUENCE_FOLDER" \
+  --output_folder "$OMOP_FOLDER" \
+  --concept_path "$OMOP_FOLDER/concept" \
+  --buffer_size 1280 \
+  --cpu_cores 10
+
+# Create observation_period
+echo "Reconstructing observation_period in $OMOP_FOLDER"
+python -u -m cehrgpt.omop.observation_period \
+  --input_folder "$OMOP_FOLDER" \
+  --output_folder "$OMOP_FOLDER" \
+  --domain_table_list "condition_occurrence drug_exposure procedure_occurrence measurement"
+
+# Create condition_era
+echo "Reconstructing condition_era in $OMOP_FOLDER"
+python -u -m cehrgpt.omop.condition_era \
+  --input_folder "$OMOP_FOLDER" \
+  --output_folder "$OMOP_FOLDER" \
+  --domain_table_list "condition_occurrence"

{cehrgpt-0.0.1 → cehrgpt-0.1.0}/src/cehrgpt/data/hf_cehrgpt_dataset.py

@@ -1,9 +1,10 @@
-from typing import Union
+from typing import Optional, Union
 
 from cehrbert.data_generators.hf_data_generator.hf_dataset import (
     FINETUNING_COLUMNS,
     apply_cehrbert_dataset_mapping,
 )
+from cehrbert.data_generators.hf_data_generator.meds_utils import CacheFileCollector
 from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments
 from datasets import Dataset, DatasetDict
 
@@ -31,16 +32,25 @@ def create_cehrgpt_pretraining_dataset(
     dataset: Union[Dataset, DatasetDict],
     cehrgpt_tokenizer: CehrGptTokenizer,
     data_args: DataTrainingArguments,
-) -> Dataset:
+    cache_file_collector: Optional[CacheFileCollector] = None,
+) -> Union[Dataset, DatasetDict]:
     required_columns = TRANSFORMER_COLUMNS + CEHRGPT_COLUMNS
+    # TODO: temp solution, this column is mixed typed and causes an issue when transforming the data
+    if not data_args.streaming:
+        if isinstance(dataset, DatasetDict):
+            all_columns = dataset["train"].column_names
+        else:
+            all_columns = dataset.column_names
+        if "visit_concept_ids" in all_columns:
+            dataset.remove_columns(["visit_concept_ids"])
     dataset = apply_cehrbert_dataset_mapping(
         dataset,
         HFCehrGptTokenizationMapping(cehrgpt_tokenizer),
         num_proc=data_args.preprocessing_num_workers,
         batch_size=data_args.preprocessing_batch_size,
         streaming=data_args.streaming,
+        cache_file_collector=cache_file_collector,
     )
-
     if not data_args.streaming:
         if isinstance(dataset, DatasetDict):
             all_columns = dataset["train"].column_names
@@ -56,8 +66,17 @@ def create_cehrgpt_finetuning_dataset(
     dataset: Union[Dataset, DatasetDict],
     cehrgpt_tokenizer: CehrGptTokenizer,
     data_args: DataTrainingArguments,
-) -> Dataset:
+    cache_file_collector: Optional[CacheFileCollector] = None,
+) -> Union[Dataset, DatasetDict]:
     required_columns = TRANSFORMER_COLUMNS + CEHRGPT_COLUMNS + FINETUNING_COLUMNS
+    # TODO: temp solution, this column is mixed typed and causes an issue when transforming the data
+    if not data_args.streaming:
+        if isinstance(dataset, DatasetDict):
+            all_columns = dataset["train"].column_names
+        else:
+            all_columns = dataset.column_names
+        if "visit_concept_ids" in all_columns:
+            dataset.remove_columns(["visit_concept_ids"])
     mapping_functions = [
         HFFineTuningMapping(cehrgpt_tokenizer),
     ]
@@ -68,6 +87,7 @@ def create_cehrgpt_finetuning_dataset(
             num_proc=data_args.preprocessing_num_workers,
             batch_size=data_args.preprocessing_batch_size,
             streaming=data_args.streaming,
+            cache_file_collector=cache_file_collector,
         )
 
     if not data_args.streaming: