Add disease history ablation modes
This commit is contained in:
@@ -30,7 +30,12 @@ from torch.nn.utils.rnn import pad_sequence
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from dataset import HealthDataset
|
||||
from dataset import (
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
HealthDataset,
|
||||
normalize_disease_history_mode,
|
||||
transform_disease_history,
|
||||
)
|
||||
from delphi2m_auc_report import (
|
||||
DEFAULT_DELPHI2M_PERIODS_YEARS,
|
||||
build_delphi2m_auc_report,
|
||||
@@ -338,6 +343,7 @@ class LandmarkDataset(Dataset):
|
||||
min_history_events: int,
|
||||
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
||||
death_token_ids: Sequence[int],
|
||||
disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
|
||||
) -> None:
|
||||
self.dataset = dataset
|
||||
self.subset_indices = np.asarray(subset_indices, dtype=np.int64)
|
||||
@@ -348,6 +354,16 @@ class LandmarkDataset(Dataset):
|
||||
"model_target_mode must be next_token or all_future, got "
|
||||
f"{self.model_target_mode!r}"
|
||||
)
|
||||
self.disease_history_mode = normalize_disease_history_mode(
|
||||
disease_history_mode
|
||||
)
|
||||
if (
|
||||
self.model_target_mode != "all_future"
|
||||
and self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
|
||||
):
|
||||
raise ValueError(
|
||||
"ordered/set disease history is only supported for all_future models"
|
||||
)
|
||||
self.min_history_events = int(min_history_events)
|
||||
|
||||
self.first_occurrence_by_token = first_occurrence_by_token
|
||||
@@ -428,10 +444,16 @@ class LandmarkDataset(Dataset):
|
||||
readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
|
||||
readout_mask[-1] = True
|
||||
else:
|
||||
event_seq_landmark = prefix_events.astype(
|
||||
np.int64, copy=False)
|
||||
time_seq_landmark = prefix_times.astype(
|
||||
np.float32, copy=False)
|
||||
(
|
||||
event_seq_landmark,
|
||||
time_seq_landmark,
|
||||
model_t_query,
|
||||
) = transform_disease_history(
|
||||
event_seq=prefix_events,
|
||||
actual_time_seq=prefix_times,
|
||||
actual_t_query=landmark_age,
|
||||
disease_history_mode=self.disease_history_mode,
|
||||
)
|
||||
landmark_pos = int(len(event_seq_landmark) - 1)
|
||||
readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
|
||||
|
||||
@@ -444,7 +466,11 @@ class LandmarkDataset(Dataset):
|
||||
"followup_end_time": np.float32(followup_end),
|
||||
"death_time": np.float32(self.patient_death_time[patient_id]),
|
||||
"landmark_pos": landmark_pos,
|
||||
"t_query": np.float32(landmark_age),
|
||||
"t_query": (
|
||||
np.float32(landmark_age)
|
||||
if self.model_target_mode == "next_token"
|
||||
else model_t_query
|
||||
),
|
||||
"event_seq": event_seq_landmark,
|
||||
"time_seq": time_seq_landmark,
|
||||
"readout_mask": readout_mask,
|
||||
@@ -1137,6 +1163,9 @@ def main() -> None:
|
||||
f"got {model_target_mode!r}"
|
||||
)
|
||||
dist_mode_cfg = str(cfg.get("dist_mode", "exponential"))
|
||||
disease_history_mode = normalize_disease_history_mode(
|
||||
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
||||
)
|
||||
|
||||
output_path = Path(
|
||||
cfg_get(args, cfg, "output_path", None)
|
||||
@@ -1162,6 +1191,7 @@ def main() -> None:
|
||||
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
||||
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
||||
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
@@ -1282,6 +1312,7 @@ def main() -> None:
|
||||
min_history_events=min_history_events,
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=death_token_ids,
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
@@ -1322,6 +1353,7 @@ def main() -> None:
|
||||
print(f"Number of selected patients: {len(subset_indices)}")
|
||||
print(f"No-event support: {bool(has_no_event)}")
|
||||
print(f"Model target mode: {model_target_mode}")
|
||||
print(f"Disease history mode: {disease_history_mode}")
|
||||
print(f"Landmark query mode: {landmark_query_mode}")
|
||||
print(
|
||||
"Landmark token mode: no_event"
|
||||
|
||||
Reference in New Issue
Block a user