Refactor DeepHealth model and related components

- Removed BaselineEncoder and CrossAttention classes from models.py.
- Introduced OtherInfoTokenizer for handling additional token types.
- Updated DeepHealth class to integrate OtherInfoTokenizer and manage extra pooling logic.
- Added support for extra_pool_reduce parameter to control pooling behavior.
- Modified forward methods to return structured output using DeepHealthOutput dataclass.
- Updated training scripts to accommodate changes in model architecture and output handling.
- Enhanced error handling and validation for input shapes and types.

models.py

@@ -1,16 +1,144 @@
+from dataclasses import dataclass
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 from backbones import (
     AgeSinusoidalEncoding,
-    BaselineEncoder,
-    CrossAttention,
     GPTBlock,
     GaussianRBFTimeBasis,
     TimeRoPE,
+    TokenAutoDiscretization,
 )
-from targets import CHECKUP_IDX, PAD_IDX
+from targets import PAD_IDX
+
+
+@dataclass
+class DeepHealthOutput:
+    hidden: torch.Tensor
+    time_seq: torch.Tensor
+    padding_mask: torch.Tensor
+    event_len: int
+
+
+class OtherInfoTokenizer(nn.Module):
+    PAD_KIND = 0
+    CONT_KIND = 1
+    CATE_KIND = 2
+
+    def __init__(
+        self,
+        n_embd: int,
+        n_types: int,
+        n_cont_types: int,
+        n_categories: int,
+        cont_type_ids: list[int],
+        n_value_kinds: int = 3,
+        n_bins: int = 16,
+    ):
+        super().__init__()
+        if len(cont_type_ids) != n_cont_types:
+            raise ValueError(
+                "cont_type_ids length must match n_cont_types, got "
+                f"{len(cont_type_ids)} vs {n_cont_types}"
+            )
+        if n_types <= 0:
+            raise ValueError(f"n_types must include PAD and be > 0, got {n_types}")
+        if n_categories <= 0:
+            raise ValueError(
+                f"n_categories must include PAD and be > 0, got {n_categories}"
+            )
+        if n_value_kinds <= self.CATE_KIND:
+            raise ValueError(
+                f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
+            )
+
+        self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
+        self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
+        self.cont_value_encoder = (
+            TokenAutoDiscretization(
+                n_cont_types=n_cont_types,
+                n_bins=n_bins,
+                n_embd=n_embd,
+            )
+            if n_cont_types > 0
+            else None
+        )
+        self.cate_value_emb = nn.Embedding(
+            n_categories,
+            n_embd,
+            padding_idx=0,
+        )
+
+        cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
+        for idx, type_id in enumerate(cont_type_ids):
+            if type_id <= 0 or type_id >= n_types:
+                raise ValueError(
+                    f"continuous type id {type_id} must be in [1, {n_types})"
+                )
+            cont_type_index[type_id] = idx
+        self.register_buffer(
+            "cont_type_index",
+            cont_type_index,
+            persistent=False,
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.type_emb.weight[0])
+        nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.kind_emb.weight[0])
+        nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.cate_value_emb.weight[0])
+
+    def forward(
+        self,
+        other_type: torch.LongTensor,
+        other_value: torch.Tensor,
+        other_value_kind: torch.LongTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if other_type.shape != other_value.shape:
+            raise ValueError(
+                "other_type and other_value must have the same shape, got "
+                f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}"
+            )
+        if other_type.shape != other_value_kind.shape:
+            raise ValueError(
+                "other_type and other_value_kind must have the same shape, got "
+                f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}"
+            )
+
+        other_valid = other_type > 0
+        type_emb = self.type_emb(other_type)
+        kind_emb = self.kind_emb(other_value_kind)
+        value_emb = torch.zeros_like(type_emb)
+
+        cont_pos = other_valid & (other_value_kind == self.CONT_KIND)
+        if cont_pos.any():
+            if self.cont_value_encoder is None:
+                raise ValueError("continuous tokens found but n_cont_types is 0")
+            cont_idx = self.cont_type_index[other_type[cont_pos]]
+            if (cont_idx < 0).any():
+                bad_type = other_type[cont_pos][cont_idx < 0][0].item()
+                raise ValueError(
+                    f"type_id={bad_type} is marked continuous but is not in "
+                    "cont_type_ids"
+                )
+            value_emb[cont_pos] = self.cont_value_encoder(
+                cont_type_idx=cont_idx,
+                value=other_value[cont_pos].to(type_emb.dtype),
+            )
+
+        cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
+        if cate_pos.any():
+            cate_id = other_value[cate_pos].long()
+            value_emb[cate_pos] = self.cate_value_emb(cate_id)
+
+        out = type_emb + kind_emb + value_emb
+        out = out * other_valid.unsqueeze(-1).to(out.dtype)
+        return out, other_valid
 
 
 class DeepHealth(nn.Module):
