Feat/vals

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(

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

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())

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
 

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()
+            all_tokenized.append(subject)
 
     return all_tokenized
 

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

bonsai/modules/datasets/FinetuneDataset.py

@@ -42,6 +42,8 @@ def __getitem__(self, index: int) -> dict:
 
         subject["segment"] = normalize_segments(subject["segment"])
         subject["attention_mask"] = torch.ones(len(subject["code"]), dtype=torch.bool)
+        if "numeric_value" in subject:
+            subject["numeric_value"] = subject["numeric_value"].float()
 
         return subject
 

bonsai/modules/datasets/PretrainDataset.py

@@ -70,14 +70,28 @@ 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
+        numeric_values = subject.get("numeric_value")
+        masked_codes, masked_numeric_values, target, numeric_target = (
+            self.mask_patient_codes(subject["code"], numeric_values)
+        )
+        # Collate / model read "code"; mask_patient_codes returns the masked sequence.
+        subject["code"] = masked_codes
         subject["target"] = target
+        if numeric_target is not None:
+            subject["numeric_value"] = masked_numeric_values
+            subject["numeric_target"] = numeric_target
         return subject
 
     def mask_patient_codes(
-        self, codes: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self,
+        codes: torch.Tensor,
+        numeric_values: Optional[torch.Tensor] = None,
+    ) -> Tuple[
+        torch.Tensor,
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[torch.Tensor],
+    ]:
         target = codes.clone()
         probability_vector = torch.full(target.shape, self.masking_select_ratio)
 
@@ -94,8 +108,18 @@ def mask_patient_codes(
             torch.bernoulli(torch.full(target.shape, self.masking_mask_ratio)).bool()
             & selected_indices
         )
+        codes = codes.clone()
         codes[indices_mask] = self.vocabulary["[MASK]"]
 
+        if numeric_values is not None:
+            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")
+        else:
+            numeric_target = None
+
         # Replace with random word and Account for already masked tokens
         random_ratio = self.masking_random_ratio / (1 - self.masking_mask_ratio)
         indicies_random = (
@@ -110,7 +134,7 @@ def mask_patient_codes(
             dtype=codes.dtype,
         )
         codes[indicies_random] = random_words[indicies_random]
-        return codes, target
+        return codes, numeric_values, target, numeric_target
 
 
 class ARPretrainDataset(PretrainDataset):
@@ -129,6 +153,16 @@ 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)
+
+        if "numeric_value" in subject:
+            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

bonsai/modules/lightningmodules/PretrainModule.py

@@ -1,4 +1,5 @@
 import lightning as L
+import torch
 from torch import nn
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import LinearLR
@@ -45,8 +46,12 @@ def __init__(
                 "scheduler_warmup_epochs": scheduler_warmup_epochs,
             }
         )
+        hparams.update(self.add_hparams())
         self.save_hyperparameters(hparams)
 
+    def add_hparams(self) -> dict:
+        return {}
+
     def configure_metrics(self, prefix: str):
         return MetricCollection(
             {
@@ -69,15 +74,18 @@ def configure_metrics(self, prefix: str):
             ],
         )
 
-    def training_step(self, batch, batch_idx):
+    def compute_loss(self, batch, concept_loss_fn):
         logits, labels = self.model(batch)
-        loss = self.train_loss(logits.view(-1, logits.size(-1)), labels.view(-1))
+        loss = concept_loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
+        return loss, logits, labels
+
+    def training_step(self, batch, batch_idx):
+        loss, _, _ = self.compute_loss(batch, self.train_loss)
         self.log("train/loss", loss, prog_bar=True)
         return loss
 
     def validation_step(self, batch, batch_idx):
-        logits, labels = self.model(batch)
-        loss = self.val_loss(logits.view(-1, logits.size(-1)), labels.view(-1))
+        loss, logits, labels = self.compute_loss(batch, self.val_loss)
         self.log("val/loss", loss, prog_bar=True)
         self.val_metrics.update(logits, labels)
         self.log_dict(self.val_metrics)
@@ -106,3 +114,58 @@ 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,
+    ):
+        self.value_loss_weight = value_loss_weight
+        super().__init__(
+            model=model,
+            compile_mode=compile_mode,
+            learning_rate=learning_rate,
+            optimizer_epsilon=optimizer_epsilon,
+            scheduler_warmup_epochs=scheduler_warmup_epochs,
+        )
+        self.value_loss_fn = nn.MSELoss()
+
+    def add_hparams(self) -> dict:
+        return {"value_loss_weight": self.value_loss_weight}
+
+    def compute_loss(self, batch, concept_loss_fn):
+        logits, labels, val_logits, val_labels = self.model(batch)
+        concept_loss = 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, logits, labels
+
+    def training_step(self, batch, batch_idx):
+        loss, concept_loss, value_loss, _, _ = self.compute_loss(batch, self.train_loss)
+        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):
+        loss, concept_loss, value_loss, logits, labels = self.compute_loss(
+            batch, self.val_loss
+        )
+        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