JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove extra info token pathway

Browse Source
This commit is contained in:
Jiarui Li
2026-07-07 16:57:49 +08:00
parent 6dfeb5a696
commit a0379daf29
13 changed files with 18 additions and 1390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

246
models.py
View File

@@ -10,7 +10,6 @@ from backbones import (
GaussianRBFTimeBasis,
TimesNetExposureEncoder,
TimeRoPE,
TokenAutoDiscretization,
)
from targets import PAD_IDX
@@ -23,125 +22,6 @@ class DeepHealthOutput:
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):
def __init__(
self,
@@ -149,17 +29,9 @@ class DeepHealth(nn.Module):
n_embd: int,
n_head: int,
n_hist_layer: int,
n_tab_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,
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,
use_exposure_encoder: bool = False,
exposure_daily_input_dim: int = 4,
@@ -182,24 +54,12 @@ 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.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,
)
self.target_mode = target_mode
self.time_mode = time_mode
self.dist_mode = dist_mode
self.extra_pool_reduce = extra_pool_reduce
self.use_exposure_encoder = use_exposure_encoder
self.n_embd = n_embd
self.vocab_size = vocab_size
@@ -283,69 +143,6 @@ class DeepHealth(nn.Module):
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 _encode_event_exposure(
self,
exposure_daily: torch.Tensor | None,
@@ -442,10 +239,6 @@ class DeepHealth(nn.Module):
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,
exposure_daily: torch.Tensor | None = None,
exposure_monthly: torch.Tensor | None = None,
exposure_daily_mask: torch.Tensor | None = None,
@@ -460,16 +253,6 @@ class DeepHealth(nn.Module):
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
@@ -489,24 +272,6 @@ class DeepHealth(nn.Module):
h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype)
h_disease = h_disease + h_exposure
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":
@@ -571,15 +336,10 @@ class DeepHealth(nn.Module):
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),
hidden=h_event,
time_seq=t_event,
padding_mask=event_mask,
event_len=event_len,
)
return h_disease[:, :event_len, :]