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

cehrgpt/rl_finetune/ppo_finetune_v2.py

@@ -0,0 +1,373 @@
+import datetime
+import os
+import pickle
+from collections import Counter, defaultdict
+from functools import partial
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+from cehrbert.models.hf_models.tokenization_utils import agg_helper
+from cehrbert.runners.runner_util import load_parquet_as_dataset
+from tqdm import tqdm
+from transformers.utils import is_flash_attn_2_available, logging
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, create_reference_model
+
+from cehrgpt.cehrgpt_args import create_inference_base_arg_parser
+from cehrgpt.generation.generate_batch_hf_gpt_sequence import generate_single_batch
+from cehrgpt.gpt_utils import get_cehrgpt_output_folder
+from cehrgpt.models.hf_cehrgpt import CEHRGPT2LMHeadModel
+from cehrgpt.models.tokenization_hf_cehrgpt import CehrGptTokenizer
+from cehrgpt.rl_finetune.cehrgpt_ppo_trainer import (
+    CehrGptPPODataCollator,
+    CehrGptPPOTrainer,
+)
+
+LOG = logging.get_logger("transformers")
+
+
+def extract_concept_frequency(records: Dict[str, Any]) -> Dict[str, int]:
+    batched_concept_ids = records["concept_ids"]
+    outputs = defaultdict(int)
+    for concept_ids in batched_concept_ids:
+        for concept_id, cnt in dict(Counter(concept_ids[4:])).items():
+            outputs[concept_id] += cnt
+    return outputs
+
+
+def main(args):
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(args.tokenizer_folder)
+    model_folder_name = os.path.join(
+        args.output_folder, get_cehrgpt_output_folder(args, cehrgpt_tokenizer), "model"
+    )
+
+    if not os.path.exists(model_folder_name):
+        os.makedirs(model_folder_name)
+
+    if args.restore_from_checkpoint:
+        try:
+            cehrgpt_model = CEHRGPT2LMHeadModel.from_pretrained(
+                model_folder_name,
+                attn_implementation=(
+                    "flash_attention_2" if is_flash_attn_2_available() else "eager"
+                ),
+                torch_dtype=(
+                    torch.bfloat16 if is_flash_attn_2_available() else torch.float32
+                ),
+            )
+        except Exception:
+            LOG.warning(
+                "Checkpoint does not exist in %s, loading from the %s",
+                model_folder_name,
+                args.model_folder,
+            )
+            cehrgpt_model = CEHRGPT2LMHeadModel.from_pretrained(
+                args.model_folder,
+                attn_implementation=(
+                    "flash_attention_2" if is_flash_attn_2_available() else "eager"
+                ),
+                torch_dtype=(
+                    torch.bfloat16 if is_flash_attn_2_available() else torch.float32
+                ),
+            )
+    else:
+        cehrgpt_model = CEHRGPT2LMHeadModel.from_pretrained(
+            args.model_folder,
+            attn_implementation=(
+                "flash_attention_2" if is_flash_attn_2_available() else "eager"
+            ),
+            torch_dtype=(
+                torch.bfloat16 if is_flash_attn_2_available() else torch.float32
+            ),
+        )
+
+    cehrgpt_model.generation_config.pad_token_id = cehrgpt_tokenizer.pad_token_id
+    cehrgpt_model.generation_config.eos_token_id = cehrgpt_tokenizer.end_token_id
+    cehrgpt_model.generation_config.bos_token_id = cehrgpt_tokenizer.end_token_id
+    model = AutoModelForCausalLMWithValueHead(cehrgpt_model).to(device)
+    model.is_peft_model = False
+    ref_model = create_reference_model(model).to(device)
+
+    # create a ppo trainer
+    ppo_trainer = CehrGptPPOTrainer(
+        config=PPOConfig(
+            batch_size=args.batch_size,
+            mini_batch_size=args.mini_batch_size,
+            init_kl_coef=args.init_kl_coef,
+            vf_coef=args.vf_coef,
+            kl_penalty=args.kl_penalty,
+            gamma=args.gamma,
+            use_score_scaling=args.use_score_scaling,
+        ),
+        model=model,
+        ref_model=ref_model,
+        tokenizer=cehrgpt_tokenizer,
+        training_data_collator=CehrGptPPODataCollator(
+            cehrgpt_tokenizer, max_length=args.context_window
+        ),
+    )
+
+    LOG.info(f"Loading tokenizer at {args.model_folder}")
+    LOG.info(f"Loading model at {args.model_folder}")
+    LOG.info(f"Will save the fine-tuned model at {model_folder_name}")
+    LOG.info(f"Context window {args.context_window}")
+    LOG.info(f"Temperature {args.temperature}")
+    LOG.info(f"Repetition Penalty {args.repetition_penalty}")
+    LOG.info(f"Sampling Strategy {args.sampling_strategy}")
+    LOG.info(f"Num beam {args.num_beams}")
+    LOG.info(f"Num beam groups {args.num_beam_groups}")
+    LOG.info(f"Epsilon cutoff {args.epsilon_cutoff}")
+    LOG.info(f"Top P {args.top_p}")
+    LOG.info(f"Top K {args.top_k}")
+    LOG.info(f"Loading demographic_info at {args.demographic_data_path}")
+
+    dataset = load_parquet_as_dataset(args.demographic_data_path).filter(
+        lambda batched: [
+            model.config.n_positions >= num_of_concepts > args.min_num_tokens
+            for num_of_concepts in batched["num_of_concepts"]
+        ],
+        batched=True,
+    )
+    parts = dataset.map(
+        partial(agg_helper, map_func=extract_concept_frequency),
+        batched=True,
+        batch_size=1000,
+        num_proc=args.num_proc,
+        remove_columns=dataset.column_names,
+    )
+
+    concept_stats = defaultdict(float)
+    for stat in tqdm(parts, desc="Aggregating the concept counts"):
+        fixed_stat = pickle.loads(stat["data"])
+        for concept_id, count in fixed_stat.items():
+            concept_stats[concept_id] += count
+    total_sum = sum(concept_stats.values())
+    for concept_id, count in concept_stats.items():
+        concept_stats[concept_id] = count / total_sum
+
+    logs = []
+    device = ppo_trainer.current_device
+    total_rows = len(dataset)
+    num_of_micro_batches = args.batch_size // args.mini_batch_size
+    for i in tqdm(range(args.num_of_steps)):
+        LOG.info(f"{datetime.datetime.now()}: Batch {i} started")
+        random_prompts = []
+        batched_sequences = []
+        batched_values = []
+        batched_value_indicators = []
+        for _ in range(num_of_micro_batches):
+            random_indices = np.random.randint(0, total_rows, args.mini_batch_size)
+            random_prompts_micro_batch = [
+                record["concept_ids"][:4] for record in dataset.select(random_indices)
+            ]
+            random_prompts.extend(random_prompts_micro_batch)
+            micro_batched_prompts = [
+                cehrgpt_tokenizer.encode(random_prompt)
+                for random_prompt in random_prompts_micro_batch
+            ]
+
+            micro_batched_sequences = generate_single_batch(
+                cehrgpt_model,
+                cehrgpt_tokenizer,
+                micro_batched_prompts,
+                max_new_tokens=args.context_window,
+                mini_num_of_concepts=args.min_num_of_concepts,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                temperature=args.temperature,
+                repetition_penalty=args.repetition_penalty,
+                num_beams=args.num_beams,
+                num_beam_groups=args.num_beam_groups,
+                epsilon_cutoff=args.epsilon_cutoff,
+                device=device,
+            )
+            # Clear the cache
+            torch.cuda.empty_cache()
+            batched_sequences.extend(micro_batched_sequences["sequences"])
+            batched_values.extend(micro_batched_sequences["values"])
+            batched_value_indicators.extend(micro_batched_sequences["value_indicators"])
+
+        LOG.info(f"{datetime.datetime.now()}: Batch {i} sequence generated")
+        reward = compute_marginal_dist_reward(
+            batched_sequences, concept_stats, cehrgpt_tokenizer
+        )
+        LOG.info(f"{datetime.datetime.now()}: Batch {i} KL divergence reward: {reward}")
+        query_tensors = []
+        response_tensors = []
+        value_tensors = []
+        value_indicator_tensors = []
+        rewards = []
+        for sequence, values, value_indicators in zip(
+            batched_sequences, batched_values, batched_value_indicators
+        ):
+            # Convert sequence to a NumPy array if it's not already one
+            sequence_array = np.asarray(sequence)
+            # Find the end token
+            condition_array = sequence_array == cehrgpt_tokenizer.end_token
+            end_index = (
+                np.argmax(condition_array)
+                if condition_array.any()
+                else len(sequence_array) - 1
+            )
+
+            sequence = sequence[: end_index + 1]
+            values = values[: end_index + 1]
+            value_indicators = value_indicators[: end_index + 1]
+
+            query_tensors.append(torch.tensor(cehrgpt_tokenizer.encode(sequence[:4])))
+            response_tensors.append(
+                torch.tensor(cehrgpt_tokenizer.encode(sequence[4:]))
+            )
+            value_tensors.append(torch.tensor(cehrgpt_tokenizer.encode_value(values)))
+            value_indicator_tensors.append(torch.tensor(value_indicators))
+            rewards.append(reward)
+
+        train_stats = ppo_trainer.step(
+            query_tensors,
+            response_tensors,
+            rewards,
+            value_tensors,
+            value_indicator_tensors,
+        )
+        LOG.info(f"{datetime.datetime.now()}: Batch {i} stats: {train_stats}")
+        logs.append(reward)
+        ppo_trainer.log_stats(stats=train_stats, batch={}, rewards=rewards)
+    ppo_trainer.save_pretrained(model_folder_name)
+    with open(os.path.join(model_folder_name, "ppo_finetune_stats.pkl"), "wb") as f:
+        pickle.dump(logs, f)
+
+
+def compute_marginal_dist_reward(
+    batched_sequences: List[List[str]],
+    expected_concept_dist: Dict[str, float],
+    tokenizer: CehrGptTokenizer,
+) -> torch.Tensor:
+    actual_concept_dist = dict(
+        Counter(
+            [
+                concept_id
+                for sequence in batched_sequences
+                for concept_id in sequence[4:]
+            ]
+        )
+    )
+    total_count = sum(actual_concept_dist.values())
+    for concept_id in actual_concept_dist.keys():
+        actual_concept_dist[concept_id] /= total_count
+    # Translate the concept ids to token ids
+    actual_dist = np.zeros(tokenizer.vocab_size)
+    actual_dist[tokenizer.encode(list(actual_concept_dist.keys()))] = list(
+        actual_concept_dist.values()
+    )
+    # Add a small epsilon to avoid log(0)
+    epsilon = 1e-10
+    logprob_dist = torch.tensor(np.log(actual_dist + epsilon))
+    # Translate the concept ids to token ids
+    ref_dist = np.zeros(tokenizer.vocab_size)
+    ref_dist[tokenizer.encode(list(expected_concept_dist.keys()))] = list(
+        expected_concept_dist.values()
+    )
+    ref_logprob_dist = torch.tensor(np.log(ref_dist + epsilon))
+
+    # Flip is required due to this issue? :https://github.com/pytorch/pytorch/issues/57459
+    return torch.exp(
+        -torch.nn.functional.kl_div(
+            ref_logprob_dist, logprob_dist, log_target=True, reduction="none"
+        ).sum(-1)
+    )
+
+
+def create_arg_parser():
+    base_arg_parser = create_inference_base_arg_parser(
+        description="Arguments for finetuning cehr-gpt using PPO"
+    )
+    base_arg_parser.add_argument(
+        "--mini_batch_size",
+        dest="mini_batch_size",
+        action="store",
+        type=int,
+        required=True,
+    )
+    base_arg_parser.add_argument(
+        "--init_kl_coef",
+        dest="init_kl_coef",
+        action="store",
+        type=float,
+        required=False,
+        default=0.1,
+    )
+    base_arg_parser.add_argument(
+        "--vf_coef",
+        dest="vf_coef",
+        action="store",
+        type=float,
+        required=False,
+        default=0.1,
+    )
+    base_arg_parser.add_argument(
+        "--kl_penalty",
+        dest="kl_penalty",
+        action="store",
+        choices=["kl", "abs", "mse", "full"],
+        required=False,
+        default="kl",
+    )
+    base_arg_parser.add_argument(
+        "--gamma",
+        dest="gamma",
+        action="store",
+        type=float,
+        required=False,
+        default=0.99,
+    )
+    base_arg_parser.add_argument(
+        "--num_proc",
+        dest="num_proc",
+        action="store",
+        type=int,
+        default=4,
+        required=False,
+    )
+    base_arg_parser.add_argument(
+        "--num_of_steps",
+        dest="num_of_steps",
+        action="store",
+        type=int,
+        default=1028,
+        required=False,
+    )
+    base_arg_parser.add_argument(
+        "--min_num_tokens",
+        dest="min_num_tokens",
+        action="store",
+        type=int,
+        default=4,
+        required=False,
+    )
+    base_arg_parser.add_argument(
+        "--demographic_data_path",
+        dest="demographic_data_path",
+        action="store",
+        help="The path for your concept_path",
+        required=True,
+    )
+    base_arg_parser.add_argument(
+        "--restore_from_checkpoint",
+        dest="restore_from_checkpoint",
+        action="store_true",
+    )
+    base_arg_parser.add_argument(
+        "--use_score_scaling",
+        dest="use_score_scaling",
+        action="store_true",
+    )
+    return base_arg_parser
+
+
+if __name__ == "__main__":
+    main(create_arg_parser().parse_args())

