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

Add training scripts for all-future and next-step supervision with DeepHealth

Browse Source 
- Implement `train_all_future.py` for training with query-conditioned all-future supervision.
- Implement `train_next_step.py` for training with next-token/next-time-point supervision.
- Introduce `train_util.py` for shared utility functions including logging, dataset splitting, and model checkpointing.
- Enhance argument parsing for both training scripts to accommodate new parameters.
- Update loss functions and model configurations to support the new training paradigms.
This commit is contained in:
Jiarui Li
2026-06-13 11:42:04 +08:00
parent 034d8065a7
commit 46a3dfe628
12 changed files with 1927 additions and 1273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
eval_data.py Normal file
View File

@@ -0,0 +1,163 @@
from __future__ import annotations
from typing import Any, Dict, Iterable, List
import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence
from dataset import AllFutureHealthDataset, HealthDataset
from targets import PAD_IDX
class AllFutureSequenceEvalDataset:
"""
Eval-only sequence view for all-future checkpoints.
All-future training uses pure disease histories, so token-level and landmark
evaluation should not reuse the next-step dataset view that contains
imputed <NO_EVENT> gap tokens.
"""
def __init__(
self,
data_prefix: str,
labels_file: str,
min_history_events: int = 1,
min_future_events: int = 1,
extra_info_types: Iterable[int] | None = None,
) -> None:
base = AllFutureHealthDataset(
data_prefix=data_prefix,
labels_file=labels_file,
split="train",
min_history_events=min_history_events,
min_future_events=min_future_events,
extra_info_types=extra_info_types,
)
self.base = base
self.label_code_to_id = base.label_code_to_id
self.label_id_to_code = base.label_id_to_code
self.vocab_size = base.vocab_size
self.n_types = base.n_types
self.n_cont_types = base.n_cont_types
self.n_categories = base.n_categories
self.cont_type_ids = base.cont_type_ids
self.extra_info_types = base.extra_info_types
self.samples: List[Dict[str, Any]] = []
for patient in base.patients:
labels = np.asarray(patient["labels"], dtype=np.int64)
times = np.asarray(patient["times"], dtype=np.float32)
if labels.size < 2:
continue
input_len = int(labels.size - 1)
self.samples.append(
{
"eid": int(patient["eid"]),
"event_seq": labels[:-1],
"time_seq": times[:-1],
"target_event_seq": labels[1:],
"target_time_seq": times[1:],
"readout_mask": np.ones(input_len, dtype=bool),
"sex": int(patient["sex"]),
"other_type": np.asarray(patient["other_type"], dtype=np.int64),
"other_value": np.asarray(patient["other_value"], dtype=np.float32),
"other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
"other_time": np.asarray(patient["other_time"], dtype=np.float32),
}
)
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
s = self.samples[idx]
return {
"event_seq": torch.from_numpy(s["event_seq"]).long(),
"time_seq": torch.from_numpy(s["time_seq"]).float(),
"target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
"readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
"sex": torch.tensor(s["sex"], dtype=torch.long),
"other_type": torch.from_numpy(s["other_type"]).long(),
"other_value": torch.from_numpy(s["other_value"]).float(),
"other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
"other_time": torch.from_numpy(s["other_time"]).float(),
}
def load_sequence_eval_dataset(
*,
model_target_mode: str,
data_prefix: str,
labels_file: str,
no_event_interval_years: float,
include_no_event_in_uts_target: bool,
min_history_events: int,
min_future_events: int,
extra_info_types: Iterable[int] | None,
):
mode = str(model_target_mode).lower()
if mode == "next_token":
return HealthDataset(
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
extra_info_types=extra_info_types,
)
if mode == "all_future":
return AllFutureSequenceEvalDataset(
data_prefix=data_prefix,
labels_file=labels_file,
min_history_events=min_history_events,
min_future_events=min_future_events,
extra_info_types=extra_info_types,
)
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
event_seq = pad_sequence(
[s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
)
time_seq = pad_sequence(
[s["time_seq"] for s in batch], batch_first=True, padding_value=0.0
)
target_event_seq = pad_sequence(
[s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
)
target_time_seq = pad_sequence(
[s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0
)
readout_mask = pad_sequence(
[s["readout_mask"] for s in batch], batch_first=True, padding_value=False
)
other_type = pad_sequence(
[s["other_type"] for s in batch], batch_first=True, padding_value=0
)
other_value = pad_sequence(
[s["other_value"] for s in batch], batch_first=True, padding_value=0.0
)
other_value_kind = pad_sequence(
[s["other_value_kind"] for s in batch], batch_first=True, padding_value=0
)
other_time = pad_sequence(
[s["other_time"] for s in batch], batch_first=True, padding_value=0.0
)
return {
"event_seq": event_seq,
"time_seq": time_seq,
"padding_mask": event_seq > PAD_IDX,
"target_event_seq": target_event_seq,
"target_time_seq": target_time_seq,
"readout_mask": readout_mask,
"sex": torch.stack([s["sex"] for s in batch]),
"other_type": other_type,
"other_value": other_value,
"other_value_kind": other_value_kind,
"other_time": other_time,
}