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

Add exposure cache and keep absolute time only

Browse Source
This commit is contained in:
Jiarui Li
2026-07-07 17:21:52 +08:00
parent a0379daf29
commit 45a857d1a6
9 changed files with 690 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
backbones.py
View File

@@ -5,114 +5,24 @@ import torch.nn as nn
import torch.nn.functional as F
class TimeRoPE(nn.Module):
def __init__(self, dim: int, base: float = 10000.0):
super().__init__()
assert dim % 2 == 0, "RoPE dim must be even"
self.dim = dim
# inv_freq is not trainable, but should move with device.
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
def precompute_cache(self, tau: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
t = tau.unsqueeze(-1) # (B, L, 1)
angles = t * self.inv_freq # (B, L, dim//2)
# Pre-expand for heads and interleave once (avoids N_layers repeats)
cos = angles.cos().unsqueeze(1).repeat_interleave(2, dim=-1)
sin = angles.sin().unsqueeze(1).repeat_interleave(2, dim=-1)
return cos, sin # (B, 1, L, dim)
@staticmethod
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
"""Rotate pairs: ``[-x2, x1, -x4, x3, ...]``."""
x1 = x[..., 0::2]
x2 = x[..., 1::2]
return torch.stack((-x2, x1), dim=-1).flatten(-2)
@staticmethod
def apply_from_cache(
q: torch.Tensor,
k: torch.Tensor,
rope_cache: tuple[torch.Tensor, torch.Tensor],
) -> tuple[torch.Tensor, torch.Tensor]:
cos, sin = rope_cache # each (B, 1, L, dim)
q_rot = q * cos + TimeRoPE._rotate_half(q) * sin
k_rot = k * cos + TimeRoPE._rotate_half(k) * sin
return q_rot, k_rot
def forward(
self,
tau: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
cache = self.precompute_cache(tau)
return self.apply_from_cache(q, k, cache)
class GaussianRBFTimeBasis(nn.Module):
def __init__(
self,
n_bases: int = 16,
max_time_diff: float = 40.0,
):
super().__init__()
self.n_bases = n_bases
# Evenly spaced RBF centres for non-negative linear time differences.
# Causal masking enforces query_time >= key_time, so diff is >= 0.
centers = torch.linspace(0.0, max_time_diff, n_bases)
self.register_buffer("centers", centers,
persistent=False) # (n_bases,)
# Learnable log-widths (initialized to center spacing on linear scale).
init_width = max(max_time_diff / max(n_bases - 1, 1), 1e-3)
init_log_width = math.log(init_width)
self.log_widths = nn.Parameter(torch.full((n_bases,), init_log_width))
def precompute_cache(self, tau: torch.Tensor) -> torch.Tensor:
time_coord = tau.float() # (B, L)
# Pairwise signed difference: query_i - key_j.
diff = time_coord.unsqueeze(
2) - time_coord.unsqueeze(1) # (B, L_q, L_k)
# Gaussian RBF: exp(-0.5 * ((diff - c) / w)^2)
diff = diff.unsqueeze(-1) # (B, L, L, 1)
widths = self.log_widths.exp() # (n_bases,)
rbf_acts = torch.exp(
-0.5 * ((diff - self.centers) / widths).square()
# (B, L, L, n_bases)
)
return rbf_acts
class TemporalAttention(nn.Module):
def __init__(
self,
n_embd: int,
n_head: int,
n_rbf_bases: int = 16,
dropout: float = 0.0,
use_time_rope: bool = True,
use_rbf_bias: bool = True,
):
super().__init__()
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
self.n_head = n_head
self.d_head = n_embd // n_head
self.scale = 1.0 / math.sqrt(self.d_head)
self.use_time_rope = use_time_rope
self.use_rbf_bias = use_rbf_bias
# QKV projection (fused for efficiency)
self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
# Output projection
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
# Layer-specific projection from shared RBF basis activations to per-head attention bias.
self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
self.resid_drop = nn.Dropout(dropout)
self.reset_parameters()
@@ -120,20 +30,12 @@ class TemporalAttention(nn.Module):
"""Match the previous version's GPT-style weight initialization."""
nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
nn.init.zeros_(self.rbf_proj.weight)
def forward(
self,
x: torch.Tensor,
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
rbf_cache: torch.Tensor | None = None,
attn_mask: torch.Tensor | None = None,
) -> torch.Tensor:
if self.use_time_rope:
assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
if self.use_rbf_bias:
assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"
B, L, _ = x.shape
H, D = self.n_head, self.d_head
@@ -141,31 +43,11 @@ class TemporalAttention(nn.Module):
qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
q, k, v = qkv.unbind(0) # each (B, H, L, D)
# --- Apply RoPE (from shared cache) --------------------------------
if self.use_time_rope:
q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)
# Build additive attention bias mask: time bias + causal/padding mask.
time_bias = None
if self.use_rbf_bias:
time_bias = self.rbf_proj(rbf_cache).permute(
0, 3, 1, 2) # (B, H, L, L)
time_bias = self.time_bias_scale.tanh() * time_bias
if time_bias is not None and attn_mask is not None:
attn_bias = time_bias + attn_mask.to(time_bias.dtype)
elif time_bias is not None:
attn_bias = time_bias
elif attn_mask is not None:
attn_bias = attn_mask
else:
attn_bias = None
out = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=attn_bias,
attn_mask=attn_mask,
dropout_p=0.0,
is_causal=False,
scale=self.scale,
@@ -218,18 +100,12 @@ class GPTBlock(nn.Module):
attn_dropout: float = 0.0,
mlp_dropout: float = 0.0,
use_time_rope: bool = False,
use_rbf_bias: bool = False,
n_rbf_bases: int = 16,
):
super().__init__()
self.attn = TemporalAttention(
n_embd=n_embd,
n_head=n_head,
n_rbf_bases=n_rbf_bases,
dropout=attn_dropout,
use_time_rope=use_time_rope,
use_rbf_bias=use_rbf_bias,
)
self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
self.ln1 = nn.LayerNorm(n_embd)
@@ -238,11 +114,9 @@ class GPTBlock(nn.Module):
def forward(
self,
x: torch.Tensor,
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
rbf_cache: torch.Tensor | None = None,
attn_mask: torch.Tensor | None = None,
) -> torch.Tensor:
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
x = x + self.attn(self.ln1(x), attn_mask)
x = x + self.mlp(self.ln2(x))
return x

