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

cehrgpt/runners/sample_packing_trainer.py

@@ -0,0 +1,168 @@
+from typing import Optional, Union
+
+from datasets import Dataset
+from torch.utils.data import DataLoader
+from transformers import Trainer
+from transformers.trainer_utils import has_length
+from transformers.utils import import_utils, logging
+
+from cehrgpt.data.sample_packing_sampler import SamplePackingBatchSampler
+
+DEFAULT_MAX_TOKENS_PER_BATCH = 16384
+
+LOG = logging.get_logger("transformers")
+
+
+class SamplePackingTrainer(Trainer):
+    def __init__(self, *args, **kwargs):
+        if "max_tokens_per_batch" in kwargs:
+            self.max_tokens_per_batch = kwargs.pop("max_tokens_per_batch")
+            LOG.info("max_tokens_per_batch: %s", self.max_tokens_per_batch)
+        else:
+            self.max_tokens_per_batch = DEFAULT_MAX_TOKENS_PER_BATCH
+            LOG.info(
+                "max_tokens_per_batch is not provided to SamplePackingTrainer and will default to %s",
+                DEFAULT_MAX_TOKENS_PER_BATCH,
+            )
+
+        if "max_position_embeddings" in kwargs:
+            self.max_position_embeddings = kwargs.pop("max_position_embeddings")
+            LOG.info("max_position_embeddings: %s", self.max_position_embeddings)
+        else:
+            self.max_position_embeddings = self.max_tokens_per_batch
+            LOG.info(
+                "max_position_embeddings is not provided to SamplePackingTrainer and will default to %s",
+                self.max_tokens_per_batch,
+            )
+
+        self.train_lengths = kwargs.pop("train_lengths", None)
+        self.validation_lengths = kwargs.pop("validation_lengths", None)
+        super().__init__(*args, **kwargs)
+        self.accelerator.even_batches = False
+
+    def num_examples(self, dataloader: DataLoader) -> int:
+        if has_length(dataloader):
+            return len(dataloader)
+        raise RuntimeError("DataLoader in SamplePackingTrainer must have length")
+
+    def get_train_dataloader(self) -> DataLoader:
+        """Returns the training dataloader with our custom batch sampler."""
+        train_dataset = self.train_dataset
+
+        if self.train_lengths is None:
+            LOG.info("Started computing lengths for the train dataset")
+            # Calculate lengths of all sequences in dataset
+            if "num_of_concepts" in train_dataset.column_names:
+                lengths = train_dataset["num_of_concepts"]
+            else:
+                lengths = [len(sample["input_ids"]) for sample in train_dataset]
+
+            LOG.info("Finished computing lengths for the train dataset")
+        else:
+            lengths = self.train_lengths
+
+        data_collator = self.data_collator
+        if import_utils.is_datasets_available() and isinstance(train_dataset, Dataset):
+            train_dataset = self._remove_unused_columns(
+                train_dataset, description="training"
+            )
+        else:
+            data_collator = self._get_collator_with_removed_columns(
+                data_collator, description="training"
+            )
+        # Create our custom batch sampler
+        batch_sampler = SamplePackingBatchSampler(
+            lengths=lengths,
+            max_tokens_per_batch=self.max_tokens_per_batch,
+            max_position_embeddings=self.max_position_embeddings,
+            drop_last=self.args.dataloader_drop_last,
+            seed=self.args.seed,
+        )
+        dataloader_params = {
+            "collate_fn": data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+            "persistent_workers": self.args.dataloader_persistent_workers,
+            "batch_sampler": batch_sampler,
+        }
+        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))
+
+    def get_eval_dataloader(
+        self, eval_dataset: Optional[Union[str, Dataset]] = None
+    ) -> DataLoader:
+        """
+        Returns the evaluation [`~torch.utils.data.DataLoader`].
+
+        Subclass and override this method if you want to inject some custom behavior.
+
+        Args:
+            eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
+                If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
+        """
+        if eval_dataset is None and self.eval_dataset is None:
+            raise ValueError("Trainer: evaluation requires an eval_dataset.")
+
+        # If we have persistent workers, don't do a fork bomb especially as eval datasets
+        # don't change during training
+        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
+        if (
+            hasattr(self, "_eval_dataloaders")
+            and dataloader_key in self._eval_dataloaders
+            and self.args.dataloader_persistent_workers
+        ):
+            return self.accelerator.prepare(self._eval_dataloaders[dataloader_key])
+
+        eval_dataset = (
+            self.eval_dataset[eval_dataset]
+            if isinstance(eval_dataset, str)
+            else eval_dataset if eval_dataset is not None else self.eval_dataset
+        )
+
+        if self.validation_lengths is None:
+            LOG.info("Started computing lengths for the train dataset")
+            # Calculate lengths of all sequences in dataset
+            if "num_of_concepts" in eval_dataset.column_names:
+                lengths = eval_dataset["num_of_concepts"]
+            else:
+                lengths = [len(sample["input_ids"]) for sample in eval_dataset]
+
+            LOG.info("Finished computing lengths for the train dataset")
+        else:
+            lengths = self.validation_lengths
+
+        data_collator = self.data_collator
+
+        if import_utils.is_datasets_available() and isinstance(eval_dataset, Dataset):
+            eval_dataset = self._remove_unused_columns(
+                eval_dataset, description="evaluation"
+            )
+        else:
+            data_collator = self._get_collator_with_removed_columns(
+                data_collator, description="evaluation"
+            )
+
+        # Create our custom batch sampler
+        batch_sampler = SamplePackingBatchSampler(
+            lengths=lengths,
+            max_tokens_per_batch=self.max_tokens_per_batch,
+            max_position_embeddings=self.max_position_embeddings,
+            drop_last=self.args.dataloader_drop_last,
+            seed=self.args.seed,
+        )
+        dataloader_params = {
+            "collate_fn": data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+            "persistent_workers": self.args.dataloader_persistent_workers,
+            "batch_sampler": batch_sampler,
+        }
+        # accelerator.free_memory() will destroy the references, so
+        # we need to store the non-prepared version
+        eval_dataloader = DataLoader(eval_dataset, **dataloader_params)
+        if self.args.dataloader_persistent_workers:
+            if hasattr(self, "_eval_dataloaders"):
+                self._eval_dataloaders[dataloader_key] = eval_dataloader
+            else:
+                self._eval_dataloaders = {dataloader_key: eval_dataloader}
+
+        return self.accelerator.prepare(eval_dataloader)

