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

cehrgpt/omop/ontology.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import collections
+import os
+from typing import Any, Dict, Iterable, Optional, Set, Union
+
+import polars as pl
+from datasets import Dataset
+
+
+# Adapted from femr.ontology
+def _get_all_codes_map(batch: Dict[str, Any]) -> Dict[str, Any]:
+    result = set()
+    for concept_ids in batch["concept_ids"]:
+        for concept_id in concept_ids:
+            if concept_id.isnumeric():
+                result.add(concept_id)
+    return {"unique_concept_ids": list(result)}
+
+
+class Ontology:
+    def __init__(self, vocab_path: str):
+        """Create an Ontology from an Athena download and an optional meds Code Metadata structure.
+
+        NOTE: This is an expensive operation.
+        It is recommended to create an ontology once and then save/load it as necessary.
+        """
+        # Load from code metadata
+        self.parents_map: Dict[str, Set[str]] = collections.defaultdict(set)
+        self.concept_vocabulary_map: Dict[str, str] = collections.defaultdict(str)
+        self.concept_domain_map: Dict[str, str] = collections.defaultdict(str)
+
+        # Load from the athena path ...
+        concept = pl.scan_parquet(os.path.join(vocab_path, "concept/*parquet"))
+        vocabulary_id_col = pl.col("vocabulary_id")
+        concept_id_col = pl.col("concept_id").cast(pl.String)
+        domain_id_col = pl.col("domain_id").cast(pl.String)
+
+        processed_concepts = (
+            concept.select(
+                concept_id_col,
+                domain_id_col,
+                vocabulary_id_col,
+                pl.col("standard_concept").is_null(),
+            )
+            .collect()
+            .rows()
+        )
+
+        non_standard_concepts = set()
+
+        for concept_id, domain_id, vocabulary_id, is_non_standard in processed_concepts:
+            # We don't want to override code metadata
+            if concept_id not in self.concept_vocabulary_map:
+                self.concept_vocabulary_map[concept_id] = vocabulary_id
+
+            if concept_id not in self.concept_domain_map:
+                self.concept_domain_map[concept_id] = domain_id
+            # We don't want to override code metadata
+            if is_non_standard:
+                non_standard_concepts.add(concept_id)
+
+        relationship = pl.scan_parquet(
+            os.path.join(vocab_path, "concept_relationship/*parquet")
+        )
+        relationship_id = pl.col("relationship_id")
+        relationship = relationship.filter(
+            relationship_id == "Maps to",
+            pl.col("concept_id_1") != pl.col("concept_id_2"),
+        )
+        for concept_id_1, concept_id_2 in (
+            relationship.select(
+                pl.col("concept_id_1").cast(pl.String),
+                pl.col("concept_id_2").cast(pl.String),
+            )
+            .collect()
+            .rows()
+        ):
+            if concept_id_1 in non_standard_concepts:
+                self.parents_map[concept_id_1].add(concept_id_2)
+
+        ancestor = pl.scan_parquet(
+            os.path.join(vocab_path, "concept_ancestor/*parquet")
+        )
+        ancestor = ancestor.filter(pl.col("min_levels_of_separation") == 1)
+        for concept_id, parent_concept_id in (
+            ancestor.select(
+                pl.col("descendant_concept_id").cast(pl.String),
+                pl.col("ancestor_concept_id").cast(pl.String),
+            )
+            .collect()
+            .rows()
+        ):
+            self.parents_map[concept_id].add(parent_concept_id)
+        self.all_parents_map: Dict[str, Set[str]] = {}
+
+    def get_domain(self, concept_id: Union[str, int]) -> Optional[str]:
+        return self.concept_domain_map.get(str(concept_id), None)
+
+    def prune_to_dataset(
+        self,
+        dataset: Dataset,
+        remove_ontologies: Set[str] = set(),
+        num_proc: int = 4,
+        batch_size: int = 1024,
+    ) -> None:
+        mapped_dataset = dataset.map(
+            _get_all_codes_map,
+            batched=True,
+            batch_size=batch_size,
+            remove_columns=dataset.column_names,
+            num_proc=num_proc,
+        )
+        valid_concept_ids = set(mapped_dataset["unique_concept_ids"])
+        all_parents = set()
+        for concept_id in valid_concept_ids:
+            all_parents |= self.get_all_parents(concept_id)
+
+        def is_valid(c: str):
+            ontology = self.concept_vocabulary_map.get(c, "")
+            return (c in valid_concept_ids) or (
+                (ontology not in remove_ontologies) and (c in all_parents)
+            )
+
+        concept_ids = set(self.parents_map.keys())
+        for concept_id in concept_ids:
+            m: Any
+            if is_valid(concept_id):
+                for m in (self.parents_map, self.concept_vocabulary_map):
+                    m[concept_id] = {a for a in m[concept_id] if is_valid(a)}
+            else:
+                for m in (self.parents_map, self.concept_vocabulary_map):
+                    if concept_id in m:
+                        del m[concept_id]
+
+        self.all_parents_map = {}
+
+        # Prime the pump
+        for concept_id in self.parents_map.keys():
+            self.get_all_parents(concept_id)
+
+    def get_parents(self, code: str) -> Iterable[str]:
+        """Get the parents for a given code."""
+        return self.parents_map.get(code, set())
+
+    def get_all_parents(self, code: str) -> Set[str]:
+        """Get all parents, including through the ontology."""
+        if code not in self.all_parents_map:
+            result = {code}
+            for parent in self.parents_map.get(code, set()):
+                result |= self.get_all_parents(parent)
+            self.all_parents_map[code] = result
+
+        return self.all_parents_map[code]

