cehrgpt 0.0.2__py3-none-any.whl → 0.1.1__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,12 +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_mapping import MedToCehrGPTDatasetMapping
+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,
@@ -51,9 +50,16 @@ 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 (
+    extract_cohort_sequences,
+    get_torch_dtype,
+    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")
 
@@ -140,11 +146,10 @@ def load_finetuned_model(
         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
+    torch_dtype = get_torch_dtype(model_args.torch_dtype)
     # Try to create a new model based on the base model
     try:
         return finetune_model_cls.from_pretrained(
@@ -156,148 +161,25 @@ 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,
+    cehrgpt_args: CehrGPTArguments,
     tokenizer: CehrGptTokenizer,
 ):
     model = load_finetuned_model(
         model_args, training_args, model_args.model_name_or_path
     )
+
+    if cehrgpt_args.class_weights:
+        model.config.class_weights = cehrgpt_args.class_weights
+        LOG.info(f"Setting class_weights to {model.config.class_weights}")
+
+    # Enable position embeddings when position embeddings are disabled in pre-training
+    if not model_args.exclude_position_ids and model.cehrgpt.exclude_position_ids:
+        LOG.info(f"Enable the position_embeddings")
+        model.cehrgpt.enable_position_embeddings()
+
     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}"
@@ -307,9 +189,6 @@ def model_init(
     # 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)
@@ -327,6 +206,7 @@ def model_init(
             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:
@@ -364,16 +244,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. "
@@ -383,101 +263,86 @@ 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,
-            )
-            try:
-                LOG.info(
-                    f"Trying to load the MEDS extension from disk at {meds_extension_path}..."
+        if is_main_process(training_args.local_rank):
+            # If the full dataset has been tokenized, we don't want to tokenize the cohort containing
+            # the subset of the data. We should slice out the portion of the tokenized sequences for each sample
+            if cehrgpt_args.tokenized_full_dataset_path is not None:
+                processed_dataset = extract_cohort_sequences(
+                    data_args, cehrgpt_args, cache_file_collector
                 )
-                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:
+                final_splits = prepare_finetune_dataset(
+                    data_args, training_args, cehrgpt_args, cache_file_collector
+                )
+                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:
-                        dataset = dataset.to_iterable_dataset(
-                            num_shards=training_args.dataloader_num_workers
+                        # 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
                         )
-            except Exception as e:
-                LOG.exception(e)
-                dataset = create_dataset_from_meds_reader(
-                    data_args=data_args,
-                    dataset_mappings=[
-                        MedToCehrGPTDatasetMapping(
+                        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,
-                            is_pretraining=False,
-                            include_inpatient_hour_token=cehrgpt_args.include_inpatient_hour_token,
+                            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:
-                    dataset.save_to_disk(str(meds_extension_path))
-                    stats = dataset.cleanup_cache_files()
-                    LOG.info(
-                        "Clean up the cached files for the cehrgpt dataset transformed from the MEDS: %s",
-                        stats,
-                    )
-                    dataset = load_from_disk(str(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}
-        )
+                    all_columns = final_splits["train"].column_names
+                    if "visit_concept_ids" in all_columns:
+                        final_splits = final_splits.remove_columns(
+                            ["visit_concept_ids"]
+                        )
 
-        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(
+                processed_dataset = create_cehrgpt_finetuning_dataset(
+                    dataset=final_splits,
                     cehrgpt_tokenizer=tokenizer,
-                    dataset=final_splits["train"],
                     data_args=data_args,
-                    concept_name_mapping={},
-                    pretrained_concept_embedding_model=pretrained_concept_embedding_model,
+                    cache_file_collector=cache_file_collector,
                 )
-                tokenizer.save_pretrained(os.path.expanduser(training_args.output_dir))
+            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,
+                )
+
+            # Remove any cached files if there are any
+            cache_file_collector.remove_cache_files()
 
-        processed_dataset = create_cehrgpt_finetuning_dataset(
-            dataset=final_splits, cehrgpt_tokenizer=tokenizer, data_args=data_args
+        # 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(str(prepared_ds_path))
-            stats = processed_dataset.cleanup_cache_files()
-            LOG.info(
-                "Clean up the cached files for the  cehrgpt finetuning dataset : %s",
-                stats,
-            )
-            processed_dataset = load_from_disk(str(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
@@ -497,40 +362,76 @@ 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,
+            negative_sampling_probability=cehrgpt_args.negative_sampling_probability,
+        )
+        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,
+            add_end_token_in_sample_packing=cehrgpt_args.add_end_token_in_sample_packing,
+        )
+    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,
         include_ttv_prediction=False,
         use_sub_time_tokenization=False,
         include_demographics=cehrgpt_args.include_demographics,
+        add_linear_prob_token=True,
     )
 
     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),
+                trainer_class,
+                partial(model_init, model_args, training_args, cehrgpt_args, tokenizer),
                 processed_dataset,
                 data_collator,
                 training_args,
                 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,
+                cehrgpt_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),
+            trainer = trainer_class(
+                model=model_init(model_args, training_args, cehrgpt_args, tokenizer),
                 data_collator=data_collator,
                 args=training_args,
                 train_dataset=processed_dataset["train"],
@@ -552,47 +453,34 @@ 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,
+    cehrgpt_args: CehrGPTArguments,
     tokenizer: CehrGptTokenizer,
     dataset: DatasetDict,
     data_collator: CehrGptDataCollator,
@@ -607,9 +495,11 @@ 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.
+        cehrgpt_args (CehrGPTArguments): CehrGPT specific 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.
@@ -628,8 +518,8 @@ def retrain_with_full_set(
     # 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),
+    final_trainer = trainer_class(
+        model=model_init(model_args, training_args, cehrgpt_args, tokenizer),
         data_collator=data_collator,
         args=training_args,
         train_dataset=full_dataset,
@@ -683,15 +573,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)
@@ -699,17 +589,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,
                 }
             )
@@ -723,7 +621,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)