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.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.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.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
self.cont_value_encoder = ( self.cont_value_encoder = (
@@ -376,6 +377,7 @@ class BaselineEncoder(nn.Module):
self.reset_parameters() self.reset_parameters()
def reset_parameters(self) -> None: 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.normal_(self.type_emb.weight, mean=0.0, std=0.02)
nn.init.zeros_(self.type_emb.weight[0]) nn.init.zeros_(self.type_emb.weight[0])
nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02) 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 = type_emb + kind_emb + value_emb
f = f * other_valid.unsqueeze(-1).to(f.dtype) f = f * other_valid.unsqueeze(-1).to(f.dtype)
if f.size(1) == 0: cls = self.cls_token.expand(f.size(0), -1, -1)
return f, other_valid 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: for block in self.blocks:
f = block(f, attn_mask=attn_mask) 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 = self.ln(f)
h = h * other_valid.unsqueeze(-1).to(h.dtype) h = h * full_valid.unsqueeze(-1).to(h.dtype)
return h, other_valid 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): class CrossAttention(nn.Module):

16
dataset.py
View File

@@ -110,7 +110,7 @@ def _cache_file_path(
selected.append(type_id) selected.append(type_id)
selected_types = ",".join(str(t) for t in selected) selected_types = ",".join(str(t) for t in selected)
signature_parts = [ signature_parts = [
"deephealthnew_dataset_cache_v2", "deephealthnew_dataset_cache_v3_checkup_state",
dataset_kind, dataset_kind,
split or "", split or "",
event_path, event_path,
@@ -320,9 +320,6 @@ class _ExpoBaseDataset(Dataset):
times_days_raw = rows[:, 1].astype(np.float32) times_days_raw = rows[:, 1].astype(np.float32)
labels_raw = rows[:, 2].astype(np.int64) labels_raw = rows[:, 2].astype(np.int64)
disease_mask = labels_raw != CHECKUP_IDX
times_days_raw = times_days_raw[disease_mask]
labels_raw = labels_raw[disease_mask]
if len(labels_raw) == 0: if len(labels_raw) == 0:
yield eid, times_days_raw, labels_raw yield eid, times_days_raw, labels_raw
continue continue
@@ -392,7 +389,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
- UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot
""" """
CACHE_VERSION = 2 CACHE_VERSION = 3
def __init__( def __init__(
self, self,
@@ -488,7 +485,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
time range, with at least one future event after every query. time range, with at least one future event after every query.
""" """
CACHE_VERSION = 4 CACHE_VERSION = 5
def __init__( def __init__(
self, self,
@@ -582,8 +579,13 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
def _is_valid_query(self, patient: Dict, t_query: float) -> bool: def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
times = patient["times"] times = patient["times"]
labels = patient["labels"]
real_event_mask = ~np.isin(
labels,
np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
)
n_hist = int((times <= t_query).sum()) n_hist = int((times <= t_query).sum())
n_future = int((times > t_query).sum()) n_future = int(((times > t_query) & real_event_mask).sum())
return ( return (
n_hist >= self.min_history_events n_hist >= self.min_history_events
and n_future >= self.min_future_events and n_future >= self.min_future_events

6
eval_data.py
View File

@@ -14,9 +14,9 @@ class AllFutureSequenceEvalDataset:
""" """
Eval-only sequence view for all-future checkpoints. Eval-only sequence view for all-future checkpoints.
All-future training uses pure disease histories, so token-level and landmark All-future training uses the observed history, including CHECKUP state
evaluation should not reuse the next-step dataset view that contains tokens, without reusing the next-step view that contains imputed
imputed <NO_EVENT> gap tokens. <NO_EVENT> gap tokens.
""" """
def __init__( def __init__(

5
evaluate_auc.py
View File

@@ -386,10 +386,7 @@ def load_model_state(
state = state_dict if state_dict is not None else load_checkpoint_state_dict( state = state_dict if state_dict is not None else load_checkpoint_state_dict(
checkpoint_path, map_location=device) checkpoint_path, map_location=device)
missing, unexpected = model.load_state_dict(state, strict=False) model.load_state_dict(state, strict=True)
if missing or unexpected:
print(
f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}")
def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]: def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]:

9
evaluate_auc_v2.py
View File

@@ -192,10 +192,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None: def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
missing, unexpected = model.load_state_dict(state_dict, strict=False) model.load_state_dict(state_dict, strict=True)
if missing or unexpected:
print(
f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}")
def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None: def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
@@ -1495,8 +1492,8 @@ def main() -> None:
load_model_state(model, state_dict) load_model_state(model, state_dict)
except RuntimeError as exc: except RuntimeError as exc:
raise RuntimeError( raise RuntimeError(
"Checkpoint vocabulary shape is incompatible with the no-event dataset/model setup. " "Checkpoint vocabulary shape is incompatible with the dataset/model setup. "
"Please ensure this run was trained with the current no-event vocabulary and matching labels file." "Please ensure this run was trained with the same special-token vocabulary and labels file."
) from exc ) from exc
if model.token_embedding.num_embeddings != dataset.vocab_size or model.risk_head.out_features != dataset.vocab_size: if model.token_embedding.num_embeddings != dataset.vocab_size or model.risk_head.out_features != dataset.vocab_size:

65
evaluate_doa_auc.py
View File

@@ -6,10 +6,9 @@ full observed record. Patients prevalent at/before DOA or incident after the
horizon are not used for that disease-horizon AUC. horizon are not used for that disease-horizon AUC.
The script adapts automatically to checkpoint target mode: The script adapts automatically to checkpoint target mode:
- next_token: use the DOA token position, inserting <NO_EVENT> at DOA when no - next_token: use the CHECKUP token position at DOA;
real disease token exists at DOA; - all_future: query the model directly with t_query=DOA. The history includes
- all_future: query the model directly with t_query=DOA, allowing empty the CHECKUP token at DOA.
disease history because other-info tokens still describe the DOA state.
""" """
from __future__ import annotations from __future__ import annotations
@@ -163,17 +162,20 @@ class DOAStatusDataset(_ExpoBaseDataset):
doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32))) doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32)))
doa_years = np.float32(doa_days / DAYS_PER_YEAR) doa_years = np.float32(doa_days / DAYS_PER_YEAR)
disease_rows = rows[rows[:, 2].astype(np.int64) != CHECKUP_IDX] raw_times = rows[:, 1].astype(np.float32) / DAYS_PER_YEAR
disease_times = disease_rows[:, 1].astype(np.float32) / DAYS_PER_YEAR raw_labels = rows[:, 2].astype(np.int64)
disease_labels_raw = disease_rows[:, 2].astype(np.int64) shifted_labels = np.where(
disease_labels = np.where( raw_labels >= NO_EVENT_IDX,
disease_labels_raw >= NO_EVENT_IDX, raw_labels + 1,
disease_labels_raw + 1, raw_labels,
disease_labels_raw,
).astype(np.int64) ).astype(np.int64)
order = np.lexsort((disease_labels, disease_times)) order = np.lexsort((shifted_labels, raw_times))
disease_times = disease_times[order].astype(np.float32) event_times = raw_times[order].astype(np.float32)
disease_labels = disease_labels[order].astype(np.int64) event_labels = shifted_labels[order].astype(np.int64)
disease_mask = event_labels != CHECKUP_IDX
disease_times = event_times[disease_mask]
disease_labels = event_labels[disease_mask]
patient_id = len(self.records) patient_id = len(self.records)
for token in np.unique(disease_labels).tolist(): for token in np.unique(disease_labels).tolist():
@@ -186,25 +188,20 @@ class DOAStatusDataset(_ExpoBaseDataset):
(patient_id, float(disease_times[int(hit[0])])) (patient_id, float(disease_times[int(hit[0])]))
) )
hist = disease_times <= doa_years hist = event_times <= doa_years
hist_events = disease_labels[hist] hist_events = event_labels[hist]
hist_times = disease_times[hist] hist_times = event_times[hist]
if self.model_target_mode == "next_token": if self.model_target_mode == "next_token":
at_doa = np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6) checkup_at_doa = (
if hist_events.size == 0 or not np.any(at_doa): (hist_events == CHECKUP_IDX)
event_seq = np.concatenate([ & np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6)
hist_events, )
np.array([NO_EVENT_IDX], dtype=np.int64), if not np.any(checkup_at_doa):
]) raise RuntimeError(f"Missing CHECKUP token at DOA for eid={eid}")
time_seq = np.concatenate([ event_seq = hist_events
hist_times, time_seq = hist_times
np.array([doa_years], dtype=np.float32), readout_pos = int(np.where(checkup_at_doa)[0][-1])
])
else:
event_seq = hist_events
time_seq = hist_times
readout_pos = int(len(event_seq) - 1)
else: else:
event_seq = hist_events event_seq = hist_events
time_seq = hist_times time_seq = hist_times
@@ -682,10 +679,10 @@ def main() -> None:
model.eval() model.eval()
if model_target_mode == "next_token" and ( if model_target_mode == "next_token" and (
model.token_embedding.num_embeddings <= NO_EVENT_IDX model.token_embedding.num_embeddings <= CHECKUP_IDX
or model.risk_head.out_features <= NO_EVENT_IDX or model.risk_head.out_features <= CHECKUP_IDX
): ):
raise RuntimeError("Next-token DOA evaluation requires <NO_EVENT> in the model vocabulary.") raise RuntimeError("Next-token DOA evaluation requires <CHECKUP> in the model vocabulary.")
eval_dataset = Subset(dataset, eval_indices) eval_dataset = Subset(dataset, eval_indices)
loader = DataLoader( loader = DataLoader(

66
models.py
View File

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

6
train_next_step.py
View File

@@ -1,9 +1,9 @@
""" """
Train DeepHealth with next-token / next-time-point supervision. Train DeepHealth with next-token / next-time-point supervision.
The dataset remains the current next-step construction: pure disease events plus The next-step dataset uses observed event histories, including CHECKUP state
optional gap <NO_EVENT> imputation are shifted into autoregressive inputs and tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
targets. UTS training reads out only same-time group ends. same-time group ends.
""" """
from __future__ import annotations from __future__ import annotations