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

cehrgpt/runners/data_utils.py

@@ -1,4 +1,10 @@
+import os
+from datetime import datetime
+from typing import Dict, List, Optional, Union
+
 import numpy as np
+import polars as pl
+import torch
 from cehrbert.data_generators.hf_data_generator.cache_util import CacheFileCollector
 from cehrbert.data_generators.hf_data_generator.meds_utils import (
     create_dataset_from_meds_reader,
@@ -12,17 +18,36 @@ from datasets import DatasetDict, concatenate_datasets, load_from_disk
 from transformers import TrainingArguments
 from transformers.utils import logging
 
-from cehrgpt.data.hf_cehrgpt_dataset_mapping import MedToCehrGPTDatasetMapping
+from cehrgpt.data.hf_cehrgpt_dataset_mapping import (
+    ExtractTokenizedSequenceDataMapping,
+    MedToCehrGPTDatasetMapping,
+)
 from cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
 
 LOG = logging.get_logger("transformers")
 
 
+def get_torch_dtype(torch_dtype: Optional[str] = None) -> Union[torch.dtype, str]:
+    if torch_dtype and hasattr(torch, torch_dtype):
+        return getattr(torch, torch_dtype)
+    return torch.float
+
+
+def data_collate_fn(features, model_type: torch.dtype, collator):
+    batch = collator(features)
+    if model_type != torch.float32:
+        for key, value in batch.items():
+            # Only convert float32 tensors to bfloat16
+            if isinstance(value, torch.Tensor) and value.dtype == torch.float32:
+                batch[key] = value.to(model_type)
+    return batch
+
+
 def prepare_finetune_dataset(
     data_args: DataTrainingArguments,
     training_args: TrainingArguments,
     cehrgpt_args: CehrGPTArguments,
-    cache_file_collector: CacheFileCollector,
+    cache_file_collector: Optional[CacheFileCollector] = None,
 ) -> DatasetDict:
     # 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:
@@ -66,8 +91,9 @@ def prepare_finetune_dataset(
                     "Clean up the cached files for the cehrgpt dataset transformed from the MEDS: %s",
                     stats,
                 )
-                # Clean up the files created from the data generator
-                cache_file_collector.remove_cache_files()
+                if cache_file_collector:
+                    # Clean up the files created from the data generator
+                    cache_file_collector.remove_cache_files()
                 dataset = load_from_disk(str(meds_extension_path))
 
         train_set = dataset["train"]
@@ -219,7 +245,7 @@ def create_dataset_splits(data_args: DataTrainingArguments, seed: int):
             )
 
         # Generate splits
-        train_set = filter_by_patient_ids(train_patient_ids)
+        train_set = filter_by_patient_ids(train_patient_ids).shuffle(seed=seed)
         validation_set = filter_by_patient_ids(val_patient_ids)
         if test_set is None:
             test_set = filter_by_patient_ids(test_patient_ids)
@@ -241,3 +267,103 @@ def create_dataset_splits(data_args: DataTrainingArguments, seed: int):
             )
 
     return train_set, validation_set, test_set
