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

cehrgpt/data/hf_cehrgpt_dataset_collator.py

@@ -162,6 +162,22 @@ class CehrGptDataCollator:
             f"batch['input_ids']: {batch['input_ids']} "
         )
 
+        if "epoch_times" in examples[0]:
+            batch_epoch_times = [
+                self._try_reverse_tensor(
+                    self._convert_to_tensor(example["epoch_times"])
+                )
+                for example in examples
+            ]
+            # Pad sequences to the max length in the batch
+            batch["epoch_times"] = self._try_reverse_tensor(
+                pad_sequence(
+                    batch_epoch_times,
+                    batch_first=True,
+                    padding_value=0,
+                ).to(torch.float32)
+            )
+
         if "position_ids" in examples[0]:
             batch_position_ids = [
                 self._try_reverse_tensor(
@@ -663,7 +679,9 @@ class CehrGptDataCollator:
 
         # Subtract one for the [END] token when sample_packing is not enabled
         new_max_length = (
-            max_length_allowed if sample_packing else max_length_allowed - 1
+            max_length_allowed - 1
+            if not sample_packing and self.pretraining
+            else max_length_allowed
         )
 
         if self.include_ttv_prediction:
@@ -685,13 +703,20 @@ class CehrGptDataCollator:
 
         # Return the record directly if the actual sequence length is less than the max sequence
         if seq_length <= new_max_length:
-            if not sample_packing:
+            if not sample_packing and self.pretraining:
                 record["input_ids"] = torch.concat(
                     [
                         self._convert_to_tensor(record["input_ids"]),
                         self._convert_to_tensor([eos_token]),
                     ]
                 )
+                if "epoch_times" in record:
+                    record["epoch_times"] = torch.concat(
+                        [
+                            self._convert_to_tensor(record["epoch_times"]),
+                            self._convert_to_tensor([record["epoch_times"][-1]]),
+                        ]
+                    )
                 if self.include_values:
                     record["value_indicators"] = torch.concat(
                         [
@@ -727,6 +752,10 @@ class CehrGptDataCollator:
                     record["input_ids"] = self._convert_to_tensor(
                         record["input_ids"][start_index : end_index + 1]
                     )
+                    if "epoch_times" in record:
+                        record["epoch_times"] = self._convert_to_tensor(
+                            record["epoch_times"][start_index : end_index + 1]
+                        )
                     if self.include_values:
                         record["value_indicators"] = self._convert_to_tensor(
                             record["value_indicators"][start_index : end_index + 1]
@@ -760,6 +789,11 @@ class CehrGptDataCollator:
             if sample_packing and "position_ids" in record:
                 record["position_ids"] = record["position_ids"][0:end_index]
 
+            if "epoch_times" in record:
+                record["epoch_times"] = self._convert_to_tensor(
+                    record["epoch_times"][0:end_index]
+                )
+
             if self.include_values:
                 record["value_indicators"] = self._convert_to_tensor(
                     record["value_indicators"][0:end_index]
@@ -792,6 +826,17 @@ class CehrGptDataCollator:
                                 ),
                             ]
                         )
+                        if "epoch_times" in record:
+                            record["epoch_times"] = torch.concat(
+                                [
+                                    torch.zeros(
+                                        [record["epoch_times"][0]], dtype=torch.float32
+                                    ),
+                                    self._convert_to_tensor(
+                                        record["epoch_times"][token_index:seq_length]
+                                    ),
+                                ]
+                            )
                         if self.include_values:
                             record["value_indicators"] = torch.concat(
                                 [
@@ -830,7 +875,7 @@ class CehrGptDataCollator:
                             )
                         break
             else:
-                start_index = seq_length - new_max_length
+                start_index = max(seq_length - new_max_length, 0)
                 end_index = seq_length
                 for i in range(start_index, end_index):
                     current_token = record["input_ids"][i]
@@ -842,6 +887,11 @@ class CehrGptDataCollator:
                             ]
                         if sample_packing and "position_ids" in record:
                             record["position_ids"] = record["position_ids"][i:end_index]
+
+                        if "epoch_times" in record:
+                            record["epoch_times"] = self._convert_to_tensor(
+                                record["epoch_times"][i:end_index]
+                            )
                         if self.include_values:
                             record["value_indicators"] = record["value_indicators"][
                                 i:end_index
@@ -863,6 +913,10 @@ class CehrGptDataCollator:
                     ]
                 if sample_packing and "position_ids" in record:
                     record["position_ids"] = record["position_ids"][-new_max_length:]
+                if "epoch_times" in record:
+                    record["epoch_times"] = self._convert_to_tensor(
+                        record["epoch_times"][-new_max_length:]
+                    )
                 if self.include_values:
                     record["value_indicators"] = record["value_indicators"][
                         -new_max_length:
@@ -873,36 +927,6 @@ class CehrGptDataCollator:
                         -new_max_length:
                     ]
 
-            if not sample_packing:
-                # Finally we add the end token to the end of the sequence
-                record["input_ids"] = torch.concat(
-                    [
-                        self._convert_to_tensor(record["input_ids"]),
-                        self._convert_to_tensor([eos_token]),
-                    ]
-                )
-                if self.include_values:
-                    record["value_indicators"] = torch.concat(
-                        [
-                            self._convert_to_tensor(record["value_indicators"]),
-                            self._convert_to_tensor([False]),
-                        ]
-                    ).to(torch.bool)
-                    record["values"] = torch.concat(
-                        [
-                            self._convert_to_tensor(record["values"]),
-                            self._convert_to_tensor(
-                                [self.tokenizer.pad_value_token_id]
-                            ),
-                        ]
-                    )
-                if self.include_ttv_prediction:
-                    record["time_to_visits"] = torch.concat(
-                        [
-                            record["time_to_visits"],
-                            self._convert_to_tensor([-100.0]),
-                        ]
-                    )
             return record
 
 

cehrgpt/data/hf_cehrgpt_dataset_mapping.py

@@ -21,7 +21,6 @@ from cehrbert_data.const.artificial_tokens import (
     DISCHARGE_UNKNOWN_TOKEN,
     GENDER_UNKNOWN_TOKEN,
     RACE_UNKNOWN_TOKEN,
-    VISIT_UNKNOWN_TOKEN,
 )
 from cehrbert_data.const.common import NA
 from cehrbert_data.decorators.patient_event_decorator_base import get_att_function
@@ -47,10 +46,16 @@ CEHRGPT_COLUMNS = [
 ]
 
 
-def convert_date_to_posix_time(index_date: datetime.date) -> float:
-    return datetime.datetime.combine(
-        index_date, datetime.datetime.min.time()
-    ).timestamp()
+def convert_date_to_posix_time(index_date: Union[datetime.date, int, float]) -> float:
+    if isinstance(index_date, datetime.date):
+        return (
+            datetime.datetime.combine(index_date, datetime.datetime.min.time())
+            .replace(tzinfo=datetime.timezone.utc)
+            .timestamp()
+        )
+    elif isinstance(index_date, datetime.datetime):
+        return index_date.replace(tzinfo=datetime.timezone.utc).timestamp()
+    return index_date
 
 
 class DatasetMappingDecorator(DatasetMapping):
@@ -128,7 +133,9 @@ class MedToCehrGPTDatasetMapping(DatasetMappingDecorator):
         cehrgpt_record["concept_as_values"].append(concept_as_value)
         cehrgpt_record["units"].append(unit)
         cehrgpt_record["is_numeric_types"].append(is_numeric_type)
-        cehrgpt_record["epoch_times"].append(time.timestamp())
+        cehrgpt_record["epoch_times"].append(
+            time.replace(tzinfo=datetime.timezone.utc).timestamp()
+        )
 
     def transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
         cehrgpt_record = {
@@ -360,7 +367,9 @@ class MedToCehrGPTDatasetMapping(DatasetMappingDecorator):
         cehrgpt_record["num_of_visits"] = len(visits)
 
         if record.get("index_date", None) is not None:
-            cehrgpt_record["index_date"] = record["index_date"]
+            cehrgpt_record["index_date"] = (
+                record["index_date"].replace(tzinfo=datetime.timezone.utc).timestamp()
+            )
         if record.get("label", None) is not None:
             cehrgpt_record["label"] = record["label"]
         if record.get("age_at_index", None) is not None:
@@ -529,9 +538,13 @@ class ExtractTokenizedSequenceDataMapping:
         prediction_start_end_times = [
             (
                 self._calculate_prediction_start_time(
-                    prediction_time_label_map["index_date"].timestamp()
+                    prediction_time_label_map["index_date"]
+                    .replace(tzinfo=datetime.timezone.utc)
+                    .timestamp()
                 ),
-                prediction_time_label_map["index_date"].timestamp(),
+                prediction_time_label_map["index_date"]
+                .replace(tzinfo=datetime.timezone.utc)
+                .timestamp(),
                 prediction_time_label_map["label"],
             )
             for prediction_time_label_map in prediction_times

cehrgpt/generation/cehrgpt_conditional_generation.py

@@ -0,0 +1,314 @@
+import datetime
+import os
+import random
+import shutil
+from pathlib import Path
+from typing import Any, Dict
+
+import numpy as np
+import polars as pl
+import torch
+import torch.distributed as dist
+from cehrbert.runners.runner_util import generate_prepared_ds_path
+from datasets import load_from_disk
+from meds import held_out_split, train_split, tuning_split
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+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.generation.generate_batch_hf_gpt_sequence import (
+    generate_single_batch,
+    normalize_value,
+)
+from cehrgpt.gpt_utils import (
+    extract_time_interval_in_days,
+    extract_time_interval_in_hours,
+    is_att_token,
+    is_inpatient_hour_token,
+    is_visit_end,
+    is_visit_start,
+)
+from cehrgpt.models.hf_cehrgpt import CEHRGPT2LMHeadModel
+from cehrgpt.models.tokenization_hf_cehrgpt import CehrGptTokenizer
+from cehrgpt.runners.data_utils import (
+    extract_cohort_sequences,
+    prepare_finetune_dataset,
+)
+from cehrgpt.runners.gpt_runner_util import parse_runner_args
+from cehrgpt.runners.hf_cehrgpt_pretrain_runner import tokenizer_exists
+
+LOG = logging.get_logger("transformers")
+
+
+def map_data_split_name(split: str) -> str:
+    if split == "train":
+        return train_split
+    elif split == "validation":
+        return tuning_split
+    elif split == "test":
+        return held_out_split
+    raise ValueError(f"Unknown split: {split}")
+
+
+def seed_all(seed: int = 42):
+    """Set seed for Python, NumPy, and PyTorch (CPU & CUDA)."""
+    random.seed(seed)  # Python random
+    np.random.seed(seed)  # NumPy
+    torch.manual_seed(seed)  # PyTorch CPU
+    torch.cuda.manual_seed(seed)  # Current GPU
+    torch.cuda.manual_seed_all(seed)  # All GPUs
+
+    # For reproducibility in dataloader workers
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+
+def generate_trajectories_per_batch(
+    batch: Dict[str, Any],
+    cehrgpt_tokenizer: CehrGptTokenizer,
+    cehrgpt_model: CEHRGPT2LMHeadModel,
+    device,
+    data_output_path: Path,
+    max_length: int,
+):
+    subject_ids = batch["person_id"].squeeze().detach().cpu().tolist()
+    prediction_times = batch["index_date"].squeeze().detach().cpu().tolist()
+    batched_epoch_times = batch["epoch_times"].detach().cpu().tolist()
+    batched_input_ids = batch["input_ids"]
+    batched_value_indicators = batch["value_indicators"]
+    batched_values = batch["values"]
+    # Make sure the batch does not exceed batch_size
+    batch_sequences = generate_single_batch(
+        cehrgpt_model,
+        cehrgpt_tokenizer,
+        batched_input_ids,
+        values=batched_values,
+        value_indicators=batched_value_indicators,
+        max_length=max_length,
+        top_p=1.0,
+        top_k=cehrgpt_tokenizer.vocab_size,
+        device=device,
+    )
+    # Clear the cache
+    torch.cuda.empty_cache()
+
+    trajectories = []
+    for sample_i, (concept_ids, value_indicators, values) in enumerate(
+        zip(
+            batch_sequences["sequences"],
+            batch_sequences["value_indicators"],
+            batch_sequences["values"],
+        )
+    ):
+        (
+            concept_ids,
+            is_numeric_types,
+            number_as_values,
+            concept_as_values,
+            units,
+        ) = normalize_value(concept_ids, values, cehrgpt_tokenizer)
+
+        epoch_times = batched_epoch_times[sample_i]
+        input_length = len(epoch_times)
+        # Getting the last observed event time from the token before the prediction time
+        window_last_observed = epoch_times[input_length - 1]
+        current_cursor = epoch_times[-1]
+        generated_epoch_times = []
+        valid_indices = []
+
+        for i in range(input_length, len(concept_ids)):
+            concept_id = concept_ids[i]
+            # We use the left padding strategy in the data collator
+            if concept_id in [cehrgpt_tokenizer.pad_token, cehrgpt_tokenizer.end_token]:
+                continue
+            # We need to construct the time stamp
+            if is_att_token(concept_id):
+                current_cursor += extract_time_interval_in_days(concept_id) * 24 * 3600
+            elif is_inpatient_hour_token(concept_id):
+                current_cursor += extract_time_interval_in_hours(concept_id) * 3600
+            elif is_visit_start(concept_id) or is_visit_end(concept_id):
+                continue
+            else:
+                valid_indices.append(i)
+                generated_epoch_times.append(
+                    datetime.datetime.utcfromtimestamp(current_cursor).replace(
+                        tzinfo=None
+                    )
+                )
+
+        trajectories.append(
+            {
+                "subject_id": subject_ids[sample_i],
+                "prediction_time": datetime.datetime.utcfromtimestamp(
+                    prediction_times[sample_i]
+                ).replace(tzinfo=None),
+                "window_last_observed_time": datetime.datetime.utcfromtimestamp(
+                    window_last_observed
+                ).replace(tzinfo=None),
+                "times": generated_epoch_times,
+                "concept_ids": np.asarray(concept_ids)[valid_indices].tolist(),
+                "numeric_values": np.asarray(number_as_values)[valid_indices].tolist(),
+                "text_value": np.asarray(concept_as_values)[valid_indices].tolist(),
+                "units": np.asarray(units)[valid_indices].tolist(),
+            }
+        )
+
+    trajectories = (
+        pl.DataFrame(trajectories)
+        .explode(["times", "concept_ids", "numeric_values", "text_value", "units"])
+        .rename(
+            {
+                "times": "time",
+                "concept_ids": "code",
+                "numeric_values": "numeric_value",
+                "units": "unit",
+            }
+        )
+        .select(
+            "subject_id",
+            "prediction_time",
+            "window_last_observed_time",
+            "time",
+            "code",
+            "numeric_value",
+            "text_value",
+            "unit",
+        )
+    )
+    trajectories.write_parquet(data_output_path)
+
+
+def main():
+    cehrgpt_args, data_args, model_args, training_args = parse_runner_args()
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+    cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(
+        model_args.tokenizer_name_or_path
+    )
+    cehrgpt_model = (
+        CEHRGPT2LMHeadModel.from_pretrained(
+            model_args.model_name_or_path,
+            attn_implementation=(
+                "flash_attention_2" if is_flash_attn_2_available() else "eager"
+            ),
+        )
+        .eval()
+        .to(device)
+    )
+    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
+
+    if not os.path.exists(training_args.output_dir):
+        os.makedirs(training_args.output_dir)
+
+    prepared_ds_path = generate_prepared_ds_path(
+        data_args, model_args, data_folder=data_args.cohort_folder
+    )
+
+    processed_dataset = None
+    if any(prepared_ds_path.glob("*")):
+        LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
+        processed_dataset = load_from_disk(str(prepared_ds_path))
+        LOG.info("Prepared dataset loaded from disk...")
+        if cehrgpt_args.expand_tokenizer:
+            if tokenizer_exists(training_args.output_dir):
+                cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(
+                    training_args.output_dir
+                )
+            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. "
+                    f"Please delete the processed dataset at {prepared_ds_path}."
+                )
+                processed_dataset = None
+                shutil.rmtree(prepared_ds_path)
+
+    if processed_dataset is None and 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)
+        else:
+            # Organize them into a single DatasetDict
+            final_splits = prepare_finetune_dataset(
+                data_args, training_args, cehrgpt_args
+            )
+            # 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=cehrgpt_tokenizer,
+                data_args=data_args,
+            )
+        if not data_args.streaming:
+            processed_dataset.save_to_disk(prepared_ds_path)
+            processed_dataset.cleanup_cache_files()
+
+    # After main-process-only operations, synchronize all processes to ensure consistency
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+
+    # We suppress the additional learning objectives in fine-tuning
+    data_collator = CehrGptDataCollator(
+        tokenizer=cehrgpt_tokenizer,
+        max_length=cehrgpt_args.generation_input_length,
+        include_values=cehrgpt_model.config.include_values,
+        pretraining=False,
+        include_ttv_prediction=False,
+        use_sub_time_tokenization=False,
+        include_demographics=False,
+        add_linear_prob_token=False,
+    )
+
+    LOG.info(
+        "Generating %s trajectories per sample",
+        cehrgpt_args.num_of_trajectories_per_sample,
+    )
+    for sample_i in range(cehrgpt_args.num_of_trajectories_per_sample):
+        for split, dataset in processed_dataset.items():
+            meds_split = map_data_split_name(split)
+            dataloader = DataLoader(
+                dataset=dataset,
+                batch_size=training_args.per_device_eval_batch_size,
+                num_workers=training_args.dataloader_num_workers,
+                collate_fn=data_collator,
+                pin_memory=training_args.dataloader_pin_memory,
+            )
+            sample_output_dir = (
+                Path(training_args.output_dir) / meds_split / f"{sample_i}"
+            )
+            sample_output_dir.mkdir(exist_ok=True, parents=True)
+            for batch_i, batch in tqdm(
+                enumerate(dataloader),
+                desc=f"Generating Trajectories for split {meds_split} with trajectory {sample_i + 1}",
+            ):
+                output_parquet_file = sample_output_dir / f"{batch_i}.parquet"
+                if output_parquet_file.exists():
+                    LOG.info("%s already exists, skip...", output_parquet_file)
+                    continue
+
+                generate_trajectories_per_batch(
+                    batch,
+                    cehrgpt_tokenizer,
+                    cehrgpt_model,
+                    device,
+                    sample_output_dir / f"{batch_i}.parquet",
+                    cehrgpt_args.generation_max_new_tokens
+                    + cehrgpt_args.generation_input_length,
+                )
+
+
+if __name__ == "__main__":
+    # ✅ Call first thing inside main()
+    seed_all(42)
+    main()

cehrgpt/generation/generate_batch_hf_gpt_sequence.py

@@ -74,7 +74,10 @@ def generate_single_batch(
     model: CEHRGPT2LMHeadModel,
     tokenizer: CehrGptTokenizer,
     prompts: List[List[int]],
-    max_new_tokens=512,
+    max_length: int,
+    values: Optional[torch.Tensor] = None,
+    value_indicators: Optional[torch.Tensor] = None,
+    max_new_tokens: Optional[int] = None,
     mini_num_of_concepts=1,
     top_p=0.95,
     top_k=50,
@@ -88,7 +91,8 @@ def generate_single_batch(
     with torch.no_grad():
         generation_config = GenerationConfig(
             repetition_penalty=repetition_penalty,
-            max_length=max_new_tokens,
+            max_new_tokens=max_new_tokens,
+            max_length=max_length,
             min_length=mini_num_of_concepts,
             temperature=temperature,
             top_p=top_p,
@@ -107,9 +111,17 @@ def generate_single_batch(
             num_beam_groups=num_beam_groups,
             epsilon_cutoff=epsilon_cutoff,
         )
+
         batched_prompts = torch.tensor(prompts).to(device)
+        if values is not None:
+            values = values.to(device)
+        if value_indicators is not None:
+            value_indicators = value_indicators.to(device)
+
         results = model.generate(
             inputs=batched_prompts,
+            values=values,
+            value_indicators=value_indicators,
             generation_config=generation_config,
             lab_token_ids=tokenizer.lab_token_ids,
         )
@@ -226,7 +238,7 @@ def main(args):
             cehrgpt_model,
             cehrgpt_tokenizer,
             random_prompts[: args.batch_size],
-            max_new_tokens=args.context_window,
+            max_length=args.context_window,
             mini_num_of_concepts=args.min_num_of_concepts,
             top_p=args.top_p,
             top_k=args.top_k,

cehrgpt/models/hf_cehrgpt.py

@@ -102,7 +102,9 @@ def is_sample_pack(attention_mask: torch.Tensor) -> bool:
         attention_mask = attention_mask.flip(dims=[1])
 
     nonzero_counts = attention_mask.sum(dim=1)
-    max_token_positions = torch.argmax(attention_mask.flip(dims=[1]), dim=1)
+    max_token_positions = torch.argmax(
+        attention_mask.to(torch.int32).flip(dims=[1]), dim=1
+    )
     max_indices = attention_mask.shape[1] - 1 - max_token_positions
     return torch.any(nonzero_counts < (max_indices + 1)).item()
 
@@ -1848,6 +1850,7 @@ class CEHRGPT2LMHeadModel(CEHRGPTPreTrainedModel):
 
         # keep track of which sequences are already finished
         batch_size, cur_len = input_ids.shape
+        model_kwargs["attention_mask"] = input_ids != pad_token_id
         if "inputs_embeds" in model_kwargs:
             cur_len = model_kwargs["inputs_embeds"].shape[1]
         this_peer_finished = False
@@ -1866,11 +1869,19 @@ class CEHRGPT2LMHeadModel(CEHRGPTPreTrainedModel):
                 [] if self.config.lab_token_ids is None else self.config.lab_token_ids,
                 dtype=torch.int32,
             )
-        value_indicators = torch.zeros_like(input_ids).to(torch.bool)
-        values = torch.zeros_like(
-            input_ids,
-            dtype=torch.int32,
-        )
+
+        if model_kwargs.get("value_indicators", None) is not None:
+            value_indicators = model_kwargs.get("value_indicators")
+        else:
+            value_indicators = torch.zeros_like(input_ids).to(torch.bool)
+
+        if model_kwargs.get("values", None) is not None:
+            values = model_kwargs.get("values")
+        else:
+            values = torch.zeros_like(
+                input_ids,
+                dtype=torch.int32,
+            )
         # Generate initial random_vectors
         if self.cehrgpt.config.causal_sfm:
             model_kwargs["random_vectors"] = torch.rand(

cehrgpt/runners/data_utils.py

@@ -47,7 +47,7 @@ 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:
@@ -91,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"]
@@ -271,7 +272,7 @@ def create_dataset_splits(data_args: DataTrainingArguments, seed: int):
 def extract_cohort_sequences(
     data_args: DataTrainingArguments,
     cehrgpt_args: CehrGPTArguments,
-    cache_file_collector: CacheFileCollector,
+    cache_file_collector: Optional[CacheFileCollector] = None,
 ) -> DatasetDict:
     """
     Extracts and processes cohort-specific tokenized sequences from a pre-tokenized dataset,.
@@ -309,9 +310,18 @@ def extract_cohort_sequences(
             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(
@@ -353,6 +363,7 @@ def extract_cohort_sequences(
         num_proc=data_args.preprocessing_num_workers,
         remove_columns=filtered_tokenized_dataset["train"].column_names,
     )
-    cache_file_collector.add_cache_files(filtered_tokenized_dataset)
-    cache_file_collector.add_cache_files(processed_dataset)
+    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

@@ -580,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_gpt_runner_argument_dataclass.py

@@ -229,3 +229,15 @@ class CehrGPTArguments:
             "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/tools/linear_prob/compute_cehrgpt_features.py

@@ -1,15 +1,14 @@
+import datetime
 import glob
 import os
 import shutil
 import uuid
-from datetime import datetime
 from functools import partial
 from pathlib import Path
 from typing import Optional, Union
 
 import numpy as np
 import pandas as pd
-import polars as pl
 import torch
 import torch.distributed as dist
 from cehrbert.data_generators.hf_data_generator.meds_utils import CacheFileCollector
@@ -25,7 +24,6 @@ from cehrgpt.data.hf_cehrgpt_dataset_collator import (
     CehrGptDataCollator,
     SamplePackingCehrGptDataCollator,
 )
-from cehrgpt.data.hf_cehrgpt_dataset_mapping import ExtractTokenizedSequenceDataMapping
 from cehrgpt.data.sample_packing_sampler import SamplePackingBatchSampler
 from cehrgpt.models.hf_cehrgpt import (
     CEHRGPT2Model,
@@ -159,24 +157,7 @@ def main():
                 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):
-                        cehrgpt_tokenizer = CehrGptTokenizer.from_pretrained(
-                            new_tokenizer_path
-                        )
-                    else:
-                        cehrgpt_tokenizer = CehrGptTokenizer.expand_trained_tokenizer(
-                            cehrgpt_tokenizer=cehrgpt_tokenizer,
-                            dataset=final_splits["train"],
-                            data_args=data_args,
-                            concept_name_mapping={},
-                        )
-                        cehrgpt_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
+                # 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:
@@ -238,10 +219,6 @@ def main():
             len(processed_dataset["test"]),
         )
 
-    LOG.info(f"cehrgpt_model.config.vocab_size: {cehrgpt_model.config.vocab_size}")
-    LOG.info(f"cehrgpt_tokenizer.vocab_size: {cehrgpt_tokenizer.vocab_size}")
-    if cehrgpt_model.config.vocab_size < cehrgpt_tokenizer.vocab_size:
-        cehrgpt_model.resize_token_embeddings(cehrgpt_tokenizer.vocab_size)
     if (
         cehrgpt_model.config.max_position_embeddings
         < model_args.max_position_embeddings
@@ -339,10 +316,12 @@ def main():
                 for data_dir in [data_args.data_folder, data_args.test_data_folder]
             ]
         )
-        # This is a pre-caution in case the index_date is not a datetime type
-        demographics_df["index_date"] = pd.to_datetime(
-            demographics_df["index_date"]
-        ).dt.date
+
+        demographics_df["index_date"] = (
+            demographics_df["index_date"].dt.tz_localize("UTC")
+            - datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)
+        ).dt.total_seconds()
+
         demographics_dict = {
             (row["person_id"], row["index_date"]): {
                 "gender_concept_id": row["gender_concept_id"],
@@ -379,9 +358,16 @@ def main():
                 prediction_time_posix = batch.pop("index_date").numpy().squeeze()
                 if prediction_time_posix.ndim == 0:
                     prediction_time_posix = np.asarray([prediction_time_posix])
+
                 prediction_time = list(
-                    map(datetime.fromtimestamp, prediction_time_posix)
+                    map(
+                        lambda posix_time: datetime.datetime.utcfromtimestamp(
+                            posix_time
+                        ).replace(tzinfo=None),
+                        prediction_time_posix,
+                    )
                 )
+
                 labels = (
                     batch.pop("classifier_label")
                     .float()
@@ -393,6 +379,10 @@ def main():
                 if labels.ndim == 0:
                     labels = np.asarray([labels])
 
+                # Right now the model does not support this column, we need to pop it
+                if "epoch_times" in batch:
+                    batch.pop("epoch_times")
+
                 batch = {k: v.to(device) for k, v in batch.items()}
                 # Forward pass
                 cehrgpt_output = cehrgpt_model(

{cehrgpt-0.1.1.dist-info → cehrgpt-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cehrgpt
-Version: 0.1.1
+Version: 0.1.2
 Summary: CEHR-GPT: Generating Electronic Health Records with Chronological Patient Timelines
 Author-email: Chao Pang <chaopang229@gmail.com>, Xinzhuo Jiang <xj2193@cumc.columbia.edu>, Krishna Kalluri <kk3326@cumc.columbia.edu>, Elise Minto <em3697@cumc.columbia.edu>, Jason Patterson <jp3477@cumc.columbia.edu>, Nishanth Parameshwar Pavinkurve <np2689@cumc.columbia.edu>, Karthik Natarajan <kn2174@cumc.columbia.edu>
 License: MIT License
@@ -105,6 +105,100 @@ sh scripts/omop_pipeline.sh \
   $OMOP_VOCAB_DIR
 ```
 
+# MEDS Support
+
+This section demonstrates how to pretrain CEHR-GPT using MIMIC-IV data in the MEDS (Medical Event Data Standard) format.
+
+## Prerequisites
+
+Set up the required environment variables before beginning:
+
+```bash
+export CEHR_GPT_MODEL_DIR=""    # Path to CEHR-GPT model directory
+export MEDS_DIR=""              # Path to MEDS data directory
+export MEDS_READER_DIR=""       # Path to MEDS reader output directory
+```
+
+## Step 1: Create MIMIC MEDS Data
+
+Transform your MIMIC files into MEDS format by following the instructions in the [MEDS_transforms](https://github.com/mmcdermott/MEDS_transforms/) repository.
+
+## Step 2: Create the MEDS Reader
+
+Convert the MEDS data for use with CEHR-GPT:
+
+```bash
+meds_reader_convert $MEDS_DIR $MEDS_READER_DIR --num_threads 10
+```
+
+## Step 3: Pretrain CEHR-GPT
+
+Run the pretraining process using the prepared MEDS data:
+
+```bash
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $CEHR_GPT_MODEL_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --do_train true --seed 42 \
+  --dataloader_num_workers 16 --dataloader_prefetch_factor 8 \
+  --hidden_size 768 --num_hidden_layers 14 --max_position_embeddings 8192 \
+  --evaluation_strategy epoch --save_strategy epoch \
+  --sample_packing --max_tokens_per_batch 16384 \
+  --warmup_steps 500 --weight_decay 0.01 \
+  --num_train_epochs 50 --learning_rate 0.0002 \
+  --use_early_stopping --early_stopping_threshold 0.001 \
+  --is_data_in_meds --inpatient_att_function_type day \
+  --att_function_type day --include_inpatient_hour_token \
+  --include_auxiliary_token --include_demographic_prompt \
+  --meds_to_cehrbert_conversion_type "MedsToBertMimic4"
+```
+
+## Step 4: Generate MEDS Trajectories
+
+### Environment Setup for Trajectory Generation
+
+Configure additional environment variables for trajectory generation with task labels (`subject_id`, `prediction_time`, `boolean_value` [optional]):
+
+```bash
+# MEDS_LABEL_COHORT_DIR must contain a set of parquet files
+export MEDS_LABEL_COHORT_DIR=""     # Path to cohort labels directory
+export MEDS_TRAJECTORY_DIR=""       # Path for trajectory output
+```
+
+### Generate Trajectories
+
+Create synthetic patient trajectories using the trained model:
+
+> **Important:** The total sequence length (`generation_input_length` + `generation_max_new_tokens`) cannot exceed the `max_position_embeddings` value (8192) defined during pretraining.
+
+```bash
+python -u -m cehrgpt.generation.cehrgpt_conditional_generation \
+  --cohort_folder $MEDS_LABEL_COHORT_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $MEDS_TRAJECTORY_DIR \
+  --per_device_eval_batch_size 16 \
+  --num_of_trajectories_per_sample 2 \
+  --generation_input_length 4096 \
+  --generation_max_new_tokens 4096 \
+  --is_data_in_meds \
+  --att_function_type day --inpatient_att_function_type day \
+  --meds_to_cehrbert_conversion_type MedsToBertMimic4 \
+  --include_auxiliary_token --include_demographic_prompt \
+  --include_inpatient_hour_token
+```
+
+### Parameters Explanation
+
+- `generation_input_length`: Controls the length of input context for generation
+- `generation_max_new_tokens`: Maximum number of new tokens to generate
+- `num_of_trajectories_per_sample`: Number of trajectories to generate per patient sample
+
 ## Citation
 ```
 @article{cehrgpt2024,

{cehrgpt-0.1.1.dist-info → cehrgpt-0.1.2.dist-info}/RECORD

@@ -13,17 +13,18 @@ cehrgpt/analysis/privacy/reid_inference.py,sha256=Pypd3QJXQNY8VljpnIEa5zeAbTZHMj
 cehrgpt/analysis/privacy/utils.py,sha256=CRA4H9mPLBjMQGKzZ_x_3ro3tMap-NjsMDVqSOjHSVQ,8226
 cehrgpt/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cehrgpt/data/hf_cehrgpt_dataset.py,sha256=hwJlGW7XiJIr6cXtmwvReQf9yLZJPD-dvJGvRg5ERqU,3755
-cehrgpt/data/hf_cehrgpt_dataset_collator.py,sha256=ACMXiaYnR3bKD5dRleL0_siEvhL-2HAFcy5eBgvxnH4,44412
-cehrgpt/data/hf_cehrgpt_dataset_mapping.py,sha256=KU0WMjc2vT1zBAl7JJkOc8dgGxsL1uFDy4dDrv-RkII,25668
+cehrgpt/data/hf_cehrgpt_dataset_collator.py,sha256=juM5HeZScgj8w15Bl1qC83Swld4gY6avh0QkSWLqITA,45465
+cehrgpt/data/hf_cehrgpt_dataset_mapping.py,sha256=_QDX9NXfmQ_S3kOf3yndb3AhoEeFiSzAOv836uYW0AY,26230
 cehrgpt/data/sample_packing_sampler.py,sha256=vovGMtmhG70DRkSCeiaDEJ_rjKZ38y-YLaI1kkhFEkI,6747
 cehrgpt/generation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cehrgpt/generation/cehrgpt_conditional_generation.py,sha256=AM76yaPyw1B-bcdei24HO0uspGZWHGKWpYpHywotTIQ,11972
 cehrgpt/generation/chatgpt_generation.py,sha256=SrnLwHLdNtnAOEg36gNjqfoT9yd12iyPgpZffL2AFJo,4428
-cehrgpt/generation/generate_batch_hf_gpt_sequence.py,sha256=uSEh8aMmPD61nGewIaPSkIqm-2AxDjCBiu4cBfxHxU4,11503
+cehrgpt/generation/generate_batch_hf_gpt_sequence.py,sha256=P8al4-zqymqEkCHCCu2sqz_45akcKF2o_AtQIjJdVmQ,11919
 cehrgpt/generation/omop_converter_batch.py,sha256=LUmCD-t_6ZP1YfNDZCqYewl-XIIaIgRZ_dAxuR_VdCQ,26275
 cehrgpt/generation/omop_entity.py,sha256=Q5Sr0AlyuPAm1FRPfnJO13q-u1fqRgYVHXruZ9g4xNE,19400
 cehrgpt/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cehrgpt/models/config.py,sha256=nOAKgH5420HLCcy7n1hE7MbqR861Iq4DTutKoAd25tg,11090
-cehrgpt/models/hf_cehrgpt.py,sha256=77CAkdMPgxD4xSpFU7gYGzRn6_Iv-4q7FnHpnZGsKxw,92450
+cehrgpt/models/hf_cehrgpt.py,sha256=3P7bOLDr7NMSedGszhmlJJN4Mhpd_65-x6uzwvSjigE,92837
 cehrgpt/models/hf_modeling_outputs.py,sha256=5X4WEYKqT37phv_e5ZAv3A_N0wqdAUJLJRm6TxS6dDQ,10356
 cehrgpt/models/pretrained_embeddings.py,sha256=vLLVs17TLpXRqCVEWQxGGwPHkUJUO7laNTeBuyBK_yk,3238
 cehrgpt/models/special_tokens.py,sha256=lrw45B4tea4Dsajn09Cz6w5D2TfHmYXikZkgwnstu_o,521
@@ -38,11 +39,11 @@ cehrgpt/omop/queries/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 cehrgpt/omop/queries/condition_era.py,sha256=LFB6vBAvshHJxtYIRkl7cfrF0kf7ay0piBKpmHBwrpE,2578
 cehrgpt/omop/queries/observation_period.py,sha256=fpzr5DMNw-QLoSwp2Iatfch88E3hyhZ75usiIdG3A0U,6410
 cehrgpt/runners/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cehrgpt/runners/data_utils.py,sha256=I6k1TkiiZR8ggw3eVO16g2lVPY-Hu3b-nbrIOKlFIO0,15528
+cehrgpt/runners/data_utils.py,sha256=i-krtBx_6rvPYtdLdDoWwOTtJcaovd0wH8gBYmgN2l4,16013
 cehrgpt/runners/gpt_runner_util.py,sha256=YJQSRW9Mo4TjXSOUOTf6BUFcs1MGFiXU5T4ztKZcYhU,3485
-cehrgpt/runners/hf_cehrgpt_finetune_runner.py,sha256=GVbHHqf5TWGbVWlQG-XurgYH8pKRjTk8ug_ib9L9U7E,28118
+cehrgpt/runners/hf_cehrgpt_finetune_runner.py,sha256=1OgxLm4T7iHv5pKi2QaSdaz9ogWo2n3sSUGp6cHDF9s,28309
 cehrgpt/runners/hf_cehrgpt_pretrain_runner.py,sha256=ERSnvB38fPYVghtKQeNTZ8VfeXnoRcCHB0cWISWaZ84,26523
-cehrgpt/runners/hf_gpt_runner_argument_dataclass.py,sha256=ejAFLM9g765p1fyeF5MITsiIeWHKkz9wTeFDeVgxSto,8851
+cehrgpt/runners/hf_gpt_runner_argument_dataclass.py,sha256=fJR4RHPqal1YI6_KUH-WlkoQLSZuBT5bKUGfPHDFrWI,9350
 cehrgpt/runners/hyperparameter_search_util.py,sha256=YWdFQ1igQs-G_wqWUrUzYraGiz8OSpSYyvid-I5nhWA,9262
 cehrgpt/runners/sample_packing_trainer.py,sha256=Zb7Aqwnk8-VqrjEKUVeg5XzZWmHxXOU2sDn1YURS-FU,7960
 cehrgpt/simulations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -65,10 +66,10 @@ cehrgpt/tools/generate_pretrained_embeddings.py,sha256=lhFSacGv8bMld6qigKZN8Op8e
 cehrgpt/tools/merge_synthetic_real_dataasets.py,sha256=O1dbQ32Le0t15fwymwAh9mfNVLEWuFwW53DNvESrWbY,7589
 cehrgpt/tools/upload_omop_tables.py,sha256=vdBAbkeAsGPA4NsyhNjelPVj3gS8yzmS1sKNM1Qk96g,3791
 cehrgpt/tools/linear_prob/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cehrgpt/tools/linear_prob/compute_cehrgpt_features.py,sha256=q0rmlBWDDEkjHjwcTouGUhCYa32a1vRicaDOAMsdW0I,20741
+cehrgpt/tools/linear_prob/compute_cehrgpt_features.py,sha256=Hpx7WvAWm2WwPHFfimCADXh019I7bwdzJ4_5_YCxQzU,19817
 cehrgpt/tools/linear_prob/train_with_cehrgpt_features.py,sha256=w0UvzMKYGenN_KDVnbzutmy8IPLUxW5hPvpKKxDSL5U,5820
-cehrgpt-0.1.1.dist-info/licenses/LICENSE,sha256=LOfC32zkfUIdGm8e_098jPbt8OHKtNWymDzxn2pA9Zk,1093
-cehrgpt-0.1.1.dist-info/METADATA,sha256=VnXH74vJQZaV7VxGiIvJnFhQA0jzJQNx86yHFkygobM,4922
-cehrgpt-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cehrgpt-0.1.1.dist-info/top_level.txt,sha256=akNCJBbMSLV8nkOzdVzdy13hMJ5CIQURnAS_YYEDVwA,17
-cehrgpt-0.1.1.dist-info/RECORD,,
+cehrgpt-0.1.2.dist-info/licenses/LICENSE,sha256=LOfC32zkfUIdGm8e_098jPbt8OHKtNWymDzxn2pA9Zk,1093
+cehrgpt-0.1.2.dist-info/METADATA,sha256=D7gGKrQThiLivViFeNm711NCP8J-wXfkueMGb6RKqV0,8481
+cehrgpt-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cehrgpt-0.1.2.dist-info/top_level.txt,sha256=akNCJBbMSLV8nkOzdVzdy13hMJ5CIQURnAS_YYEDVwA,17
+cehrgpt-0.1.2.dist-info/RECORD,,