223
dataset.py
View File

@@ -1,7 +1,8 @@
# dataset.py
from __future__ import annotations
from typing import Dict, List, Literal, Optional, Tuple
from pathlib import Path
from typing import Dict, Iterable, List, Literal, Optional, Tuple
import numpy as np
import pandas as pd
@@ -19,6 +20,92 @@ from targets import (
ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
DAILY_EXPOSURE_SHAPE = (1826, 4)
MONTHLY_EXPOSURE_SHAPE = (241, 2)
class ExposureCache:
"""Random-access view over files produced by prepare_exposure_cache.py."""
def __init__(self, cache_dir: str | Path):
cache_dir = Path(cache_dir)
self.cache_dir = cache_dir
eid_path = cache_dir / "exposure_eid.npy"
token_path = cache_dir / "exposure_token.npy"
onset_date_path = cache_dir / "exposure_onset_date.npy"
if not (eid_path.is_file() and token_path.is_file() and onset_date_path.is_file()):
raise FileNotFoundError(
"Exposure cache must contain exposure_eid.npy, "
"exposure_token.npy, and exposure_onset_date.npy. "
"Regenerate it with the current prepare_exposure_cache.py."
)
self.eids = np.load(eid_path, mmap_mode="r")
self.raw_tokens = np.load(token_path, mmap_mode="r")
self.onset_dates = np.load(onset_date_path, mmap_mode="r")
self.daily = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r")
self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="r")
quality_path = cache_dir / "exposure_quality.npy"
self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None
if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE:
raise ValueError(
f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, "
f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}"
)
if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE:
raise ValueError(
f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, "
f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}"
)
n_rows = len(self.eids)
if (
len(self.raw_tokens) != n_rows
or len(self.onset_dates) != n_rows
or self.daily.shape[0] != n_rows
or self.monthly.shape[0] != n_rows
):
raise ValueError("Exposure cache metadata/daily/monthly row counts do not match")
self._key_to_index: dict[tuple[int, int, int], int] | None = None
def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
keys: dict[tuple[int, int, int], int] = {}
eids = np.asarray(self.eids, dtype=np.int64)
tokens = np.asarray(self.raw_tokens, dtype=np.int64)
onset_dates = np.asarray(self.onset_dates, dtype="datetime64[D]")
for idx, (eid, token, onset_date) in enumerate(zip(eids, tokens, onset_dates)):
birth_date = birth_date_by_eid.get(int(eid))
if birth_date is None or np.isnat(onset_date) or np.isnat(birth_date):
continue
age_days = int((onset_date - birth_date).astype("timedelta64[D]").astype(np.int64))
if age_days < 0:
continue
keys[(int(eid), int(token), age_days)] = idx
self._key_to_index = keys
def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray:
if self._key_to_index is None:
raise RuntimeError("ExposureCache.build_age_index must be called before lookup")
out = np.full(len(raw_tokens), -1, dtype=np.int64)
real = raw_tokens > 1
if not np.any(real):
return out
real_pos = np.nonzero(real)[0]
out[real_pos] = [
self._key_to_index.get((int(eid), int(raw_tokens[pos]), int(round(float(age_days[pos])))), -1)
for pos in real_pos
]
return out
def daily_window(self, index: int) -> np.ndarray:
if index < 0:
return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
return np.asarray(self.daily[index], dtype=np.float32)
def monthly_window(self, index: int) -> np.ndarray:
if index < 0:
return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
return np.asarray(self.monthly[index], dtype=np.float32)
def load_label_vocab(
@@ -87,11 +174,19 @@ class _ExpoBaseDataset(Dataset):
labels_file: str = "labels.csv",
no_event_interval_years: float = 5.0,
include_no_event_in_uts_target: bool = False,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
) -> None:
self.data_prefix = data_prefix
self.labels_file = labels_file
self.no_event_interval_years = float(no_event_interval_years)
self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
self.exposure_cache = (
ExposureCache(exposure_cache_dir)
if exposure_cache_dir is not None
else None
)
self.mask_onset_exposure = bool(mask_onset_exposure)
self.label_code_to_id, self.label_id_to_code = load_label_vocab(
labels_file,
@@ -112,6 +207,9 @@ class _ExpoBaseDataset(Dataset):
basic_table = basic_table.loc[unique_eids]
self._prepare_sex(basic_table, unique_eids)
self._prepare_birth_dates(basic_table, unique_eids)
if self.exposure_cache is not None:
self.exposure_cache.build_age_index(self.birth_date_mapping)
max_id_in_vocab = max(self.label_id_to_code.keys())
max_id_in_data = int(self.event_data[:, 2].max()) if len(self.event_data) > 0 else 0
@@ -140,6 +238,25 @@ class _ExpoBaseDataset(Dataset):
)
self.sex_mapping = {int(eid): int(s) for eid, s in zip(unique_eids, sex01)}
def _prepare_birth_dates(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None:
if "date_of_birth" not in basic_table.columns:
if self.exposure_cache is None:
self.birth_date_mapping = {}
return
raise ValueError(
"Exposure alignment requires ukb_basic_info.csv to contain "
"'date_of_birth'. Regenerate it with the current prepare_data.py."
)
birth = pd.to_datetime(basic_table["date_of_birth"], errors="coerce")
if birth.isna().any() and self.exposure_cache is not None:
raise ValueError("date_of_birth contains missing or invalid values")
birth_np = birth.to_numpy(dtype="datetime64[D]")
self.birth_date_mapping = {
int(eid): np.datetime64(date, "D")
for eid, date in zip(unique_eids, birth_np)
if not np.isnat(date)
}
def _iter_patient_events(
self,
*,
@@ -176,6 +293,46 @@ class _ExpoBaseDataset(Dataset):
"sex": self.sex_mapping[eid],
}
def _raw_tokens_from_model_tokens(self, model_tokens: np.ndarray) -> np.ndarray:
raw_tokens = np.full(len(model_tokens), -1, dtype=np.int64)
real = model_tokens > NO_EVENT_IDX
raw_tokens[real] = model_tokens[real].astype(np.int64) - 1
return raw_tokens
def _exposure_indices_for_inputs(
self,
eid: int,
input_events: np.ndarray,
input_times_days: np.ndarray,
) -> np.ndarray | None:
if self.exposure_cache is None:
return None
raw_tokens = self._raw_tokens_from_model_tokens(input_events)
return self.exposure_cache.lookup_indices(
eid=eid,
raw_tokens=raw_tokens,
age_days=input_times_days,
)
def _load_exposure_windows(self, exposure_index: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
if self.exposure_cache is None:
raise RuntimeError("Exposure cache is not enabled")
daily = np.stack(
[self.exposure_cache.daily_window(int(idx)) for idx in exposure_index],
axis=0,
).astype(np.float32, copy=True)
monthly = np.stack(
[self.exposure_cache.monthly_window(int(idx)) for idx in exposure_index],
axis=0,
).astype(np.float32, copy=True)
if self.mask_onset_exposure:
daily[:, 0, :] = np.nan
monthly[:, 0, :] = np.nan
return torch.from_numpy(daily).float(), torch.from_numpy(monthly).float()
class NextStepHealthDataset(_ExpoBaseDataset):
"""
Dataset for next-token and next-time-point losses with unified other-info
@@ -194,12 +351,16 @@ class NextStepHealthDataset(_ExpoBaseDataset):
labels_file: str = "labels.csv",
no_event_interval_years: float = 5.0,
include_no_event_in_uts_target: bool = False,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
) -> None:
super().__init__(
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,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
)
self.samples: List[Dict] = []
@@ -221,7 +382,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
require_sorted=True,
)
self.samples.append({
sample = {
"eid": eid,
"event_seq": target_pack.next_token.input_events,
"time_seq": target_pack.next_token.input_times_years,
@@ -231,14 +392,22 @@ class NextStepHealthDataset(_ExpoBaseDataset):
"target_dt_unique": target_pack.unique_time_set.target_dt_unique,
"target_multi_hot": target_pack.unique_time_set.target_multi_hot,
**features,
})
}
exposure_index = self._exposure_indices_for_inputs(
eid=eid,
input_events=target_pack.next_token.input_events,
input_times_days=times_days[:-1],
)
if exposure_index is not None:
sample["exposure_index"] = exposure_index
self.samples.append(sample)
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> Dict:
s = self.samples[idx]
return {
out = {
"event_seq": torch.from_numpy(s["event_seq"]).long(),
"time_seq": torch.from_numpy(s["time_seq"]).float(),
"sex": torch.tensor(s["sex"], dtype=torch.long),
@@ -248,6 +417,11 @@ class NextStepHealthDataset(_ExpoBaseDataset):
"target_dt_unique": torch.from_numpy(s["target_dt_unique"]).float(),
"target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(),
}
if "exposure_index" in s:
daily, monthly = self._load_exposure_windows(s["exposure_index"])
out["exposure_daily"] = daily
out["exposure_monthly"] = monthly
return out
class AllFutureHealthDataset(_ExpoBaseDataset):
@@ -273,6 +447,8 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
min_history_events: int = 1,
min_future_events: int = 1,
validation_query_seed: int = 42,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
) -> None:
if split not in {"train", "valid", "test"}:
raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -282,6 +458,8 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
)
self.split = split
@@ -310,6 +488,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
patient = {
"eid": eid,
"times": times_years,
"times_days": times_days.astype(np.float32),
"labels": labels.astype(np.int64),
"t_obs": float(times_years.max()),
**features,
@@ -406,11 +585,12 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
def _build_item(self, patient: Dict, t_query: float) -> Dict:
times = patient["times"]
times_days = patient["times_days"]
labels = patient["labels"]
hist = times <= t_query
fut = times > t_query
return {
out = {
"event_seq": torch.from_numpy(labels[hist]).long(),
"time_seq": torch.from_numpy(times[hist]).float(),
"t_query": torch.tensor(t_query, dtype=torch.float32),
@@ -419,6 +599,17 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
"exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32),
"sex": torch.tensor(patient["sex"], dtype=torch.long),
}
if self.exposure_cache is not None:
exposure_index = self._exposure_indices_for_inputs(
eid=int(patient["eid"]),
input_events=labels[hist],
input_times_days=times_days[hist],
)
if exposure_index is not None:
daily, monthly = self._load_exposure_windows(exposure_index)
out["exposure_daily"] = daily
out["exposure_monthly"] = monthly
return out
def __len__(self) -> int:
if self.split == "train":
@@ -441,6 +632,22 @@ def _collate_common_static(batch: List[Dict]) -> Dict:
}
def _pad_exposure(batch: List[Dict], key: str, shape: tuple[int, int]) -> torch.Tensor:
max_len = max(int(s["event_seq"].numel()) for s in batch)
out = torch.full(
(len(batch), max_len, shape[0], shape[1]),
float("nan"),
dtype=torch.float32,
)
for idx, sample in enumerate(batch):
value = sample.get(key)
if value is None:
continue
seq_len = int(value.size(0))
out[idx, :seq_len] = value
return out
def next_step_collate_fn(batch: List[Dict]) -> Dict:
event_seq = pad_sequence(
[s["event_seq"] for s in batch],
@@ -489,6 +696,9 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict:
"target_multi_hot": target_multi_hot,
}
out.update(_collate_common_static(batch))
if any("exposure_daily" in s for s in batch):
out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
return out
@@ -524,6 +734,9 @@ def all_future_collate_fn(batch: List[Dict]) -> Dict:
"exposure": torch.stack([s["exposure"] for s in batch]),
}
out.update(_collate_common_static(batch))
if any("exposure_daily" in s for s in batch):
out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
return out

69
eval_data.py
View File

@@ -1,12 +1,18 @@
from __future__ import annotations
from pathlib import Path
from typing import Any, Dict, List
import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence
from dataset import AllFutureHealthDataset, HealthDataset
from dataset import (
DAILY_EXPOSURE_SHAPE,
MONTHLY_EXPOSURE_SHAPE,
AllFutureHealthDataset,
HealthDataset,
)
from targets import PAD_IDX
@@ -25,6 +31,8 @@ class AllFutureSequenceEvalDataset:
labels_file: str,
min_history_events: int = 1,
min_future_events: int = 1,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
) -> None:
base = AllFutureHealthDataset(
data_prefix=data_prefix,
@@ -32,6 +40,8 @@ class AllFutureSequenceEvalDataset:
split="train",
min_history_events=min_history_events,
min_future_events=min_future_events,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
)
self.base = base
@@ -43,11 +53,12 @@ class AllFutureSequenceEvalDataset:
for patient in base.patients:
labels = np.asarray(patient["labels"], dtype=np.int64)
times = np.asarray(patient["times"], dtype=np.float32)
times_days = np.asarray(patient["times_days"], dtype=np.float32)
if labels.size < 2:
continue
input_len = int(labels.size - 1)
self.samples.append(
{
sample := {
"eid": int(patient["eid"]),
"event_seq": labels[:-1],
"time_seq": times[:-1],
@@ -57,13 +68,21 @@ class AllFutureSequenceEvalDataset:
"sex": int(patient["sex"]),
}
)
if base.exposure_cache is not None:
exposure_index = base._exposure_indices_for_inputs(
eid=int(patient["eid"]),
input_events=labels[:-1],
input_times_days=times_days[:-1],
)
if exposure_index is not None:
sample["exposure_index"] = exposure_index
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
s = self.samples[idx]
return {
out = {
"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(),
@@ -71,6 +90,11 @@ class AllFutureSequenceEvalDataset:
"readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
"sex": torch.tensor(s["sex"], dtype=torch.long),
}
if "exposure_index" in s:
daily, monthly = self.base._load_exposure_windows(s["exposure_index"])
out["exposure_daily"] = daily
out["exposure_monthly"] = monthly
return out
def load_sequence_eval_dataset(
@@ -82,6 +106,8 @@ def load_sequence_eval_dataset(
include_no_event_in_uts_target: bool,
min_history_events: int,
min_future_events: int,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
):
mode = str(model_target_mode).lower()
if mode == "next_token":
@@ -90,6 +116,8 @@ def load_sequence_eval_dataset(
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
)
if mode == "all_future":
return AllFutureSequenceEvalDataset(
@@ -97,6 +125,8 @@ def load_sequence_eval_dataset(
labels_file=labels_file,
min_history_events=min_history_events,
min_future_events=min_future_events,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
)
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
@@ -117,7 +147,7 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
readout_mask = pad_sequence(
[s["readout_mask"] for s in batch], batch_first=True, padding_value=False
)
return {
out = {
"event_seq": event_seq,
"time_seq": time_seq,
"padding_mask": event_seq > PAD_IDX,
@@ -126,3 +156,34 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
"readout_mask": readout_mask,
"sex": torch.stack([s["sex"] for s in batch]),
}
if any("exposure_daily" in s for s in batch):
out["exposure_daily"] = _pad_eval_exposure(
batch,
"exposure_daily",
DAILY_EXPOSURE_SHAPE,
)
out["exposure_monthly"] = _pad_eval_exposure(
batch,
"exposure_monthly",
MONTHLY_EXPOSURE_SHAPE,
)
return out
def _pad_eval_exposure(
batch: List[Dict[str, torch.Tensor]],
key: str,
shape: tuple[int, int],
) -> torch.Tensor:
max_len = max(int(s["event_seq"].numel()) for s in batch)
out = torch.full(
(len(batch), max_len, shape[0], shape[1]),
float("nan"),
dtype=torch.float32,
)
for idx, sample in enumerate(batch):
value = sample.get(key)
if value is None:
continue
out[idx, : int(value.size(0))] = value
return out

23
evaluate_auc.py
View File

@@ -321,9 +321,16 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
n_head=int(cfg_get(args, cfg, "n_head", 10)),
n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
target_mode=model_target_mode,
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)),
exposure_d_model=cfg_get(args, cfg, "exposure_d_model", None),
exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 3)),
exposure_n_convnext_blocks=int(cfg_get(args, cfg, "exposure_n_convnext_blocks", 2)),
exposure_conv_kernel_size=int(cfg_get(args, cfg, "exposure_conv_kernel_size", 7)),
exposure_mlp_ratio=float(cfg_get(args, cfg, "exposure_mlp_ratio", 4.0)),
exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)),
)
@@ -524,6 +531,16 @@ def infer_readout_hidden(
sex=batch_dev["sex"][active],
padding_mask=padding_mask[active],
t_query=time_seq[active, pos],
exposure_daily=(
batch_dev["exposure_daily"][active]
if "exposure_daily" in batch_dev
else None
),
exposure_monthly=(
batch_dev["exposure_monthly"][active]
if "exposure_monthly" in batch_dev
else None
),
target_mode="all_future",
)
hidden[active, pos, :] = hidden_pos.float()
@@ -534,6 +551,8 @@ def infer_readout_hidden(
time_seq=time_seq,
sex=batch_dev["sex"],
padding_mask=padding_mask,
exposure_daily=batch_dev.get("exposure_daily"),
exposure_monthly=batch_dev.get("exposure_monthly"),
target_mode="next_token",
)
ro = readout(
@@ -1317,6 +1336,8 @@ def main() -> None:
include_no_event_in_uts_target=include_no_event,
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
exposure_cache_dir=cfg.get("exposure_cache_dir", None),
mask_onset_exposure=bool(cfg.get("mask_onset_exposure", False)),
)
validate_dataset_metadata(dataset, cfg)

33
models.py
View File

@@ -7,9 +7,7 @@ import torch.nn.functional as F
from backbones import (
AgeSinusoidalEncoding,
GPTBlock,
GaussianRBFTimeBasis,
TimesNetExposureEncoder,
TimeRoPE,
)
from targets import PAD_IDX
@@ -30,7 +28,6 @@ class DeepHealth(nn.Module):
n_head: int,
n_hist_layer: int,
target_mode: str = "next_token", # "next_token" or "all_future"
time_mode: str = "relative", # "relative" or "absolute"
dist_mode: str = "exponential", # "exponential", "weibull" or "mixed"
dropout: float = 0.0,
use_exposure_encoder: bool = False,
@@ -48,9 +45,6 @@ class DeepHealth(nn.Module):
if target_mode not in ["next_token", "all_future"]:
raise ValueError(
"target_mode must be either 'next_token' or 'all_future'")
if time_mode not in ["relative", "absolute"]:
raise ValueError(
"time_mode must be either 'relative' or 'absolute'")
if dist_mode not in ["exponential", "weibull", "mixed"]:
raise ValueError(
"dist_mode must be either 'exponential', 'weibull' or 'mixed'")
@@ -58,7 +52,6 @@ class DeepHealth(nn.Module):
self.gender_embedding = nn.Embedding(
2, n_embd) # Assuming binary gender
self.target_mode = target_mode
self.time_mode = time_mode
self.dist_mode = dist_mode
self.use_exposure_encoder = use_exposure_encoder
self.n_embd = n_embd
@@ -94,32 +87,14 @@ class DeepHealth(nn.Module):
nn.init.zeros_(self.rho_death_head.weight)
nn.init.constant_(self.rho_death_head.bias, 0.5413)
if time_mode == "absolute":
self.age_encoding = AgeSinusoidalEncoding(n_embd)
self.blocks = nn.ModuleList([
GPTBlock(
n_embd=n_embd,
n_head=n_head,
use_time_rope=False,
use_rbf_bias=False,
mlp_dropout=dropout,
) for _ in range(n_hist_layer)
])
self.rope = None
self.rbf = None
elif time_mode == "relative":
self.age_encoding = None
self.blocks = nn.ModuleList([
GPTBlock(
n_embd=n_embd,
n_head=n_head,
use_time_rope=True,
use_rbf_bias=True,
mlp_dropout=dropout,
) for _ in range(n_hist_layer)
])
self.rope = TimeRoPE(n_embd // n_head)
self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
self.final_ln = nn.LayerNorm(n_embd)
self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
@@ -291,14 +266,8 @@ class DeepHealth(nn.Module):
h_disease = h_disease + sex_emb
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
rope_cache = None
rbf_cache = None
if self.time_mode == "absolute":
h_disease = h_disease + self.age_encoding(t_disease)
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
elif self.time_mode == "relative":
rope_cache = self.rope.precompute_cache(t_disease)
rbf_cache = self.rbf.precompute_cache(t_disease)
attn_mask = self._make_history_attn_mask(
padding_mask=padding_mask,
@@ -308,8 +277,6 @@ class DeepHealth(nn.Module):
for block in self.blocks:
h_disease = block(
h_disease,
rope_cache=rope_cache,
rbf_cache=rbf_cache,
attn_mask=attn_mask,
)
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)

21
prepare_data.py
View File

@@ -35,6 +35,17 @@ import pandas as pd # Pandas for data manipulation
import tqdm # Progress bar for chunk processing
def _unique_preserve_order(values):
"""Return unique values while preserving first-seen order."""
seen = set()
out = []
for value in values:
if value not in seen:
seen.add(value)
out.append(value)
return out
# CSV mapping field IDs to human-readable names
field_map_file = "field_ids_enriched.csv"
@@ -72,7 +83,7 @@ basic_info_fields = [
# Fields needed for tabular extraction from raw CSV.
tabular_fields = _unique_preserve_order(
basic_info_fields + assessment_fields + exposure_fields
basic_info_fields + assessment_fields + exposure_fields + ["date_of_birth"]
)
# TSV mapping field IDs to ICD10-related date columns
@@ -144,6 +155,7 @@ for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
),
errors="coerce",
)
ukb_chunk["date_of_birth"] = dob.dt.strftime("%Y-%m-%d")
# Use only date variables that actually exist in the current chunk
present_date_vars = [c for c in date_vars if c in ukb_chunk.columns]
@@ -253,8 +265,11 @@ data = data[np.isin(data[:, 0], valid_eids)]
final_tabular = final_tabular.loc[valid_eids]
final_tabular = final_tabular.convert_dtypes()
# Save basic sex information separately.
basic_info = final_tabular[["sex"]].copy()
# Save basic information needed by the model and exposure-date alignment.
basic_cols = ["sex"]
if "date_of_birth" in final_tabular.columns:
basic_cols.append("date_of_birth")
basic_info = final_tabular[basic_cols].copy()
basic_info.to_csv("ukb_basic_info.csv")
# Save event data

300
prepare_exposure_cache.py Normal file
View File

@@ -0,0 +1,300 @@
"""Build a random-access exposure cache from disease-level parquet files.
The README-described exposure dataset is stored as one daily and one monthly
parquet file per disease. That layout is good for disease-specific analysis but
too expensive for mini-batch training, where we need exposure windows aligned
to arbitrary event sequences.
This script converts those parquet files into a compact directory:
exposure_keys.npy uint64 legacy keys, key = (eid << 16) | raw_token
exposure_eid.npy int64 eid per exposure row
exposure_token.npy int32 raw disease token per exposure row
exposure_onset_date.npy datetime64[D] onset date per exposure row
exposure_daily.npy float32 memmap, shape (N, 1826, 4)
channels: tmean, tmax, tmin, rhmean
exposure_monthly.npy float32 memmap, shape (N, 241, 2)
channels: tmean, rhmean
exposure_quality.npy float32 memmap, shape (N, 4)
n_days, n_rh_days, n_months, n_rh_months
exposure_manifest.json metadata
The raw token convention follows the exposure README: padding=0, checkup=1,
and the first row of labels.csv is token=2. The model dataset inserts
<NO_EVENT> at token 2 and shifts real disease tokens by +1 internally; dataset
lookup converts back to these raw tokens before reading this cache. Dataset
alignment uses (eid, raw_token, onset_date - date_of_birth) so that raw
calendar dates in the exposure files match the age-day event times used by the
model.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Iterable
import numpy as np
import pandas as pd
DAILY_LENGTH = 1826
MONTHLY_LENGTH = 241
DAILY_CHANNELS = ("tmean", "tmax", "tmin", "rhmean")
MONTHLY_CHANNELS = ("tmean", "rhmean")
QUALITY_COLUMNS = (
"n_days_nonmissing",
"n_rh_days_nonmissing",
"n_months_nonmissing",
"n_rh_months_nonmissing",
)
def encode_exposure_key(eid: np.ndarray, raw_token: np.ndarray) -> np.ndarray:
eid_u64 = np.asarray(eid, dtype=np.uint64)
token_u64 = np.asarray(raw_token, dtype=np.uint64)
if np.any(token_u64 >= (1 << 16)):
raise ValueError("raw_token must fit in 16 bits")
return (eid_u64 << np.uint64(16)) | token_u64
def _daily_columns() -> list[str]:
cols: list[str] = []
for name in DAILY_CHANNELS:
cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_LENGTH))
return cols
def _monthly_columns() -> list[str]:
cols: list[str] = []
for name in MONTHLY_CHANNELS:
cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_LENGTH))
return cols
def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
"""Return the subset of requested columns present in a parquet file."""
try:
import pyarrow.parquet as pq
except ImportError as exc:
raise ImportError(
"prepare_exposure_cache.py requires pyarrow. Install requirements "
"or run `pip install pyarrow`."
) from exc
schema_names = set(pq.ParquetFile(path).schema.names)
return [col for col in columns if col in schema_names]
def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame:
return pd.read_parquet(path, columns=columns)
def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: int) -> np.ndarray:
arr = df.reindex(columns=cols).to_numpy(dtype=np.float32, copy=True)
return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1)
def _count_rows(summary: pd.DataFrame) -> int:
if "n_cases" in summary.columns:
return int(summary["n_cases"].sum())
return int(sum(pd.read_parquet(path, columns=["eid"]).shape[0] for path in summary["daily_path"]))
def build_exposure_cache(
*,
exposure_dir: str | Path,
output_dir: str | Path,
summary_file: str = "summary.csv",
overwrite: bool = False,
) -> int:
exposure_dir = Path(exposure_dir)
output_dir = Path(output_dir)
summary_path = exposure_dir / summary_file
if not summary_path.is_file():
raise FileNotFoundError(f"summary.csv not found: {summary_path}")
output_dir.mkdir(parents=True, exist_ok=True)
keys_path = output_dir / "exposure_keys.npy"
eid_path = output_dir / "exposure_eid.npy"
token_path = output_dir / "exposure_token.npy"
onset_date_path = output_dir / "exposure_onset_date.npy"
daily_path = output_dir / "exposure_daily.npy"
monthly_path = output_dir / "exposure_monthly.npy"
quality_path = output_dir / "exposure_quality.npy"
manifest_path = output_dir / "exposure_manifest.json"
outputs = [
keys_path,
eid_path,
token_path,
onset_date_path,
daily_path,
monthly_path,
quality_path,
manifest_path,
]
if any(path.exists() for path in outputs) and not overwrite:
raise FileExistsError(
f"{output_dir} already contains exposure cache files; pass --overwrite"
)
summary = pd.read_csv(summary_path)
required = {"label_code", "daily_file", "monthly_file"}
missing = required - set(summary.columns)
if missing:
raise ValueError(f"{summary_path} is missing columns: {sorted(missing)}")
summary = summary.copy()
summary["daily_path"] = summary["daily_file"].map(lambda name: exposure_dir / str(name))
summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name))
n_rows = _count_rows(summary)
keys = np.lib.format.open_memmap(keys_path, mode="w+", dtype=np.uint64, shape=(n_rows,))
eids_mm = np.lib.format.open_memmap(eid_path, mode="w+", dtype=np.int64, shape=(n_rows,))
tokens_mm = np.lib.format.open_memmap(token_path, mode="w+", dtype=np.int32, shape=(n_rows,))
onset_dates_mm = np.lib.format.open_memmap(
onset_date_path,
mode="w+",
dtype="datetime64[D]",
shape=(n_rows,),
)
daily_mm = np.lib.format.open_memmap(
daily_path,
mode="w+",
dtype=np.float32,
shape=(n_rows, DAILY_LENGTH, len(DAILY_CHANNELS)),
)
monthly_mm = np.lib.format.open_memmap(
monthly_path,
mode="w+",
dtype=np.float32,
shape=(n_rows, MONTHLY_LENGTH, len(MONTHLY_CHANNELS)),
)
quality_mm = np.lib.format.open_memmap(
quality_path,
mode="w+",
dtype=np.float32,
shape=(n_rows, len(QUALITY_COLUMNS)),
)
daily_cols = _daily_columns()
monthly_cols = _monthly_columns()
offset = 0
for row in summary.itertuples(index=False):
daily_file = Path(row.daily_path)
monthly_file = Path(row.monthly_path)
if not daily_file.is_file():
raise FileNotFoundError(f"Missing daily parquet: {daily_file}")
if not monthly_file.is_file():
raise FileNotFoundError(f"Missing monthly parquet: {monthly_file}")
daily_read_cols = [
"eid",
"onset_date",
"token",
*_safe_columns(daily_file, daily_cols),
*_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]),
]
monthly_read_cols = [
"eid",
"onset_date",
"token",
*_safe_columns(monthly_file, monthly_cols),
*_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]),
]
daily_df = _read_parquet_columns(daily_file, daily_read_cols)
monthly_df = _read_parquet_columns(monthly_file, monthly_read_cols)
if len(daily_df) != len(monthly_df):
raise ValueError(
f"Daily/monthly row count mismatch for {row.label_code}: "
f"{len(daily_df)} vs {len(monthly_df)}"
)
monthly_df = monthly_df.set_index(["eid", "onset_date", "token"]).reindex(
pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
).reset_index()
n = len(daily_df)
end = offset + n
if end > n_rows:
raise RuntimeError("Exposure cache row count exceeded preallocated size")
keys[offset:end] = encode_exposure_key(
daily_df["eid"].to_numpy(dtype=np.int64),
daily_df["token"].to_numpy(dtype=np.int64),
)
eids_mm[offset:end] = daily_df["eid"].to_numpy(dtype=np.int64)
tokens_mm[offset:end] = daily_df["token"].to_numpy(dtype=np.int32)
onset_dates_mm[offset:end] = pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).to_numpy(dtype="datetime64[D]")
daily_mm[offset:end] = _reshape_window(
daily_df,
daily_cols,
DAILY_LENGTH,
len(DAILY_CHANNELS),
)
monthly_mm[offset:end] = _reshape_window(
monthly_df,
monthly_cols,
MONTHLY_LENGTH,
len(MONTHLY_CHANNELS),
)
quality_mm[offset:end, 0] = daily_df.get("n_days_nonmissing", np.nan)
quality_mm[offset:end, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
quality_mm[offset:end, 2] = monthly_df.get("n_months_nonmissing", np.nan)
quality_mm[offset:end, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
offset = end
if offset != n_rows:
keys.flush()
eids_mm.flush()
tokens_mm.flush()
onset_dates_mm.flush()
daily_mm.flush()
monthly_mm.flush()
quality_mm.flush()
keys = np.lib.format.open_memmap(keys_path, mode="r+", dtype=np.uint64, shape=(offset,))
raise RuntimeError(
f"Expected {n_rows} rows from summary but wrote {offset}. "
"Regenerate summary.csv or remove n_cases before building."
)
manifest = {
"source_dir": str(exposure_dir),
"n_rows": int(n_rows),
"legacy_key": "(eid << 16) | raw_token",
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
"requires_basic_info_column": "date_of_birth",
"daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
"daily_channels": list(DAILY_CHANNELS),
"monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
"monthly_channels": list(MONTHLY_CHANNELS),
"quality_columns": list(QUALITY_COLUMNS),
"raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",
}
manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
return int(n_rows)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--exposure-dir", required=True)
parser.add_argument("--output-dir", default="ukb_exposure_cache")
parser.add_argument("--summary-file", default="summary.csv")
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args()
n_rows = build_exposure_cache(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
summary_file=args.summary_file,
overwrite=args.overwrite,
)
print(f"Wrote {n_rows:,} exposure rows to {args.output_dir}")
if __name__ == "__main__":
main()

1
requirements.txt
View File

@@ -3,3 +3,4 @@ pandas
torch
tqdm
scikit-learn
pyarrow

68
train_next_step.py
View File

@@ -47,6 +47,11 @@ MODEL_INPUT_KEYS = (
"padding_mask",
)
EXPOSURE_INPUT_KEYS = (
"exposure_daily",
"exposure_monthly",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
@@ -68,9 +73,16 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--n_embd", type=int, default=120)
parser.add_argument("--n_head", type=int, default=10)
parser.add_argument("--n_hist_layer", type=int, default=12)
parser.add_argument("--time_mode", type=str, default="relative",
choices=["relative", "absolute"])
parser.add_argument("--dropout", type=float, default=0.0)
parser.add_argument("--exposure_cache_dir", type=str, default=None)
parser.add_argument("--mask_onset_exposure", action="store_true")
parser.add_argument("--exposure_d_model", type=int, default=None)
parser.add_argument("--exposure_n_layers", type=int, default=2)
parser.add_argument("--exposure_top_k", type=int, default=3)
parser.add_argument("--exposure_n_convnext_blocks", type=int, default=2)
parser.add_argument("--exposure_conv_kernel_size", type=int, default=7)
parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0)
parser.add_argument("--no_exposure_gate", action="store_true")
parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "uts"])
@@ -137,9 +149,16 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
n_head=args.n_head,
n_hist_layer=args.n_hist_layer,
target_mode="next_token",
time_mode=args.time_mode,
dist_mode="exponential",
dropout=args.dropout,
use_exposure_encoder=args.exposure_cache_dir is not None,
exposure_d_model=args.exposure_d_model,
exposure_n_layers=args.exposure_n_layers,
exposure_top_k=args.exposure_top_k,
exposure_n_convnext_blocks=args.exposure_n_convnext_blocks,
exposure_conv_kernel_size=args.exposure_conv_kernel_size,
exposure_mlp_ratio=args.exposure_mlp_ratio,
exposure_use_gate=not args.no_exposure_gate,
)
@@ -201,18 +220,24 @@ def compute_next_step_loss(
device: torch.device,
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
batch_cpu = batch
input_keys = list(MODEL_INPUT_KEYS)
input_keys.extend(key for key in EXPOSURE_INPUT_KEYS if key in batch_cpu)
batch = move_batch_to_device(
{key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
{key: batch_cpu[key] for key in input_keys},
device,
)
model_out = model(
event_seq=batch["event_seq"],
time_seq=batch["time_seq"],
sex=batch["sex"],
padding_mask=batch["padding_mask"],
target_mode="next_token",
return_output=True,
)
model_kwargs = {
"event_seq": batch["event_seq"],
"time_seq": batch["time_seq"],
"sex": batch["sex"],
"padding_mask": batch["padding_mask"],
"target_mode": "next_token",
"return_output": True,
}
if "exposure_daily" in batch:
model_kwargs["exposure_daily"] = batch["exposure_daily"]
model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
model_out = model(**model_kwargs)
if not isinstance(model_out, DeepHealthOutput):
raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
targets = build_augmented_next_step_targets(
@@ -329,6 +354,16 @@ def build_metadata(
"dataset_metadata": {
"vocab_size": int(dataset.vocab_size),
},
"use_exposure_encoder": args.exposure_cache_dir is not None,
"exposure_cache_dir": args.exposure_cache_dir,
"mask_onset_exposure": bool(args.mask_onset_exposure),
"exposure_d_model": args.exposure_d_model,
"exposure_n_layers": int(args.exposure_n_layers),
"exposure_top_k": int(args.exposure_top_k),
"exposure_n_convnext_blocks": int(args.exposure_n_convnext_blocks),
"exposure_conv_kernel_size": int(args.exposure_conv_kernel_size),
"exposure_mlp_ratio": float(args.exposure_mlp_ratio),
"exposure_use_gate": not bool(args.no_exposure_gate),
"split_sizes": {
"train": int(len(train_subset)),
"val": int(len(val_subset)),
@@ -347,8 +382,10 @@ def main() -> None:
run_dir, run_name = create_unique_run_dir(
lambda timestamp: (
f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
f"gap_{args.no_event_interval_years:g}y_{timestamp}"
f"absolute_exponential_next_token_{args.target_mode}_"
f"gap_{args.no_event_interval_years:g}y_"
f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_"
f"{timestamp}"
)
)
logger = setup_logging(run_dir)
@@ -356,12 +393,15 @@ def main() -> None:
logger.info(f"Starting next-step training run: {run_name}")
logger.info(f"Device: {device}")
logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
dataset = HealthDataset(
data_prefix=args.data_prefix,
labels_file=args.labels_file,
no_event_interval_years=args.no_event_interval_years,
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
exposure_cache_dir=args.exposure_cache_dir,
mask_onset_exposure=args.mask_onset_exposure,
)
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
train_subset, val_subset, test_subset = split_dataset_by_eid_files(