cehrgpt 0.0.1__py3-none-any.whl

cehrgpt/runners/hf_cehrgpt_pretrain_runner.py

@@ -0,0 +1,370 @@
+import os
+from typing import Optional, Union
+
+import torch
+from cehrbert.data_generators.hf_data_generator.meds_utils import (
+    create_dataset_from_meds_reader,
+)
+from cehrbert.runners.hf_runner_argument_dataclass import (
+    DataTrainingArguments,
+    ModelArguments,
+)
+from cehrbert.runners.runner_util import (
+    generate_prepared_ds_path,
+    get_last_hf_checkpoint,
+    get_meds_extension_path,
+    load_parquet_as_dataset,
+)
+from datasets import Dataset, DatasetDict, IterableDatasetDict, load_from_disk
+from transformers import AutoConfig, Trainer, TrainingArguments, set_seed
+from transformers.utils import is_flash_attn_2_available, logging
+
+from cehrgpt.data.hf_cehrgpt_dataset import create_cehrgpt_pretraining_dataset
+from cehrgpt.data.hf_cehrgpt_dataset_collator import CehrGptDataCollator
+from cehrgpt.models.config import CEHRGPTConfig
+from cehrgpt.models.hf_cehrgpt import CEHRGPT2LMHeadModel
+from cehrgpt.models.pretrained_embeddings import PretrainedEmbeddings
+from cehrgpt.models.tokenization_hf_cehrgpt import CehrGptTokenizer
+from cehrgpt.runners.gpt_runner_util import parse_runner_args
+from src.cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
+
+LOG = logging.get_logger("transformers")
+
+
+def tokenizer_exists(tokenizer_name_or_path: str) -> bool:
+    # Try to load the pretrained tokenizer
+    try:
+        CehrGptTokenizer.from_pretrained(os.path.abspath(tokenizer_name_or_path))
+        return True
+    except Exception:
+        LOG.info(f"The tokenizer does not exist at {tokenizer_name_or_path}")
+        return False
+
+
+def load_and_create_tokenizer(
+    data_args: DataTrainingArguments,
+    model_args: ModelArguments,
+    cehrgpt_args: CehrGPTArguments,
+    dataset: Optional[Union[Dataset, DatasetDict]] = None,
+) -> CehrGptTokenizer:
+    # Try to load the pretrained tokenizer
+    tokenizer_abspath = os.path.expanduser(model_args.tokenizer_name_or_path)
+    try:
+        tokenizer = CehrGptTokenizer.from_pretrained(tokenizer_abspath)
+    except Exception as e:
+        LOG.warning(e)
+        if dataset is None:
+            raise RuntimeError(
+                f"Failed to load the tokenizer from {tokenizer_abspath} with the error \n{e}\n"
+                f"Tried to create the tokenizer, however the dataset is not provided."
+            )
+        tokenizer = CehrGptTokenizer.train_tokenizer(
+            dataset,
+            {},
+            data_args,
+            PretrainedEmbeddings(cehrgpt_args.pretrained_embedding_path),
+        )
+        tokenizer.save_pretrained(tokenizer_abspath)
+
+    return tokenizer
+
+
+def load_and_create_model(
+    model_args: ModelArguments,
+    cehrgpt_args: CehrGPTArguments,
+    training_args: TrainingArguments,
+    tokenizer: CehrGptTokenizer,
+) -> CEHRGPT2LMHeadModel:
+    attn_implementation = (
+        "flash_attention_2" if is_flash_attn_2_available() else "eager"
+    )
+    torch_dtype = torch.bfloat16 if training_args.bf16 else torch.float32
+    model_abspath = os.path.expanduser(model_args.model_name_or_path)
+    if cehrgpt_args.continue_pretrain:
+        try:
+            return CEHRGPT2LMHeadModel.from_pretrained(
+                model_abspath,
+                attn_implementation=attn_implementation,
+                torch_dtype=torch_dtype,
+            )
+        except Exception as e:
+            LOG.error(
+                f"When continue_pretrain is set to True, it assumes that CEHR-GPT has been trained "
+                f"and will be used to pretrain on new datasets. The CEHR-GPT checkpoint must exist at {model_abspath}"
+            )
+            raise e
+    try:
+        model_config = AutoConfig.from_pretrained(
+            model_abspath, attn_implementation=attn_implementation
+        )
+    except Exception as e:
+        LOG.warning(e)
+        if cehrgpt_args.causal_sfm:
+            model_args.max_position_embeddings += 1
+        if len(tokenizer.pretrained_token_ids) > 0:
+            pretrained_embedding_dim = tokenizer.pretrained_embeddings.shape[1]
+        else:
+            pretrained_embedding_dim = model_args.hidden_size
+        model_config = CEHRGPTConfig(
+            vocab_size=tokenizer.vocab_size,
+            value_vocab_size=tokenizer.value_vocab_size,
+            time_token_vocab_size=tokenizer.time_token_vocab_size,
+            bos_token_id=tokenizer.end_token_id,
+            eos_token_id=tokenizer.end_token_id,
+            lab_token_ids=tokenizer.lab_token_ids,
+            token_to_time_token_mapping=tokenizer.token_to_time_token_mapping,
+            attn_implementation=attn_implementation,
+            causal_sfm=cehrgpt_args.causal_sfm,
+            demographics_size=cehrgpt_args.demographics_size,
+            lab_token_penalty=cehrgpt_args.lab_token_penalty,
+            lab_token_loss_weight=cehrgpt_args.lab_token_loss_weight,
+            entropy_penalty=cehrgpt_args.entropy_penalty,
+            entropy_penalty_alpha=cehrgpt_args.entropy_penalty_alpha,
+            n_pretrained_embeddings_layers=cehrgpt_args.n_pretrained_embeddings_layers,
+            use_pretrained_embeddings=len(tokenizer.pretrained_token_ids) > 0,
+            pretrained_embedding_dim=pretrained_embedding_dim,
+            **model_args.as_dict(),
+        )
+    model = CEHRGPT2LMHeadModel(model_config)
+    if tokenizer.pretrained_token_ids:
+        model.cehrgpt.update_pretrained_embeddings(
+            tokenizer.pretrained_token_ids,
+            tokenizer.pretrained_embeddings,
+        )
+    if model.config.torch_dtype == torch.bfloat16:
+        return model.bfloat16()
+    elif model.config.torch_dtype == torch.float16:
+        return model.half()
+    return model
+
+
+def main():
+    cehrgpt_args, data_args, model_args, training_args = parse_runner_args()
+
+    if data_args.streaming:
+        # This is for disabling the warning message https://github.com/huggingface/transformers/issues/5486
+        # This happens only when streaming is enabled
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        # The iterable dataset doesn't have sharding implemented, so the number of works has to be set to 0
+        # Otherwise the trainer will throw an error
+        training_args.dataloader_num_workers = 0
+        training_args.dataloader_prefetch_factor = 0
+
+    prepared_ds_path = generate_prepared_ds_path(data_args, model_args)
+    if os.path.exists(os.path.join(data_args.data_folder, "dataset_dict.json")):
+        LOG.info(f"Loading prepared dataset from disk at {data_args.data_folder}...")
+        processed_dataset = load_from_disk(data_args.data_folder)
+        # If the data has been processed in the past, it's assume the tokenizer has been created before.
+        # we load the CEHR-GPT tokenizer from the output folder, otherwise an exception will be raised.
+        tokenizer_name_or_path = os.path.expanduser(
+            training_args.output_dir
+            if cehrgpt_args.expand_tokenizer
+            else model_args.tokenizer_name_or_path
+        )
+        if not tokenizer_exists(tokenizer_name_or_path):
+            raise RuntimeError(
+                f"The dataset has been tokenized but the corresponding tokenizer: "
+                f"{model_args.tokenizer_name_or_path} does not exist"
+            )
+        cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(tokenizer_name_or_path)
+    elif any(prepared_ds_path.glob("*")):
+        LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
+        processed_dataset = load_from_disk(str(prepared_ds_path))
+        LOG.info("Prepared dataset loaded from disk...")
+        # If the data has been processed in the past, it's assume the tokenizer has been created before.
+        # we load the CEHR-GPT tokenizer from the output folder, otherwise an exception will be raised.
+        tokenizer_name_or_path = os.path.expanduser(
+            training_args.output_dir
+            if cehrgpt_args.expand_tokenizer
+            else model_args.tokenizer_name_or_path
+        )
+        if not tokenizer_exists(tokenizer_name_or_path):
+            raise RuntimeError(
+                f"The dataset has been tokenized but the corresponding tokenizer: "
+                f"{model_args.tokenizer_name_or_path} does not exist"
+            )
+        cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(tokenizer_name_or_path)
+    else:
+        # If the data is in the MEDS format, we need to convert it to the CEHR-BERT format
+        if data_args.is_data_in_meds:
+            meds_extension_path = get_meds_extension_path(
+                data_folder=data_args.data_folder,
+                dataset_prepared_path=data_args.dataset_prepared_path,
+            )
+            try:
+                LOG.info(
+                    "Trying to load the MEDS extension from disk at %s...",
+                    meds_extension_path,
+                )
+                dataset = load_from_disk(meds_extension_path)
+                if data_args.streaming:
+                    if isinstance(dataset, DatasetDict):
+                        dataset = {
+                            k: v.to_iterable_dataset(
+                                num_shards=training_args.dataloader_num_workers
+                            )
+                            for k, v in dataset.items()
+                        }
+                    else:
+                        dataset = dataset.to_iterable_dataset(
+                            num_shards=training_args.dataloader_num_workers
+                        )
+            except FileNotFoundError as e:
+                LOG.exception(e)
+                dataset = create_dataset_from_meds_reader(
+                    data_args, is_pretraining=True
+                )
+                if not data_args.streaming:
+                    dataset.save_to_disk(meds_extension_path)
+        else:
+            # Load the dataset from the parquet files
+            dataset = load_parquet_as_dataset(
+                data_args.data_folder, split="train", streaming=data_args.streaming
+            )
+            # If streaming is enabled, we need to manually split the data into train/val
+            if data_args.streaming and data_args.validation_split_num:
+                dataset = dataset.shuffle(buffer_size=10_000, seed=training_args.seed)
+                train_set = dataset.skip(data_args.validation_split_num)
+                val_set = dataset.take(data_args.validation_split_num)
+                dataset = DatasetDict({"train": train_set, "test": val_set})
+            elif data_args.validation_split_percentage:
+                dataset = dataset.train_test_split(
+                    test_size=data_args.validation_split_percentage,
+                    seed=training_args.seed,
+                )
+            else:
+                raise RuntimeError(
+                    f"Can not split the data. If streaming is enabled, validation_split_num needs to be "
+                    f"defined, otherwise validation_split_percentage needs to be provided. "
+                    f"The current values are:\n"
+                    f"validation_split_percentage: {data_args.validation_split_percentage}\n"
+                    f"validation_split_num: {data_args.validation_split_num}\n"
+                    f"streaming: {data_args.streaming}"
+                )
+
+        # Create the CEHR-GPT tokenizer if it's not available in the output folder
+        cehrgpt_tokenizer = load_and_create_tokenizer(
+            data_args=data_args,
+            model_args=model_args,
+            cehrgpt_args=cehrgpt_args,
+            dataset=dataset,
+        )
+        # Retrain the tokenizer in case we want to pretrain the model further using different datasets
+        if cehrgpt_args.expand_tokenizer:
+            new_tokenizer_path = os.path.expanduser(training_args.output_dir)
+            try:
+                cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(new_tokenizer_path)
+            except Exception:
+                cehrgpt_tokenizer = CehrGptTokenizer.expand_trained_tokenizer(
+                    cehrgpt_tokenizer=cehrgpt_tokenizer,
+                    dataset=dataset["train"],
+                    data_args=data_args,
+                    concept_name_mapping={},
+                    pretrained_concept_embedding_model=PretrainedEmbeddings(
+                        cehrgpt_args.pretrained_embedding_path
+                    ),
+                )
+                cehrgpt_tokenizer.save_pretrained(
+                    os.path.expanduser(training_args.output_dir)
+                )
+
+        # sort the patient features chronologically and tokenize the data
+        processed_dataset = create_cehrgpt_pretraining_dataset(
+            dataset=dataset, cehrgpt_tokenizer=cehrgpt_tokenizer, data_args=data_args
+        )
+        # only save the data to the disk if it is not streaming
+        if not data_args.streaming:
+            processed_dataset.save_to_disk(prepared_ds_path)
+
+    def filter_func(examples):
+        if cehrgpt_args.drop_long_sequences:
+            return [
+                model_args.max_position_embeddings >= _ >= data_args.min_num_tokens
+                for _ in examples["num_of_concepts"]
+            ]
+        else:
+            return [_ >= data_args.min_num_tokens for _ in examples["num_of_concepts"]]
+
+    # Create the args for batched filtering
+    filter_args = {"batched": True, "batch_size": data_args.preprocessing_batch_size}
+    # If the dataset is not in a streaming mode, we could add num_proc to enable parallelization
+    if not data_args.streaming:
+        filter_args["num_proc"] = data_args.preprocessing_num_workers
+
+    # The filter can't be applied to a DatasetDict of IterableDataset (in case of streaming)
+    # we need to iterate through all the datasets and apply the filter separately
+    if isinstance(processed_dataset, DatasetDict) or isinstance(
+        processed_dataset, IterableDatasetDict
+    ):
+        for key in processed_dataset.keys():
+            processed_dataset[key] = processed_dataset[key].filter(
+                filter_func, **filter_args
+            )
+    else:
+        processed_dataset = processed_dataset.filter(filter_func, **filter_args)
+
+    model = load_and_create_model(
+        model_args, cehrgpt_args, training_args, cehrgpt_tokenizer
+    )
+
+    # Expand tokenizer to adapt to the new pretraining dataset
+    if model.config.vocab_size < cehrgpt_tokenizer.vocab_size:
+        model.resize_token_embeddings(cehrgpt_tokenizer.vocab_size)
+        # Update the pretrained embedding weights if they are available
+        if model.config.use_pretrained_embeddings:
+            model.cehrgpt.update_pretrained_embeddings(
+                cehrgpt_tokenizer.pretrained_token_ids,
+                cehrgpt_tokenizer.pretrained_embeddings,
+            )
+        elif cehrgpt_tokenizer.pretrained_token_ids:
+            model.config.pretrained_embedding_dim = (
+                cehrgpt_tokenizer.pretrained_embeddings.shape[1]
+            )
+            model.config.use_pretrained_embeddings = True
+            model.cehrgpt.initialize_pretrained_embeddings()
+            model.cehrgpt.update_pretrained_embeddings(
+                cehrgpt_tokenizer.pretrained_token_ids,
+                cehrgpt_tokenizer.pretrained_embeddings,
+            )
+
+    # Detecting last checkpoint.
+    last_checkpoint = get_last_hf_checkpoint(training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    if not data_args.streaming:
+        processed_dataset.set_format("pt")
+
+    trainer = Trainer(
+        model=model,
+        data_collator=CehrGptDataCollator(
+            tokenizer=cehrgpt_tokenizer,
+            max_length=model_args.max_position_embeddings,
+            shuffle_records=data_args.shuffle_records,
+            include_ttv_prediction=model_args.include_ttv_prediction,
+            use_sub_time_tokenization=model_args.use_sub_time_tokenization,
+            include_values=model_args.include_values,
+        ),
+        train_dataset=processed_dataset["train"],
+        eval_dataset=processed_dataset["test"],
+        args=training_args,
+    )
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    train_result = trainer.train(resume_from_checkpoint=checkpoint)
+    trainer.save_model()  # Saves the tokenizer too for easy upload
+    metrics = train_result.metrics
+
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+    trainer.save_state()
+
+
+if __name__ == "__main__":
+    main()

cehrgpt/runners/hf_gpt_runner_argument_dataclass.py

@@ -0,0 +1,137 @@
+import dataclasses
+from typing import List, Optional
+
+
+@dataclasses.dataclass
+class CehrGPTArguments:
+    """Arguments pertaining to what data we are going to input our model for training and eval."""
+
+    include_demographics: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to always include the demographics for the long sequences that are longer than the model context window."
+        },
+    )
+    continue_pretrain: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to continue to pretrain cehrgpt on the new dataset"
+        },
+    )
+    pretrained_embedding_path: Optional[str] = dataclasses.field(
+        default=None,
+        metadata={"help": "The path to the concept pretrained embeddings"},
+    )
+    retrain_with_full: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to retrain the model on the full set after early stopping"
+        },
+    )
+    expand_tokenizer: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to expand the tokenizer for fine-tuning."
+        },
+    )
+    few_shot_predict: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to use a few shots to train the model"
+        },
+    )
+    n_shots: Optional[int] = dataclasses.field(
+        default=128,
+        metadata={"help": "The number of examples from the training set."},
+    )
+    hyperparameter_tuning_percentage: Optional[float] = dataclasses.field(
+        default=0.1,
+        metadata={
+            "help": "The percentage of the train/val will be use for hyperparameter tuning."
+        },
+    )
+    n_trials: Optional[int] = dataclasses.field(
+        default=10,
+        metadata={
+            "help": "The number of trails will be use for hyperparameter tuning."
+        },
+    )
+    hyperparameter_tuning: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={"help": "A flag to indicate if we want to do hyperparameter tuning."},
+    )
+    hyperparameter_batch_sizes: Optional[List[int]] = dataclasses.field(
+        default_factory=lambda: [4, 8, 16],
+        metadata={"help": "Hyperparameter search batch sizes"},
+    )
+    hyperparameter_num_train_epochs: Optional[List[int]] = dataclasses.field(
+        default_factory=lambda: [10],
+        metadata={"help": "Hyperparameter search num_train_epochs"},
+    )
+    lr_low: Optional[float] = dataclasses.field(
+        default=1e-5,
+        metadata={
+            "help": "The lower bound of the learning rate range for hyperparameter tuning."
+        },
+    )
+    lr_high: Optional[float] = dataclasses.field(
+        default=5e-5,
+        metadata={
+            "help": "The upper bound of the learning rate range for hyperparameter tuning."
+        },
+    )
+    weight_decays_low: Optional[float] = dataclasses.field(
+        default=1e-3,
+        metadata={
+            "help": "The lower bound of the weight decays range for hyperparameter tuning."
+        },
+    )
+    weight_decays_high: Optional[float] = dataclasses.field(
+        default=1e-2,
+        metadata={
+            "help": "The upper bound of the weight decays range for hyperparameter tuning."
+        },
+    )
+    causal_sfm: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether the GPT conforms to the causal Standard Fairness Model"
+        },
+    )
+    demographics_size: Optional[int] = dataclasses.field(
+        default=4,
+        metadata={
+            "help": "The number of demographics tokens in the patient sequence "
+            "It defaults to 4, assuming the demographics tokens follow this pattern [Year][Age][Gender][Race]"
+        },
+    )
+    drop_long_sequences: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "The lower bound of the learning rate range for hyperparameter tuning."
+        },
+    )
+    lab_token_penalty: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to use lab token loss penalty."
+        },
+    )
+    lab_token_loss_weight: Optional[float] = dataclasses.field(
+        default=1.0,
+        metadata={"help": "lab_token_loss_weight penalty co-efficient"},
+    )
+    entropy_penalty: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={"help": "A flag to indicate whether we want to use entropy penalty."},
+    )
+    entropy_penalty_alpha: Optional[float] = dataclasses.field(
+        default=0.01,
+        metadata={"help": "Entropy penalty co-efficient"},
+    )
+    n_pretrained_embeddings_layers: Optional[int] = dataclasses.field(
+        default=2,
+        metadata={
+            "help": "The number of feed forward layers for transforming pretrained embeddings to internal embeddings"
+        },
+    )