Add target construction and training script for DeepHealth model
- Implemented target construction in `targets.py` for next-token and unique-time set supervision. - Added validation functions and utility methods for target building. - Created a comprehensive training script in `train.py` that includes data loading, model building, optimizer setup, and training loop with early stopping and logging. - Integrated loss functions and readout mechanisms based on target modes. - Established dataset splitting and DataLoader configurations for training, validation, and testing.
This commit is contained in:
394
targets.py
Normal file
394
targets.py
Normal file
@@ -0,0 +1,394 @@
|
||||
# targets.py
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
PAD_IDX = 0
|
||||
CHECKUP_IDX = 1
|
||||
NO_EVENT_IDX = 2
|
||||
DAYS_PER_YEAR = 365.25
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NextTokenTargets:
|
||||
"""
|
||||
Delphi2M-style next-token supervision targets.
|
||||
|
||||
Shapes:
|
||||
input_events: (L,)
|
||||
input_times_years: (L,)
|
||||
target_events: (L,)
|
||||
target_times_years:(L,)
|
||||
|
||||
where L = N - 1.
|
||||
"""
|
||||
input_events: np.ndarray
|
||||
input_times_years: np.ndarray
|
||||
target_events: np.ndarray
|
||||
target_times_years: np.ndarray
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UniqueTimeSetTargets:
|
||||
"""
|
||||
Unique-time set supervision targets.
|
||||
|
||||
Shapes:
|
||||
readout_mask: (L,)
|
||||
target_dt_unique: (L,)
|
||||
target_multi_hot: (L, vocab_size)
|
||||
|
||||
where L = N - 1.
|
||||
|
||||
Only group-end positions can have readout_mask=True.
|
||||
target_dt_unique is measured in years.
|
||||
"""
|
||||
readout_mask: np.ndarray
|
||||
target_dt_unique: np.ndarray
|
||||
target_multi_hot: np.ndarray
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TargetPack:
|
||||
"""
|
||||
Combined target package for one patient sequence.
|
||||
|
||||
Contains both next-token targets and unique-time-set targets.
|
||||
The training pipeline decides which one to use.
|
||||
"""
|
||||
next_token: NextTokenTargets
|
||||
unique_time_set: UniqueTimeSetTargets
|
||||
|
||||
|
||||
def _as_numpy_1d(
|
||||
x: np.ndarray,
|
||||
name: str,
|
||||
dtype: np.dtype | type | None = None,
|
||||
) -> np.ndarray:
|
||||
arr = np.asarray(x)
|
||||
if arr.ndim != 1:
|
||||
raise ValueError(f"{name} must be 1D, got shape {arr.shape}")
|
||||
if dtype is not None:
|
||||
arr = arr.astype(dtype)
|
||||
return arr
|
||||
|
||||
|
||||
def validate_event_sequence(
|
||||
labels: np.ndarray,
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
require_sorted: bool = True,
|
||||
) -> None:
|
||||
"""
|
||||
Validate one patient's event sequence.
|
||||
|
||||
labels:
|
||||
1D integer label ids.
|
||||
|
||||
times_days:
|
||||
1D event times in days.
|
||||
|
||||
require_sorted:
|
||||
If True, raises when times_days is not non-decreasing.
|
||||
"""
|
||||
labels = _as_numpy_1d(labels, "labels")
|
||||
times_days = _as_numpy_1d(times_days, "times_days")
|
||||
|
||||
if len(labels) != len(times_days):
|
||||
raise ValueError(
|
||||
f"labels and times_days must have the same length, "
|
||||
f"got {len(labels)} and {len(times_days)}"
|
||||
)
|
||||
|
||||
if len(labels) == 0:
|
||||
raise ValueError("Empty event sequence is not valid.")
|
||||
|
||||
if np.any(labels < 0):
|
||||
raise ValueError("labels contains negative ids.")
|
||||
|
||||
if not np.all(np.isfinite(times_days)):
|
||||
raise ValueError("times_days contains non-finite values.")
|
||||
|
||||
if require_sorted and len(times_days) > 1:
|
||||
if np.any(np.diff(times_days) < 0):
|
||||
raise ValueError("times_days must be non-decreasing.")
|
||||
|
||||
|
||||
def build_next_token_targets(
|
||||
labels: np.ndarray,
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
require_sorted: bool = True,
|
||||
) -> NextTokenTargets:
|
||||
"""
|
||||
Build Delphi2M-style autoregressive next-token targets.
|
||||
|
||||
Given full sequence:
|
||||
|
||||
labels: [x0, x1, x2, ..., xN-1]
|
||||
times_days: [t0, t1, t2, ..., tN-1]
|
||||
|
||||
returns:
|
||||
|
||||
input_events: [x0, x1, ..., xN-2]
|
||||
input_times_years: [t0, t1, ..., tN-2] / 365.25
|
||||
target_events: [x1, x2, ..., xN-1]
|
||||
target_times_years: [t1, t2, ..., tN-1] / 365.25
|
||||
|
||||
This function does not ignore PAD/CHECKUP/NO_EVENT. Ignoring belongs to
|
||||
the loss function because different objectives may use different ignore ids.
|
||||
"""
|
||||
labels = _as_numpy_1d(labels, "labels", np.int64)
|
||||
times_days = _as_numpy_1d(times_days, "times_days", np.float32)
|
||||
validate_event_sequence(labels, times_days, require_sorted=require_sorted)
|
||||
|
||||
if len(labels) < 2:
|
||||
raise ValueError(
|
||||
"Need at least two events to build next-token targets."
|
||||
)
|
||||
|
||||
input_events = labels[:-1].astype(np.int64)
|
||||
input_times_years = (times_days[:-1] / DAYS_PER_YEAR).astype(np.float32)
|
||||
|
||||
target_events = labels[1:].astype(np.int64)
|
||||
target_times_years = (times_days[1:] / DAYS_PER_YEAR).astype(np.float32)
|
||||
|
||||
return NextTokenTargets(
|
||||
input_events=input_events,
|
||||
input_times_years=input_times_years,
|
||||
target_events=target_events,
|
||||
target_times_years=target_times_years,
|
||||
)
|
||||
|
||||
|
||||
def build_unique_time_set_targets(
|
||||
labels: np.ndarray,
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
vocab_size: int,
|
||||
ignored_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
require_sorted: bool = True,
|
||||
) -> UniqueTimeSetTargets:
|
||||
"""
|
||||
Build next-unique-time set targets.
|
||||
|
||||
This is the target construction used by your UTS / default mode.
|
||||
|
||||
For each input position i:
|
||||
- only if i is the last token of its timestamp group;
|
||||
- find the next distinct timestamp group;
|
||||
- target is the set of valid event labels at that next timestamp.
|
||||
|
||||
Example:
|
||||
|
||||
t=49: X
|
||||
t=50: A, B, C
|
||||
t=51: D, E
|
||||
|
||||
Supervises:
|
||||
|
||||
X@49 -> {A, B, C}@50
|
||||
group_end@50 -> {D, E}@51
|
||||
|
||||
It does NOT supervise:
|
||||
|
||||
A@50 -> B@50
|
||||
B@50 -> C@50
|
||||
|
||||
Parameters
|
||||
----------
|
||||
labels:
|
||||
Full event sequence labels, shape (N,).
|
||||
|
||||
times_days:
|
||||
Full event sequence times in days, shape (N,).
|
||||
|
||||
vocab_size:
|
||||
Size of output vocabulary.
|
||||
|
||||
ignored_target_ids:
|
||||
Label ids that should not enter target_multi_hot.
|
||||
Usually:
|
||||
no no-event: {0, 1}
|
||||
with no-event: {0, 1, 2}
|
||||
For UTS, I recommend ignoring <NO_EVENT> unless explicitly testing it
|
||||
as an event target.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UniqueTimeSetTargets
|
||||
"""
|
||||
labels = _as_numpy_1d(labels, "labels", np.int64)
|
||||
times_days = _as_numpy_1d(times_days, "times_days", np.float32)
|
||||
validate_event_sequence(labels, times_days, require_sorted=require_sorted)
|
||||
|
||||
if vocab_size <= 0:
|
||||
raise ValueError(f"vocab_size must be positive, got {vocab_size}")
|
||||
|
||||
if len(labels) < 2:
|
||||
raise ValueError(
|
||||
"Need at least two events to build unique-time-set targets."
|
||||
)
|
||||
|
||||
input_len = len(labels) - 1
|
||||
|
||||
readout_mask = np.zeros(input_len, dtype=bool)
|
||||
target_dt_unique = np.zeros(input_len, dtype=np.float32)
|
||||
target_multi_hot = np.zeros((input_len, vocab_size), dtype=bool)
|
||||
|
||||
ignored = {int(x) for x in ignored_target_ids}
|
||||
|
||||
unique_times = np.unique(times_days)
|
||||
time_to_group_idx = {t: i for i, t in enumerate(unique_times)}
|
||||
group_indices = np.array([time_to_group_idx[t]
|
||||
for t in times_days], dtype=np.int64)
|
||||
|
||||
for i in range(input_len):
|
||||
current_group = group_indices[i]
|
||||
|
||||
is_last_in_group = (
|
||||
i == input_len - 1
|
||||
or group_indices[i + 1] != current_group
|
||||
)
|
||||
if not is_last_in_group:
|
||||
continue
|
||||
|
||||
next_group_idx = current_group + 1
|
||||
if next_group_idx >= len(unique_times):
|
||||
continue
|
||||
|
||||
next_time = unique_times[next_group_idx]
|
||||
next_labels = labels[group_indices == next_group_idx]
|
||||
|
||||
valid_next_labels: list[int] = []
|
||||
for lab in next_labels:
|
||||
lab_int = int(lab)
|
||||
if lab_int in ignored:
|
||||
continue
|
||||
if lab_int < 0 or lab_int >= vocab_size:
|
||||
continue
|
||||
valid_next_labels.append(lab_int)
|
||||
|
||||
# If next timestamp contains only technical tokens, do not supervise UTS.
|
||||
if len(valid_next_labels) == 0:
|
||||
continue
|
||||
|
||||
readout_mask[i] = True
|
||||
target_dt_unique[i] = float(next_time - times_days[i]) / DAYS_PER_YEAR
|
||||
target_multi_hot[i, valid_next_labels] = True
|
||||
|
||||
return UniqueTimeSetTargets(
|
||||
readout_mask=readout_mask,
|
||||
target_dt_unique=target_dt_unique.astype(np.float32),
|
||||
target_multi_hot=target_multi_hot,
|
||||
)
|
||||
|
||||
|
||||
def build_all_targets(
|
||||
labels: np.ndarray,
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
vocab_size: int,
|
||||
ignored_uts_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
require_sorted: bool = True,
|
||||
) -> TargetPack:
|
||||
"""
|
||||
Build both next-token targets and unique-time-set targets for one patient.
|
||||
|
||||
This is the function dataset.py should usually call during initialization.
|
||||
|
||||
The dataset can then store:
|
||||
event_seq = target_pack.next_token.input_events
|
||||
time_seq = target_pack.next_token.input_times_years
|
||||
|
||||
target_event_seq = target_pack.next_token.target_events
|
||||
target_time_seq = target_pack.next_token.target_times_years
|
||||
|
||||
readout_mask = target_pack.unique_time_set.readout_mask
|
||||
target_dt_unique = target_pack.unique_time_set.target_dt_unique
|
||||
target_multi_hot = target_pack.unique_time_set.target_multi_hot
|
||||
"""
|
||||
next_token = build_next_token_targets(
|
||||
labels=labels,
|
||||
times_days=times_days,
|
||||
require_sorted=require_sorted,
|
||||
)
|
||||
|
||||
unique_time_set = build_unique_time_set_targets(
|
||||
labels=labels,
|
||||
times_days=times_days,
|
||||
vocab_size=vocab_size,
|
||||
ignored_target_ids=ignored_uts_target_ids,
|
||||
require_sorted=require_sorted,
|
||||
)
|
||||
|
||||
return TargetPack(
|
||||
next_token=next_token,
|
||||
unique_time_set=unique_time_set,
|
||||
)
|
||||
|
||||
|
||||
def get_group_end_mask_from_times(
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
input_len: int | None = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Convenience utility for debugging.
|
||||
|
||||
Returns a bool mask indicating the last token of each same-time group
|
||||
within the input sequence.
|
||||
|
||||
If input_len is None, uses len(times_days) - 1, matching model input length.
|
||||
"""
|
||||
times_days = _as_numpy_1d(times_days, "times_days", np.float32)
|
||||
|
||||
if input_len is None:
|
||||
input_len = len(times_days) - 1
|
||||
|
||||
if input_len < 0 or input_len > len(times_days):
|
||||
raise ValueError(
|
||||
f"Invalid input_len={input_len} for sequence length {len(times_days)}"
|
||||
)
|
||||
|
||||
out = np.zeros(input_len, dtype=bool)
|
||||
|
||||
for i in range(input_len):
|
||||
is_last_in_group = (
|
||||
i == input_len - 1
|
||||
or times_days[i + 1] != times_days[i]
|
||||
)
|
||||
out[i] = is_last_in_group
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def summarize_targets(
|
||||
target_pack: TargetPack,
|
||||
) -> dict[str, int | float]:
|
||||
"""
|
||||
Small debugging helper for logging.
|
||||
"""
|
||||
nt = target_pack.next_token
|
||||
uts = target_pack.unique_time_set
|
||||
|
||||
n_tokens = int(len(nt.input_events))
|
||||
n_readout = int(uts.readout_mask.sum())
|
||||
n_positive_labels = int(uts.target_multi_hot.sum())
|
||||
|
||||
mean_set_size = (
|
||||
float(n_positive_labels / n_readout)
|
||||
if n_readout > 0
|
||||
else 0.0
|
||||
)
|
||||
|
||||
return {
|
||||
"n_input_tokens": n_tokens,
|
||||
"n_uts_readouts": n_readout,
|
||||
"n_uts_positive_labels": n_positive_labels,
|
||||
"mean_uts_set_size": mean_set_size,
|
||||
}
|
||||
Reference in New Issue
Block a user