[Bounty] Adacare PyHealth 2.0 (#613)

* Add DeepNestedSequenceProcessor and DeepNestedFloatsProcessor for handling deeply nested sequences

* Add unit tests for DeepNestedSequenceProcessor and DeepNestedFloatsProcessor

* Refactor EmbeddingModel to support DeepNestedSequenceProcessor and DeepNestedFloatsProcessor, and support mask outputs

* Refactor AdaCare model to integrate new dataset and processor classes, update input handling, and forward propagation logic

* Remove output mask handling for passthrough tensors in EmbeddingModel forward method

* Add unit tests for AdaCare model including initialization, forward and backward passes, loss checks, and output shapes

* Update mortality prediction example script for MIMIC-III using adacare

* Remove property decorator from size function in StageNetTensorProcessor class

Author: Razin93937

examples/mortality_mimic3_adacare.ipynb

@@ -13,68 +13,62 @@
     StageNetTensorProcessor,
     TensorProcessor,
     TimeseriesProcessor,
+    DeepNestedSequenceProcessor,
+    DeepNestedFloatsProcessor,
 )
 from .base_model import BaseModel
 
-
 class EmbeddingModel(BaseModel):
     """
     EmbeddingModel is responsible for creating embedding layers for different types of input data.
 
     This model automatically creates appropriate embedding transformations based on the processor type:
 
-    - SequenceProcessor: Creates nn.Embedding for categorical sequences (e.g., diagnosis codes)
-      Input: (batch, seq_len) with integer indices
-      Output: (batch, seq_len, embedding_dim)
+    - SequenceProcessor: nn.Embedding
+        Input: (batch, seq_len)
+        Output: (batch, seq_len, embedding_dim)
 
-    - TimeseriesProcessor: Creates nn.Linear for time series features
-      Input: (batch, seq_len, num_features)
-      Output: (batch, seq_len, embedding_dim)
+    - NestedSequenceProcessor: nn.Embedding
+        Input: (batch, num_visits, max_codes_per_visit)
+        Output: (batch, num_visits, max_codes_per_visit, embedding_dim)
 
-    - TensorProcessor: Creates nn.Linear for fixed-size numerical features
-      Input: (batch, feature_size)
-      Output: (batch, embedding_dim)
+    - DeepNestedSequenceProcessor: nn.Embedding
+        Input: (batch, num_groups, num_visits, max_codes_per_visit)
+        Output: (batch, num_groups, num_visits, max_codes_per_visit, embedding_dim)
 
-    - MultiHotProcessor: Creates nn.Linear for multi-hot encoded categorical features
-      Input: (batch, num_categories) binary tensor
-      Output: (batch, embedding_dim)
-      Note: Converts sparse categorical representations to dense embeddings
+    - TimeseriesProcessor / NestedFloatsProcessor / DeepNestedFloatsProcessor / StageNetTensorProcessor:
+        nn.Linear over the last dimension
+        Input: (..., size)
+        Output: (..., embedding_dim)
 
-    - Other processors with size(): Creates nn.Linear if processor reports a positive size
-      Input: (batch, size)
-      Output: (batch, embedding_dim)
+    - TensorProcessor: nn.Linear (size inferred from first sample)
 
-    Attributes:
-        dataset (SampleDataset): The dataset containing input processors.
-        embedding_layers (nn.ModuleDict): A dictionary of embedding layers for each input field.
-        embedding_dim (int): The target embedding dimension for all features.
+    - MultiHotProcessor: nn.Linear over multi-hot vector
     """
 
     def __init__(self, dataset: SampleDataset, embedding_dim: int = 128):
-        """
-        Initializes the EmbeddingModel with the given dataset and embedding dimension.
-
-        Args:
-            dataset (SampleDataset): The dataset containing input processors.
-            embedding_dim (int): The dimension of the embedding space. Default is 128.
-        """
         super().__init__(dataset)
         self.embedding_dim = embedding_dim
         self.embedding_layers = nn.ModuleDict()
+
         for field_name, processor in self.dataset.input_processors.items():
+            # Deep categorical: use special module that collapses last dim to embedding_dim
+
+            # Regular categorical sequences -> nn.Embedding (adds embedding dim)
             if isinstance(
                 processor,
                 (
                     SequenceProcessor,
                     StageNetProcessor,
                     NestedSequenceProcessor,
+                    DeepNestedSequenceProcessor
                 ),
             ):
-                # Categorical codes -> use nn.Embedding
                 vocab_size = len(processor.code_vocab)