cehrgpt/simulations/generate_plots.py

@@ -0,0 +1,95 @@
+import json
+import os
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def main(output_dir: str):
+    with open(os.path.join(output_dir, "time_embedding_metrics.json"), "r") as f:
+        time_embedding_metrics = json.load(f)
+    with open(os.path.join(output_dir, "time_token_metrics.json"), "r") as f:
+        time_token_metrics = json.load(f)
+
+    common_steps = list(
+        set(time_embedding_metrics["steps"]) & set(time_token_metrics["steps"])
+    )
+
+    time_embedding_aucs = []
+    time_embedding_accuracies = []
+    for step, roc_auc, accuracy in zip(
+        time_embedding_metrics["steps"],
+        time_embedding_metrics["roc_auc"],
+        time_embedding_metrics["accuracy"],
+    ):
+        if step in common_steps:
+            time_embedding_aucs.append(roc_auc)
+            time_embedding_accuracies.append(accuracy)
+
+    time_token_aucs = []
+    time_token_accuracies = []
+    for step, roc_auc, accuracy in zip(
+        time_token_metrics["steps"],
+        time_token_metrics["roc_auc"],
+        time_token_metrics["accuracy"],
+    ):
+        if step in common_steps:
+            time_token_aucs.append(roc_auc)
+            time_token_accuracies.append(accuracy)
+
+    # Create the accuracy plot
+    plt.figure(figsize=(8, 5))  # Define figure size
+    plt.plot(
+        common_steps,
+        time_embedding_accuracies,
+        linestyle="-",
+        color="b",
+        label="Time Embedding",
+        lw=1,
+    )
+    plt.plot(
+        common_steps,
+        time_token_accuracies,
+        linestyle="--",
+        color="r",
+        label="Time Token",
+        lw=1,
+    )
+    plt.title("Accuracy Comparison Over Time")
+    plt.xlabel("Training Steps")
+    plt.ylabel("Accuracy")
+    plt.legend()
+    plt.grid(False)
+    plt.savefig(os.path.join(output_dir, "accuracy_comparison.png"))
+
+    # Create the ROC AUC plot
+    plt.figure(figsize=(8, 5))  # Define figure size
+    plt.plot(
+        common_steps,
+        time_embedding_aucs,
+        linestyle="-",
+        color="b",
+        label="Time Embedding",
+        lw=1,
+    )
+    plt.plot(
+        common_steps,
+        time_token_aucs,
+        linestyle="--",
+        color="r",
+        label="Time Token",
+        lw=1,
+    )
+    plt.title("ROC AUC Comparison Over Time")
+    plt.xlabel("Training Steps")
+    plt.ylabel("ROC AUC")
+    plt.legend()
+    plt.grid(False)
+    plt.savefig(
+        os.path.join(output_dir, "roc_auc_comparison.png")
+    )  # Save the plot as a PNG file
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])

