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

24
backbones.py
View File

@@ -333,6 +333,7 @@ class BaselineEncoder(nn.Module):
)
self.n_embd = n_embd
self.cls_token = nn.Parameter(torch.zeros(1, 1, n_embd))
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 = (
@@ -376,6 +377,7 @@ class BaselineEncoder(nn.Module):
self.reset_parameters()
def reset_parameters(self) -> None:
nn.init.normal_(self.cls_token, mean=0.0, std=0.02)
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)
@@ -439,17 +441,27 @@ class BaselineEncoder(nn.Module):
f = type_emb + kind_emb + value_emb
f = f * other_valid.unsqueeze(-1).to(f.dtype)
if f.size(1) == 0:
return f, other_valid
cls = self.cls_token.expand(f.size(0), -1, -1)
f = torch.cat([cls, f], dim=1)
cls_valid = torch.ones(
other_valid.size(0),
1,
device=other_valid.device,
dtype=torch.bool,
)
full_valid = torch.cat([cls_valid, other_valid], dim=1)
attn_mask = self._make_attn_mask(other_valid, f.dtype)
attn_mask = self._make_attn_mask(full_valid, f.dtype)
for block in self.blocks:
f = block(f, attn_mask=attn_mask)
f = f * other_valid.unsqueeze(-1).to(f.dtype)
f = f * full_valid.unsqueeze(-1).to(f.dtype)
h = self.ln(f)
h = h * other_valid.unsqueeze(-1).to(h.dtype)
return h, other_valid
h = h * full_valid.unsqueeze(-1).to(h.dtype)
cls_summary = h[:, 0, :]
token_h = h[:, 1:, :]
token_h = token_h * other_valid.unsqueeze(-1).to(token_h.dtype)
return token_h, other_valid, cls_summary
class CrossAttention(nn.Module):