cehrgpt/runners/hf_cehrgpt_dpo_runner.py

@@ -0,0 +1,119 @@
+from cehrbert.data_generators.hf_data_generator.hf_dataset import (
+    apply_cehrbert_dataset_mapping,
+)
+from cehrbert.runners.runner_util import (
+    generate_prepared_ds_path,
+    get_last_hf_checkpoint,
+    load_parquet_as_dataset,
+)
+from datasets import DatasetDict, load_from_disk
+from transformers import set_seed
+from transformers.utils import is_flash_attn_2_available, logging
+
+from cehrgpt.data.hf_cehrgpt_dpo_collator import CehrGptDPODataCollator
+from cehrgpt.data.hf_cehrgpt_dpo_dataset_mapping import HFCehrGptDPOTokenizationMapping
+from cehrgpt.models.hf_cehrgpt import CEHRGPT2LMHeadModel
+from cehrgpt.rl_finetune.cehrgpt_dpo_trainer import CehrGptDPOTrainer
+from cehrgpt.runners.gpt_runner_util import parse_dpo_runner_args
+from cehrgpt.runners.hf_cehrgpt_finetune_runner import load_pretrained_tokenizer
+
+LOG = logging.get_logger("transformers")
+
+
+def main():
+    cehrgpt_args, data_args, model_args, dpo_config = parse_dpo_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
+    )
+    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...")
+    else:
+        dataset = load_parquet_as_dataset(data_args.data_folder)
+        # Random split
+        dataset = dataset.train_test_split(
+            test_size=data_args.validation_split_percentage, seed=dpo_config.seed
+        )
+        processed_dataset = apply_cehrbert_dataset_mapping(
+            dataset,
+            mapping_function=HFCehrGptDPOTokenizationMapping(tokenizer),
+            batch_size=data_args.preprocessing_batch_size,
+            num_proc=data_args.preprocessing_num_workers,
+            streaming=data_args.streaming,
+        )
+
+        processed_dataset = processed_dataset.filter(
+            lambda batch: [
+                len(chosen_concept_ids) < model_args.max_position_embeddings
+                for chosen_concept_ids in batch["chosen_concept_ids"]
+            ],
+            batched=True,
+            batch_size=data_args.preprocessing_batch_size,
+            num_proc=data_args.preprocessing_num_workers,
+        ).filter(
+            lambda batch: [
+                len(rejected_concept_ids) < model_args.max_position_embeddings
+                for rejected_concept_ids in batch["rejected_concept_ids"]
+            ],
+            batched=True,
+            batch_size=data_args.preprocessing_batch_size,
+            num_proc=data_args.preprocessing_num_workers,
+        )
+        processed_dataset.save_to_disk(prepared_ds_path)
+
+    # Set seed before initializing model.
+    set_seed(dpo_config.seed)
+    processed_dataset.set_format("pt")
+
+    # A hacky way to prevent the training from removing unmatched inputs
+    dpo_config.label_names = [
+        "chosen_input_ids",
+        "rejected_input_ids",
+        "chosen_concept_values",
+        "rejected_concept_values",
+        "chosen_concept_value_masks",
+        "rejected_concept_value_masks",
+    ]
+
+    attn_implementation = (
+        "flash_attention_2" if is_flash_attn_2_available() else "eager"
+    )
+    model = CEHRGPT2LMHeadModel.from_pretrained(
+        model_args.model_name_or_path,
+        attn_implementation=attn_implementation,
+    )
+    ref_model = CEHRGPT2LMHeadModel.from_pretrained(
+        model_args.model_name_or_path,
+        attn_implementation=attn_implementation,
+    )
+
+    # Initialize Trainer for final training on the combined train+val set
+    trainer = CehrGptDPOTrainer(
+        model=model,
+        ref_model=ref_model,
+        args=dpo_config,
+        tokenizer=tokenizer,
+        train_dataset=processed_dataset["train"],
+        eval_dataset=processed_dataset["test"],
+        data_collator=CehrGptDPODataCollator(
+            tokenizer=tokenizer,
+            max_length=model_args.max_position_embeddings,
+            pretraining=False,
+            include_ttv_prediction=False,
+            use_sub_time_tokenization=False,
+        ),
+    )
+    # Train the model on the combined train + val set
+    checkpoint = get_last_hf_checkpoint(dpo_config)
+    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_cehrgpt_finetune_runner.py