-                # For NestedSequenceProcessor, don't use padding_idx
-                # because empty visits need non-zero embeddings
-                if isinstance(processor, NestedSequenceProcessor):
+
+                # For NestedSequenceProcessor and DeepNestedSequenceProcessor, don't use padding_idx
+                # because empty visits/groups need non-zero embeddings.
+                if isinstance(processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor)):
                     self.embedding_layers[field_name] = nn.Embedding(
                         num_embeddings=vocab_size,
                         embedding_dim=embedding_dim,
@@ -86,22 +80,25 @@ def __init__(self, dataset: SampleDataset, embedding_dim: int = 128):
                         embedding_dim=embedding_dim,
                         padding_idx=0,
                     )
+
+            # Numeric features (including deep nested floats) -> nn.Linear over last dim
             elif isinstance(
                 processor,
                 (
                     TimeseriesProcessor,
                     StageNetTensorProcessor,
                     NestedFloatsProcessor,
+                    DeepNestedFloatsProcessor,
                 ),
             ):
-                # Numeric features -> use nn.Linear
-                # Both processors have .size attribute
+                # Assuming processor.size() returns the last-dim size
+                in_features = processor.size()
                 self.embedding_layers[field_name] = nn.Linear(
-                    in_features=processor.size, out_features=embedding_dim
+                    in_features=in_features, out_features=embedding_dim
                 )
+
             elif isinstance(processor, TensorProcessor):
-                # For tensor processor, we need to determine the input size
-                # from the first sample in the dataset
+                # Infer size from first sample
                 sample_tensor = None
                 for sample in dataset.samples:
                     if field_name in sample:
@@ -114,43 +111,51 @@ def __init__(self, dataset: SampleDataset, embedding_dim: int = 128):
                     self.embedding_layers[field_name] = nn.Linear(
                         in_features=input_size, out_features=embedding_dim
                     )
+
             elif isinstance(processor, MultiHotProcessor):
-                # MultiHotProcessor produces fixed-size binary vectors
-                # Use processor.size() to get the vocabulary size (num_categories)
                 num_categories = processor.size()
                 self.embedding_layers[field_name] = nn.Linear(
                     in_features=num_categories, out_features=embedding_dim
                 )
+
             else:
                 print(
                     "Warning: No embedding created for field due to lack of compatible processor:",
                     field_name,
                 )
 
-    def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        """
-        Forward pass to compute embeddings for the input data.
-
-        Args:
-            inputs (Dict[str, torch.Tensor]): A dictionary of input tensors.
-
-        Returns:
-            Dict[str, torch.Tensor]: A dictionary of embedded tensors.
-        """
-        embedded = {}
+    def forward(self,
+                inputs: Dict[str, torch.Tensor],
+                output_mask: bool = False
+                ) -> Dict[str, torch.Tensor] | tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        
+        embedded: Dict[str, torch.Tensor] = {}
+        masks: Dict[str, torch.Tensor] = {} if output_mask else None
+        
         for field_name, tensor in inputs.items():
-            if field_name in self.embedding_layers:
-                tensor = tensor.to(self.device)
-                embedded[field_name] = self.embedding_layers[field_name](tensor)
-            else:
-                embedded[field_name] = tensor  # passthrough for continuous features
-        return embedded
+            processor = self.dataset.input_processors.get(field_name, None)
+            
+            if field_name not in self.embedding_layers:
+                # No embedding layer -> passthrough
+                embedded[field_name] = tensor
+                continue
+            
+            tensor = tensor.to(self.device)            
+            embedded[field_name] = self.embedding_layers[field_name](tensor)
+            
+            if output_mask:
+                # Generate a mask for this field
+                if hasattr(processor, "code_vocab"):
+                    pad_idx = processor.code_vocab.get("<pad>", 0)
+                else:
+                    pad_idx = 0
+                    
+                masks[field_name] = (tensor != pad_idx)
+        
+        if output_mask:
+            return embedded, masks
+        else:
+            return embedded
 
     def __repr__(self) -> str:
-        """
-        Returns a string representation of the EmbeddingModel.
-
-        Returns:
-            str: A string representation of the model.
-        """
         return f"EmbeddingModel(embedding_layers={self.embedding_layers})"

examples/mortality_mimic3_adacare.py

@@ -30,6 +30,10 @@ def get_processor(name: str):
     NestedFloatsProcessor,
     NestedSequenceProcessor,
 )
+from .deep_nested_sequence_processor import (
+    DeepNestedFloatsProcessor,
+    DeepNestedSequenceProcessor,
+)
 from .raw_processor import RawProcessor
 from .sequence_processor import SequenceProcessor
 from .signal_processor import SignalProcessor