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

Enhance DeepHealth model to incorporate CHECKUP state tokens in next-step training and evaluation, update dataset cache versioning, and improve handling of observed event histories.

Browse Source
This commit is contained in:
Jiarui Li
2026-06-15 14:10:09 +08:00
parent 593ecd2e71
commit c3e49db859
8 changed files with 111 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
models.py
View File

@@ -10,6 +10,7 @@ from backbones import (
GaussianRBFTimeBasis,
TimeRoPE,
)
from targets import CHECKUP_IDX, PAD_IDX
class DeepHealth(nn.Module):
@@ -128,12 +129,24 @@ class DeepHealth(nn.Module):
dtype=dtype,
).masked_fill(~valid, -1e4)[:, None, :, :]
def _insert_baseline_summary(
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 _encode_other_tokens(
self,
other_type: torch.LongTensor,
other_value: torch.Tensor,
other_value_kind: torch.LongTensor,
) -> tuple[torch.Tensor, torch.Tensor]:
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return self.token_encoder(
other_type=other_type,
other_value=other_value,
@@ -173,12 +186,40 @@ class DeepHealth(nn.Module):
)
if padding_mask is None:
padding_mask = event_seq > 0
padding_mask = event_seq > PAD_IDX
else:
padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
h_disease = self.token_embedding(event_seq)
t_disease = time_seq
h_token, token_mask, baseline_summary = self._encode_other_tokens(
other_type=other_type,
other_value=other_value,
other_value_kind=other_value_kind,
)
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)}"
)
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 = 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)
@@ -222,27 +263,6 @@ class DeepHealth(nn.Module):
h_disease = self.final_ln(h_disease)
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
h_token, token_mask = self._encode_other_tokens(
other_type=other_type,
other_value=other_value,
other_value_kind=other_value_kind,
)
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)}"
)
token_time = other_time.to(device=h_token.device, dtype=t_disease.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)
if mode == "all_future":
return h_disease[:, -1, :]
return h_disease