From 45a857d1a613836be60ce0b9a8c271872e4db198 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Tue, 7 Jul 2026 17:21:52 +0800 Subject: [PATCH] Add exposure cache and keep absolute time only --- backbones.py | 130 +---------------- dataset.py | 223 +++++++++++++++++++++++++++- eval_data.py | 69 ++++++++- evaluate_auc.py | 23 ++- models.py | 53 ++----- prepare_data.py | 21 ++- prepare_exposure_cache.py | 300 ++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + train_next_step.py | 68 +++++++-- 9 files changed, 690 insertions(+), 198 deletions(-) create mode 100644 prepare_exposure_cache.py diff --git a/backbones.py b/backbones.py index 66f3bd7..678e834 100644 --- a/backbones.py +++ b/backbones.py @@ -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 diff --git a/dataset.py b/dataset.py index 0691fdf..fb6fe29 100644 --- a/dataset.py +++ b/dataset.py @@ -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 diff --git a/eval_data.py b/eval_data.py index 7b45a48..af585c8 100644 --- a/eval_data.py +++ b/eval_data.py @@ -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 diff --git a/evaluate_auc.py b/evaluate_auc.py index d8a0a69..93a7d1f 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -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) diff --git a/models.py b/models.py index 9495980..97543be 100644 --- a/models.py +++ b/models.py @@ -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.age_encoding = AgeSinusoidalEncoding(n_embd) + self.blocks = nn.ModuleList([ + GPTBlock( + n_embd=n_embd, + n_head=n_head, + mlp_dropout=dropout, + ) for _ in range(n_hist_layer) + ]) 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) + h_disease = h_disease + self.age_encoding(t_disease) + h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) 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) diff --git a/prepare_data.py b/prepare_data.py index b100767..e458c71 100644 --- a/prepare_data.py +++ b/prepare_data.py @@ -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 diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py new file mode 100644 index 0000000..cc36e4d --- /dev/null +++ b/prepare_exposure_cache.py @@ -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 + 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() diff --git a/requirements.txt b/requirements.txt index a4871bd..9e01347 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ pandas torch tqdm scikit-learn +pyarrow diff --git a/train_next_step.py b/train_next_step.py index 87521ea..b19570a 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -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(