diff --git a/bonsai/functional/censoring.py b/bonsai/functional/censoring.py index 9cfd30dcb..a8bfcdc25 100644 --- a/bonsai/functional/censoring.py +++ b/bonsai/functional/censoring.py @@ -18,6 +18,8 @@ def censor_subject( # Slice everything up to idx for embed_name in ["code", "abspos", "segment", "age"]: subject[embed_name] = subject[embed_name][:idx] + if "numeric_value" in subject: + subject["numeric_value"] = subject["numeric_value"][:idx] if predict_token_id is not None: subject = append_predict_token(subject, censor_date_abspos, predict_token_id) @@ -49,6 +51,13 @@ def append_predict_token( ), ) ) + if "numeric_value" in subject: + subject["numeric_value"] = torch.cat( + ( + subject["numeric_value"], + torch.tensor([float("nan")], dtype=subject["numeric_value"].dtype), + ) + ) age_in_years = float((censor_date_abspos - subject["abspos"][0]) / (365.25 * 24)) subject["age"] = torch.cat( diff --git a/bonsai/functional/collate.py b/bonsai/functional/collate.py index c1eeda163..846526c90 100644 --- a/bonsai/functional/collate.py +++ b/bonsai/functional/collate.py @@ -16,6 +16,14 @@ def dynamic_padding(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Ten output["target"] = torch.nn.utils.rnn.pad_sequence( collected["target"], batch_first=True, padding_value=-100 ) + if "numeric_value" in collected: + output["numeric_value"] = torch.nn.utils.rnn.pad_sequence( + collected["numeric_value"], batch_first=True, padding_value=float("nan") + ) + if "numeric_target" in collected: + output["numeric_target"] = torch.nn.utils.rnn.pad_sequence( + collected["numeric_target"], batch_first=True, padding_value=float("nan") + ) output["subject_id"] = torch.tensor(collected["subject_id"]) return output diff --git a/bonsai/functional/create_data.py b/bonsai/functional/create_data.py index e2f85b194..08ae442a2 100644 --- a/bonsai/functional/create_data.py +++ b/bonsai/functional/create_data.py @@ -5,13 +5,16 @@ from bonsai.functional.features import create_features -def drop_duplicates(df: pl.DataFrame) -> pl.DataFrame: +def drop_duplicates( + df: pl.DataFrame, numeric_column: Optional[str] = None +) -> pl.DataFrame: pre = len(df) - df = df.unique(subset=["subject_id", "code", "time"], maintain_order=True) + dup_cols = ["subject_id", "code", "time"] + if numeric_column is not None: + dup_cols.append(numeric_column) + df = df.unique(subset=dup_cols, maintain_order=True) if pre != len(df): - logging.info( - f"Dropped {pre - len(df)} duplicate rows based on subject_id, code, and time" - ) + logging.info(f"Dropped {pre - len(df)} duplicate rows based on {dup_cols}") return df @@ -21,6 +24,7 @@ def process_split( path_output_dir: Path, tokenizer, exclude_regex: Optional[str] = None, + numeric_column: Optional[str] = None, ): path_output_dir_split = path_output_dir / split path_output_dir_split.mkdir(parents=True, exist_ok=True) @@ -38,11 +42,15 @@ def process_split( logging.info(f"Processing shard {shard_idx}/{len(shards)}: {shard}") # Load - shard_df = pl.read_parquet(shard) + load_cols = ["subject_id", "code", "time"] + if numeric_column is not None: + load_cols.append(numeric_column) + shard_df = pl.read_parquet(shard, columns=load_cols) + data_counts["loaded"] += len(shard_df) # Drop duplicates - shard_df = drop_duplicates(shard_df) + shard_df = drop_duplicates(shard_df, numeric_column=numeric_column) data_counts["after_duplicates"] += len(shard_df) # Optional: Exclude codes based on regex @@ -51,20 +59,23 @@ def process_split( data_counts["after_exclusion"] += len(shard_df) # Create features - features = create_features(shard_df) + features = create_features(shard_df, numeric_column=numeric_column) data_counts["after_features"] += len(features) # Tokenize tokenized = tokenizer(features) # Cast to correct dtypes - tokenized = tokenized.select( + cols = [ pl.col("subject_id").cast(pl.Int64), pl.col("code").cast(pl.Int64), pl.col("age").cast(pl.Float32), pl.col("abspos").cast(pl.Float32), pl.col("segment").cast(pl.Int32), - ) + ] + if numeric_column is not None: + cols.append(pl.col("numeric_value").cast(pl.Float32)) + tokenized = tokenized.select(cols) tokenized.write_parquet(path_output_dir_split / f"{shard.stem}.parquet") ids.extend(tokenized["subject_id"].unique()) diff --git a/bonsai/functional/features.py b/bonsai/functional/features.py index d48bf9aa2..4ef420b04 100644 --- a/bonsai/functional/features.py +++ b/bonsai/functional/features.py @@ -1,10 +1,12 @@ import polars as pl import logging from datetime import datetime -from typing import Tuple, Union +from typing import Tuple, Union, Optional -def create_features(df: pl.DataFrame) -> pl.DataFrame: +def create_features( + df: pl.DataFrame, numeric_column: Optional[str] = None +) -> pl.DataFrame: """ Create background, age, absolute position, and segment features. TODO: Death? @@ -39,7 +41,10 @@ def create_features(df: pl.DataFrame) -> pl.DataFrame: ) ) - features = features.select("subject_id", "code", "age", "abspos", "segment") + select_cols = ["subject_id", "code", "age", "abspos", "segment"] + if numeric_column is not None: + select_cols.append(pl.col(numeric_column).alias("numeric_value")) + features = features.select(select_cols) return features diff --git a/bonsai/functional/subject_data.py b/bonsai/functional/subject_data.py index 38e3f0867..5578c0130 100644 --- a/bonsai/functional/subject_data.py +++ b/bonsai/functional/subject_data.py @@ -10,18 +10,20 @@ def prepare_subject_data(split_path: Path) -> List[Dict[str, torch.Tensor]]: all_tokenized = [] for shard in split_path.glob("*.parquet"): tokenized_data = pl.read_parquet(shard) + has_numeric = "numeric_value" in tokenized_data.columns # Convert to training format for subject_id, group in tokenized_data.group_by("subject_id"): - all_tokenized.append( - { - "subject_id": subject_id[0], - "code": group["code"].to_torch(), - "abspos": group["abspos"].to_torch(), - "segment": group["segment"].to_torch(), - "age": group["age"].to_torch(), - } - ) + subject = { + "subject_id": subject_id[0], + "code": group["code"].to_torch(), + "abspos": group["abspos"].to_torch(), + "segment": group["segment"].to_torch(), + "age": group["age"].to_torch(), + } + if has_numeric: + subject["numeric_value"] = group["numeric_value"].to_torch().float() + all_tokenized.append(subject) return all_tokenized diff --git a/bonsai/functional/truncation.py b/bonsai/functional/truncation.py index 3d2c69dba..67e33b910 100644 --- a/bonsai/functional/truncation.py +++ b/bonsai/functional/truncation.py @@ -15,4 +15,6 @@ def truncate_subject( for embed_name in ["code", "abspos", "segment", "age"]: subject[embed_name] = subject[embed_name][idxs] + if "numeric_value" in subject: + subject["numeric_value"] = subject["numeric_value"][idxs] return subject diff --git a/bonsai/modules/datamodules/PretrainDataModule.py b/bonsai/modules/datamodules/PretrainDataModule.py index 9617f904b..90c6e4adc 100644 --- a/bonsai/modules/datamodules/PretrainDataModule.py +++ b/bonsai/modules/datamodules/PretrainDataModule.py @@ -6,8 +6,8 @@ from bonsai.functional.collate import dynamic_padding from bonsai.functional.subject_data import filter_subject_data from bonsai.modules.datasets.PretrainDataset import ( - MLMPretrainDataset, ARPretrainDataset, + MLMPretrainDataset, ) diff --git a/bonsai/modules/datasets/PretrainDataset.py b/bonsai/modules/datasets/PretrainDataset.py index 5c8a37ffc..eff8610ee 100644 --- a/bonsai/modules/datasets/PretrainDataset.py +++ b/bonsai/modules/datasets/PretrainDataset.py @@ -71,7 +71,7 @@ def __init__( def __getitem__(self, index: int) -> dict: subject = super().__getitem__(index) masked_codes, target = self.mask_patient_codes(subject["code"]) - subject["concept"] = masked_codes + subject["code"] = masked_codes subject["target"] = target return subject @@ -81,22 +81,19 @@ def mask_patient_codes( target = codes.clone() probability_vector = torch.full(target.shape, self.masking_select_ratio) - # Ignore special tokens special_token_mask = codes < self.masking_n_special_tokens probability_vector.masked_fill_(special_token_mask, value=0.0) - # Get MLM mask selected_indices = torch.bernoulli(probability_vector).bool() target[~selected_indices] = -100 - # Replace with [MASK] indices_mask = ( torch.bernoulli(torch.full(target.shape, self.masking_mask_ratio)).bool() & selected_indices ) + codes = codes.clone() codes[indices_mask] = self.vocabulary["[MASK]"] - # Replace with random word and Account for already masked tokens random_ratio = self.masking_random_ratio / (1 - self.masking_mask_ratio) indicies_random = ( torch.bernoulli(torch.full(target.shape, random_ratio)).bool() @@ -113,6 +110,61 @@ def mask_patient_codes( return codes, target +class ValueMLMPretrainDataset(MLMPretrainDataset): + def __getitem__(self, index: int) -> dict: + subject = super(MLMPretrainDataset, self).__getitem__(index) + masked_codes, masked_numeric_values, target, numeric_target = ( + self.mask_patient_codes(subject["code"], subject["numeric_value"]) + ) + subject["code"] = masked_codes + subject["target"] = target + subject["numeric_value"] = masked_numeric_values + subject["numeric_target"] = numeric_target + return subject + + def mask_patient_codes( + self, + codes: torch.Tensor, + numeric_values: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + target = codes.clone() + probability_vector = torch.full(target.shape, self.masking_select_ratio) + + special_token_mask = codes < self.masking_n_special_tokens + probability_vector.masked_fill_(special_token_mask, value=0.0) + + selected_indices = torch.bernoulli(probability_vector).bool() + target[~selected_indices] = -100 + + indices_mask = ( + torch.bernoulli(torch.full(target.shape, self.masking_mask_ratio)).bool() + & selected_indices + ) + codes = codes.clone() + codes[indices_mask] = self.vocabulary["[MASK]"] + + numeric_values = numeric_values.clone() + numeric_target = numeric_values.clone() + # Hide input values at [MASK]; keep targets only at those positions. + numeric_values[indices_mask] = float("nan") + numeric_target[~indices_mask] = float("nan") + + random_ratio = self.masking_random_ratio / (1 - self.masking_mask_ratio) + indicies_random = ( + torch.bernoulli(torch.full(target.shape, random_ratio)).bool() + & selected_indices + & ~indices_mask + ) + random_words = torch.randint( + self.masking_n_special_tokens, + len(self.vocabulary), + target.shape, + dtype=codes.dtype, + ) + codes[indicies_random] = random_words[indicies_random] + return codes, numeric_values, target, numeric_target + + class ARPretrainDataset(PretrainDataset): def __init__( self, @@ -129,6 +181,25 @@ def __getitem__(self, index: int) -> dict: subject = super().__getitem__(index) subject["target"] = subject["code"][1:] subject["target"] = subject["target"].masked_fill(subject["target"] == 0, -100) + + for key in ["code", "abspos", "segment", "age", "attention_mask"]: + subject[key] = subject[key][:-1] + return subject + + +class ValueARPretrainDataset(ARPretrainDataset): + def __getitem__(self, index: int) -> dict: + subject = super(ARPretrainDataset, self).__getitem__(index) + subject["target"] = subject["code"][1:] + subject["target"] = subject["target"].masked_fill(subject["target"] == 0, -100) + + values = subject["numeric_value"] + subject["numeric_target"] = values[1:].clone() + subject["numeric_value"] = values[:-1] + subject["numeric_target"] = subject["numeric_target"].masked_fill( + subject["target"] == -100, float("nan") + ) + for key in ["code", "abspos", "segment", "age", "attention_mask"]: subject[key] = subject[key][:-1] return subject diff --git a/bonsai/modules/lightningmodules/FinetuneModule.py b/bonsai/modules/lightningmodules/FinetuneModule.py index 89479c69f..2cbaffb23 100644 --- a/bonsai/modules/lightningmodules/FinetuneModule.py +++ b/bonsai/modules/lightningmodules/FinetuneModule.py @@ -1,6 +1,7 @@ from os.path import join from pathlib import Path from typing import Optional +import logging import lightning as L import polars as pl @@ -49,6 +50,18 @@ def __init__( ) self.save_hyperparameters(hparams) + def on_load_checkpoint(self, checkpoint: dict) -> None: + """Warn on value_embedding_mode changes vs the pretrain ckpt.""" + ckpt_mode = checkpoint.get("hyper_parameters", {}).get("value_embedding_mode") + ft_mode = self.model.hparams.get("value_embedding_mode") + if ckpt_mode != ft_mode: + logging.warning( + "value_embedding_mode changed between pretrain and finetune: " + "checkpoint=%r, finetune model=%r. ", + ckpt_mode, + ft_mode, + ) + def configure_metrics(self, prefix: str): return MetricCollection( { diff --git a/bonsai/modules/lightningmodules/PretrainModule.py b/bonsai/modules/lightningmodules/PretrainModule.py index 197f96d9f..33d6ae6ac 100644 --- a/bonsai/modules/lightningmodules/PretrainModule.py +++ b/bonsai/modules/lightningmodules/PretrainModule.py @@ -15,6 +15,28 @@ pass +class ConceptValueLoss(nn.Module): + """Concept CE + weighted value MSE; empty value batches contribute 0 value loss.""" + + def __init__(self, concept_loss_fn: nn.Module, value_loss_weight: float = 1.0): + super().__init__() + self.concept_loss_fn = concept_loss_fn + self.value_loss_fn = nn.MSELoss() + self.value_loss_weight = value_loss_weight + + def forward(self, logits, labels, val_logits, val_labels): + concept_loss = self.concept_loss_fn( + logits.view(-1, logits.size(-1)), labels.view(-1) + ) + # MLM may mask no numeric tokens in a batch → empty tensors → NaN mean MSE + if val_logits.numel() == 0: + value_loss = concept_loss.new_zeros(()) + else: + value_loss = self.value_loss_fn(val_logits, val_labels) + loss = concept_loss + self.value_loss_weight * value_loss + return loss, concept_loss, value_loss + + class PretrainModule(L.LightningModule): def __init__( self, @@ -106,3 +128,47 @@ def configure_optimizers(self): "frequency": 1, } return [optimizer], [scheduler_config] + + +class ValuePretrainModule(PretrainModule): + def __init__( + self, + model: nn.Module, + compile_mode: str = None, + learning_rate: float = 5e-4, + optimizer_epsilon: float = 1e-6, + scheduler_warmup_epochs: int = 0, + value_loss_weight: float = 1.0, + ): + super().__init__( + model=model, + compile_mode=compile_mode, + learning_rate=learning_rate, + optimizer_epsilon=optimizer_epsilon, + scheduler_warmup_epochs=scheduler_warmup_epochs, + ) + self.train_loss = ConceptValueLoss(self.train_loss, value_loss_weight) + self.val_loss = ConceptValueLoss(self.val_loss, value_loss_weight) + self.save_hyperparameters({"value_loss_weight": value_loss_weight}) + + def training_step(self, batch, batch_idx): + logits, labels, val_logits, val_labels = self.model(batch) + loss, concept_loss, value_loss = self.train_loss( + logits, labels, val_logits, val_labels + ) + self.log("train/loss", loss, prog_bar=True) + self.log("train/concept_loss", concept_loss, prog_bar=True) + self.log("train/value_loss", value_loss, prog_bar=True) + return loss + + def validation_step(self, batch, batch_idx): + logits, labels, val_logits, val_labels = self.model(batch) + loss, concept_loss, value_loss = self.val_loss( + logits, labels, val_logits, val_labels + ) + self.log("val/loss", loss, prog_bar=True) + self.log("val/concept_loss", concept_loss, prog_bar=True) + self.log("val/value_loss", value_loss, prog_bar=True) + self.val_metrics.update(logits, labels) + self.log_dict(self.val_metrics) + return loss diff --git a/bonsai/modules/networks/bonsai_nets.py b/bonsai/modules/networks/bonsai_nets.py index 42be3cb17..2d0381101 100644 --- a/bonsai/modules/networks/bonsai_nets.py +++ b/bonsai/modules/networks/bonsai_nets.py @@ -1,8 +1,12 @@ import importlib.util +import torch import torch.nn as nn -from bonsai.modules.networks.components.embeddings import EhrEmbeddings +from bonsai.modules.networks.components.embeddings import ( + EhrEmbeddings, + EhrValueEmbeddings, +) from bonsai.modules.networks.components.layer import TransformerLayer _FLASH_ATTENTION_AVAILABLE = importlib.util.find_spec("flash_attn") is not None @@ -63,16 +67,20 @@ def __init__( "dropout": dropout, "causal": causal, "attn_type": attn_type, + "value_embedding_mode": None, } - def forward(self, batch): - x = self.embeddings( + def embed(self, batch: dict) -> torch.Tensor: + return self.embeddings( code=batch["code"], age=batch["age"], abspos=batch["abspos"], segment=batch["segment"], ) + def forward(self, batch: dict) -> torch.Tensor: + x = self.embed(batch) + # SDPA requires a broadcasted attention mask if self.hparams["attn_type"] == "sdpa" and not self.hparams["causal"]: attn_mask = batch["attention_mask"][:, None, None, :] @@ -87,21 +95,62 @@ def forward(self, batch): cu_seqlens=batch.get("cu_seqlens"), ) x = self.layernorm(x) - return x +class BonsaiValueBase(BonsaiBase): + def __init__( + self, + vocab_size, + max_seqlen, + hidden_size, + num_layers, + num_attention_heads, + bias, + dropout, + attention_dropout, + causal, + attn_type, + value_embedding_mode, + ): + super().__init__( + vocab_size=vocab_size, + max_seqlen=max_seqlen, + hidden_size=hidden_size, + num_layers=num_layers, + num_attention_heads=num_attention_heads, + bias=bias, + dropout=dropout, + attention_dropout=attention_dropout, + causal=causal, + attn_type=attn_type, + ) + self.embeddings = EhrValueEmbeddings( + vocab_size=vocab_size, + hidden_size=hidden_size, + max_seqlen=max_seqlen, + value_embedding_mode=value_embedding_mode, + ) + self.hparams["value_embedding_mode"] = value_embedding_mode + + def embed(self, batch: dict) -> torch.Tensor: + return self.embeddings( + code=batch["code"], + age=batch["age"], + abspos=batch["abspos"], + segment=batch["segment"], + numeric_value=batch["numeric_value"], + ) + + class BonsaiPretrain(BonsaiBase): def __init__( self, - # Embedding / vocab vocab_size, max_seqlen, - # Model dimensions hidden_size, num_layers, num_attention_heads, - # Attention / behavior bias, dropout, attention_dropout, @@ -121,7 +170,6 @@ def __init__( attn_type=attn_type, ) self.pretrain_head = nn.Linear(hidden_size, vocab_size, bias=bias) - # Weight tying (shares weights from code embedding to pretrain head) self.pretrain_head.weight = self.embeddings.code_embedding.weight @@ -129,32 +177,74 @@ def forward(self, batch: dict): last_hidden_state = super().forward(batch) labels = batch["target"] - # Predicts only on the non-masked tokens mask = labels != -100 last_hidden_state = last_hidden_state[mask] labels = labels[mask] - logits = self.pretrain_head(last_hidden_state) return logits, labels +class BonsaiValuePretrain(BonsaiValueBase): + def __init__( + self, + vocab_size, + max_seqlen, + hidden_size, + num_layers, + num_attention_heads, + bias, + dropout, + attention_dropout, + causal, + attn_type, + value_embedding_mode, + ): + super().__init__( + vocab_size=vocab_size, + max_seqlen=max_seqlen, + hidden_size=hidden_size, + num_layers=num_layers, + num_attention_heads=num_attention_heads, + bias=bias, + dropout=dropout, + attention_dropout=attention_dropout, + causal=causal, + attn_type=attn_type, + value_embedding_mode=value_embedding_mode, + ) + self.pretrain_head = nn.Linear(hidden_size, vocab_size, bias=bias) + self.pretrain_head.weight = self.embeddings.code_embedding.weight + self.pretrain_head_value = nn.Linear(hidden_size, 1, bias=bias) + + def forward(self, batch: dict): + last_hidden_state = super().forward(batch) + labels = batch["target"] + + mask = labels != -100 + last_hidden_state = last_hidden_state[mask] + labels = labels[mask] + logits = self.pretrain_head(last_hidden_state) + + val_labels = batch["numeric_target"][mask] + val_mask = ~torch.isnan(val_labels) + val_logits = self.pretrain_head_value(last_hidden_state[val_mask]).squeeze(-1) + val_labels = val_labels[val_mask] + return logits, labels, val_logits, val_labels + + class BonsaiFinetune(BonsaiBase): def __init__( self, - # Embedding / vocab vocab_size, max_seqlen, - # Model dimensions hidden_size, num_layers, num_attention_heads, - # Attention / behavior bias, dropout, attention_dropout, causal, attn_type, - # Misc predict_token_id, ): super().__init__( @@ -174,11 +264,45 @@ def __init__( def forward(self, batch: dict): last_hidden_state = super().forward(batch) - - # Extracts the hidden states corresponding to the predict token for each subject pred_tokens = batch["code"] == self.hparams["predict_token_id"] last_hidden_state = last_hidden_state[pred_tokens] + return self.finetune_head(last_hidden_state) - logits = self.finetune_head(last_hidden_state) - return logits +class BonsaiValueFinetune(BonsaiValueBase): + def __init__( + self, + vocab_size, + max_seqlen, + hidden_size, + num_layers, + num_attention_heads, + bias, + dropout, + attention_dropout, + causal, + attn_type, + predict_token_id, + value_embedding_mode, + ): + super().__init__( + vocab_size=vocab_size, + max_seqlen=max_seqlen, + hidden_size=hidden_size, + num_layers=num_layers, + num_attention_heads=num_attention_heads, + bias=bias, + dropout=dropout, + attention_dropout=attention_dropout, + causal=causal, + attn_type=attn_type, + value_embedding_mode=value_embedding_mode, + ) + self.hparams["predict_token_id"] = predict_token_id + self.finetune_head = nn.Linear(hidden_size, 1, bias=bias) + + def forward(self, batch: dict): + last_hidden_state = super().forward(batch) + pred_tokens = batch["code"] == self.hparams["predict_token_id"] + last_hidden_state = last_hidden_state[pred_tokens] + return self.finetune_head(last_hidden_state) diff --git a/bonsai/modules/networks/components/embeddings.py b/bonsai/modules/networks/components/embeddings.py index a3385490d..452d5a6ec 100644 --- a/bonsai/modules/networks/components/embeddings.py +++ b/bonsai/modules/networks/components/embeddings.py @@ -26,8 +26,8 @@ def forward( abspos: torch.Tensor, segment: torch.LongTensor, ) -> torch.Tensor: - embeddings = self.code_embedding(code) + embeddings = self.code_embedding(code) embeddings += self.age_embedding(age) embeddings += self.abspos_embedding(abspos) embeddings += self.segment_embedding(segment) @@ -35,6 +35,36 @@ def forward( return embeddings +class EhrValueEmbeddings(EhrEmbeddings): + def __init__( + self, + vocab_size: int, + hidden_size: int, + max_seqlen: int, + value_embedding_mode: str, + ): + super().__init__(vocab_size, hidden_size, max_seqlen) + self.value_embedding_mode = value_embedding_mode + self.numeric_value_embedding = ContinuousEmbedding( + hidden_size, value_embedding_mode + ) + + def forward( + self, + code: torch.LongTensor, + age: torch.Tensor, + abspos: torch.Tensor, + segment: torch.LongTensor, + numeric_value: torch.Tensor, + ) -> torch.Tensor: + embeddings = self.code_embedding(code) + embeddings = self.numeric_value_embedding(numeric_value, embeddings) + embeddings += self.age_embedding(age) + embeddings += self.abspos_embedding(abspos) + embeddings += self.segment_embedding(segment) + return embeddings + + class Time2Vec(nn.Module): """Time2Vec embedding layer that combines linear and periodic components. @@ -96,3 +126,37 @@ def forward(self, tau: torch.Tensor) -> torch.Tensor: periodic = self.f(linear_2 + self.phi) return torch.cat((linear_1, periodic), dim=-1) + + +class ContinuousEmbedding(nn.Module): + def __init__(self, hidden_size: int, value_embedding_mode: str): + super().__init__() + self.value_embedding_mode = value_embedding_mode + self.hidden_size = hidden_size + + self.value_proj = nn.Sequential( + nn.Linear(1, hidden_size), nn.ReLU(), nn.Linear(hidden_size, hidden_size) + ) + + if self.value_embedding_mode == "film": + self.gamma_layer = nn.Linear(hidden_size, hidden_size) + self.beta_layer = nn.Linear(hidden_size, hidden_size) + else: + raise ValueError( + f"Unknown value_embedding_mode: {self.value_embedding_mode}" + ) + + def forward( + self, values: torch.Tensor, concept_embeds: torch.Tensor + ) -> torch.Tensor: + mask = (~torch.isnan(values)).float().unsqueeze(-1) + values_safe = torch.where(torch.isnan(values), torch.zeros_like(values), values) + value_embed = self.value_proj(values_safe.unsqueeze(-1)) * mask + + if self.value_embedding_mode == "film": + gamma = self.gamma_layer(concept_embeds) + beta = self.beta_layer(concept_embeds) + fused = (gamma * value_embed + beta).to(dtype=concept_embeds.dtype) + return fused * mask + concept_embeds * (1 - mask) + + raise ValueError(f"Unknown value_embedding_mode: {self.value_embedding_mode}") diff --git a/bonsai/run/create_data.py b/bonsai/run/create_data.py index e0cfd70e7..fcb50de46 100644 --- a/bonsai/run/create_data.py +++ b/bonsai/run/create_data.py @@ -51,6 +51,7 @@ def main(cfg: DictConfig) -> None: path_output_dir=path_output_dir, tokenizer=tokenizer, exclude_regex=cfg.exclude_regex, + numeric_column=cfg.numeric_column, ) ids.extend(split_ids) tokenizer.freeze_vocabulary() # freeze after first split (train) to prevent data leakage diff --git a/bonsai/run/finetune.py b/bonsai/run/finetune.py index c27049a76..875e20d54 100644 --- a/bonsai/run/finetune.py +++ b/bonsai/run/finetune.py @@ -6,6 +6,7 @@ import torch from dotenv import load_dotenv from hydra.core.hydra_config import HydraConfig +from hydra.utils import instantiate from lightning.pytorch.callbacks import ModelCheckpoint from lightning.pytorch.loggers import CSVLogger from omegaconf import DictConfig, OmegaConf @@ -18,7 +19,6 @@ from bonsai.functional.versioning import generate_unused_run_id from bonsai.modules.datamodules.FinetuneDataModule import FinetuneDataModule from bonsai.modules.lightningmodules.FinetuneModule import FinetuneModule -from bonsai.modules.networks.bonsai_nets import BonsaiFinetune from bonsai.paths import get_config_path OmegaConf.register_new_resolver( @@ -76,16 +76,14 @@ def main(cfg: DictConfig) -> None: ), ) - model = BonsaiFinetune( + model = instantiate( + cfg.model, vocab_size=len(vocab), max_seqlen=pretrain_cfg["max_seqlen"], hidden_size=pretrain_cfg["hidden_size"], num_layers=pretrain_cfg["num_layers"], num_attention_heads=pretrain_cfg["num_attention_heads"], bias=pretrain_cfg["bias"], - dropout=cfg.model.dropout, - attention_dropout=cfg.model.attention_dropout, - causal=cfg.model.causal, attn_type=pretrain_cfg["attn_type"], predict_token_id=vocab["[CLS]"], ) diff --git a/bonsai/run/pretrain.py b/bonsai/run/pretrain.py index d05f89ef1..e5a0feff5 100644 --- a/bonsai/run/pretrain.py +++ b/bonsai/run/pretrain.py @@ -2,7 +2,7 @@ import lightning as L from dotenv import load_dotenv from hydra.core.hydra_config import HydraConfig -from hydra.utils import get_class +from hydra.utils import get_class, instantiate from lightning.pytorch.callbacks import ModelCheckpoint from lightning.pytorch.loggers import CSVLogger from omegaconf import DictConfig, OmegaConf @@ -10,8 +10,11 @@ from bonsai.functional.pathing import get_experiment_output_path from bonsai.functional.versioning import generate_unused_run_id from bonsai.modules.datamodules.PretrainDataModule import PretrainDataModule -from bonsai.modules.lightningmodules.PretrainModule import PretrainModule -from bonsai.modules.networks.bonsai_nets import BonsaiPretrain +from bonsai.modules.lightningmodules.PretrainModule import ( + PretrainModule, + ValuePretrainModule, +) +from bonsai.modules.networks.bonsai_nets import BonsaiValuePretrain from bonsai.paths import get_config_path OmegaConf.register_new_resolver( @@ -47,18 +50,21 @@ def main(cfg: DictConfig) -> None: max_len=cfg.training.max_len, ) - model = BonsaiPretrain( - vocab_size=len(data_module.vocabulary), - max_seqlen=cfg.model.max_seqlen, - hidden_size=cfg.model.hidden_size, - num_layers=cfg.model.num_layers, - num_attention_heads=cfg.model.num_attention_heads, - bias=cfg.model.bias, - dropout=cfg.model.dropout, - attention_dropout=cfg.model.attention_dropout, - causal=cfg.model.causal, - attn_type=cfg.model.attn_type, + model = instantiate(cfg.model, vocab_size=len(data_module.vocabulary)) + module_kwargs = dict( + model=model, + compile_mode=cfg.hardware.compile_mode, + learning_rate=cfg.training.learning_rate, + optimizer_epsilon=cfg.training.optimizer_epsilon, + scheduler_warmup_epochs=cfg.training.scheduler_warmup_epochs, ) + if isinstance(model, BonsaiValuePretrain): + lightning_module = ValuePretrainModule( + **module_kwargs, + value_loss_weight=cfg.training.get("value_loss_weight", 1.0), + ) + else: + lightning_module = PretrainModule(**module_kwargs) ckpt_callback = ModelCheckpoint( dirpath=model_save_dir, @@ -70,14 +76,6 @@ def main(cfg: DictConfig) -> None: save_last=True, ) - lightning_module = PretrainModule( - model=model, - compile_mode=cfg.hardware.compile_mode, - learning_rate=cfg.training.learning_rate, - optimizer_epsilon=cfg.training.optimizer_epsilon, - scheduler_warmup_epochs=cfg.training.scheduler_warmup_epochs, - ) - trainer = L.Trainer( accelerator=cfg.hardware.accelerator, accumulate_grad_batches=cfg.training.accumulate_grad_batches, diff --git a/bonsai/run/train.py b/bonsai/run/train.py index 6e1dcd386..543852b8c 100644 --- a/bonsai/run/train.py +++ b/bonsai/run/train.py @@ -5,6 +5,7 @@ import polars as pl import torch from dotenv import load_dotenv +from hydra.utils import instantiate from lightning.pytorch.callbacks import ModelCheckpoint from lightning.pytorch.loggers import CSVLogger from omegaconf import DictConfig, OmegaConf @@ -17,7 +18,6 @@ from bonsai.functional.versioning import generate_unused_run_id from bonsai.modules.datamodules.FinetuneDataModule import FinetuneDataModule from bonsai.modules.lightningmodules.FinetuneModule import FinetuneModule -from bonsai.modules.networks.bonsai_nets import BonsaiFinetune from bonsai.paths import get_config_path OmegaConf.register_new_resolver( @@ -68,17 +68,9 @@ def main(cfg: DictConfig) -> None: ), ) - model = BonsaiFinetune( + model = instantiate( + cfg.model, vocab_size=len(vocab), - max_seqlen=cfg.model.max_seqlen, - hidden_size=cfg.model.hidden_size, - num_layers=cfg.model.num_layers, - num_attention_heads=cfg.model.num_attention_heads, - bias=cfg.model.bias, - dropout=cfg.model.dropout, - attention_dropout=cfg.model.attention_dropout, - causal=cfg.model.causal, - attn_type=cfg.model.attn_type, predict_token_id=vocab["[CLS]"], ) diff --git a/configs/data_creation/default_create_data.yaml b/configs/data_creation/default_create_data.yaml index 0377dbb67..74763c75b 100644 --- a/configs/data_creation/default_create_data.yaml +++ b/configs/data_creation/default_create_data.yaml @@ -15,7 +15,9 @@ paths: exclude_regex: +numeric_column: "numeric_value_normalized" + tokenizer: vocabulary: cutoffs: - sep_tokens: true + sep_tokens: true \ No newline at end of file diff --git a/configs/examples/example_data.yaml b/configs/examples/example_data.yaml index 8a8399cc3..9f6e28684 100644 --- a/configs/examples/example_data.yaml +++ b/configs/examples/example_data.yaml @@ -16,7 +16,9 @@ paths: exclude_regex: ^(?:LAB).* # example regex to exclude all features that start with 'LAB' +numeric_column: + tokenizer: vocabulary: cutoffs: - sep_tokens: true + sep_tokens: true \ No newline at end of file diff --git a/configs/examples/example_finetune.yaml b/configs/examples/example_finetune.yaml index 1cac449ab..5da968368 100644 --- a/configs/examples/example_finetune.yaml +++ b/configs/examples/example_finetune.yaml @@ -9,5 +9,4 @@ training: batch_size: 2 epochs: 5 limit_val_batches: 10 - limit_train_batches: 10 - \ No newline at end of file + limit_train_batches: 10 \ No newline at end of file diff --git a/configs/examples/example_finetune_val.yaml b/configs/examples/example_finetune_val.yaml new file mode 100644 index 000000000..043f38828 --- /dev/null +++ b/configs/examples/example_finetune_val.yaml @@ -0,0 +1,18 @@ +# @package _global_ +defaults: + - ../finetune@ + - ../hardware/cpu@hardware + - _self_ + +training: + max_len: 512 # must match (or be ≤) pretrained max_seqlen + batch_size: 64 # val has ~150 subjects; 256 + drop_last left 0 val batches + epochs: 5 + limit_val_batches: 1.0 + limit_train_batches: 1.0 + learning_rate: 1e-3 + scheduler_warmup_epochs: 1 + +model: + _target_: bonsai.modules.networks.bonsai_nets.BonsaiValueFinetune + value_embedding_mode: film \ No newline at end of file diff --git a/configs/examples/example_outcome1.yaml b/configs/examples/example_outcome1.yaml index 85098aa09..ec06ac675 100644 --- a/configs/examples/example_outcome1.yaml +++ b/configs/examples/example_outcome1.yaml @@ -40,4 +40,4 @@ outcome: conditions: # exposure censor: # ALWAYS RELATIVE TO INDEX - relative_hour_shift: -1000 # 0 sets index=censor + relative_hour_shift: -1000 # 0 sets index=censor \ No newline at end of file diff --git a/configs/examples/example_outcome_val.yaml b/configs/examples/example_outcome_val.yaml new file mode 100644 index 000000000..40957bd1b --- /dev/null +++ b/configs/examples/example_outcome_val.yaml @@ -0,0 +1,39 @@ +# @package _global_ +defaults: + - ../core/base_create@ + - _self_ + +dataset: correlated_MEDS_data + +paths: + input_dir: example_data/${dataset} + save_path: ${oc.env:BONSAI_PROCESSED_DATA}/${dataset}/outcomes/${hydra:job.config_name}.parquet + +splits: + - train + - tuning + - held_out + + + +outcome: + exclude: # OPTIONAL + + outcome: + dependence: independent # independent/dependent + conditions: + - col: code + vals: [ SYN//Negative ] + + index: + type: absolute # absolute/relative/exposure + absolute_date: # absolute + year: 2023 + month: 1 + day: 1 + relative_hour_shift: 0 + dependence: # exposure + conditions: # exposure + + censor: # ALWAYS RELATIVE TO INDEX + relative_hour_shift: 48 # 0 sets index=censor diff --git a/configs/examples/example_pretrain_val.yaml b/configs/examples/example_pretrain_val.yaml new file mode 100644 index 000000000..ecfe636ce --- /dev/null +++ b/configs/examples/example_pretrain_val.yaml @@ -0,0 +1,20 @@ +# @package _global_ +defaults: + - ../pretrain@ + - ../hardware/1gpu6cpu@hardware + - _self_ + +training: + max_len: 512 + batch_size: 2 + epochs: 10 + limit_val_batches: 10 + limit_train_batches: 10 + value_loss_weight: 1.0 + +paths: + dataset_class: bonsai.modules.datasets.PretrainDataset.ValueARPretrainDataset + +model: + _target_: bonsai.modules.networks.bonsai_nets.BonsaiValuePretrain + value_embedding_mode: film diff --git a/configs/examples/example_value_data.yaml b/configs/examples/example_value_data.yaml new file mode 100644 index 000000000..b53ca72e6 --- /dev/null +++ b/configs/examples/example_value_data.yaml @@ -0,0 +1,24 @@ +# @package _global_ +defaults: + - ../core/base_create@ + - _self_ + +splits: + - train + - tuning + - held_out + +dataset: lab_correlated_MEDS_data_not_fused + +paths: + input_dir: example_data/${dataset} + output_dir: ${oc.env:BONSAI_PROCESSED_DATA}/${dataset} + +exclude_regex: ^(?:D//).* # example regex to exclude all features that start with 'D' + +numeric_column: "numeric_value_normalized" + +tokenizer: + vocabulary: + cutoffs: + sep_tokens: true diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 700478b41..ef6f05d25 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -43,6 +43,8 @@ training: eval_monitor_metric: val/loss model: + # For numeric values, switch to BonsaiValueFinetune and set value_embedding_mode. + _target_: bonsai.modules.networks.bonsai_nets.BonsaiFinetune max_seqlen: ${training.max_len} hidden_size: 64 num_layers: 4 @@ -50,5 +52,5 @@ model: dropout: 0.1 attention_dropout: 0.0 bias: false - causal: false + causal: true attn_type: sdpa diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 89f9c02c3..90cf2ba8f 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -10,7 +10,7 @@ paths: train_split: ${paths.dir}/subject_data_train.pt val_split: ${paths.dir}/subject_data_tuning.pt vocab: ${paths.dir}/vocabulary.pt - dataset_class: bonsai.modules.datasets.PretrainDataset.ARPretrainDataset # or MLMPretrainDataset + dataset_class: bonsai.modules.datasets.PretrainDataset.ARPretrainDataset # or MLMPretrainDataset / ValueARPretrainDataset / ValueMLMPretrainDataset population: ${paths.dir}/population_full.csv ckpt_path: @@ -24,6 +24,7 @@ training: scheduler_warmup_epochs: 1 limit_val_batches: 1.0 limit_train_batches: 1.0 + value_loss_weight: 1.0 ckpt_every_n_epoch: 1 cutoff_date: # Optional # day: 1 @@ -31,12 +32,14 @@ training: # year: 2015 # MLM specific # masking: - # masking_select_ratio: 1. + # masking_select_ratio: 0.2 # masking_mask_ratio: .8 # masking_random_ratio: .1 # masking_ignore_special_tokens: true model: + # For numeric values, switch to BonsaiValuePretrain and set value_embedding_mode. + _target_: bonsai.modules.networks.bonsai_nets.BonsaiPretrain max_seqlen: ${training.max_len} hidden_size: 64 num_layers: 4 @@ -44,7 +47,7 @@ model: dropout: 0.1 attention_dropout: 0.0 bias: false - causal: false + causal: true attn_type: sdpa diff --git a/example_data/example_MEDS_data/held_out/0.parquet b/example_data/example_MEDS_data/held_out/0.parquet deleted file mode 100644 index 68888fc02..000000000 Binary files a/example_data/example_MEDS_data/held_out/0.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data/train/0.parquet b/example_data/example_MEDS_data/train/0.parquet deleted file mode 100644 index fe53592f2..000000000 Binary files a/example_data/example_MEDS_data/train/0.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data/tuning/0.parquet b/example_data/example_MEDS_data/tuning/0.parquet deleted file mode 100644 index 01ac3d376..000000000 Binary files a/example_data/example_MEDS_data/tuning/0.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/0.parquet b/example_data/example_MEDS_data_wo_held_out/train/0.parquet deleted file mode 100644 index b2f5c8cdf..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/0.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/1.parquet b/example_data/example_MEDS_data_wo_held_out/train/1.parquet deleted file mode 100644 index 9c814a181..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/1.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/2.parquet b/example_data/example_MEDS_data_wo_held_out/train/2.parquet deleted file mode 100644 index d1af61df7..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/2.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/3.parquet b/example_data/example_MEDS_data_wo_held_out/train/3.parquet deleted file mode 100644 index e81fa4fff..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/3.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/4.parquet b/example_data/example_MEDS_data_wo_held_out/train/4.parquet deleted file mode 100644 index 3f667138d..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/4.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/5.parquet b/example_data/example_MEDS_data_wo_held_out/train/5.parquet deleted file mode 100644 index 3fcdd924e..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/5.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/6.parquet b/example_data/example_MEDS_data_wo_held_out/train/6.parquet deleted file mode 100644 index 4f1889797..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/6.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/train/7.parquet b/example_data/example_MEDS_data_wo_held_out/train/7.parquet deleted file mode 100644 index 73caef40c..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/train/7.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/tuning/0.parquet b/example_data/example_MEDS_data_wo_held_out/tuning/0.parquet deleted file mode 100644 index 0308ea1af..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/tuning/0.parquet and /dev/null differ diff --git a/example_data/example_MEDS_data_wo_held_out/tuning/1.parquet b/example_data/example_MEDS_data_wo_held_out/tuning/1.parquet deleted file mode 100644 index 5fb43403a..000000000 Binary files a/example_data/example_MEDS_data_wo_held_out/tuning/1.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/hash_to_integer_map.pt b/example_data/lab_correlated_MEDS_data/hash_to_integer_map.pt deleted file mode 100644 index 5b6022474..000000000 Binary files a/example_data/lab_correlated_MEDS_data/hash_to_integer_map.pt and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/held_out/0.parquet b/example_data/lab_correlated_MEDS_data/held_out/0.parquet deleted file mode 100644 index 068b634f9..000000000 Binary files a/example_data/lab_correlated_MEDS_data/held_out/0.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/held_out/1.parquet b/example_data/lab_correlated_MEDS_data/held_out/1.parquet deleted file mode 100644 index fc3d80693..000000000 Binary files a/example_data/lab_correlated_MEDS_data/held_out/1.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/0.parquet b/example_data/lab_correlated_MEDS_data/train/0.parquet deleted file mode 100644 index e882e2e2d..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/0.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/1.parquet b/example_data/lab_correlated_MEDS_data/train/1.parquet deleted file mode 100644 index aa1b0ee18..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/1.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/10.parquet b/example_data/lab_correlated_MEDS_data/train/10.parquet deleted file mode 100644 index 7f00d8c02..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/10.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/11.parquet b/example_data/lab_correlated_MEDS_data/train/11.parquet deleted file mode 100644 index 3bbfeb285..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/11.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/12.parquet b/example_data/lab_correlated_MEDS_data/train/12.parquet deleted file mode 100644 index 1e286a666..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/12.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/13.parquet b/example_data/lab_correlated_MEDS_data/train/13.parquet deleted file mode 100644 index fcc192e0b..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/13.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/14.parquet b/example_data/lab_correlated_MEDS_data/train/14.parquet deleted file mode 100644 index 8d996fca9..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/14.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/15.parquet b/example_data/lab_correlated_MEDS_data/train/15.parquet deleted file mode 100644 index 526f8b540..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/15.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/2.parquet b/example_data/lab_correlated_MEDS_data/train/2.parquet deleted file mode 100644 index 347e38bc0..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/2.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/3.parquet b/example_data/lab_correlated_MEDS_data/train/3.parquet deleted file mode 100644 index 047478ec8..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/3.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/4.parquet b/example_data/lab_correlated_MEDS_data/train/4.parquet deleted file mode 100644 index bb187696f..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/4.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/5.parquet b/example_data/lab_correlated_MEDS_data/train/5.parquet deleted file mode 100644 index 63628eb7f..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/5.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/6.parquet b/example_data/lab_correlated_MEDS_data/train/6.parquet deleted file mode 100644 index 9f8ca2bf9..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/6.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/7.parquet b/example_data/lab_correlated_MEDS_data/train/7.parquet deleted file mode 100644 index 1f72768fc..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/7.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/8.parquet b/example_data/lab_correlated_MEDS_data/train/8.parquet deleted file mode 100644 index 1e464cfb5..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/8.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/train/9.parquet b/example_data/lab_correlated_MEDS_data/train/9.parquet deleted file mode 100644 index 7825a1306..000000000 Binary files a/example_data/lab_correlated_MEDS_data/train/9.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/tuning/0.parquet b/example_data/lab_correlated_MEDS_data/tuning/0.parquet deleted file mode 100644 index 1993f588e..000000000 Binary files a/example_data/lab_correlated_MEDS_data/tuning/0.parquet and /dev/null differ diff --git a/example_data/lab_correlated_MEDS_data/tuning/1.parquet b/example_data/lab_correlated_MEDS_data/tuning/1.parquet deleted file mode 100644 index 08b5d8071..000000000 Binary files a/example_data/lab_correlated_MEDS_data/tuning/1.parquet and /dev/null differ