+
+
+def extract_cohort_sequences(
+    data_args: DataTrainingArguments,
+    cehrgpt_args: CehrGPTArguments,
+    cache_file_collector: Optional[CacheFileCollector] = None,
+) -> DatasetDict:
+    """
+    Extracts and processes cohort-specific tokenized sequences from a pre-tokenized dataset,.
+
+    based on the provided cohort Parquet files and observation window constraints.
+
+    This function performs the following steps:
+    1. Loads cohort definitions from Parquet files located in `data_args.cohort_folder`.
+    2. Renames relevant columns if the data originates from a Meds format.
+    3. Filters a pre-tokenized dataset (loaded from `cehrgpt_args.tokenized_full_dataset_path`)
+       to include only patients present in the cohort.
+    4. Aggregates each person's index date and label into a mapping.
+    5. Checks for consistency to ensure all cohort person_ids are present in the tokenized dataset.
+    6. Applies a transformation (`ExtractTokenizedSequenceDataMapping`) to generate
+       observation-window-constrained patient sequences.
+    7. Caches both the filtered and processed datasets using the provided `cache_file_collector`.
+
+    Args:
+        data_args (DataTrainingArguments): Configuration parameters for data processing,
+            including cohort folder, observation window, batch size, and parallelism.
+        cehrgpt_args (CehrGPTArguments): Contains paths to pre-tokenized datasets and CEHR-GPT-specific arguments.
+        cache_file_collector (CacheFileCollector): Utility to register and manage dataset cache files.
+
+    Returns:
+        DatasetDict: A Hugging Face `DatasetDict` containing the processed datasets (e.g., train/validation/test),
+                     where each entry includes sequences filtered and truncated by the observation window.
+
+    Raises:
+        RuntimeError: If any `person_id` in the cohort is missing from the tokenized dataset.
+    """
+
+    cohort = pl.read_parquet(os.path.join(data_args.cohort_folder, "*.parquet"))
+    if data_args.is_data_in_meds:
+        cohort = cohort.rename(
+            mapping={
+                "prediction_time": "index_date",
+                "subject_id": "person_id",
+                "boolean_value": "label",
+            }
+        )
+    all_person_ids = cohort["person_id"].unique().to_list()
+    # In case the label column does not exist, we add a fake column to the dataframe so subsequent process can work
+    if "label" not in cohort.columns:
+        cohort = cohort.with_columns(
+            pl.Series(
+                name="label", values=np.zeros_like(cohort["person_id"].to_numpy())
+            )
+        )
+
+    # data_args.observation_window
+    tokenized_dataset = load_from_disk(cehrgpt_args.tokenized_full_dataset_path)
+    filtered_tokenized_dataset = tokenized_dataset.filter(
+        lambda batch: [person_id in all_person_ids for person_id in batch["person_id"]],
+        batched=True,
+        batch_size=data_args.preprocessing_batch_size,
+        num_proc=data_args.preprocessing_num_workers,
+    )
+    person_index_date_agg = cohort.group_by("person_id").agg(
+        pl.struct("index_date", "label").alias("index_date_label")
+    )
+    # Convert to dictionary
+    person_index_date_map: Dict[int, List[datetime]] = dict(
+        zip(
+            person_index_date_agg["person_id"].to_list(),
+            person_index_date_agg["index_date_label"].to_list(),
+        )
+    )
+    LOG.info(f"person_index_date_agg: {person_index_date_agg}")
+    tokenized_person_ids = []
+    for _, dataset in filtered_tokenized_dataset.items():
+        tokenized_person_ids.extend(dataset["person_id"])
+    missing_person_ids = [
+        person_id
+        for person_id in person_index_date_map.keys()
+        if person_id not in tokenized_person_ids
+    ]
+    if missing_person_ids:
+        raise RuntimeError(
+            f"There are {len(missing_person_ids)} missing in the tokenized dataset. "
+            f"The list contains: {missing_person_ids}"
+        )
+    processed_dataset = filtered_tokenized_dataset.map(
+        ExtractTokenizedSequenceDataMapping(
+            person_index_date_map, data_args.observation_window
+        ).batch_transform,
+        batched=True,
+        batch_size=data_args.preprocessing_batch_size,
+        num_proc=data_args.preprocessing_num_workers,
+        remove_columns=filtered_tokenized_dataset["train"].column_names,
+    )
+    if cache_file_collector:
+        cache_file_collector.add_cache_files(filtered_tokenized_dataset)
+        cache_file_collector.add_cache_files(processed_dataset)
+    return processed_dataset

cehrgpt/runners/hf_cehrgpt_finetune_runner.py

@@ -50,7 +50,11 @@ 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.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
@@ -142,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(
@@ -161,11 +164,22 @@ def load_finetuned_model(
 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}"
@@ -175,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)
@@ -195,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:
@@ -252,46 +264,55 @@ def main():
 
     if processed_dataset is None:
         if is_main_process(training_args.local_rank):
-            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:
-                    # 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
+            # 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
+                )
+            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:
+                        # 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
                         )
-                    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 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)
+                        )
+
+                # 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()
@@ -350,8 +371,7 @@ def main():
             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"],
+            negative_sampling_probability=cehrgpt_args.negative_sampling_probability,
         )
         training_args.per_device_train_batch_size = 1
         training_args.per_device_eval_batch_size = 1
@@ -359,6 +379,7 @@ def main():
             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
