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

Add exposure cache and keep absolute time only

Browse Source
This commit is contained in:
Jiarui Li
2026-07-07 17:21:52 +08:00
parent a0379daf29
commit 45a857d1a6
9 changed files with 690 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
models.py
View File

@@ -7,9 +7,7 @@ import torch.nn.functional as F
from backbones import (
AgeSinusoidalEncoding,
GPTBlock,
GaussianRBFTimeBasis,
TimesNetExposureEncoder,
TimeRoPE,
)
from targets import PAD_IDX
@@ -30,7 +28,6 @@ class DeepHealth(nn.Module):
n_head: int,
n_hist_layer: int,
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"
dropout: float = 0.0,
use_exposure_encoder: bool = False,
@@ -48,9 +45,6 @@ class DeepHealth(nn.Module):
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 dist_mode not in ["exponential", "weibull", "mixed"]:
raise ValueError(
"dist_mode must be either 'exponential', 'weibull' or 'mixed'")
@@ -58,7 +52,6 @@ class DeepHealth(nn.Module):
self.gender_embedding = nn.Embedding(
2, n_embd) # Assuming binary gender
self.target_mode = target_mode
self.time_mode = time_mode
self.dist_mode = dist_mode
self.use_exposure_encoder = use_exposure_encoder
self.n_embd = n_embd
@@ -94,32 +87,14 @@ class DeepHealth(nn.Module):
nn.init.zeros_(self.rho_death_head.weight)
nn.init.constant_(self.rho_death_head.bias, 0.5413)
if time_mode == "absolute":
self.age_encoding = AgeSinusoidalEncoding(n_embd)
self.blocks = nn.ModuleList([
GPTBlock(
n_embd=n_embd,
n_head=n_head,
use_time_rope=False,
use_rbf_bias=False,
mlp_dropout=dropout,
) for _ in range(n_hist_layer)
])
self.rope = None
self.rbf = None
elif time_mode == "relative":
self.age_encoding = None
self.blocks = nn.ModuleList([
GPTBlock(
n_embd=n_embd,
n_head=n_head,
use_time_rope=True,
use_rbf_bias=True,
mlp_dropout=dropout,
) for _ in range(n_hist_layer)
])
self.rope = TimeRoPE(n_embd // n_head)
self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
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)
@@ -291,14 +266,8 @@ class DeepHealth(nn.Module):
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)
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,
@@ -308,8 +277,6 @@ class DeepHealth(nn.Module):
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)