132 lines
3.4 KiB
Python
132 lines
3.4 KiB
Python
# targets.py
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
|
|
|
|
PAD_IDX = 0
|
|
RESERVED_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
|
|
|
|
|
|
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/RESERVED/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,
|
|
)
|