cehrgpt/runners/hf_cehrgpt_finetune_runner.py

@@ -1,3 +1,4 @@
+import glob
 import json
 import os
 import random
@@ -175,11 +176,6 @@ def model_init(
         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}"
@@ -379,7 +375,6 @@ 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
@@ -406,8 +401,9 @@ def main():
     )
 
     if training_args.do_train:
+        output_dir = training_args.output_dir
         if cehrgpt_args.hyperparameter_tuning:
-            training_args = perform_hyperparameter_search(
+            training_args, run_id = perform_hyperparameter_search(
                 trainer_class,
                 partial(model_init, model_args, training_args, cehrgpt_args, tokenizer),
                 processed_dataset,
@@ -416,18 +412,28 @@ def main():
                 model_args,
                 cehrgpt_args,
             )
-
-        if cehrgpt_args.retrain_with_full:
-            # Always retrain with the full set when hyperparameter tuning is set to true
-            retrain_with_full_set(
-                trainer_class,
-                model_args,
-                training_args,
-                cehrgpt_args,
-                tokenizer,
-                processed_dataset,
-                data_collator,
+            # We enforce retraining if cehrgpt_args.hyperparameter_tuning_percentage < 1.0
+            cehrgpt_args.retrain_with_full |= (
+                cehrgpt_args.hyperparameter_tuning_percentage < 1.0
             )
+            output_dir = os.path.join(training_args.output_dir, f"run-{run_id}")
+
+        if cehrgpt_args.hyperparameter_tuning and not cehrgpt_args.retrain_with_full:
+            folders = glob.glob(os.path.join(output_dir, "checkpoint-*"))
+            if len(folders) == 0:
+                raise RuntimeError(
+                    f"There must be a checkpoint folder under {output_dir}"
+                )
+            checkpoint_dir = folders[0]
+            LOG.info("Best trial checkpoint folder: %s", checkpoint_dir)
+            for file_name in os.listdir(checkpoint_dir):
+                try:
+                    full_file_name = os.path.join(checkpoint_dir, file_name)
+                    destination = os.path.join(training_args.output_dir, file_name)
+                    if os.path.isfile(full_file_name):
+                        shutil.copy2(full_file_name, destination)
+                except Exception as e:
+                    LOG.error("Failed to copy %s: %s", file_name, str(e))
         else:
             # Initialize Trainer for final training on the combined train+val set
             trainer = trainer_class(
@@ -476,63 +482,6 @@ def main():
         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,
-) -> 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:
-        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.
-
-    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_class(
-        model=model_init(model_args, training_args, cehrgpt_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,
@@ -652,9 +601,6 @@ def load_lora_model(
     # 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

cehrgpt/runners/hf_cehrgpt_pretrain_runner.py

@@ -2,7 +2,9 @@ import os
 from functools import partial
 from typing import Optional, Union
 
+import datasets
 import numpy as np
+import pandas as pd
 import torch
 import torch.distributed as dist
 from cehrbert.data_generators.hf_data_generator.meds_utils import (
@@ -20,7 +22,7 @@ from cehrbert.runners.runner_util import (
     load_parquet_as_dataset,
 )
 from datasets import Dataset, DatasetDict, IterableDatasetDict, load_from_disk
-from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed
+from transformers import EarlyStoppingCallback, Trainer, set_seed
 from transformers.trainer_utils import is_main_process
 from transformers.utils import is_flash_attn_2_available, logging
 
@@ -34,6 +36,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.omop.ontology import Ontology
 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
@@ -70,68 +73,64 @@ def load_and_create_tokenizer(
     data_args: DataTrainingArguments,
     model_args: ModelArguments,
     cehrgpt_args: CehrGPTArguments,
-    dataset: Optional[Union[Dataset, DatasetDict]] = None,
+    dataset: Union[Dataset, DatasetDict],
 ) -> 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:
-        tokenizer = CehrGptTokenizer.from_pretrained(tokenizer_abspath)
-    except Exception as e:
-        LOG.warning(e)
-        if dataset is None:
+    if not tokenizer_exists(tokenizer_abspath):
+        if cehrgpt_args.include_motor_time_to_event and not cehrgpt_args.vocab_dir:
             raise RuntimeError(
-                f"Failed to load the tokenizer from {tokenizer_abspath} with the error \n{e}\n"
-                f"Tried to create the tokenizer, however the dataset is not provided."
+                "motor_vocab_dir must be specified if include_motor_time_to_event is True"
+            )
+        ontology: Optional[Ontology] = None
+        concept_name_mapping = {}
+        if cehrgpt_args.vocab_dir:
+            LOG.info("Loading concept data from disk at %s", cehrgpt_args.vocab_dir)
+            concept_pd = pd.read_parquet(
+                os.path.join(cehrgpt_args.vocab_dir, "concept")
             )
+            for row in concept_pd.itertuples():
+                concept_name_mapping[str(getattr(row, "concept_id"))] = getattr(
+                    row, "concept_name"
+                )
+
+            if cehrgpt_args.motor_use_ontology:
+                LOG.info("Creating ontology for MOTOR TTE predictions")
+                ontology = Ontology(cehrgpt_args.vocab_dir)
+                train_val_dataset = datasets.concatenate_datasets(
+                    [dataset["train"], dataset["validation"]]
+                )
+                ontology.prune_to_dataset(
+                    train_val_dataset,
+                    num_proc=data_args.preprocessing_num_workers,
+                    remove_ontologies={"SPL", "HemOnc", "LOINC"},
+                )
+
         LOG.info("Started training the tokenizer ...")
+        train_val_dataset = datasets.concatenate_datasets(
+            [dataset["train"], dataset["validation"]]
+        )
         tokenizer = CehrGptTokenizer.train_tokenizer(
-            dataset,
+            train_val_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,
-            (
+            num_motor_tasks=(
                 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,
+            ontology=ontology,
         )
         LOG.info("Finished training the tokenizer ...")
         tokenizer.save_pretrained(tokenizer_abspath)
         LOG.info("Saved the tokenizer to %s", tokenizer_abspath)
-
+    else:
+        LOG.info("The tokenizer exists and will be loaded from %s", tokenizer_abspath)
+        tokenizer = CehrGptTokenizer.from_pretrained(tokenizer_abspath)
     return tokenizer
 
 
@@ -187,7 +186,10 @@ def load_and_create_model(
 
         model_args_cehrgpt = model_args.as_dict()
         model_args_cehrgpt.pop("attn_implementation")
+        # CEHR-GPT does not support this anymore
+        model_args_cehrgpt.pop("exclude_position_ids")
         model_config = CEHRGPTConfig(
+            activation_function=cehrgpt_args.activation_function,
             vocab_size=tokenizer.vocab_size,
             value_vocab_size=tokenizer.value_vocab_size,
             time_token_vocab_size=tokenizer.time_token_vocab_size,
@@ -207,6 +209,7 @@ def load_and_create_model(
             n_pretrained_embeddings_layers=cehrgpt_args.n_pretrained_embeddings_layers,
             use_pretrained_embeddings=len(tokenizer.pretrained_token_ids) > 0,
             pretrained_embedding_dim=pretrained_embedding_dim,
+            apply_rotary=cehrgpt_args.apply_rotary,
             sample_packing_max_positions=(
                 cehrgpt_args.max_tokens_per_batch
                 if cehrgpt_args.sample_packing
@@ -217,6 +220,8 @@ def load_and_create_model(
             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,
+            n_inner=cehrgpt_args.inner_dim,
+            decoder_mlp=cehrgpt_args.decoder_mlp,
             **model_args_cehrgpt,
         )
 
@@ -235,7 +240,6 @@ def load_and_create_model(
 
 def main():
     cehrgpt_args, data_args, model_args, training_args = parse_runner_args()
-
     if cehrgpt_args.sample_packing and data_args.streaming:
         raise RuntimeError(
             f"sample_packing is not supported when streaming is enabled, please set streaming to False"
@@ -530,7 +534,6 @@ def main():
             SamplePackingCehrGptDataCollator,
             cehrgpt_args.max_tokens_per_batch,
             model_args.max_position_embeddings,
-            add_end_token_in_sample_packing=cehrgpt_args.add_end_token_in_sample_packing,
         )
     else:
         trainer_class = Trainer
@@ -552,6 +555,7 @@ def main():
             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,
+            motor_sampling_probability=cehrgpt_args.motor_sampling_probability,
         ),
         train_dataset=processed_dataset["train"],
         eval_dataset=(

cehrgpt/runners/hf_gpt_runner_argument_dataclass.py

@@ -1,5 +1,7 @@
 import dataclasses
-from typing import List, Optional
+from typing import List, Literal, Optional
+
+from cehrgpt.models.gpt2 import ACT2FN
 
 
 @dataclasses.dataclass
@@ -12,6 +14,14 @@ class CehrGPTArguments:
             "help": "The path to the tokenized dataset created for the full population"
         },
     )
+    activation_function: Literal[tuple(ACT2FN.keys())] = dataclasses.field(
+        default="gelu_new",
+        metadata={"help": "The activation function to use"},
+    )
+    decoder_mlp: Literal["GPT2MLP", "LlamaMLP"] = dataclasses.field(
+        default="GPT2MLP",
+        metadata={"help": "The decoder MLP architecture"},
+    )
     include_inpatient_hour_token: Optional[bool] = dataclasses.field(
         default=True,
         metadata={"help": "Include inpatient hour token"},
@@ -54,6 +64,14 @@ class CehrGPTArguments:
         default=128,
         metadata={"help": "The number of examples from the training set."},
     )
+    hyperparameter_tuning: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={"help": "A flag to indicate if we want to do hyperparameter tuning."},
+    )
+    hyperparameter_tuning_is_grid: Optional[bool] = dataclasses.field(
+        default=True,
+        metadata={"help": "A flag to indicate if we want to do hyperparameter tuning."},
+    )
     hyperparameter_tuning_percentage: Optional[float] = dataclasses.field(
         default=0.1,
         metadata={
@@ -66,10 +84,6 @@ class CehrGPTArguments:
             "help": "The number of trails will be use for hyperparameter tuning."
         },
     )
-    hyperparameter_tuning: Optional[bool] = dataclasses.field(
-        default=False,
-        metadata={"help": "A flag to indicate if we want to do hyperparameter tuning."},
-    )
     hyperparameter_batch_sizes: Optional[List[int]] = dataclasses.field(
         default_factory=lambda: [4, 8, 16],
         metadata={"help": "Hyperparameter search batch sizes"},
@@ -78,29 +92,13 @@ class CehrGPTArguments:
         default_factory=lambda: [10],
         metadata={"help": "Hyperparameter search num_train_epochs"},
     )
-    lr_low: Optional[float] = dataclasses.field(
-        default=1e-5,
-        metadata={
-            "help": "The lower bound of the learning rate range for hyperparameter tuning."
-        },
-    )
-    lr_high: Optional[float] = dataclasses.field(
-        default=5e-5,
-        metadata={
-            "help": "The upper bound of the learning rate range for hyperparameter tuning."
-        },
-    )
-    weight_decays_low: Optional[float] = dataclasses.field(
-        default=1e-3,
-        metadata={
-            "help": "The lower bound of the weight decays range for hyperparameter tuning."
-        },
+    hyperparameter_learning_rates: Optional[List[int]] = dataclasses.field(
+        default_factory=lambda: [1e-5],
+        metadata={"help": "Hyperparameter search learning rates"},
     )
-    weight_decays_high: Optional[float] = dataclasses.field(
-        default=1e-2,
-        metadata={
-            "help": "The upper bound of the weight decays range for hyperparameter tuning."
-        },
+    hyperparameter_weight_decays: Optional[List[int]] = dataclasses.field(
+        default_factory=lambda: [1e-2],
+        metadata={"help": "Hyperparameter search learning rates"},
     )
     causal_sfm: Optional[bool] = dataclasses.field(
         default=False,
@@ -168,6 +166,16 @@ class CehrGPTArguments:
             "help": "A threshold to denote how much the specified metric must improve to satisfy early stopping conditions."
         },
     )
+    inner_dim: Optional[int] = dataclasses.field(
+        default=None,
+        metadata={"help": "The dimensionality of the hidden layer"},
+    )
+    apply_rotary: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to use rotary encoder layers"
+        },
+    )
     sample_packing: Optional[bool] = dataclasses.field(
         default=False,
         metadata={
@@ -177,12 +185,6 @@ class CehrGPTArguments:
     max_tokens_per_batch: int = dataclasses.field(
         default=16384, metadata={"help": "Maximum number of tokens in each batch"}
     )
-    add_end_token_in_sample_packing: Optional[bool] = dataclasses.field(
-        default=False,
-        metadata={
-            "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={
@@ -203,7 +205,17 @@ class CehrGPTArguments:
             "help": "The number of times each motor_num_time_pieces piece has to be"
         },
     )
-    concept_dir: Optional[str] = dataclasses.field(
+    motor_use_ontology: Optional[bool] = dataclasses.field(
+        default=False,
+        metadata={
+            "help": "A flag to indicate whether we want to use motor_use_ontology"
+        },
+    )
+    motor_sampling_probability: Optional[float] = dataclasses.field(
+        default=0.0,
+        metadata={"help": "A flag to indicate whether we want to use sample packing"},
+    )
+    vocab_dir: Optional[str] = dataclasses.field(
         default=None,
         metadata={"help": "The directory where the concept data is stored."},
     )