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

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