cehrgpt/simulations/run_simulation.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+# This script runs various Python simulations and generates plots
+# It accepts three parameters: output directory, number of steps, and number of samples
+
+# Check if all arguments are provided
+if [ "$#" -ne 3 ]; then
+    echo "Usage: $0 <output_dir> <n_steps> <n_samples>"
+    exit 1
+fi
+
+# Assigning command line arguments to variables
+OUTPUT_DIR="$1"
+N_STEPS="$2"
+N_SAMPLES="$3"
+
+# Run time token simulation
+python -u -m cehrgpt.simulations.time_token_simulation --output_dir "$OUTPUT_DIR" --n_steps "$N_STEPS" --n_samples "$N_SAMPLES"
+
+# Run time embedding simulation
+python -u -m cehrgpt.simulations.time_embedding_simulation --output_dir "$OUTPUT_DIR" --n_steps "$N_STEPS" --n_samples "$N_SAMPLES"
+
+# Generate plots
+python -u -m cehrgpt.simulations.generate_plots "$OUTPUT_DIR"

cehrgpt/simulations/time_embedding_simulation.py

@@ -0,0 +1,250 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.optim as optim
+from sklearn.metrics import accuracy_score, roc_auc_score
+from torch.nn import CrossEntropyLoss
+from transformers import BertConfig, BertModel
+
+
+class ModelTimeEmbedding(torch.nn.Module):
+    def __init__(self, vocab_size: int):
+        super(ModelTimeEmbedding, self).__init__()
+        self.embedding = torch.nn.Embedding(vocab_size, 16)
+        self.bert = BertModel(
+            BertConfig(
+                vocab_size=vocab_size,
+                hidden_size=16,
+                num_attention_heads=2,
+                num_hidden_layers=2,
+                intermediate_size=32,
+                hidden_dropout_prob=0.0,
+                attention_probs_dropout_prob=0.0,
+                max_position_embeddings=2,
+            ),
+            add_pooling_layer=False,
+        )
+        self.linear = torch.nn.Linear(32, 2)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        time_stamps: torch.LongTensor,
+        labels: Optional[torch.LongTensor] = None,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        bz = input_ids.shape[0]
+        x = self.embedding(input_ids)
+        t = self.embedding(time_stamps)
+        x = x + t
+        bert_output = self.bert.forward(inputs_embeds=x, return_dict=True)
+        output = bert_output.last_hidden_state.reshape((bz, 32))
+        y = self.linear(output)
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(y, labels)
+        return loss, y
+
+
+def generate_simulation_data(sample_size: int = 1000, seed: int = 42) -> np.ndarray:
+    np.random.seed(seed)  # Set the seed for reproducibility
+
+    # Define input values and time stamps
+    x_values = [0, 1]
+    time_stamp_values = list(range(0, 21))
+
+    # Generate random choices for features and time stamps
+    x1 = np.random.choice(x_values, size=sample_size)
+    x2 = np.random.choice(x_values, size=sample_size)
+    t1 = np.random.choice(time_stamp_values, size=sample_size)
+    t2 = t1 + np.random.choice(time_stamp_values, size=sample_size)
+
+    # Define conditions based on time differences
+    time_diff = t2 - t1
+    # Complex condition involving modulo operation
+    is_custom_func_1 = (x1 == 1) & (time_diff % 4 == 0)
+    is_custom_func_2 = (x1 == 0) & (time_diff % 3 == 0)
+    is_xor = time_diff <= 7
+    is_and = (time_diff > 7) & (time_diff <= 14)
+    is_or = (time_diff > 14) & (time_diff <= 21)
+
+    # Logical operations based on x1 and x2
+    xor = (x2 != x1).astype(int)
+    logical_and = (x2 & x1).astype(int)
+    logical_or = (x2 | x1).astype(int)
+    # Additional complexity: introduce a new rule based on a more complex condition
+    custom_func_1_result = (x2 == 0).astype(int)  # For example, use a different rule
+    custom_func_2_result = (x2 == 1).astype(int)  # For example, use a different rule
+
+    # Determine output based on multiple conditions
+    y = np.where(
+        is_custom_func_1,
+        custom_func_1_result,
+        np.where(
+            is_custom_func_2,
+            custom_func_2_result,
+            np.where(
+                is_xor,
+                xor,
+                np.where(is_and, logical_and, np.where(is_or, logical_or, 0)),
+            ),
+        ),
+    )
+
+    # Return the data as a single numpy array with features and output
+    return np.column_stack((x1, x2, t1, t2, y))
+
+
+def create_time_embedding_tokenizer(simulated_data):
+    vocab = []
+    for row in simulated_data:
+        x1, x2, t1, t2, y = row
+        x1 = f"c-{x1}"
+        x2 = f"c-{x2}"
+        t1 = f"t-{t1}"
+        t2 = f"t-{t2}"
+        if x1 not in vocab:
+            vocab.append(x1)
+        if x2 not in vocab:
+            vocab.append(x2)
+        if t1 not in vocab:
+            vocab.append(t1)
+        if t2 not in vocab:
+            vocab.append(t2)
+    return {c: i + 1 for i, c in enumerate(vocab)}
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def eval_step(
+    simulated_data,
+    time_embedding_tokenizer,
+    time_embedding_model,
+    time_embedding_optimizer,
+):
+    time_embedding_optimizer.zero_grad()
+    time_embedding_model.eval()
+    eval_input_ids = []
+    eval_time_stamps = []
+    eval_y = []
+    for row in simulated_data:
+        x1, x2, t1, t2, y = row
+        x1 = f"c-{x1}"
+        x2 = f"c-{x2}"
+        t1 = f"t-{t1}"
+        t2 = f"t-{t2}"
+        eval_input_ids.append(
+            [time_embedding_tokenizer[x1], time_embedding_tokenizer[x2]]
+        )
+        eval_time_stamps.append(
+            [time_embedding_tokenizer[t1], time_embedding_tokenizer[t2]]
+        )
+        eval_y.append(y)
+    eval_input_ids = torch.tensor(eval_input_ids, dtype=torch.long).to(device)
+    eval_time_stamps = torch.tensor(eval_time_stamps, dtype=torch.long).to(device)
+    eval_y = np.asarray(eval_y)
+    with torch.no_grad():
+        # Compute loss and forward pass
+        _, y_pred = time_embedding_model(eval_input_ids, eval_time_stamps)
+        y_probs = torch.nn.functional.softmax(y_pred, dim=1)
+        y_probs = y_probs.detach().cpu().numpy()
+        # print(np.concatenate((y_probs, batched_y[:, None]), axis=1))
+        roc_auc = roc_auc_score(eval_y, y_probs[:, 1])
+        accuracy = accuracy_score(eval_y, y_probs[:, 1] > y_probs[:, 0])
+        print(f"ROC AUC: {roc_auc}")
+        print(f"Accuracy: {accuracy}")
+    return roc_auc, accuracy
+
+
+def train_step(
+    simulated_data,
+    time_embedding_tokenizer,
+    time_embedding_model,
+    time_embedding_optimizer,
+):
+    batched_input_ids = []
+    batched_time_stamps = []
+    batched_y = []
+    indices = np.random.choice(simulated_data.shape[0], size=8, replace=False)
+    for row in simulated_data[indices, :]:
+        x1, x2, t1, t2, y = row
+        x1 = f"c-{x1}"
+        x2 = f"c-{x2}"
+        t1 = f"t-{t1}"
+        t2 = f"t-{t2}"
+        batched_input_ids.append(
+            [time_embedding_tokenizer[x1], time_embedding_tokenizer[x2]]
+        )
+        batched_time_stamps.append(
+            [time_embedding_tokenizer[t1], time_embedding_tokenizer[t2]]
+        )
+        batched_y.append(y)
+    batched_input_ids = torch.tensor(batched_input_ids, dtype=torch.long).to(device)
+    batched_time_stamps = torch.tensor(batched_time_stamps, dtype=torch.long).to(device)
+    batched_y = torch.tensor(batched_y, dtype=torch.long).to(device)
+    # Zero the gradients
+    time_embedding_optimizer.zero_grad()
+    # Compute loss and forward pass
+    loss, _ = time_embedding_model(batched_input_ids, batched_time_stamps, batched_y)
+    # Backward pass (compute gradients)
+    loss.backward()
+    # Update model parameters
+    time_embedding_optimizer.step()
+    return loss
+
+
+def main(args):
+    simulated_data = generate_simulation_data(args.n_samples)
+    time_embedding_tokenizer = create_time_embedding_tokenizer(simulated_data)
+    time_embedding_model = ModelTimeEmbedding(len(time_embedding_tokenizer) + 1).to(
+        device
+    )
+    time_embedding_optimizer = optim.Adam(time_embedding_model.parameters(), lr=0.001)
+    steps = []
+    roc_aucs = []
+    accuracies = []
+    for step in range(args.n_steps):
+        loss = train_step(
+            simulated_data,
+            time_embedding_tokenizer,
+            time_embedding_model,
+            time_embedding_optimizer,
+        )
+        print(f"Step {step}: Loss = {loss.item()}")
+        # Evaluation
+        if (
+            args.n_steps % args.eval_frequency == 0
+            and args.n_steps > args.eval_frequency
+        ):
+            # Zero the gradients
+            roc_auc, accuracy = eval_step(
+                simulated_data,
+                time_embedding_tokenizer,
+                time_embedding_model,
+                time_embedding_optimizer,
+            )
+            steps.append(step)
+            roc_aucs.append(roc_auc)
+            accuracies.append(accuracy)
+    return {"steps": steps, "roc_auc": roc_aucs, "accuracy": accuracies}
+
+
+if __name__ == "__main__":
+    import argparse
+    import json
+    from pathlib import Path
+
+    parser = argparse.ArgumentParser("Model with time embedding simulation")
+    parser.add_argument("--output_dir", type=str, required=True)
+    parser.add_argument("--n_steps", type=int, default=10000)
+    parser.add_argument("--n_samples", type=int, default=1000)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--eval_frequency", type=int, default=100)
+    args = parser.parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(exist_ok=True, parents=True)
+    metrics = main(args)
+    with open(output_dir / "time_embedding_metrics.json", "w") as f:
+        json.dump(metrics, f)