refactor: isolate Delphi2M next-token pipeline
This commit is contained in:
263
targets.py
263
targets.py
@@ -2,8 +2,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@@ -32,38 +30,6 @@ class NextTokenTargets:
|
||||
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,
|
||||
@@ -163,232 +129,3 @@ def build_next_token_targets(
|
||||
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