cehrgpt 0.0.1__py3-none-any.whl

cehrgpt/runners/hf_cehrgpt_finetune_runner.py

@@ -0,0 +1,746 @@
+import json
+import os
+import random
+import shutil
+from datetime import datetime
+from functools import partial
+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,
+)
+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
+from scipy.special import expit as sigmoid
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import (
+    EarlyStoppingCallback,
+    Trainer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+    set_seed,
+)
+from transformers.tokenization_utils_base import LARGE_INTEGER
+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.models.hf_cehrgpt import (
+    CEHRGPTConfig,
+    CehrGptForClassification,
+    CEHRGPTPreTrainedModel,
+)
+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 cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
+from cehrgpt.runners.hyperparameter_search_util import perform_hyperparameter_search
+
+LOG = logging.get_logger("transformers")
+
+
+class UpdateNumEpochsBeforeEarlyStoppingCallback(TrainerCallback):
+    """
+    Callback to update metrics with the number of epochs completed before early stopping.
+
+    based on the best evaluation metric (e.g., eval_loss).
+    """
+
+    def __init__(self, model_folder: str):
+        self._model_folder = model_folder
+        self._metrics_path = os.path.join(
+            model_folder, "num_epochs_trained_before_early_stopping.json"
+        )
+        self._num_epochs_before_early_stopping = 0
+        self._best_val_loss = float("inf")
+
+    @property
+    def num_epochs_before_early_stopping(self):
+        return self._num_epochs_before_early_stopping
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        if os.path.exists(self._metrics_path):
+            with open(self._metrics_path, "r") as f:
+                metrics = json.load(f)
+            self._num_epochs_before_early_stopping = metrics[
+                "num_epochs_before_early_stopping"
+            ]
+            self._best_val_loss = metrics["best_val_loss"]
+
+    def on_evaluate(self, args, state, control, **kwargs):
+        # Ensure metrics is available in kwargs
+        metrics = kwargs.get("metrics")
+        if metrics is not None and "eval_loss" in metrics:
+            # Check and update if a new best metric is achieved
+            if metrics["eval_loss"] < self._best_val_loss:
+                self._num_epochs_before_early_stopping = round(state.epoch)
+                self._best_val_loss = metrics["eval_loss"]
+
+    def on_save(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        with open(self._metrics_path, "w") as f:
+            json.dump(
+                {
+                    "num_epochs_before_early_stopping": self._num_epochs_before_early_stopping,
+                    "best_val_loss": self._best_val_loss,
+                },
+                f,
+            )
+
+
+def load_pretrained_tokenizer(
+    model_args,
+) -> CehrGptTokenizer:
+    try:
+        return CehrGptTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
+    except Exception:
+        raise ValueError(
+            f"Can not load the pretrained tokenizer from {model_args.tokenizer_name_or_path}"
+        )
+
+
+def load_finetuned_model(
+    model_args: ModelArguments,
+    training_args: TrainingArguments,
+    model_name_or_path: str,
+) -> CEHRGPTPreTrainedModel:
+    if model_args.finetune_model_type == FineTuneModelType.POOLING.value:
+        finetune_model_cls = CehrGptForClassification
+    else:
+        raise ValueError(
+            f"finetune_model_type can be one of the following types {FineTuneModelType.POOLING.value}"
+        )
+
+    attn_implementation = (
+        "flash_attention_2" if is_flash_attn_2_available() else "eager"
+    )
+    torch_dtype = torch.bfloat16 if training_args.bf16 else torch.float32
+    # Try to create a new model based on the base model
+    try:
+        return finetune_model_cls.from_pretrained(
+            model_name_or_path,
+            attn_implementation=attn_implementation,
+            torch_dtype=torch_dtype,
+        )
+    except ValueError:
+        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,
+    tokenizer: CehrGptTokenizer,
+):
+    model = load_finetuned_model(
+        model_args, training_args, model_args.model_name_or_path
+    )
+    if model.config.max_position_embeddings < model_args.max_position_embeddings:
+        LOG.info(
+            f"Increase model.config.max_position_embeddings to {model_args.max_position_embeddings}"
+        )
+        model.config.max_position_embeddings = model_args.max_position_embeddings
+        model.resize_position_embeddings(model_args.max_position_embeddings)
+    # Enable include_values when include_values is set to be False during pre-training
+    if model_args.include_values and not model.cehrgpt.include_values:
+        model.cehrgpt.include_values = True
+    # Enable position embeddings when position embeddings are disabled in pre-training
+    if not model_args.exclude_position_ids and model.cehrgpt.exclude_position_ids:
+        model.cehrgpt.exclude_position_ids = False
+    # Expand tokenizer to adapt to the finetuning dataset
+    if model.config.vocab_size < tokenizer.vocab_size:
+        model.resize_token_embeddings(tokenizer.vocab_size)
+        # Update the pretrained embedding weights if they are available
+        if model.config.use_pretrained_embeddings:
+            model.cehrgpt.update_pretrained_embeddings(
+                tokenizer.pretrained_token_ids, tokenizer.pretrained_embeddings
+            )
+        elif tokenizer.pretrained_token_ids:
+            model.config.pretrained_embedding_dim = (
+                tokenizer.pretrained_embeddings.shape[1]
+            )
+            model.config.use_pretrained_embeddings = True
+            model.cehrgpt.initialize_pretrained_embeddings()
+            model.cehrgpt.update_pretrained_embeddings(
+                tokenizer.pretrained_token_ids, tokenizer.pretrained_embeddings
+            )
+    # Expand value tokenizer to adapt to the fine-tuning dataset
+    if model.config.include_values:
+        if model.config.value_vocab_size < tokenizer.value_vocab_size:
+            model.resize_value_embeddings(tokenizer.value_vocab_size)
+    # If lora is enabled, we add LORA adapters to the model
+    if model_args.use_lora:
+        # When LORA is used, the trainer could not automatically find this label,
+        # therefore we need to manually set label_names to "classifier_label" so the model
+        # can compute the loss during the evaluation
+        if training_args.label_names:
+            training_args.label_names.append("classifier_label")
+        else:
+            training_args.label_names = ["classifier_label"]
+
+        if model_args.finetune_model_type == FineTuneModelType.POOLING.value:
+            config = LoraConfig(
+                r=model_args.lora_rank,
+                lora_alpha=model_args.lora_alpha,
+                target_modules=model_args.target_modules,
+                lora_dropout=model_args.lora_dropout,
+                bias="none",
+                modules_to_save=["classifier", "age_batch_norm", "dense_layer"],
+            )
+            model = get_peft_model(model, config)
+        else:
+            raise ValueError(
+                f"The LORA adapter is not supported for {model_args.finetune_model_type}"
+            )
+    return model
+
+
+def main():
+    cehrgpt_args, data_args, model_args, training_args = parse_runner_args()
+    tokenizer = load_pretrained_tokenizer(model_args)
+    prepared_ds_path = generate_prepared_ds_path(
+        data_args, model_args, data_folder=data_args.cohort_folder
+    )
+
+    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:
+                tokenizer = CehrGptTokenizer.from_pretrained(training_args.output_dir)
+            except Exception:
+                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. "
+                    f"Please delete the processed dataset at {prepared_ds_path}."
+                )
+                processed_dataset = None
+                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,
+            )
+            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
+                        )
+            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.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
+        )
+        if not data_args.streaming:
+            processed_dataset.save_to_disk(prepared_ds_path)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    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
+        num_shots = max(cehrgpt_args.n_shots, 2)
+        random_train_indices = random.sample(
+            range(len(processed_dataset["train"])), k=num_shots
+        )
+        test_size = max(int(num_shots * data_args.validation_split_percentage), 1)
+        few_shot_train_val_set = processed_dataset["train"].select(random_train_indices)
+        train_val = few_shot_train_val_set.train_test_split(
+            test_size=test_size, seed=training_args.seed
+        )
+        few_shot_train_set, few_shot_val_set = train_val["train"], train_val["test"]
+        processed_dataset["train"] = few_shot_train_set
+        processed_dataset["validation"] = few_shot_val_set
+
+    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
+    # We suppress the additional learning objectives in fine-tuning
+    data_collator = CehrGptDataCollator(
+        tokenizer=tokenizer,
+        max_length=(
+            config.max_position_embeddings - 1
+            if config.causal_sfm
+            else config.max_position_embeddings
+        ),
+        include_values=model_args.include_values,
+        pretraining=False,
+        include_ttv_prediction=False,
+        use_sub_time_tokenization=False,
+        include_demographics=cehrgpt_args.include_demographics,
+    )
+
+    if training_args.do_train:
+        if cehrgpt_args.hyperparameter_tuning:
+            model_args.early_stopping_patience = LARGE_INTEGER
+            training_args = perform_hyperparameter_search(
+                partial(model_init, model_args, training_args, tokenizer),
+                processed_dataset,
+                data_collator,
+                training_args,
+                model_args,
+                cehrgpt_args,
+            )
+            # 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
+            )
+        else:
+            # Initialize Trainer for final training on the combined train+val set
+            trainer = Trainer(
+                model=model_init(model_args, training_args, tokenizer),
+                data_collator=data_collator,
+                args=training_args,
+                train_dataset=processed_dataset["train"],
+                eval_dataset=processed_dataset["validation"],
+                callbacks=[
+                    EarlyStoppingCallback(model_args.early_stopping_patience),
+                    UpdateNumEpochsBeforeEarlyStoppingCallback(
+                        training_args.output_dir
+                    ),
+                ],
+                tokenizer=tokenizer,
+            )
+            # Train the model on the combined train + val set
+            checkpoint = get_last_hf_checkpoint(training_args)
+            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()
+
+            # 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:
+        test_dataloader = DataLoader(
+            dataset=processed_dataset["test"],
+            batch_size=training_args.per_device_eval_batch_size,
+            num_workers=training_args.dataloader_num_workers,
+            collate_fn=data_collator,
+            pin_memory=training_args.dataloader_pin_memory,
+        )
+        do_predict(test_dataloader, model_args, training_args, cehrgpt_args)
+
+
+def retrain_with_full_set(
+    model_args: ModelArguments,
+    training_args: TrainingArguments,
+    tokenizer: CehrGptTokenizer,
+    dataset: DatasetDict,
+    data_collator: CehrGptDataCollator,
+) -> None:
+    """
+    Retrains a model on the full training and validation dataset for final performance evaluation.
+
+    This function consolidates the training and validation datasets into a single
+    dataset for final model training, updates the output directory for the final model,
+    and disables evaluation during training. It resumes from the latest checkpoint if available,
+    trains the model on the combined dataset, and saves the model along with training metrics
+    and state information.
+
+    Args:
+        model_args (ModelArguments): Model configuration and hyperparameters.
+        training_args (TrainingArguments): Training configuration, including output directory,
+                                           evaluation strategy, and other training parameters.
+        tokenizer (CehrGptTokenizer): Tokenizer instance specific to CEHR-GPT.
+        dataset (DatasetDict): A dictionary containing the 'train' and 'validation' datasets.
+        data_collator (CehrGptDataCollator): Data collator for handling data batching and tokenization.
+
+    Returns:
+        None
+    """
+    # Initialize Trainer for final training on the combined train+val set
+    full_dataset = concatenate_datasets([dataset["train"], dataset["validation"]])
+    training_args.output_dir = os.path.join(training_args.output_dir, "full")
+    LOG.info(
+        "Final output_dir for final_training_args.output_dir %s",
+        training_args.output_dir,
+    )
+    Path(training_args.output_dir).mkdir(exist_ok=True)
+    # Disable evaluation
+    training_args.evaluation_strategy = "no"
+    checkpoint = get_last_hf_checkpoint(training_args)
+    final_trainer = Trainer(
+        model=model_init(model_args, training_args, tokenizer),
+        data_collator=data_collator,
+        args=training_args,
+        train_dataset=full_dataset,
+        tokenizer=tokenizer,
+    )
+    final_train_result = final_trainer.train(resume_from_checkpoint=checkpoint)
+    final_trainer.save_model()  # Saves the tokenizer too for easy upload
+    metrics = final_train_result.metrics
+    final_trainer.log_metrics("train", metrics)
+    final_trainer.save_metrics("train", metrics)
+    final_trainer.save_state()
+
+
+def do_predict(
+    test_dataloader: DataLoader,
+    model_args: ModelArguments,
+    training_args: TrainingArguments,
+    cehrgpt_args: CehrGPTArguments,
+):
+    """
+    Performs inference on the test dataset using a fine-tuned model, saves predictions and evaluation metrics.
+
+    The reason we created this custom do_predict is that there is a memory leakage for transformers trainer.predict(),
+    for large test sets, it will throw the CPU OOM error
+
+    Args:
+        test_dataloader (DataLoader): DataLoader containing the test dataset, with batches of input features and labels.
+        model_args (ModelArguments): Arguments for configuring and loading the fine-tuned model.
+        training_args (TrainingArguments): Arguments related to training, evaluation, and output directories.
+        cehrgpt_args (CehrGPTArguments):
+    Returns:
+        None. Results are saved to disk.
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load model and LoRA adapters if applicable
+    model = (
+        load_finetuned_model(model_args, training_args, training_args.output_dir)
+        if not model_args.use_lora
+        else load_lora_model(model_args, training_args, cehrgpt_args)
+    )
+
+    model = model.to(device).eval()
+
+    # Ensure prediction folder exists
+    test_prediction_folder = Path(training_args.output_dir) / "test_predictions"
+    test_prediction_folder.mkdir(parents=True, exist_ok=True)
+
+    LOG.info("Generating predictions for test set at %s", test_prediction_folder)
+
+    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
+            )
+            batch = {k: v.to(device) for k, v in batch.items()}
+            # Forward pass
+            output = model(**batch, output_attentions=False, output_hidden_states=False)
+            test_losses.append(output.loss.item())
+
+            # Collect logits and labels for prediction
+            logits = output.logits.float().cpu().numpy().squeeze()
+            labels = (
+                batch["classifier_label"].float().cpu().numpy().squeeze().astype(bool)
+            )
+            probabilities = sigmoid(logits)
+            # 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,
+                    "boolean_value": labels,
+                }
+            )
+            test_prediction_pd.to_parquet(test_prediction_folder / f"{index}.parquet")
+
+    LOG.info(
+        "Computing metrics using the test set predictions at %s", test_prediction_folder
+    )
+    # Load all predictions
+    test_prediction_pd = pd.read_parquet(test_prediction_folder)
+    # Compute metrics and save results
+    metrics = compute_metrics(
+        references=test_prediction_pd.boolean_value,
+        probs=test_prediction_pd.boolean_prediction_probability,
+    )
+    metrics["test_loss"] = np.mean(test_losses)
+
+    test_results_path = Path(training_args.output_dir) / "test_results.json"
+    with open(test_results_path, "w") as f:
+        json.dump(metrics, f, indent=4)
+
+    LOG.info("Test results: %s", metrics)
+
+
+def load_lora_model(
+    model_args: ModelArguments,
+    training_args: TrainingArguments,
+    cehrgpt_args: CehrGPTArguments,
+) -> PeftModel:
+    LOG.info("Loading base model from %s", model_args.model_name_or_path)
+    model = load_finetuned_model(
+        model_args, training_args, model_args.model_name_or_path
+    )
+    # Enable include_values when include_values is set to be False during pre-training
+    if model_args.include_values and not model.cehrgpt.include_values:
+        model.cehrgpt.include_values = True
+    # Enable position embeddings when position embeddings are disabled in pre-training
+    if not model_args.exclude_position_ids and model.cehrgpt.exclude_position_ids:
+        model.cehrgpt.exclude_position_ids = False
+    if cehrgpt_args.expand_tokenizer:
+        tokenizer = CehrGptTokenizer.from_pretrained(training_args.output_dir)
+        # Expand tokenizer to adapt to the finetuning dataset
+        if model.config.vocab_size < tokenizer.vocab_size:
+            model.resize_token_embeddings(tokenizer.vocab_size)
+        if (
+            model.config.include_values
+            and model.config.value_vocab_size < tokenizer.value_vocab_size
+        ):
+            model.resize_value_embeddings(tokenizer.value_vocab_size)
+    LOG.info("Loading LoRA adapter from %s", training_args.output_dir)
+    return PeftModel.from_pretrained(model, model_id=training_args.output_dir)
+
+
+if __name__ == "__main__":
+    main()