cehrgpt 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

cehrgpt/runners/hf_cehrgpt_finetune_runner.py

@@ -9,20 +9,16 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 import torch
-from cehrbert.data_generators.hf_data_generator.meds_utils import (
-    create_dataset_from_meds_reader,
-)
+import torch.distributed as dist
+from cehrbert.data_generators.hf_data_generator.meds_utils import CacheFileCollector
 from cehrbert.runners.hf_cehrbert_finetune_runner import compute_metrics
 from cehrbert.runners.hf_runner_argument_dataclass import (
-    DataTrainingArguments,
     FineTuneModelType,
     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 DatasetDict, concatenate_datasets, load_from_disk
 from peft import LoraConfig, PeftModel, get_peft_model
@@ -38,11 +34,15 @@ from transformers import (
     TrainingArguments,
     set_seed,
 )
-from transformers.tokenization_utils_base import LARGE_INTEGER
+from transformers.trainer_utils import is_main_process
 from transformers.utils import is_flash_attn_2_available, logging
 
 from cehrgpt.data.hf_cehrgpt_dataset import create_cehrgpt_finetuning_dataset
-from cehrgpt.data.hf_cehrgpt_dataset_collator import CehrGptDataCollator
+from cehrgpt.data.hf_cehrgpt_dataset_collator import (
+    CehrGptDataCollator,
+    SamplePackingCehrGptDataCollator,
+)
+from cehrgpt.data.sample_packing_sampler import SamplePackingBatchSampler
 from cehrgpt.models.hf_cehrgpt import (
     CEHRGPTConfig,
     CehrGptForClassification,
@@ -50,9 +50,12 @@ from cehrgpt.models.hf_cehrgpt import (
 )
 from cehrgpt.models.pretrained_embeddings import PretrainedEmbeddings
 from cehrgpt.models.tokenization_hf_cehrgpt import CehrGptTokenizer
+from cehrgpt.runners.data_utils import prepare_finetune_dataset
 from cehrgpt.runners.gpt_runner_util import parse_runner_args
+from cehrgpt.runners.hf_cehrgpt_pretrain_runner import tokenizer_exists
 from cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
 from cehrgpt.runners.hyperparameter_search_util import perform_hyperparameter_search
+from cehrgpt.runners.sample_packing_trainer import SamplePackingTrainer
 
 LOG = logging.get_logger("transformers")
 
@@ -155,140 +158,6 @@ def load_finetuned_model(
         raise ValueError(f"Can not load the finetuned model from {model_name_or_path}")
 
 
-def create_dataset_splits(data_args: DataTrainingArguments, seed: int):
-    """
-    Creates training, validation, and testing dataset splits based on specified splitting strategies.
-
-    This function splits a dataset into training, validation, and test sets, using either chronological,
-    patient-based, or random splitting strategies, depending on the parameters provided in `data_args`.
-
-    - **Chronological split**: Sorts by a specified date and splits based on historical and future data.
-    - **Patient-based split**: Splits by unique patient IDs to ensure that patients in each split are distinct.
-    - **Random split**: Performs a straightforward random split of the dataset.
-
-    If `data_args.test_data_folder` is provided, a test set is loaded directly from it. Otherwise,
-    the test set is created by further splitting the validation set based on `test_eval_ratio`.
-
-    Parameters:
-        data_args (DataTrainingArguments): A configuration object containing data-related arguments, including:
-            - `data_folder` (str): Path to the main dataset.
-            - `test_data_folder` (str, optional): Path to an optional test dataset.
-            - `chronological_split` (bool): Whether to split chronologically.
-            - `split_by_patient` (bool): Whether to split by unique patient IDs.
-            - `validation_split_percentage` (float): Percentage of data to use for validation.
-            - `test_eval_ratio` (float): Ratio of test to validation data when creating a test set from validation.
-            - `preprocessing_num_workers` (int): Number of processes for parallel data filtering.
-            - `preprocessing_batch_size` (int): Batch size for batched operations.
-        seed (int): Random seed for reproducibility of splits.
-
-    Returns:
-        Tuple[Dataset, Dataset, Dataset]: A tuple containing:
-            - `train_set` (Dataset): Training split of the dataset.
-            - `validation_set` (Dataset): Validation split of the dataset.
-            - `test_set` (Dataset): Test split of the dataset.
-
-    Raises:
-        FileNotFoundError: If `data_args.data_folder` or `data_args.test_data_folder` does not exist.
-        ValueError: If incompatible arguments are passed for splitting strategies.
-
-    Example Usage:
-        data_args = DataTrainingArguments(
-            data_folder="data/",
-            validation_split_percentage=0.1,
-            test_eval_ratio=0.2,
-            chronological_split=True
-        )
-        train_set, validation_set, test_set = create_dataset_splits(data_args, seed=42)
-    """
-    dataset = load_parquet_as_dataset(data_args.data_folder)
-    test_set = (
-        None
-        if not data_args.test_data_folder
-        else load_parquet_as_dataset(data_args.test_data_folder)
-    )
-
-    if data_args.chronological_split:
-        # Chronological split by sorting on `index_date`
-        dataset = dataset.sort("index_date")
-        total_size = len(dataset)
-        train_end = int((1 - data_args.validation_split_percentage) * total_size)
-
-        # Perform the split
-        train_set = dataset.select(range(0, train_end))
-        validation_set = dataset.select(range(train_end, total_size))
-
-        if test_set is None:
-            test_valid_split = validation_set.train_test_split(
-                test_size=data_args.test_eval_ratio, seed=seed
-            )
-            validation_set, test_set = (
-                test_valid_split["train"],
-                test_valid_split["test"],
-            )
-
-    elif data_args.split_by_patient:
-        # Patient-based split
-        LOG.info("Using the split_by_patient strategy")
-        unique_patient_ids = dataset.unique("person_id")
-        LOG.info(f"There are {len(unique_patient_ids)} patients in total")
-
-        np.random.seed(seed)
-        np.random.shuffle(unique_patient_ids)
-
-        train_end = int(
-            len(unique_patient_ids) * (1 - data_args.validation_split_percentage)
-        )
-        train_patient_ids = set(unique_patient_ids[:train_end])
-
-        if test_set is None:
-            validation_end = int(
-                train_end
-                + len(unique_patient_ids)
-                * data_args.validation_split_percentage
-                * data_args.test_eval_ratio
-            )
-            val_patient_ids = set(unique_patient_ids[train_end:validation_end])
-            test_patient_ids = set(unique_patient_ids[validation_end:])
-        else:
-            val_patient_ids, test_patient_ids = (
-                set(unique_patient_ids[train_end:]),
-                None,
-            )
-
-        # Helper function to apply patient-based filtering
-        def filter_by_patient_ids(patient_ids):
-            return dataset.filter(
-                lambda batch: [pid in patient_ids for pid in batch["person_id"]],
-                num_proc=data_args.preprocessing_num_workers,
-                batched=True,
-                batch_size=data_args.preprocessing_batch_size,
-            )
-
-        # Generate splits
-        train_set = filter_by_patient_ids(train_patient_ids)
-        validation_set = filter_by_patient_ids(val_patient_ids)
-        if test_set is None:
-            test_set = filter_by_patient_ids(test_patient_ids)
-
-    else:
-        # Random split
-        train_val = dataset.train_test_split(
-            test_size=data_args.validation_split_percentage, seed=seed
-        )
-        train_set, validation_set = train_val["train"], train_val["test"]
-
-        if test_set is None:
-            test_valid_split = validation_set.train_test_split(
-                test_size=data_args.test_eval_ratio, seed=seed
-            )
-            validation_set, test_set = (
-                test_valid_split["train"],
-                test_valid_split["test"],
-            )
-
-    return train_set, validation_set, test_set
-
-
 def model_init(
     model_args: ModelArguments,
     training_args: TrainingArguments,
@@ -363,16 +232,16 @@ def main():
     prepared_ds_path = generate_prepared_ds_path(
         data_args, model_args, data_folder=data_args.cohort_folder
     )
-
+    cache_file_collector = CacheFileCollector()
     processed_dataset = None
     if 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 cehrgpt_args.expand_tokenizer:
-            try:
+            if tokenizer_exists(training_args.output_dir):
                 tokenizer = CehrGptTokenizer.from_pretrained(training_args.output_dir)
-            except Exception:
+            else:
                 LOG.warning(
                     f"CehrGptTokenizer must exist in {training_args.output_dir} "
                     f"when the dataset has been processed and expand_tokenizer is set to True. "
@@ -382,81 +251,77 @@ def main():
                 shutil.rmtree(prepared_ds_path)
 
     if processed_dataset is None:
-        # 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.cohort_folder,
-                dataset_prepared_path=data_args.dataset_prepared_path,
+        if is_main_process(training_args.local_rank):
+            final_splits = prepare_finetune_dataset(
+                data_args, training_args, cehrgpt_args, cache_file_collector
             )
-            try:
-                LOG.info(
-                    f"Trying to load the MEDS extension from disk at {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
+            if cehrgpt_args.expand_tokenizer:
+                new_tokenizer_path = os.path.expanduser(training_args.output_dir)
+                if tokenizer_exists(new_tokenizer_path):
+                    tokenizer = CehrGptTokenizer.from_pretrained(new_tokenizer_path)
+                else:
+                    # Try to use the defined pretrained embeddings if exists, Otherwise we default to the pretrained model
+                    # embedded in the pretrained model
+                    pretrained_concept_embedding_model = PretrainedEmbeddings(
+                        cehrgpt_args.pretrained_embedding_path
+                    )
+                    if not pretrained_concept_embedding_model.exists:
+                        pretrained_concept_embedding_model = (
+                            tokenizer.pretrained_concept_embedding_model
                         )
-            except Exception as e:
-                LOG.exception(e)
-                dataset = create_dataset_from_meds_reader(
-                    data_args, is_pretraining=False
-                )
-                if not data_args.streaming:
-                    dataset.save_to_disk(meds_extension_path)
-            train_set = dataset["train"]
-            validation_set = dataset["validation"]
-            test_set = dataset["test"]
-        else:
-            train_set, validation_set, test_set = create_dataset_splits(
-                data_args=data_args, seed=training_args.seed
-            )
-        # Organize them into a single DatasetDict
-        final_splits = DatasetDict(
-            {"train": train_set, "validation": validation_set, "test": test_set}
-        )
-
-        if cehrgpt_args.expand_tokenizer:
-            new_tokenizer_path = os.path.expanduser(training_args.output_dir)
-            try:
-                tokenizer = CehrGptTokenizer.from_pretrained(new_tokenizer_path)
-            except Exception:
-                # Try to use the defined pretrained embeddings if exists,
-                # Otherwise we default to the pretrained model embedded in the pretrained model
-                pretrained_concept_embedding_model = PretrainedEmbeddings(
-                    cehrgpt_args.pretrained_embedding_path
-                )
-                if not pretrained_concept_embedding_model.exists:
-                    pretrained_concept_embedding_model = (
-                        tokenizer.pretrained_concept_embedding_model
+                    tokenizer = CehrGptTokenizer.expand_trained_tokenizer(
+                        cehrgpt_tokenizer=tokenizer,
+                        dataset=final_splits["train"],
+                        data_args=data_args,
+                        concept_name_mapping={},
+                        pretrained_concept_embedding_model=pretrained_concept_embedding_model,
                     )
-                tokenizer = CehrGptTokenizer.expand_trained_tokenizer(
-                    cehrgpt_tokenizer=tokenizer,
-                    dataset=final_splits["train"],
-                    data_args=data_args,
-                    concept_name_mapping={},
-                    pretrained_concept_embedding_model=pretrained_concept_embedding_model,
+                    tokenizer.save_pretrained(
+                        os.path.expanduser(training_args.output_dir)
+                    )
+
+            # TODO: temp solution, this column is mixed typed and causes an issue when transforming the data
+            if not data_args.streaming:
+                all_columns = final_splits["train"].column_names
+                if "visit_concept_ids" in all_columns:
+                    final_splits = final_splits.remove_columns(["visit_concept_ids"])
+
+            processed_dataset = create_cehrgpt_finetuning_dataset(
+                dataset=final_splits,
+                cehrgpt_tokenizer=tokenizer,
+                data_args=data_args,
+                cache_file_collector=cache_file_collector,
+            )
+            if not data_args.streaming:
+                processed_dataset.save_to_disk(str(prepared_ds_path))
+                stats = processed_dataset.cleanup_cache_files()
+                LOG.info(
+                    "Clean up the cached files for the  cehrgpt finetuning dataset : %s",
+                    stats,
                 )
-                tokenizer.save_pretrained(os.path.expanduser(training_args.output_dir))
 
-        processed_dataset = create_cehrgpt_finetuning_dataset(
-            dataset=final_splits, cehrgpt_tokenizer=tokenizer, data_args=data_args
+            # Remove any cached files if there are any
+            cache_file_collector.remove_cache_files()
+
+        # After main-process-only operations, synchronize all processes to ensure consistency
+        if dist.is_available() and dist.is_initialized():
+            dist.barrier()
+
+        # Loading tokenizer in all processes in torch distributed training
+        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 data_args.streaming:
-            processed_dataset.save_to_disk(prepared_ds_path)
+        tokenizer = CehrGptTokenizer.from_pretrained(tokenizer_name_or_path)
+        # Load the dataset from disk again to in torch distributed training
+        processed_dataset = load_from_disk(str(prepared_ds_path))
 
     # Set seed before initializing model.
     set_seed(training_args.seed)
 
-    processed_dataset.set_format("pt")
+    if not data_args.streaming and not cehrgpt_args.sample_packing:
+        processed_dataset.set_format("pt")
 
     if cehrgpt_args.few_shot_predict:
         # At least we need two examples to have a validation set for early stopping
@@ -476,13 +341,40 @@ def main():
     config = CEHRGPTConfig.from_pretrained(model_args.model_name_or_path)
     if config.max_position_embeddings < model_args.max_position_embeddings:
         config.max_position_embeddings = model_args.max_position_embeddings
+
+    # persist this parameter in case this is overwritten by sample packing
+    per_device_eval_batch_size = training_args.per_device_eval_batch_size
+
+    if cehrgpt_args.sample_packing:
+        trainer_class = partial(
+            SamplePackingTrainer,
+            max_tokens_per_batch=cehrgpt_args.max_tokens_per_batch,
+            max_position_embeddings=config.max_position_embeddings,
+            train_lengths=processed_dataset["train"]["num_of_concepts"],
+            validation_lengths=processed_dataset["validation"]["num_of_concepts"],
+        )
+        training_args.per_device_train_batch_size = 1
+        training_args.per_device_eval_batch_size = 1
+        data_collator_fn = partial(
+            SamplePackingCehrGptDataCollator,
+            cehrgpt_args.max_tokens_per_batch,
+            config.max_position_embeddings,
+        )
+    else:
+        trainer_class = Trainer
+        data_collator_fn = CehrGptDataCollator
+
     # We suppress the additional learning objectives in fine-tuning
-    data_collator = CehrGptDataCollator(
+    data_collator = data_collator_fn(
         tokenizer=tokenizer,
         max_length=(
-            config.max_position_embeddings - 1
-            if config.causal_sfm
-            else config.max_position_embeddings
+            cehrgpt_args.max_tokens_per_batch
+            if cehrgpt_args.sample_packing
+            else (
+                config.max_position_embeddings - 1
+                if config.causal_sfm
+                else config.max_position_embeddings
+            )
         ),
         include_values=model_args.include_values,
         pretraining=False,
@@ -493,8 +385,8 @@ def main():
 
     if training_args.do_train:
         if cehrgpt_args.hyperparameter_tuning:
-            model_args.early_stopping_patience = LARGE_INTEGER
             training_args = perform_hyperparameter_search(
+                trainer_class,
                 partial(model_init, model_args, training_args, tokenizer),
                 processed_dataset,
                 data_collator,
@@ -502,13 +394,20 @@ def main():
                 model_args,
                 cehrgpt_args,
             )
+
+        if cehrgpt_args.retrain_with_full:
             # Always retrain with the full set when hyperparameter tuning is set to true
             retrain_with_full_set(
-                model_args, training_args, tokenizer, processed_dataset, data_collator
+                trainer_class,
+                model_args,
+                training_args,
+                tokenizer,
+                processed_dataset,
+                data_collator,
             )
         else:
             # Initialize Trainer for final training on the combined train+val set
-            trainer = Trainer(
+            trainer = trainer_class(
                 model=model_init(model_args, training_args, tokenizer),
                 data_collator=data_collator,
                 args=training_args,
@@ -531,45 +430,31 @@ def main():
             trainer.save_metrics("train", metrics)
             trainer.save_state()
 
-            # Retrain the model with full set using the num of epoches before earlying stopping
-            if cehrgpt_args.retrain_with_full:
-                update_num_epoch_before_early_stopping_callback = None
-                for callback in trainer.callback_handler.callbacks:
-                    if isinstance(callback, UpdateNumEpochsBeforeEarlyStoppingCallback):
-                        update_num_epoch_before_early_stopping_callback = callback
-
-                if update_num_epoch_before_early_stopping_callback is None:
-                    raise RuntimeError(
-                        f"{UpdateNumEpochsBeforeEarlyStoppingCallback} must be included as a callback!"
-                    )
-                final_num_epochs = (
-                    update_num_epoch_before_early_stopping_callback.num_epochs_before_early_stopping
-                )
-                training_args.num_train_epochs = final_num_epochs
-                LOG.info(
-                    "Num Epochs before early stopping: %s",
-                    training_args.num_train_epochs,
-                )
-                retrain_with_full_set(
-                    model_args,
-                    training_args,
-                    tokenizer,
-                    processed_dataset,
-                    data_collator,
-                )
-
     if training_args.do_predict:
+        if cehrgpt_args.sample_packing:
+            batch_sampler = SamplePackingBatchSampler(
+                lengths=processed_dataset["test"]["num_of_concepts"],
+                max_tokens_per_batch=cehrgpt_args.max_tokens_per_batch,
+                max_position_embeddings=config.max_position_embeddings,
+                drop_last=training_args.dataloader_drop_last,
+                seed=training_args.seed,
+            )
+            per_device_eval_batch_size = 1
+        else:
+            batch_sampler = None
         test_dataloader = DataLoader(
             dataset=processed_dataset["test"],
-            batch_size=training_args.per_device_eval_batch_size,
+            batch_size=per_device_eval_batch_size,
             num_workers=training_args.dataloader_num_workers,
             collate_fn=data_collator,
             pin_memory=training_args.dataloader_pin_memory,
+            batch_sampler=batch_sampler,
         )
         do_predict(test_dataloader, model_args, training_args, cehrgpt_args)
 
 
 def retrain_with_full_set(
+    trainer_class,
     model_args: ModelArguments,
     training_args: TrainingArguments,
     tokenizer: CehrGptTokenizer,
@@ -586,6 +471,7 @@ def retrain_with_full_set(
     and state information.
 
     Args:
+        trainer_class: Trainer or its subclass
         model_args (ModelArguments): Model configuration and hyperparameters.
         training_args (TrainingArguments): Training configuration, including output directory,
                                            evaluation strategy, and other training parameters.
@@ -607,7 +493,7 @@ def retrain_with_full_set(
     # Disable evaluation
     training_args.evaluation_strategy = "no"
     checkpoint = get_last_hf_checkpoint(training_args)
-    final_trainer = Trainer(
+    final_trainer = trainer_class(
         model=model_init(model_args, training_args, tokenizer),
         data_collator=data_collator,
         args=training_args,
@@ -662,15 +548,15 @@ def do_predict(
     test_losses = []
     with torch.no_grad():
         for index, batch in enumerate(tqdm(test_dataloader, desc="Predicting")):
-            person_ids = batch.pop("person_id").numpy().squeeze().astype(int)
-            index_dates = (
-                map(
-                    datetime.fromtimestamp,
-                    batch.pop("index_date").numpy().squeeze(axis=-1).tolist(),
-                )
-                if "index_date" in batch
-                else None
-            )
+            person_ids = batch.pop("person_id").numpy().astype(int).squeeze()
+            if person_ids.ndim == 0:
+                person_ids = np.asarray([person_ids])
+
+            index_dates = batch.pop("index_date").numpy().squeeze()
+            if index_dates.ndim == 0:
+                index_dates = np.asarray([index_dates])
+            index_dates = list(map(datetime.fromtimestamp, index_dates.tolist()))
+
             batch = {k: v.to(device) for k, v in batch.items()}
             # Forward pass
             output = model(**batch, output_attentions=False, output_hidden_states=False)
@@ -678,17 +564,25 @@ def do_predict(
 
             # Collect logits and labels for prediction
             logits = output.logits.float().cpu().numpy().squeeze()
+            if logits.ndim == 0:
+                logits = np.asarray([logits])
+            probabilities = sigmoid(logits)
+
             labels = (
-                batch["classifier_label"].float().cpu().numpy().squeeze().astype(bool)
+                batch["classifier_label"].float().cpu().numpy().astype(bool).squeeze()
             )
-            probabilities = sigmoid(logits)
+            if labels.ndim == 0:
+                labels = np.asarray([labels])
+
             # Save predictions to parquet file
             test_prediction_pd = pd.DataFrame(
                 {
                     "subject_id": person_ids,
                     "prediction_time": index_dates,
-                    "boolean_prediction_probability": probabilities,
-                    "boolean_prediction": logits,
+                    "predicted_boolean_probability": probabilities,
+                    "predicted_boolean_value": pd.Series(
+                        [None] * len(person_ids), dtype=bool
+                    ),
                     "boolean_value": labels,
                 }
             )
@@ -702,7 +596,7 @@ def do_predict(
     # Compute metrics and save results
     metrics = compute_metrics(
         references=test_prediction_pd.boolean_value,
-        probs=test_prediction_pd.boolean_prediction_probability,
+        probs=test_prediction_pd.predicted_boolean_probability,
     )
     metrics["test_loss"] = np.mean(test_losses)