@@ -43,6 +43,7 @@ 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.models.hf_cehrgpt import (
     CEHRGPTConfig,
     CehrGptForClassification,
@@ -408,10 +409,24 @@ def main():
             except Exception as e:
                 LOG.exception(e)
                 dataset = create_dataset_from_meds_reader(
-                    data_args, is_pretraining=False
+                    data_args=data_args,
+                    dataset_mappings=[
+                        MedToCehrGPTDatasetMapping(
+                            data_args=data_args,
+                            is_pretraining=False,
+                            include_inpatient_hour_token=cehrgpt_args.include_inpatient_hour_token,
+                        )
+                    ],
                 )
                 if not data_args.streaming:
-                    dataset.save_to_disk(meds_extension_path)
+                    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"]
@@ -451,7 +466,13 @@ def main():
             dataset=final_splits, cehrgpt_tokenizer=tokenizer, data_args=data_args
         )
         if not data_args.streaming:
-            processed_dataset.save_to_disk(prepared_ds_path)
+            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))
 
     # Set seed before initializing model.
     set_seed(training_args.seed)

cehrgpt/runners/hf_cehrgpt_pretrain_runner.py

@@ -21,12 +21,13 @@ 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.data.hf_cehrgpt_dataset_mapping import MedToCehrGPTDatasetMapping
 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
