349 lines
14 KiB
Python
349 lines
14 KiB
Python
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 = 64,
|
|
exposure_n_layers: int = 2,
|
|
exposure_top_k: int = 2,
|
|
exposure_n_backbone_blocks: int = 1,
|
|
exposure_backbone_kernel_size: int = 5,
|
|
exposure_backbone_expansion: float = 2.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_backbone_blocks=exposure_n_backbone_blocks,
|
|
backbone_kernel_size=exposure_backbone_kernel_size,
|
|
backbone_expansion=exposure_backbone_expansion,
|
|
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
|
|
if target_mode == "all_future":
|
|
self.query_token = nn.Parameter(torch.zeros(n_embd))
|
|
nn.init.normal_(self.query_token, mean=0.0, std=0.02)
|
|
else:
|
|
self.register_parameter("query_token", None)
|
|
|
|
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":
|
|
if self.query_token is None:
|
|
raise RuntimeError(
|
|
"all_future forward requires a model initialized with "
|
|
"target_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
|