@@ -381,13 +402,14 @@ def main():
         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:
             training_args = perform_hyperparameter_search(
                 trainer_class,
-                partial(model_init, model_args, training_args, tokenizer),
+                partial(model_init, model_args, training_args, cehrgpt_args, tokenizer),
                 processed_dataset,
                 data_collator,
                 training_args,
@@ -401,6 +423,7 @@ def main():
                 trainer_class,
                 model_args,
                 training_args,
+                cehrgpt_args,
                 tokenizer,
                 processed_dataset,
                 data_collator,
@@ -408,7 +431,7 @@ def main():
         else:
             # Initialize Trainer for final training on the combined train+val set
             trainer = trainer_class(
-                model=model_init(model_args, training_args, tokenizer),
+                model=model_init(model_args, training_args, cehrgpt_args, tokenizer),
                 data_collator=data_collator,
                 args=training_args,
                 train_dataset=processed_dataset["train"],
@@ -457,6 +480,7 @@ def retrain_with_full_set(
     trainer_class,
     model_args: ModelArguments,
     training_args: TrainingArguments,
+    cehrgpt_args: CehrGPTArguments,
     tokenizer: CehrGptTokenizer,
     dataset: DatasetDict,
     data_collator: CehrGptDataCollator,
@@ -475,6 +499,7 @@ def retrain_with_full_set(
         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.
@@ -494,7 +519,7 @@ def retrain_with_full_set(
     training_args.evaluation_strategy = "no"
     checkpoint = get_last_hf_checkpoint(training_args)
     final_trainer = trainer_class(
-        model=model_init(model_args, training_args, tokenizer),
+        model=model_init(model_args, training_args, cehrgpt_args, tokenizer),
         data_collator=data_collator,
         args=training_args,
         train_dataset=full_dataset,
@@ -555,7 +580,15 @@ def do_predict(
             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()))
+
+            index_dates = list(
+                map(
+                    lambda posix_time: datetime.utcfromtimestamp(posix_time).replace(
+                        tzinfo=None
+                    ),
+                    index_dates.tolist(),
+                )
+            )
 
             batch = {k: v.to(device) for k, v in batch.items()}
             # Forward pass

cehrgpt/runners/hf_cehrgpt_pretrain_runner.py

@@ -34,6 +34,7 @@ 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.data_utils import get_torch_dtype
 from cehrgpt.runners.gpt_runner_util import parse_runner_args
 from cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
 from cehrgpt.runners.sample_packing_trainer import SamplePackingTrainer
@@ -71,6 +72,36 @@ def load_and_create_tokenizer(
     cehrgpt_args: CehrGPTArguments,
     dataset: Optional[Union[Dataset, DatasetDict]] = None,
 ) -> CehrGptTokenizer:
+
+    concept_name_mapping = {}
+    allowed_motor_codes = list()
+    if cehrgpt_args.concept_dir:
+        import pandas as pd
+        from cehrbert_data.const.artificial_tokens import DEATH_TOKEN
+        from meds.schema import death_code
+
+        LOG.info("Loading concept data from disk at %s", cehrgpt_args.concept_dir)
+        concept_pd = pd.read_parquet(cehrgpt_args.concept_dir)
+        LOG.info(
+            "Creating concept name mapping and motor_time_to_event_codes from disk at %s",
+            cehrgpt_args.concept_dir,
+        )
+        for row in concept_pd.itertuples():
+            concept_name_mapping[str(getattr(row, "concept_id"))] = getattr(
+                row, "concept_name"
+            )
+            if (
+                cehrgpt_args.include_motor_time_to_event
+                and getattr(row, "domain_id")
+                in ["Condition", "Procedure", "Drug", "Visit"]
+                and getattr(row, "standard_concept") == "S"
+            ):
+                allowed_motor_codes.append(str(getattr(row, "concept_id")))
+        LOG.info(
+            "Adding death codes for MOTOR TTE predictions: %s",
+            [DEATH_TOKEN, death_code],
+        )
+        allowed_motor_codes.extend([DEATH_TOKEN, death_code])
     # Try to load the pretrained tokenizer
     tokenizer_abspath = os.path.expanduser(model_args.tokenizer_name_or_path)
     try:
@@ -85,9 +116,17 @@ def load_and_create_tokenizer(
         LOG.info("Started training the tokenizer ...")
         tokenizer = CehrGptTokenizer.train_tokenizer(
             dataset,
-            {},
+            concept_name_mapping,
             data_args,
             PretrainedEmbeddings(cehrgpt_args.pretrained_embedding_path),
+            allowed_motor_codes if cehrgpt_args.include_motor_time_to_event else None,
+            (
+                cehrgpt_args.num_motor_tasks
+                if cehrgpt_args.include_motor_time_to_event
+                else None
+            ),
+            apply_entropy_filter=cehrgpt_args.apply_entropy_filter,
+            min_prevalence=cehrgpt_args.min_prevalence,
         )
         LOG.info("Finished training the tokenizer ...")
         tokenizer.save_pretrained(tokenizer_abspath)
@@ -99,13 +138,12 @@ def load_and_create_tokenizer(
 def load_and_create_model(
     model_args: ModelArguments,
     cehrgpt_args: CehrGPTArguments,
-    training_args: TrainingArguments,
     tokenizer: CehrGptTokenizer,
 ) -> CEHRGPT2LMHeadModel:
     attn_implementation = (
         "flash_attention_2" if is_flash_attn_2_available() else "eager"
     )
-    torch_dtype = torch.bfloat16 if training_args.bf16 else torch.float32
+    torch_dtype = get_torch_dtype(model_args.torch_dtype)
     model_abspath = os.path.expanduser(model_args.model_name_or_path)
     if cehrgpt_args.continue_pretrain:
         try:
@@ -147,6 +185,8 @@ def load_and_create_model(
         else:
             pretrained_embedding_dim = model_args.hidden_size
 
+        model_args_cehrgpt = model_args.as_dict()
+        model_args_cehrgpt.pop("attn_implementation")
         model_config = CEHRGPTConfig(
             vocab_size=tokenizer.vocab_size,
             value_vocab_size=tokenizer.value_vocab_size,
@@ -172,7 +212,12 @@ def load_and_create_model(
                 if cehrgpt_args.sample_packing
                 else model_args.max_position_embeddings
             ),
-            **model_args.as_dict(),
+            include_motor_time_to_event=cehrgpt_args.include_motor_time_to_event,
+            motor_tte_vocab_size=tokenizer.motor_tte_vocab_size,
+            motor_time_to_event_weight=cehrgpt_args.motor_time_to_event_weight,
+            motor_num_time_pieces=cehrgpt_args.motor_num_time_pieces,
+            ve_token_id=tokenizer.ve_token_id,
+            **model_args_cehrgpt,
         )
 
     model = CEHRGPT2LMHeadModel(model_config)
@@ -348,6 +393,8 @@ def main():
                         pretrained_concept_embedding_model=PretrainedEmbeddings(
                             cehrgpt_args.pretrained_embedding_path
                         ),
+                        apply_entropy_filter=cehrgpt_args.apply_entropy_filter,
+                        min_prevalence=cehrgpt_args.min_prevalence,
                     )
                     cehrgpt_tokenizer.save_pretrained(
                         os.path.expanduser(training_args.output_dir)
@@ -421,9 +468,11 @@ def main():
     else:
         processed_dataset = processed_dataset.filter(filter_func, **filter_args)
 
-    model = load_and_create_model(
-        model_args, cehrgpt_args, training_args, cehrgpt_tokenizer
-    )
+    model = load_and_create_model(model_args, cehrgpt_args, cehrgpt_tokenizer)
+
+    # Try to update motor tte vocab size if the new configuration is different from the existing one
+    if cehrgpt_args.include_motor_time_to_event:
+        model.update_motor_tte_vocab_size(cehrgpt_tokenizer.motor_tte_vocab_size)
 
     # Expand tokenizer to adapt to the new pretraining dataset
     if model.config.vocab_size < cehrgpt_tokenizer.vocab_size:
@@ -500,6 +549,9 @@ def main():
             include_ttv_prediction=model_args.include_ttv_prediction,
             use_sub_time_tokenization=model_args.use_sub_time_tokenization,
             include_values=model_args.include_values,
+            include_motor_time_to_event=cehrgpt_args.include_motor_time_to_event,
+            motor_tte_vocab_size=model.config.motor_tte_vocab_size,
+            motor_num_time_pieces=cehrgpt_args.motor_num_time_pieces,
         ),
         train_dataset=processed_dataset["train"],
         eval_dataset=(

cehrgpt/runners/hf_gpt_runner_argument_dataclass.py

@@ -6,6 +6,12 @@ from typing import List, Optional
 class CehrGPTArguments:
     """Arguments pertaining to what data we are going to input our model for training and eval."""
 
+    tokenized_full_dataset_path: Optional[str] = dataclasses.field(
+        default=None,
+        metadata={
+            "help": "The path to the tokenized dataset created for the full population"
+        },
+    )
     include_inpatient_hour_token: Optional[bool] = dataclasses.field(
         default=True,
         metadata={"help": "Include inpatient hour token"},
@@ -177,7 +183,61 @@ class CehrGPTArguments:
             "help": "A flag to indicate whether we want to add end token in sample packing"
         },
     )
+    include_motor_time_to_event: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to include the motor time to events"
+        },
+    )
+    num_motor_tasks: Optional[int] = dataclasses.field(
+        default=10000,
+        metadata={"help": "The number of max MOTOR tasks"},
+    )
+    motor_time_to_event_weight: Optional[float] = dataclasses.field(
+        default=1.0,
+        metadata={"help": "The MOTOR time to event loss weight"},
+    )
+    motor_num_time_pieces: Optional[int] = dataclasses.field(
+        default=8,
+        metadata={
+            "help": "The number of times each motor_num_time_pieces piece has to be"
+        },
+    )
+    concept_dir: Optional[str] = dataclasses.field(
+        default=None,
+        metadata={"help": "The directory where the concept data is stored."},
+    )
     average_over_sequence: bool = dataclasses.field(
         default=False,
         metadata={"help": "Whether or not to average tokens per sequence"},
     )
+    apply_entropy_filter: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={"help": "A flag to indicate whether we want to use entropy filter."},
+    )
+    min_prevalence: Optional[float] = dataclasses.field(
+        default=1 / 1000,
+        metadata={"help": "The min_prevalence to keep the concepts in the tokenizer"},
+    )
+    class_weights: Optional[List[int]] = dataclasses.field(
+        default=None,
+        metadata={"help": "The class weights for training"},
+    )
+    negative_sampling_probability: Optional[float] = dataclasses.field(
+        default=None,
+        metadata={
+            "help": "The probability of negative samples will be included in the training data"
+        },
+    )
+    num_of_trajectories_per_sample: Optional[int] = dataclasses.field(
+        default=1,
+        metadata={"help": "The number of trajectories per sample"},
+    )
+    generation_input_length: Optional[int] = dataclasses.field(
+        default=1024,
+        metadata={"help": "The length of the input sequence"},
+    )
+    generation_max_new_tokens: Optional[int] = dataclasses.field(
+        default=1024,
+        metadata={"help": "The maximum number of tokens in the generation sequence"},
+    )

cehrgpt/runners/hyperparameter_search_util.py

@@ -4,12 +4,7 @@ from typing import Callable, Tuple
 import optuna
 from cehrbert.runners.hf_runner_argument_dataclass import ModelArguments
 from datasets import Dataset, DatasetDict
-from transformers import (
-    EarlyStoppingCallback,
-    Trainer,
-    TrainerCallback,
-    TrainingArguments,
-)
+from transformers import EarlyStoppingCallback, TrainerCallback, TrainingArguments
 from transformers.utils import logging
 
 from cehrgpt.data.hf_cehrgpt_dataset_collator import CehrGptDataCollator
@@ -85,7 +80,9 @@ def hp_space(
             "per_device_train_batch_size", batch_sizes
         ),
         "weight_decay": trial.suggest_float("weight_decay", *weight_decays, log=True),
-        "num_train_epochs": trial.suggest_int("num_train_epochs", *num_train_epochs),
+        "num_train_epochs": trial.suggest_categorical(
+            "num_train_epochs", num_train_epochs
+        ),
     }
 
 
@@ -217,6 +214,8 @@ def perform_hyperparameter_search(
             backend="optuna",
             n_trials=cehrgpt_args.n_trials,
             compute_objective=lambda m: m["optuna_best_metric"],
+            # Ensure reproducibility
+            sampler=optuna.samplers.TPESampler(seed=training_args.seed),
         )
         LOG.info("Best hyperparameters: %s", best_trial.hyperparameters)
         # Update training arguments with best hyperparameters and set epochs based on adjusted effective epochs