@@ -30,6 +158,7 @@ class DeepHealth(nn.Module):
         target_mode: str = "next_token",  # "next_token" or "all_future"
         time_mode: str = "relative",  # "relative" or "absolute"
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
+        extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
     ):
         super().__init__()
@@ -42,30 +171,24 @@ class DeepHealth(nn.Module):
         if dist_mode not in ["exponential", "weibull", "mixed"]:
             raise ValueError(
                 "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
+        if extra_pool_reduce not in {"mean", "sum"}:
+            raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
         self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
         self.gender_embedding = nn.Embedding(
             2, n_embd)  # Assuming binary gender
-        self.token_encoder = BaselineEncoder(
+        self.tokenizer = OtherInfoTokenizer(
             n_embd=n_embd,
-            n_head=n_head,
             n_types=n_types,
             n_cont_types=n_cont_types,
             n_categories=n_categories,
             cont_type_ids=cont_type_ids,
             n_value_kinds=n_value_kinds,
             n_bins=n_bins,
-            n_tab_layer=n_tab_layer,
-            dropout=dropout,
-        )
-        self.cross_attention = CrossAttention(
-            n_embd=n_embd,
-            n_head=n_head,
-            dropout=dropout,
-            n_rbf_bases=16,
         )
         self.target_mode = target_mode
         self.time_mode = time_mode
         self.dist_mode = dist_mode
+        self.extra_pool_reduce = extra_pool_reduce
         self.n_embd = n_embd
         self.vocab_size = vocab_size
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
@@ -111,6 +234,8 @@ class DeepHealth(nn.Module):
 
         self.final_ln = nn.LayerNorm(n_embd)
         self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
+        if target_mode == "next_token":
+            self.risk_head.weight = self.token_embedding.weight
         self.query_token = nn.Parameter(torch.zeros(n_embd))
         nn.init.normal_(self.query_token, mean=0.0, std=0.02)
 
@@ -129,55 +254,59 @@ class DeepHealth(nn.Module):
             dtype=dtype,
         ).masked_fill(~valid, -1e4)[:, None, :, :]
 
-    def _insert_baseline_summary(
+    def _pool_other_by_time(
         self,
-        h_disease: torch.Tensor,
-        event_seq: torch.Tensor,
-        baseline_summary: torch.Tensor,
-    ) -> torch.Tensor:
-        checkup_mask = event_seq == CHECKUP_IDX
-        if not checkup_mask.any():
-            return h_disease
-        summary = baseline_summary.to(device=h_disease.device, dtype=h_disease.dtype)
-        return torch.where(checkup_mask.unsqueeze(-1), summary[:, None, :], h_disease)
-
-    def _baseline_cls_time(
-        self,
-        event_seq: torch.Tensor,
-        time_seq: torch.Tensor,
-        padding_mask: torch.Tensor,
-    ) -> torch.Tensor:
-        checkup_mask = event_seq == CHECKUP_IDX
-        inf = torch.full_like(time_seq, float("inf"))
-        first_checkup = torch.where(checkup_mask, time_seq, inf).min(dim=1).values
-        has_checkup = torch.isfinite(first_checkup)
-        fallback_time = torch.where(
-            padding_mask,
-            time_seq,
-            torch.full_like(time_seq, float("-inf")),
-        ).max(dim=1).values
-        fallback_time = torch.where(
-            torch.isfinite(fallback_time),
-            fallback_time,
-            torch.zeros_like(fallback_time),
-        )
-        return torch.where(has_checkup, first_checkup, fallback_time)
-
-    def _encode_other_tokens(
-        self,
-        other_type: torch.LongTensor,
-        other_value: torch.Tensor,
-        other_value_kind: torch.LongTensor,
+        h_other: torch.Tensor,
         other_time: torch.Tensor,
-        cls_time: torch.Tensor,
+        other_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        return self.token_encoder(
-            other_type=other_type,
-            other_value=other_value,
-            other_value_kind=other_value_kind,
-            other_time=other_time,
-            cls_time=cls_time,
+        batch_size, _n_other, n_embd = h_other.shape
+        group_counts = [
+            int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel())
+            for b in range(batch_size)
+        ]
+        max_groups = max(group_counts, default=0)
+        pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
+        pooled_time = other_time.new_zeros(batch_size, max_groups)
+        pooled_mask = torch.zeros(
+            batch_size,
+            max_groups,
+            dtype=torch.bool,
+            device=h_other.device,
         )
+        if max_groups == 0:
+            return pooled_h, pooled_time, pooled_mask
+
+        for b in range(batch_size):
+            valid_time = other_time[b, other_mask[b]]
+            if valid_time.numel() == 0:
+                continue
+            valid_h = h_other[b, other_mask[b]]
+            unique_time, inverse = torch.unique(
+                valid_time,
+                sorted=True,
+                return_inverse=True,
+            )
+            n_groups = unique_time.numel()
+            group_h = valid_h.new_zeros(n_groups, n_embd)
+            group_h.scatter_add_(
+                0,
+                inverse[:, None].expand(-1, n_embd),
+                valid_h,
+            )
+            if self.extra_pool_reduce == "mean":
+                counts = valid_h.new_zeros(n_groups, 1)
+                counts.scatter_add_(
+                    0,
+                    inverse[:, None],
+                    torch.ones_like(valid_h[:, :1]),
+                )
+                group_h = group_h / counts.clamp_min(1.0)
+
+            pooled_h[b, :n_groups] = group_h
+            pooled_time[b, :n_groups] = unique_time
+            pooled_mask[b, :n_groups] = True
+        return pooled_h, pooled_time, pooled_mask
 
     def _forward_shared(
         self,
@@ -191,6 +320,7 @@ class DeepHealth(nn.Module):
         other_value: torch.Tensor | None = None,
         other_value_kind: torch.LongTensor | None = None,
         other_time: torch.FloatTensor | None = None,
+        return_output: bool = False,
         **unused_kwargs,
     ) -> torch.Tensor:
         if unused_kwargs:
@@ -216,6 +346,7 @@ class DeepHealth(nn.Module):
         else:
             padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
 
+        event_len = event_seq.size(1)
         h_disease = self.token_embedding(event_seq)
         t_disease = time_seq
 
@@ -225,40 +356,29 @@ class DeepHealth(nn.Module):
                 f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
             )
         other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
-        cls_time = self._baseline_cls_time(
-            event_seq=event_seq,
-            time_seq=time_seq,
-            padding_mask=padding_mask,
-        )
-        h_token, token_mask, baseline_summary = self._encode_other_tokens(
+        h_other, other_mask = self.tokenizer(
             other_type=other_type,
             other_value=other_value,
             other_value_kind=other_value_kind,
+        )
+        h_other = h_other.to(device=event_seq.device)
+        other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
+        h_other, other_time, other_mask = self._pool_other_by_time(
+            h_other=h_other,
             other_time=other_time,
-            cls_time=cls_time,
+            other_mask=other_mask,
         )
-        token_time = other_time.to(device=h_token.device, dtype=time_seq.dtype)
 
-        h_disease = self.cross_attention(
-            h_disease=h_disease,
-            t_disease=t_disease,
-            h_token=h_token,
-            t_token=token_time,
-            token_mask=token_mask,
-        )
-        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-        h_disease = self._insert_baseline_summary(
-            h_disease=h_disease,
-            event_seq=event_seq,
-            baseline_summary=baseline_summary,
-        )
+        h_disease = torch.cat([h_disease, h_other], dim=1)
+        t_disease = torch.cat([t_disease, other_time], dim=1)
+        padding_mask = torch.cat([padding_mask, other_mask], dim=1)
         h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 
         if mode == "all_future":
             batch_size = event_seq.size(0)
             query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
             h_disease = torch.cat([h_disease, query], dim=1)
-            t_disease = torch.cat([time_seq, t_query[:, None]], dim=1)
+            t_disease = torch.cat([t_disease, t_query[:, None]], dim=1)
             query_mask = torch.ones(
                 batch_size,
                 1,
@@ -298,8 +418,28 @@ class DeepHealth(nn.Module):
         h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 
         if mode == "all_future":
-            return h_disease[:, -1, :]
-        return h_disease
+            hidden = h_disease[:, -1, :]
+            if return_output:
+                return DeepHealthOutput(
+                    hidden=hidden,
+                    time_seq=t_query[:, None],
+                    padding_mask=torch.ones(
+                        hidden.size(0),
+                        1,
+                        dtype=torch.bool,
+                        device=hidden.device,
+                    ),
+                    event_len=event_len,
+                )
+            return hidden
+        if return_output:
+            return DeepHealthOutput(
+                hidden=h_disease,
+                time_seq=t_disease,
+                padding_mask=padding_mask,
+                event_len=event_len,
+            )
+        return h_disease[:, :event_len, :]
 
     def forward_next_token(self, **kwargs) -> torch.Tensor:
         return self._forward_shared(mode="next_token", **kwargs)