+from cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
 
 LOG = logging.get_logger("transformers")
 
@@ -82,11 +83,25 @@ def load_and_create_model(
     model_abspath = os.path.expanduser(model_args.model_name_or_path)
     if cehrgpt_args.continue_pretrain:
         try:
-            return CEHRGPT2LMHeadModel.from_pretrained(
+            pretrained_model = CEHRGPT2LMHeadModel.from_pretrained(
                 model_abspath,
                 attn_implementation=attn_implementation,
                 torch_dtype=torch_dtype,
             )
+            if (
+                pretrained_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}"
+                )
+                pretrained_model.config.max_position_embeddings = (
+                    model_args.max_position_embeddings
+                )
+                pretrained_model.resize_position_embeddings(
+                    model_args.max_position_embeddings
+                )
+            return pretrained_model
         except Exception as e:
             LOG.error(
                 f"When continue_pretrain is set to True, it assumes that CEHR-GPT has been trained "
@@ -94,7 +109,7 @@ def load_and_create_model(
             )
             raise e
     try:
-        model_config = AutoConfig.from_pretrained(
+        model_config = CEHRGPTConfig.from_pretrained(
             model_abspath, attn_implementation=attn_implementation
         )
     except Exception as e:
@@ -148,7 +163,7 @@ def main():
         # 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
+        training_args.dataloader_prefetch_factor = None
 
     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")):
@@ -212,14 +227,29 @@ def main():
             except FileNotFoundError as e:
                 LOG.exception(e)
                 dataset = create_dataset_from_meds_reader(
-                    data_args, is_pretraining=True
+                    data_args=data_args,
+                    dataset_mappings=[
+                        MedToCehrGPTDatasetMapping(
+                            data_args=data_args,
+                            is_pretraining=True,
+                            include_inpatient_hour_token=cehrgpt_args.include_inpatient_hour_token,
+                        )
+                    ],
                 )
                 if not data_args.streaming:
-                    dataset.save_to_disk(meds_extension_path)
+                    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))
         else:
             # Load the dataset from the parquet files
             dataset = load_parquet_as_dataset(
-                data_args.data_folder, split="train", streaming=data_args.streaming
+                os.path.expanduser(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:
@@ -274,7 +304,13 @@ def main():
         )
         # 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)
+            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 pretraining dataset: %s",
+                stats,
+            )
+            processed_dataset = load_from_disk(str(prepared_ds_path))
 
     def filter_func(examples):
         if cehrgpt_args.drop_long_sequences:

cehrgpt/runners/hf_gpt_runner_argument_dataclass.py

@@ -6,6 +6,10 @@ from typing import List, Optional
 class CehrGPTArguments:
     """Arguments pertaining to what data we are going to input our model for training and eval."""
 
+    include_inpatient_hour_token: Optional[bool] = dataclasses.field(
+        default=True,
+        metadata={"help": "Include inpatient hour token"},
+    )
     include_demographics: Optional[bool] = dataclasses.field(
         default=False,
         metadata={