from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F from backbones import ( AgeSinusoidalEncoding, GPTBlock, TimesNetExposureEncoder, ) from targets import PAD_IDX @dataclass class DeepHealthOutput: hidden: torch.Tensor time_seq: torch.Tensor padding_mask: torch.Tensor event_len: int class DeepHealth(nn.Module): def __init__( self, vocab_size: int, n_embd: int, n_head: int, n_hist_layer: int, target_mode: str = "next_token", # "next_token" or "all_future" dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" dropout: float = 0.0, use_exposure_encoder: bool = False, exposure_daily_input_dim: int = 4, exposure_monthly_input_dim: int = 2, exposure_d_model: int | None = None, exposure_n_layers: int = 2, exposure_top_k: int = 3, exposure_n_convnext_blocks: int = 2, exposure_conv_kernel_size: int = 7, exposure_mlp_ratio: float = 4.0, exposure_use_gate: bool = True, ): super().__init__() if target_mode not in ["next_token", "all_future"]: raise ValueError( "target_mode must be either 'next_token' or 'all_future'") if dist_mode not in ["exponential", "weibull", "mixed"]: raise ValueError( "dist_mode must be either 'exponential', 'weibull' or 'mixed'") self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0) self.gender_embedding = nn.Embedding( 2, n_embd) # Assuming binary gender self.target_mode = target_mode self.dist_mode = dist_mode self.use_exposure_encoder = use_exposure_encoder self.n_embd = n_embd self.vocab_size = vocab_size self.exposure_encoder = ( TimesNetExposureEncoder( n_embd=n_embd, daily_input_dim=exposure_daily_input_dim, monthly_input_dim=exposure_monthly_input_dim, d_model=exposure_d_model, n_layers=exposure_n_layers, top_k=exposure_top_k, n_convnext_blocks=exposure_n_convnext_blocks, conv_kernel_size=exposure_conv_kernel_size, mlp_ratio=exposure_mlp_ratio, dropout=dropout, use_gate=exposure_use_gate, ) if use_exposure_encoder else None ) nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) nn.init.zeros_(self.token_embedding.weight[0]) nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02) if dist_mode == "weibull": self.rho_head = nn.Linear(n_embd, vocab_size) nn.init.zeros_(self.rho_head.weight) nn.init.constant_(self.rho_head.bias, 0.5413) if dist_mode == "mixed": self.death_idx = vocab_size - 1 self.rho_death_head = nn.Linear(n_embd, 1) nn.init.zeros_(self.rho_death_head.weight) nn.init.constant_(self.rho_death_head.bias, 0.5413) self.age_encoding = AgeSinusoidalEncoding(n_embd) self.blocks = nn.ModuleList([ GPTBlock( n_embd=n_embd, n_head=n_head, mlp_dropout=dropout, ) for _ in range(n_hist_layer) ]) 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) def _make_history_attn_mask( self, padding_mask: torch.Tensor, time_seq: torch.Tensor, dtype: torch.dtype, ) -> torch.Tensor: valid_key = padding_mask[:, None, :] # (B, 1, L) visible_by_time = time_seq[:, None, :] <= time_seq[:, :, None] valid = valid_key & visible_by_time return torch.zeros( valid.shape, device=valid.device, dtype=dtype, ).masked_fill(~valid, -1e4)[:, None, :, :] def _encode_event_exposure( self, exposure_daily: torch.Tensor | None, exposure_monthly: torch.Tensor | None, exposure_daily_mask: torch.Tensor | None, exposure_monthly_mask: torch.Tensor | None, event_shape: tuple[int, int], ) -> torch.Tensor | None: if self.exposure_encoder is None: if exposure_daily is not None or exposure_monthly is not None: raise ValueError( "Exposure tensors were provided but use_exposure_encoder=False" ) return None if exposure_daily is None or exposure_monthly is None: raise ValueError( "exposure_daily and exposure_monthly are required when " "use_exposure_encoder=True" ) batch_size, event_len = event_shape if exposure_daily.shape[:2] != event_shape: raise ValueError( "exposure_daily must have shape (B, L, T_daily, C_daily), got " f"{tuple(exposure_daily.shape)} for event shape {event_shape}" ) if exposure_monthly.shape[:2] != event_shape: raise ValueError( "exposure_monthly must have shape (B, L, T_monthly, C_monthly), got " f"{tuple(exposure_monthly.shape)} for event shape {event_shape}" ) if exposure_daily.dim() != 4 or exposure_monthly.dim() != 4: raise ValueError( "exposure_daily and exposure_monthly must both be 4D tensors" ) def flatten_mask(mask: torch.Tensor | None, ref: torch.Tensor) -> torch.Tensor | None: if mask is None: return None if mask.shape[:2] != event_shape: raise ValueError( "exposure mask must start with shape (B, L), got " f"{tuple(mask.shape)}" ) if mask.dim() not in {3, 4}: raise ValueError( "exposure mask must have shape (B, L, T) or (B, L, T, C), " f"got {tuple(mask.shape)}" ) if mask.shape[2] != ref.shape[2]: raise ValueError( "exposure mask time dimension does not match exposure tensor" ) if mask.dim() == 4 and mask.shape[3] != ref.shape[3]: raise ValueError( "exposure mask channel dimension does not match exposure tensor" ) return mask.reshape(batch_size * event_len, *mask.shape[2:]) daily = exposure_daily.reshape( batch_size * event_len, exposure_daily.size(2), exposure_daily.size(3), ) monthly = exposure_monthly.reshape( batch_size * event_len, exposure_monthly.size(2), exposure_monthly.size(3), ) daily_mask = flatten_mask(exposure_daily_mask, exposure_daily) monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly) exposure_device = exposure_daily.device exposure_dtype = self.token_embedding.weight.dtype daily = daily.to(device=exposure_device, dtype=exposure_dtype) monthly = monthly.to(device=exposure_device, dtype=exposure_dtype) if daily_mask is not None: daily_mask = daily_mask.to(device=exposure_device) if monthly_mask is not None: monthly_mask = monthly_mask.to(device=exposure_device) exposure_emb = self.exposure_encoder( daily=daily, monthly=monthly, daily_mask=daily_mask, monthly_mask=monthly_mask, ) return exposure_emb.reshape(batch_size, event_len, self.n_embd) def _forward_shared( self, event_seq: torch.LongTensor, time_seq: torch.FloatTensor, sex: torch.LongTensor, mode: str, padding_mask: torch.Tensor | None = None, t_query: torch.FloatTensor | None = None, exposure_daily: torch.Tensor | None = None, exposure_monthly: torch.Tensor | None = None, exposure_daily_mask: torch.Tensor | None = None, exposure_monthly_mask: torch.Tensor | None = None, return_output: bool = False, **unused_kwargs, ) -> torch.Tensor | DeepHealthOutput: if unused_kwargs: unknown = ", ".join(sorted(unused_kwargs)) raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}") if mode not in {"next_token", "all_future"}: raise ValueError("mode must be either 'next_token' or 'all_future'") if mode == "all_future" and t_query is None: raise ValueError("t_query is required when mode='all_future'") if padding_mask is None: padding_mask = event_seq > PAD_IDX 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) h_exposure = self._encode_event_exposure( exposure_daily=exposure_daily, exposure_monthly=exposure_monthly, exposure_daily_mask=exposure_daily_mask, exposure_monthly_mask=exposure_monthly_mask, event_shape=(event_seq.size(0), event_len), ) if h_exposure is not None: h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype) h_disease = h_disease + h_exposure t_disease = time_seq 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([t_disease, t_query[:, None]], dim=1) query_mask = torch.ones( batch_size, 1, dtype=torch.bool, device=event_seq.device, ) padding_mask = torch.cat([padding_mask, query_mask], dim=1) sex_emb = self.gender_embedding(sex)[:, None, :] h_disease = h_disease + sex_emb h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) h_disease = h_disease + self.age_encoding(t_disease) h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) attn_mask = self._make_history_attn_mask( padding_mask=padding_mask, time_seq=t_disease, dtype=h_disease.dtype, ) for block in self.blocks: h_disease = block( h_disease, attn_mask=attn_mask, ) h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) h_disease = self.final_ln(h_disease) h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) if mode == "all_future": 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: h_event = h_disease[:, :event_len, :] t_event = t_disease[:, :event_len] event_mask = padding_mask[:, :event_len] return DeepHealthOutput( hidden=h_event, time_seq=t_event, padding_mask=event_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) def forward_all_future(self, **kwargs) -> torch.Tensor: return self._forward_shared(mode="all_future", **kwargs) def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor: mode = self.target_mode if target_mode is None else target_mode return self._forward_shared(mode=mode, **kwargs) def calc_risk(self, x: torch.Tensor) -> torch.Tensor: return self.risk_head(x) def calc_weibull_rho(self, x: torch.Tensor) -> torch.Tensor: if self.dist_mode != "weibull": raise RuntimeError( f"calc_weibull_rho called with dist_mode={self.dist_mode!r}" ) return F.softplus(self.rho_head(x)) + 1e-6 def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor: if self.dist_mode != "mixed": raise RuntimeError( f"calc_death_rho called with dist_mode={self.dist_mode!r}" ) return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6