537 lines
21 KiB
Python
537 lines
21 KiB
Python
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from backbones import (
|
|
AgeSinusoidalEncoding,
|
|
GaussianRBFTimeBasis,
|
|
TimeRoPE,
|
|
TokenAutoDiscretization,
|
|
build_backbone_block,
|
|
)
|
|
from model_architectures import resolve_model_architecture
|
|
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,
|
|
continuous_value_center: torch.Tensor | list[float] | None = None,
|
|
continuous_value_scale: torch.Tensor | list[float] | None = None,
|
|
):
|
|
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,
|
|
)
|
|
if n_cont_types > 0:
|
|
if continuous_value_center is None or continuous_value_scale is None:
|
|
raise ValueError(
|
|
"Continuous values require train-split RobustScale center "
|
|
"and scale statistics"
|
|
)
|
|
center = self._coerce_scaler_buffer(
|
|
continuous_value_center,
|
|
n_cont_types=n_cont_types,
|
|
name="continuous_value_center",
|
|
)
|
|
scale = self._coerce_scaler_buffer(
|
|
continuous_value_scale,
|
|
n_cont_types=n_cont_types,
|
|
name="continuous_value_scale",
|
|
)
|
|
if not torch.isfinite(center).all():
|
|
raise ValueError(
|
|
"continuous_value_center must contain only finite values"
|
|
)
|
|
if not torch.isfinite(scale).all() or torch.any(scale <= 0):
|
|
raise ValueError(
|
|
"continuous_value_scale must be finite and strictly positive"
|
|
)
|
|
self.register_buffer("continuous_value_center", center)
|
|
self.register_buffer("continuous_value_scale", scale)
|
|
else:
|
|
# ``None`` buffers are omitted from state_dict. This preserves the
|
|
# exact checkpoint schema used by models trained before continuous
|
|
# value scaling was introduced.
|
|
self.register_buffer("continuous_value_center", None)
|
|
self.register_buffer("continuous_value_scale", None)
|
|
|
|
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()
|
|
|
|
@staticmethod
|
|
def _coerce_scaler_buffer(
|
|
value: torch.Tensor | list[float] | None,
|
|
*,
|
|
n_cont_types: int,
|
|
name: str,
|
|
) -> torch.Tensor:
|
|
if value is None:
|
|
raise ValueError(f"{name} is required")
|
|
tensor = torch.as_tensor(value, dtype=torch.float32).detach().clone()
|
|
if tensor.shape != (n_cont_types,):
|
|
raise ValueError(
|
|
f"{name} must have shape ({n_cont_types},), got {tuple(tensor.shape)}"
|
|
)
|
|
return tensor
|
|
|
|
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"
|
|
)
|
|
cont_value = other_value[cont_pos].to(type_emb.dtype)
|
|
if (
|
|
self.continuous_value_center is None
|
|
or self.continuous_value_scale is None
|
|
):
|
|
raise RuntimeError("RobustScale buffers are missing")
|
|
center = self.continuous_value_center[cont_idx].to(type_emb.dtype)
|
|
scale = self.continuous_value_scale[cont_idx].to(type_emb.dtype)
|
|
cont_value = (cont_value - center) / scale
|
|
value_emb[cont_pos] = self.cont_value_encoder(
|
|
cont_type_idx=cont_idx,
|
|
value=cont_value,
|
|
)
|
|
|
|
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):
|
|
def __init__(
|
|
self,
|
|
vocab_size: int,
|
|
n_embd: int,
|
|
n_head: int,
|
|
n_layer: 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,
|
|
continuous_value_center: torch.Tensor | list[float] | None = None,
|
|
continuous_value_scale: torch.Tensor | list[float] | None = None,
|
|
target_mode: str = "next_token", # "next_token" or "all_future"
|
|
time_mode: str = "absolute", # next_token requires absolute
|
|
dist_mode: str = "exponential", # "exponential" or "weibull"
|
|
extra_pool_reduce: str = "mean",
|
|
dropout: float = 0.0,
|
|
model_architecture: str | None = None,
|
|
):
|
|
super().__init__()
|
|
if target_mode not in ["next_token", "all_future"]:
|
|
raise ValueError(
|
|
"target_mode must be either 'next_token' or 'all_future'")
|
|
if time_mode not in ["relative", "absolute"]:
|
|
raise ValueError(
|
|
"time_mode must be either 'relative' or 'absolute'")
|
|
if target_mode == "next_token" and time_mode != "absolute":
|
|
raise ValueError(
|
|
"next_token is reserved for Delphi2M reproduction and "
|
|
"requires time_mode='absolute'"
|
|
)
|
|
if dist_mode not in ["exponential", "weibull"]:
|
|
raise ValueError(
|
|
"dist_mode must be either 'exponential' or 'weibull'")
|
|
if extra_pool_reduce not in {"mean", "sum"}:
|
|
raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
|
|
if n_layer < 1:
|
|
raise ValueError(f"n_layer must be >= 1, got {n_layer}")
|
|
model_architecture = resolve_model_architecture(model_architecture)
|
|
self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
|
|
self.gender_embedding = nn.Embedding(
|
|
2, n_embd) # Assuming binary gender
|
|
self.tokenizer = OtherInfoTokenizer(
|
|
n_embd=n_embd,
|
|
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,
|
|
continuous_value_center=continuous_value_center,
|
|
continuous_value_scale=continuous_value_scale,
|
|
)
|
|
self.target_mode = target_mode
|
|
self.time_mode = time_mode
|
|
self.dist_mode = dist_mode
|
|
self.extra_pool_reduce = extra_pool_reduce
|
|
self.model_architecture = model_architecture
|
|
self.n_layer = n_layer
|
|
self.n_embd = n_embd
|
|
self.vocab_size = vocab_size
|
|
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 time_mode == "absolute":
|
|
self.age_encoding = AgeSinusoidalEncoding(n_embd)
|
|
self.blocks = nn.ModuleList([
|
|
build_backbone_block(
|
|
model_architecture,
|
|
n_embd=n_embd,
|
|
n_head=n_head,
|
|
use_time_rope=False,
|
|
use_rbf_bias=False,
|
|
mlp_dropout=dropout,
|
|
) for _ in range(n_layer)
|
|
])
|
|
self.rope = None
|
|
self.rbf = None
|
|
elif time_mode == "relative":
|
|
self.age_encoding = None
|
|
self.blocks = nn.ModuleList([
|
|
build_backbone_block(
|
|
model_architecture,
|
|
n_embd=n_embd,
|
|
n_head=n_head,
|
|
use_time_rope=True,
|
|
use_rbf_bias=True,
|
|
mlp_dropout=dropout,
|
|
) for _ in range(n_layer)
|
|
])
|
|
self.rope = TimeRoPE(n_embd // n_head)
|
|
self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
|
|
|
|
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 _pool_other_by_time(
|
|
self,
|
|
h_other: torch.Tensor,
|
|
other_time: torch.Tensor,
|
|
other_mask: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
batch_size, n_other, n_embd = h_other.shape
|
|
if n_other == 0:
|
|
empty_h = h_other.new_zeros(batch_size, 0, n_embd)
|
|
empty_t = other_time.new_zeros(batch_size, 0)
|
|
empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device)
|
|
return empty_h, empty_t, empty_m
|
|
|
|
masked_time = other_time.masked_fill(~other_mask, float("inf"))
|
|
_sorted_time_with_pad, order = masked_time.sort(dim=1)
|
|
sorted_time = other_time.gather(1, order)
|
|
sorted_mask = other_mask.gather(1, order)
|
|
sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd))
|
|
|
|
group_start = torch.zeros_like(sorted_mask)
|
|
group_start[:, 0] = sorted_mask[:, 0]
|
|
group_start[:, 1:] = sorted_mask[:, 1:] & (
|
|
sorted_time[:, 1:] != sorted_time[:, :-1]
|
|
)
|
|
group_id = group_start.long().cumsum(dim=1) - 1
|
|
max_groups = int(group_start.sum(dim=1).max().item())
|
|
|
|
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
|
|
|
|
safe_group_id = group_id.clamp_min(0)
|
|
pooled_h.scatter_add_(
|
|
1,
|
|
safe_group_id.unsqueeze(-1).expand_as(sorted_h),
|
|
sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype),
|
|
)
|
|
if self.extra_pool_reduce == "mean":
|
|
counts = h_other.new_zeros(batch_size, max_groups, 1)
|
|
counts.scatter_add_(
|
|
1,
|
|
safe_group_id.unsqueeze(-1),
|
|
sorted_mask.unsqueeze(-1).to(h_other.dtype),
|
|
)
|
|
pooled_h = pooled_h / counts.clamp_min(1.0)
|
|
|
|
pooled_time.scatter_add_(
|
|
1,
|
|
safe_group_id,
|
|
sorted_time * group_start.to(sorted_time.dtype),
|
|
)
|
|
group_count = group_start.sum(dim=1)
|
|
arange_groups = torch.arange(max_groups, device=h_other.device)
|
|
pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
|
|
return pooled_h, pooled_time, pooled_mask
|
|
|
|
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,
|
|
other_type: torch.LongTensor | None = None,
|
|
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 | 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 (
|
|
other_type is None
|
|
or other_value is None
|
|
or other_value_kind is None
|
|
or other_time is None
|
|
):
|
|
raise ValueError(
|
|
"DeepHealth expects other_type, other_value, "
|
|
"other_value_kind, and other_time."
|
|
)
|
|
|
|
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)
|
|
t_disease = time_seq
|
|
|
|
if other_time.shape != other_type.shape:
|
|
raise ValueError(
|
|
"other_time must have the same shape as other_type, got "
|
|
f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
|
|
)
|
|
other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
|
|
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_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([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)
|
|
|
|
rope_cache = None
|
|
rbf_cache = None
|
|
if self.time_mode == "absolute":
|
|
h_disease = h_disease + self.age_encoding(t_disease)
|
|
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
|
elif self.time_mode == "relative":
|
|
rope_cache = self.rope.precompute_cache(t_disease)
|
|
rbf_cache = self.rbf.precompute_cache(t_disease)
|
|
|
|
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,
|
|
rope_cache=rope_cache,
|
|
rbf_cache=rbf_cache,
|
|
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]
|
|
h_extra, t_extra, extra_mask = self._pool_other_by_time(
|
|
h_other=h_disease[:, event_len:, :],
|
|
other_time=t_disease[:, event_len:],
|
|
other_mask=padding_mask[:, event_len:],
|
|
)
|
|
return DeepHealthOutput(
|
|
hidden=torch.cat([h_event, h_extra], dim=1),
|
|
time_seq=torch.cat([t_event, t_extra], dim=1),
|
|
padding_mask=torch.cat([event_mask, extra_mask], dim=1),
|
|
event_len=event_len,
|
|
)
|
|
return h_disease[:, :event_len, :]
|
|
|
|
def forward(self, **kwargs) -> torch.Tensor:
|
|
return self._forward_shared(mode=self.target_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
|