Initial DeepHealthExpo next-token codebase

2026-07-07 15:43:11 +08:00
commit e110e99f66
26 changed files with 10278 additions and 0 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -0,0 +1,37 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+
+# Environments
+.venv/
+venv/
+env/
+
+# Training outputs
+runs/
+outputs/
+checkpoints/
+*.pt
+*.pth
+
+# Large/local data
+ukb_*_data.npy
+ukb_event_data.npy
+ukb_other_info.npy
+ukb_data.csv
+ukb_basic_info.csv
+ukb_train_eid.csv
+ukb_val_eid.csv
+ukb_test_eid.csv
+UKB_*.csv
+cate_types.csv
+*.parquet
+
+# OS/editor
+.DS_Store
+Thumbs.db
+.vscode/
--- a/28
+++ b/28
@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Jiarui Li
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT INCLUDING NEGLIGENCE OR OTHERWISE ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--- a/README.md
+++ b/README.md
@@ -0,0 +1,60 @@
+# DeepHealthExpo
+
+Next-token DeepHealth training code for disease-event sequence modeling with optional extra/exposure information.
+
+This repository is a clean code-only extraction from the main DeepHealth project. It keeps the next-token training path and reusable model/data utilities, while excluding large UKB data files, trained checkpoints, result folders, and all-future training entry points.
+
+## Included
+
+- `train_next_step.py`: next-token / UTS training entry point.
+- `dataset.py`: next-step event sequence dataset with unified extra-info tokens.
+- `models.py`, `backbones.py`: DeepHealth Transformer backbone.
+- `losses.py`, `readouts.py`, `targets.py`: training targets, losses, and readout utilities.
+- `evaluate_auc.py`, `evaluate_token_auc.py`: next-token checkpoint evaluation utilities.
+- `prepare_data.py`, `prepare_event_dates.py`, `event_date_utils.py`: data preparation helpers.
+- `extra_info_types_*.txt`: reusable extra-info type selections.
+
+## Not Included
+
+The repository intentionally does not include raw or derived UKB arrays, split files, checkpoints, or run outputs.
+
+Expected local data files for training normally include:
+
+```text
+ukb_event_data.npy
+ukb_other_info.npy
+ukb_basic_info.csv
+ukb_train_eid.csv
+ukb_val_eid.csv
+ukb_test_eid.csv
+cate_types.csv
+```
+
+`labels.csv` and `field_ids_enriched.csv` are included because they define the model vocabulary and preparation metadata.
+
+## Example
+
+```bash
+python train_next_step.py \
+  --data_prefix ukb \
+  --labels_file labels.csv \
+  --extra_info_types_file extra_info_types_exposure_only.txt \
+  --target_mode uts \
+  --time_mode relative
+```
+
+For strict next-token Delphi-style training:
+
+```bash
+python train_next_step.py --target_mode delphi2m --readout_name token
+```
+
+## Exposure Modeling Direction
+
+For onset-aligned environmental exposure parquet files, the first intended extension is single-stream event enhancement:
+
+```text
+disease event token + pre-onset exposure embedding -> same next-token Transformer
+```
+
+The key constraint is that a disease event's own pre-onset exposure must not be used to predict that same disease event.
--- a/backbones.py
+++ b/backbones.py
@@ -0,0 +1,328 @@
+import math
+
+import torch
+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()
+
+    def reset_parameters(self) -> None:
+        """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
+
+        # --- QKV ----------------------------------------------------------
+        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,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=self.scale,
+        )
+
+        # --- Aggregate & project out --------------------------------------
+        out = out.transpose(1, 2).reshape(B, L, H * D)
+        return self.resid_drop(self.out_proj(out))
+
+
+class SwiGLU(nn.Module):
+    def __init__(
+        self,
+        n_embd: int,
+        hidden_dim: int | None = None,
+        dropout: float = 0.0,
+        bias: bool = True,
+    ):
+        super().__init__()
+        hidden_dim = hidden_dim if hidden_dim is not None else int(
+            n_embd * 2.5)
+
+        self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias)  # gate path
+        self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias)  # value path
+        # output projection
+        self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
+        self.drop = nn.Dropout(dropout)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """GPT-style parameter initialization for MLP paths."""
+        nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
+        if self.w1.bias is not None:
+            nn.init.zeros_(self.w1.bias)
+            nn.init.zeros_(self.w2.bias)
+            nn.init.zeros_(self.w3.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """``(B, L, n_embd) -> (B, L, n_embd)``."""
+        return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
+
+
+class GPTBlock(nn.Module):
+    def __init__(
+        self,
+        n_embd: int,
+        n_head: int,
+
+        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)
+        self.ln2 = nn.LayerNorm(n_embd)
+
+    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.mlp(self.ln2(x))
+        return x
+
+
+class TokenAutoDiscretization(nn.Module):
+    def __init__(
+        self,
+        n_cont_types: int,
+        n_bins: int,
+        n_embd: int,
+    ):
+        super().__init__()
+        if n_cont_types <= 0:
+            raise ValueError(f"n_cont_types must be > 0, got {n_cont_types}")
+        if n_bins <= 1:
+            raise ValueError(f"n_bins must be > 1, got {n_bins}")
+        if n_embd <= 0:
+            raise ValueError(f"n_embd must be > 0, got {n_embd}")
+
+        self.n_cont_types = n_cont_types
+        self.n_bins = n_bins
+        self.n_embd = n_embd
+        self.weight = nn.Parameter(torch.empty(n_cont_types, n_bins))
+        self.bias = nn.Parameter(torch.empty(n_cont_types, n_bins))
+        self.bin_emb = nn.Parameter(torch.empty(n_cont_types, n_bins, n_embd))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.bias)
+        nn.init.normal_(self.bin_emb, mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        cont_type_idx: torch.LongTensor,  # (N,)
+        value: torch.Tensor,              # (N,)
+    ) -> torch.Tensor:
+        if cont_type_idx.dim() != 1:
+            raise ValueError(
+                f"cont_type_idx must be 1D, got {tuple(cont_type_idx.shape)}"
+            )
+        if value.dim() != 1:
+            raise ValueError(f"value must be 1D, got {tuple(value.shape)}")
+        if cont_type_idx.numel() != value.numel():
+            raise ValueError("cont_type_idx and value must have the same length")
+
+        w = self.weight[cont_type_idx]      # (N, n_bins)
+        b = self.bias[cont_type_idx]        # (N, n_bins)
+        e = self.bin_emb[cont_type_idx]     # (N, n_bins, D)
+        logits = value.unsqueeze(-1) * w + b
+        probs = torch.softmax(logits, dim=-1)
+        return torch.einsum("nb,nbd->nd", probs, e)
+
+
+
+class AgeSinusoidalEncoding(nn.Module):
+
+    def __init__(self, embedding_dim: int):
+
+        super().__init__()
+        if embedding_dim % 2 != 0:
+            raise ValueError(
+                f"Embedding dimension must be an even number, but got {embedding_dim}")
+
+        self.embedding_dim = embedding_dim
+
+        i = torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
+        divisor = torch.pow(10000, i / self.embedding_dim)
+        self.register_buffer('divisor', divisor)
+        self.linear = nn.Linear(embedding_dim, embedding_dim, bias=False)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+
+        t_years = t
+        # Broadcast (B, L, 1) against (1, 1, D/2) to get (B, L, D/2)
+        args = t_years.unsqueeze(-1) / self.divisor.view(1, 1, -1)
+        # Interleave cos and sin along the last dimension
+        output = torch.zeros(t.shape[0], t.shape[1],
+                             self.embedding_dim, device=t.device)
+        output[:, :, 0::2] = torch.cos(args)
+        output[:, :, 1::2] = torch.sin(args)
+        output = self.linear(output)
+        return output
--- a/dataset.py
+++ b/dataset.py
@@ -0,0 +1,674 @@
+# dataset.py
+from __future__ import annotations
+
+from typing import Dict, Iterable, List, Literal, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Dataset
+
+from targets import (
+    CHECKUP_IDX,
+    DAYS_PER_YEAR,
+    NO_EVENT_IDX,
+    PAD_IDX,
+    build_all_targets,
+)
+
+
+ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
+
+
+def load_label_vocab(
+    labels_file: str,
+    include_no_event: bool = True,
+) -> Tuple[Dict[str, int], Dict[int, str]]:
+    label_id_to_code: Dict[int, str] = {
+        PAD_IDX: "<PAD>",
+        CHECKUP_IDX: "<CHECKUP>",
+    }
+    if include_no_event:
+        label_id_to_code[NO_EVENT_IDX] = "<NO_EVENT>"
+
+    offset = NO_EVENT_IDX + 1 if include_no_event else CHECKUP_IDX + 1
+    label_code_to_id: Dict[str, int] = {}
+    with open(labels_file, encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            parts = line.strip().split()
+            if not parts:
+                continue
+            idx = offset + i
+            code = parts[0]
+            label_code_to_id[code] = idx
+            label_id_to_code[idx] = code
+    return label_code_to_id, label_id_to_code
+
+
+def _insert_gap_no_event_tokens(
+    times_days: np.ndarray,
+    labels: np.ndarray,
+    interval_years: float = 5.0,
+) -> Tuple[np.ndarray, np.ndarray]:
+    if len(times_days) < 2:
+        return times_days, labels
+
+    step_days = interval_years * DAYS_PER_YEAR
+    unique_times = np.unique(times_days.astype(np.float64))
+    extra_times: List[float] = []
+
+    for i in range(len(unique_times) - 1):
+        t_left = float(unique_times[i])
+        t_right = float(unique_times[i + 1])
+        if t_right - t_left <= step_days:
+            continue
+        first = np.ceil((t_left + 1e-6) / step_days) * step_days
+        t = first
+        while t < t_right - 1e-6:
+            extra_times.append(t)
+            t += step_days
+
+    if not extra_times:
+        return times_days, labels
+
+    extra_arr = np.array(extra_times, dtype=np.float32)
+    no_event_labels = np.full(len(extra_arr), NO_EVENT_IDX, dtype=np.int64)
+    all_times = np.concatenate([times_days.astype(np.float32), extra_arr])
+    all_labels = np.concatenate([labels.astype(np.int64), no_event_labels])
+    order = np.lexsort((all_labels, all_times))
+    return all_times[order], all_labels[order]
+
+
+class _ExpoBaseDataset(Dataset):
+    def __init__(
+        self,
+        data_prefix: str = "ukb",
+        labels_file: str = "labels.csv",
+        no_event_interval_years: float = 5.0,
+        include_no_event_in_uts_target: bool = False,
+        extra_info_types: Iterable[int] | None = None,
+    ) -> 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.requested_extra_info_types = (
+            None
+            if extra_info_types is None
+            else list(dict.fromkeys(int(t) for t in extra_info_types))
+        )
+
+        self.label_code_to_id, self.label_id_to_code = load_label_vocab(
+            labels_file,
+            include_no_event=True,
+        )
+
+        event_data = np.load(f"{data_prefix}_event_data.npy")
+        if event_data.ndim != 2 or event_data.shape[1] < 3:
+            raise ValueError(f"event_data must have shape (N, 3+), got {event_data.shape}")
+        event_data = event_data[:, :3].copy()
+        order = np.lexsort((event_data[:, 2], event_data[:, 1], event_data[:, 0]))
+        self.event_data = event_data[order]
+
+        basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
+        other_info = np.load(f"{data_prefix}_other_info.npy")
+        if other_info.ndim != 2 or other_info.shape[1] != 5:
+            raise ValueError(
+                f"other_info must have shape (N, 5), got {other_info.shape}"
+            )
+        cate_types = pd.read_csv("cate_types.csv")
+        required_cate_cols = {"type", "name", "n_categories"}
+        missing_cate_cols = required_cate_cols - set(cate_types.columns)
+        if missing_cate_cols:
+            raise ValueError(
+                f"cate_types.csv is missing columns: {sorted(missing_cate_cols)}"
+            )
+
+        basic_table.index = basic_table.index.astype(np.int64)
+
+        unique_eids = np.unique(self.event_data[:, 0].astype(np.int64))
+        basic_table = basic_table.loc[unique_eids]
+
+        self._prepare_sex(basic_table, unique_eids)
+        self._prepare_other_info(other_info, cate_types, unique_eids)
+
+        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
+        max_id_in_data += 1
+        self.vocab_size = max(max_id_in_vocab, max_id_in_data) + 1
+
+        if not self.include_no_event_in_uts_target:
+            self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
+        else:
+            self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX}
+
+    def _prepare_sex(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None:
+        sex_values = pd.to_numeric(basic_table["sex"], errors="coerce").to_numpy()
+        if np.isnan(sex_values).any():
+            raise ValueError("sex column contains missing or non-numeric values")
+
+        sex_values = sex_values.astype(np.int64)
+        sex_unique = np.unique(sex_values)
+        if np.all(np.isin(sex_unique, [0, 1])):
+            sex01 = sex_values
+        elif np.all(np.isin(sex_unique, [1, 2])):
+            sex01 = sex_values - 1
+        else:
+            raise ValueError(
+                f"Unexpected sex values: {sex_unique.tolist()}. Expected {{0,1}} or {{1,2}}."
+            )
+        self.sex_mapping = {int(eid): int(s) for eid, s in zip(unique_eids, sex01)}
+
+    def _prepare_other_info(
+        self,
+        other_info: np.ndarray,
+        cate_types: pd.DataFrame,
+        unique_eids: np.ndarray,
+    ) -> None:
+        other_info = other_info.copy()
+        other_info[:, 0] = other_info[:, 0].astype(np.int64)
+        other_info[:, 1] = other_info[:, 1].astype(np.int64)
+        other_info[:, 3] = other_info[:, 3].astype(np.int64)
+
+        available_types = sorted(
+            int(t) for t in np.unique(other_info[:, 1]) if int(t) > 0
+        )
+        if self.requested_extra_info_types is None:
+            selected_types = available_types
+        else:
+            selected_types = self.requested_extra_info_types
+            missing = sorted(set(selected_types) - set(available_types))
+            if missing:
+                raise ValueError(f"Requested extra_info_types not found: {missing}")
+
+        keep = np.isin(other_info[:, 0].astype(np.int64), unique_eids)
+        keep &= np.isin(other_info[:, 1].astype(np.int64), selected_types)
+        other_info = other_info[keep]
+
+        cate_counts = {
+            int(row["type"]): int(row["n_categories"])
+            for _, row in cate_types.iterrows()
+        }
+        cate_offsets: Dict[int, int] = {}
+        next_offset = 0
+        for type_id in selected_types:
+            if type_id in cate_counts:
+                cate_offsets[type_id] = next_offset
+                next_offset += cate_counts[type_id]
+
+        kinds = other_info[:, 3].astype(np.int64)
+        types = other_info[:, 1].astype(np.int64)
+        cate_rows = kinds == 2
+        for type_id in np.unique(types[cate_rows]):
+            type_id = int(type_id)
+            if type_id not in cate_offsets:
+                raise ValueError(
+                    f"type {type_id} appears categorical but is missing from cate_types.csv"
+                )
+            row_mask = cate_rows & (types == type_id)
+            local_value = other_info[row_mask, 2].astype(np.int64)
+            other_info[row_mask, 2] = local_value + cate_offsets[type_id]
+
+        cont_type_ids = [
+            int(t)
+            for t in selected_types
+            if np.any((types == int(t)) & (kinds == 1))
+        ]
+
+        self.extra_info_types = selected_types
+        self.cate_type_offsets = cate_offsets
+        self.n_types = (max(selected_types) + 1) if selected_types else 1
+        self.cont_type_ids = cont_type_ids
+        self.n_cont_types = len(cont_type_ids)
+        self.n_categories = next_offset + 1
+
+        order = np.lexsort((other_info[:, 4], other_info[:, 1], other_info[:, 0]))
+        other_info = other_info[order]
+        self.other_info_by_eid: Dict[int, Dict[str, np.ndarray]] = {}
+
+        for eid in unique_eids.astype(np.int64):
+            self.other_info_by_eid[int(eid)] = {
+                "other_type": np.zeros(0, dtype=np.int64),
+                "other_value": np.zeros(0, dtype=np.float32),
+                "other_value_kind": np.zeros(0, dtype=np.int64),
+                "other_time": np.zeros(0, dtype=np.float32),
+            }
+
+        if len(other_info) == 0:
+            return
+
+        eids, starts = np.unique(other_info[:, 0].astype(np.int64), return_index=True)
+        ends = np.concatenate([starts[1:], [len(other_info)]])
+        for eid_raw, start, end in zip(eids, starts, ends):
+            rows = other_info[start:end]
+            self.other_info_by_eid[int(eid_raw)] = {
+                "other_type": rows[:, 1].astype(np.int64),
+                "other_value": rows[:, 2].astype(np.float32),
+                "other_value_kind": rows[:, 3].astype(np.int64),
+                "other_time": (rows[:, 4].astype(np.float32) / DAYS_PER_YEAR),
+            }
+
+    def _iter_patient_events(
+        self,
+        *,
+        impute_no_event_gaps: bool,
+    ) -> Iterable[tuple[int, np.ndarray, np.ndarray]]:
+        unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True)
+        ends = np.concatenate([starts[1:], [len(self.event_data)]])
+        for eid_raw, start, end in zip(unique_eids, starts, ends):
+            eid = int(eid_raw)
+            rows = self.event_data[start:end]
+            times_days_raw = rows[:, 1].astype(np.float32)
+            labels_raw = rows[:, 2].astype(np.int64)
+
+            if len(labels_raw) == 0:
+                yield eid, times_days_raw, labels_raw
+                continue
+
+            labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw)
+            if not impute_no_event_gaps:
+                yield eid, times_days_raw, labels_raw
+                continue
+
+            times_days, labels = _insert_gap_no_event_tokens(
+                times_days_raw,
+                labels_raw,
+                interval_years=self.no_event_interval_years,
+            )
+            yield eid, times_days, labels
+
+    def _split_features(self, eid: int) -> Optional[Dict]:
+        other_info = self.other_info_by_eid.get(eid)
+        if other_info is None:
+            return None
+        return {
+            "sex": self.sex_mapping[eid],
+            **other_info,
+        }
+
+class NextStepHealthDataset(_ExpoBaseDataset):
+    """
+    Dataset for next-token and next-time-point losses with unified other-info
+    tokens.
+
+    Returned targets cover both:
+        - Delphi2MLoss: target_event_seq, target_time_seq
+        - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot
+    """
+
+    CACHE_VERSION = 3
+
+    def __init__(
+        self,
+        data_prefix: str = "ukb",
+        labels_file: str = "labels.csv",
+        no_event_interval_years: float = 5.0,
+        include_no_event_in_uts_target: bool = False,
+        extra_info_types: Iterable[int] | None = None,
+    ) -> 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,
+            extra_info_types=extra_info_types,
+        )
+
+        self.samples: List[Dict] = []
+        for eid, times_days, labels in self._iter_patient_events(
+            impute_no_event_gaps=True,
+        ):
+            if len(labels) < 2:
+                continue
+
+            features = self._split_features(eid)
+            if features is None:
+                continue
+
+            target_pack = build_all_targets(
+                labels=labels,
+                times_days=times_days,
+                vocab_size=self.vocab_size,
+                ignored_uts_target_ids=self.ignored_uts_target_ids,
+                require_sorted=True,
+            )
+
+            self.samples.append({
+                "eid": eid,
+                "event_seq": target_pack.next_token.input_events,
+                "time_seq": target_pack.next_token.input_times_years,
+                "target_event_seq": target_pack.next_token.target_events,
+                "target_time_seq": target_pack.next_token.target_times_years,
+                "readout_mask": target_pack.unique_time_set.readout_mask,
+                "target_dt_unique": target_pack.unique_time_set.target_dt_unique,
+                "target_multi_hot": target_pack.unique_time_set.target_multi_hot,
+                **features,
+            })
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Dict:
+        s = self.samples[idx]
+        return {
+            "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),
+            "other_type": torch.from_numpy(s["other_type"]).long(),
+            "other_value": torch.from_numpy(s["other_value"]).float(),
+            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
+            "other_time": torch.from_numpy(s["other_time"]).float(),
+            "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
+            "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
+            "readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
+            "target_dt_unique": torch.from_numpy(s["target_dt_unique"]).float(),
+            "target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(),
+        }
+
+
+class AllFutureHealthDataset(_ExpoBaseDataset):
+    """
+    Dataset with unified other-info tokens and DeepHealthV2-style all-future
+    targets.
+
+    Train samples one query time per patient at each __getitem__ call.
+    Valid/test use random-but-fixed query points. For each patient with N real
+    disease events, N - 2 query points are sampled from the eligible observed
+    time range, with at least one future event after every query.
+    """
+
+    CACHE_VERSION = 5
+
+    def __init__(
+        self,
+        data_prefix: str = "ukb",
+        labels_file: str = "labels.csv",
+        split: Literal["train", "valid", "test"] = "train",
+        no_event_interval_years: float = 5.0,
+        include_no_event_in_uts_target: bool = False,
+        min_history_events: int = 1,
+        min_future_events: int = 1,
+        validation_query_seed: int = 42,
+        extra_info_types: Iterable[int] | None = None,
+    ) -> None:
+        if split not in {"train", "valid", "test"}:
+            raise ValueError(f"split must be train/valid/test, got {split!r}")
+
+        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,
+            extra_info_types=extra_info_types,
+        )
+
+        self.split = split
+        self.min_history_events = int(min_history_events)
+        self.min_future_events = int(min_future_events)
+        self.validation_query_seed = int(validation_query_seed)
+        self.patients: List[Dict] = []
+        self.valid_queries: List[Tuple[int, float]] = []
+        validation_rng = None
+        if split in {"valid", "test"}:
+            split_offset = 0 if split == "valid" else 1_000_003
+            validation_rng = np.random.RandomState(self.validation_query_seed + split_offset)
+
+        for eid, times_days, labels in self._iter_patient_events(
+            impute_no_event_gaps=False,
+        ):
+            times_years = (times_days / DAYS_PER_YEAR).astype(np.float32)
+            unique_times = np.unique(times_years)
+            if len(labels) < 2 or len(unique_times) < 2:
+                continue
+
+            features = self._split_features(eid)
+            if features is None:
+                continue
+
+            patient = {
+                "eid": eid,
+                "times": times_years,
+                "labels": labels.astype(np.int64),
+                "t_obs": float(times_years.max()),
+                **features,
+            }
+
+            pidx = len(self.patients)
+            self.patients.append(patient)
+
+            if split in {"valid", "test"}:
+                if validation_rng is None:
+                    raise RuntimeError("validation_rng was not initialized")
+                self.valid_queries.extend(
+                    (pidx, t_query)
+                    for t_query in self._sample_fixed_validation_queries(
+                        patient,
+                        validation_rng,
+                    )
+                )
+
+        if split in {"valid", "test"} and not self.valid_queries:
+            raise ValueError("No random-but-fixed validation query points were built.")
+
+    def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
+        times = patient["times"]
+        labels = patient["labels"]
+        real_event_mask = ~np.isin(
+            labels,
+            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
+        )
+        n_hist = int((times <= t_query).sum())
+        n_future = int(((times > t_query) & real_event_mask).sum())
+        return (
+            n_hist >= self.min_history_events
+            and n_future >= self.min_future_events
+            and patient["t_obs"] > t_query
+        )
+
+    def _sample_fixed_validation_queries(
+        self,
+        patient: Dict,
+        rng: np.random.RandomState,
+    ) -> List[float]:
+        times = np.asarray(patient["times"], dtype=np.float32)
+        labels = np.asarray(patient["labels"], dtype=np.int64)
+        real_event_mask = ~np.isin(
+            labels,
+            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
+        )
+        real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
+        n_real_events = int(real_times.size)
+        n_queries = max(0, n_real_events - 2)
+        if n_queries == 0:
+            return []
+
+        min_hist = int(self.min_history_events)
+        min_future = int(self.min_future_events)
+        if n_real_events < min_hist + min_future:
+            return []
+
+        left = float(real_times[min_hist - 1])
+        right_event_time = float(real_times[n_real_events - min_future])
+        right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
+        if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
+            return []
+
+        queries: List[float] = []
+        max_attempts = max(100, n_queries * 50)
+        for _ in range(max_attempts):
+            if len(queries) >= n_queries:
+                break
+            t_query = float(rng.uniform(left, float(right)))
+            if self._is_valid_query(patient, t_query):
+                queries.append(t_query)
+
+        return queries
+
+    def _sample_train_query(self, patient: Dict) -> float:
+        unique_times = np.unique(patient["times"])
+        if len(unique_times) < 2:
+            raise RuntimeError("Training patient has fewer than two unique times.")
+
+        j = np.random.randint(1, len(unique_times))
+        left = float(unique_times[j - 1])
+        right = float(unique_times[j])
+
+        if right - left <= ONE_DAY_YEARS:
+            t_query = right - ONE_DAY_YEARS
+        else:
+            t_query = np.random.uniform(left, right - ONE_DAY_YEARS)
+
+        if not self._is_valid_query(patient, t_query):
+            t_query = right - 1e-6
+        return float(t_query)
+
+    def _build_item(self, patient: Dict, t_query: float) -> Dict:
+        times = patient["times"]
+        labels = patient["labels"]
+        hist = times <= t_query
+        fut = times > t_query
+
+        return {
+            "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),
+            "future_targets": torch.from_numpy(labels[fut]).long(),
+            "future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(),
+            "exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32),
+            "sex": torch.tensor(patient["sex"], dtype=torch.long),
+            "other_type": torch.from_numpy(patient["other_type"]).long(),
+            "other_value": torch.from_numpy(patient["other_value"]).float(),
+            "other_value_kind": torch.from_numpy(patient["other_value_kind"]).long(),
+            "other_time": torch.from_numpy(patient["other_time"]).float(),
+        }
+
+    def __len__(self) -> int:
+        if self.split == "train":
+            return len(self.patients)
+        return len(self.valid_queries)
+
+    def __getitem__(self, idx: int) -> Dict:
+        if self.split == "train":
+            patient = self.patients[idx]
+            t_query = self._sample_train_query(patient)
+        else:
+            pidx, t_query = self.valid_queries[idx]
+            patient = self.patients[pidx]
+        return self._build_item(patient, t_query)
+
+
+def _collate_common_static(batch: List[Dict]) -> Dict:
+    return {
+        "sex": torch.stack([s["sex"] for s in batch]),
+        "other_type": pad_sequence(
+            [s["other_type"] for s in batch],
+            batch_first=True,
+            padding_value=0,
+        ),
+        "other_value": pad_sequence(
+            [s["other_value"] for s in batch],
+            batch_first=True,
+            padding_value=0.0,
+        ),
+        "other_value_kind": pad_sequence(
+            [s["other_value_kind"] for s in batch],
+            batch_first=True,
+            padding_value=0,
+        ),
+        "other_time": pad_sequence(
+            [s["other_time"] for s in batch],
+            batch_first=True,
+            padding_value=0.0,
+        ),
+    }
+
+
+def next_step_collate_fn(batch: List[Dict]) -> Dict:
+    event_seq = pad_sequence(
+        [s["event_seq"] for s in batch],
+        batch_first=True,
+        padding_value=PAD_IDX,
+    )
+    time_seq = pad_sequence(
+        [s["time_seq"] for s in batch],
+        batch_first=True,
+        padding_value=0.0,
+    )
+    target_event_seq = pad_sequence(
+        [s["target_event_seq"] for s in batch],
+        batch_first=True,
+        padding_value=PAD_IDX,
+    )
+    target_time_seq = pad_sequence(
+        [s["target_time_seq"] for s in batch],
+        batch_first=True,
+        padding_value=0.0,
+    )
+    readout_mask = pad_sequence(
+        [s["readout_mask"] for s in batch],
+        batch_first=True,
+        padding_value=False,
+    )
+    target_dt_unique = pad_sequence(
+        [s["target_dt_unique"] for s in batch],
+        batch_first=True,
+        padding_value=0.0,
+    )
+    target_multi_hot = pad_sequence(
+        [s["target_multi_hot"] for s in batch],
+        batch_first=True,
+        padding_value=False,
+    )
+
+    out = {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "target_event_seq": target_event_seq,
+        "target_time_seq": target_time_seq,
+        "readout_mask": readout_mask,
+        "target_dt_unique": target_dt_unique,
+        "target_multi_hot": target_multi_hot,
+    }
+    out.update(_collate_common_static(batch))
+    return out
+
+
+def all_future_collate_fn(batch: List[Dict]) -> Dict:
+    event_seq = pad_sequence(
+        [s["event_seq"] for s in batch],
+        batch_first=True,
+        padding_value=PAD_IDX,
+    )
+    time_seq = pad_sequence(
+        [s["time_seq"] for s in batch],
+        batch_first=True,
+        padding_value=0.0,
+    )
+    future_targets = pad_sequence(
+        [s["future_targets"] for s in batch],
+        batch_first=True,
+        padding_value=PAD_IDX,
+    )
+    future_dt = pad_sequence(
+        [s["future_dt"] for s in batch],
+        batch_first=True,
+        padding_value=0.0,
+    )
+
+    out = {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "t_query": torch.stack([s["t_query"] for s in batch]),
+        "future_targets": future_targets,
+        "future_dt": future_dt,
+        "exposure": torch.stack([s["exposure"] for s in batch]),
+    }
+    out.update(_collate_common_static(batch))
+    return out
+
+
+HealthDataset = NextStepHealthDataset
+collate_fn = next_step_collate_fn
--- a/eval_data.py
+++ b/eval_data.py
@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, List
+
+import numpy as np
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from dataset import AllFutureHealthDataset, HealthDataset
+from targets import PAD_IDX
+
+
+class AllFutureSequenceEvalDataset:
+    """
+    Eval-only sequence view for all-future checkpoints.
+
+    All-future training uses the observed history, including CHECKUP state
+    tokens, without reusing the next-step view that contains imputed
+    <NO_EVENT> gap tokens.
+    """
+
+    def __init__(
+        self,
+        data_prefix: str,
+        labels_file: str,
+        min_history_events: int = 1,
+        min_future_events: int = 1,
+        extra_info_types: Iterable[int] | None = None,
+    ) -> None:
+        base = AllFutureHealthDataset(
+            data_prefix=data_prefix,
+            labels_file=labels_file,
+            split="train",
+            min_history_events=min_history_events,
+            min_future_events=min_future_events,
+            extra_info_types=extra_info_types,
+        )
+
+        self.base = base
+        self.label_code_to_id = base.label_code_to_id
+        self.label_id_to_code = base.label_id_to_code
+        self.vocab_size = base.vocab_size
+        self.n_types = base.n_types
+        self.n_cont_types = base.n_cont_types
+        self.n_categories = base.n_categories
+        self.cont_type_ids = base.cont_type_ids
+        self.extra_info_types = base.extra_info_types
+
+        self.samples: List[Dict[str, Any]] = []
+        for patient in base.patients:
+            labels = np.asarray(patient["labels"], dtype=np.int64)
+            times = np.asarray(patient["times"], dtype=np.float32)
+            if labels.size < 2:
+                continue
+            input_len = int(labels.size - 1)
+            self.samples.append(
+                {
+                    "eid": int(patient["eid"]),
+                    "event_seq": labels[:-1],
+                    "time_seq": times[:-1],
+                    "target_event_seq": labels[1:],
+                    "target_time_seq": times[1:],
+                    "readout_mask": np.ones(input_len, dtype=bool),
+                    "sex": int(patient["sex"]),
+                    "other_type": np.asarray(patient["other_type"], dtype=np.int64),
+                    "other_value": np.asarray(patient["other_value"], dtype=np.float32),
+                    "other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
+                    "other_time": np.asarray(patient["other_time"], dtype=np.float32),
+                }
+            )
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        s = self.samples[idx]
+        return {
+            "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(),
+            "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
+            "readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
+            "sex": torch.tensor(s["sex"], dtype=torch.long),
+            "other_type": torch.from_numpy(s["other_type"]).long(),
+            "other_value": torch.from_numpy(s["other_value"]).float(),
+            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
+            "other_time": torch.from_numpy(s["other_time"]).float(),
+        }
+
+
+def load_sequence_eval_dataset(
+    *,
+    model_target_mode: str,
+    data_prefix: str,
+    labels_file: str,
+    no_event_interval_years: float,
+    include_no_event_in_uts_target: bool,
+    min_history_events: int,
+    min_future_events: int,
+    extra_info_types: Iterable[int] | None,
+):
+    mode = str(model_target_mode).lower()
+    if mode == "next_token":
+        return HealthDataset(
+            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,
+            extra_info_types=extra_info_types,
+        )
+    if mode == "all_future":
+        return AllFutureSequenceEvalDataset(
+            data_prefix=data_prefix,
+            labels_file=labels_file,
+            min_history_events=min_history_events,
+            min_future_events=min_future_events,
+            extra_info_types=extra_info_types,
+        )
+    raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
+
+
+def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    event_seq = pad_sequence(
+        [s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    time_seq = pad_sequence(
+        [s["time_seq"] for s in batch], batch_first=True, padding_value=0.0
+    )
+    target_event_seq = pad_sequence(
+        [s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    target_time_seq = pad_sequence(
+        [s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0
+    )
+    readout_mask = pad_sequence(
+        [s["readout_mask"] for s in batch], batch_first=True, padding_value=False
+    )
+    other_type = pad_sequence(
+        [s["other_type"] for s in batch], batch_first=True, padding_value=0
+    )
+    other_value = pad_sequence(
+        [s["other_value"] for s in batch], batch_first=True, padding_value=0.0
+    )
+    other_value_kind = pad_sequence(
+        [s["other_value_kind"] for s in batch], batch_first=True, padding_value=0
+    )
+    other_time = pad_sequence(
+        [s["other_time"] for s in batch], batch_first=True, padding_value=0.0
+    )
+
+    return {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "target_event_seq": target_event_seq,
+        "target_time_seq": target_time_seq,
+        "readout_mask": readout_mask,
+        "sex": torch.stack([s["sex"] for s in batch]),
+        "other_type": other_type,
+        "other_value": other_value,
+        "other_value_kind": other_value_kind,
+        "other_time": other_time,
+    }
--- a/evaluate_auc.py
+++ b/evaluate_auc.py
--- a/evaluate_token_auc.py
+++ b/evaluate_token_auc.py
@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from evaluate_auc import main
+
+
+if __name__ == "__main__":
+    main()
--- a/event_date_utils.py
+++ b/event_date_utils.py
@@ -0,0 +1,104 @@
+"""Read and query calendar-dated disease-event arrays from prepare_event_dates.py."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+import pandas as pd
+
+
+REQUIRED_FIELDS = {"eid", "event_date", "token"}
+
+
+def load_event_dates(path: str | Path) -> np.ndarray:
+    """Load and validate the structured ``.npy`` event array."""
+    events = np.load(path)
+    if events.dtype.names is None or not REQUIRED_FIELDS.issubset(events.dtype.names):
+        raise ValueError(
+            "Expected a structured .npy with eid, event_date, token fields. "
+            "Create it with prepare_event_dates.py."
+        )
+    return events
+
+
+def load_token_labels(labels_file: str | Path) -> dict[int, str]:
+    """Load token -> human-readable code using the project label convention."""
+    labels = {1: "CHECKUP"}
+    with Path(labels_file).open(encoding="utf-8") as handle:
+        for index, line in enumerate(handle):
+            code = line.strip().split(" ", maxsplit=1)[0]
+            if code:
+                labels[index + 2] = code
+    return labels
+
+
+@dataclass
+class EventDateIndex:
+    """Small in-memory query wrapper for exposure-linkage and cohort scripts."""
+
+    events: np.ndarray
+    token_labels: dict[int, str] | None = None
+
+    @classmethod
+    def from_files(
+        cls,
+        event_file: str | Path,
+        labels_file: str | Path | None = None,
+    ) -> "EventDateIndex":
+        labels = load_token_labels(labels_file) if labels_file is not None else None
+        return cls(load_event_dates(event_file), labels)
+
+    def to_frame(self, events: np.ndarray | None = None) -> pd.DataFrame:
+        """Convert records to a convenient, calendar-dated DataFrame."""
+        data = self.events if events is None else events
+        frame = pd.DataFrame(
+            {
+                "eid": data["eid"].astype("int64"),
+                "event_date": pd.to_datetime(data["event_date"]),
+                "token": data["token"].astype("int32"),
+            }
+        )
+        if self.token_labels is not None:
+            frame["label_code"] = frame["token"].map(self.token_labels).fillna("UNKNOWN")
+        return frame.sort_values(["eid", "event_date", "token"], kind="stable").reset_index(drop=True)
+
+    def for_eid(self, eid: int) -> pd.DataFrame:
+        """Return every stored disease/death event for one participant."""
+        return self.to_frame(self.events[self.events["eid"] == int(eid)])
+
+    def between(
+        self,
+        start: str | pd.Timestamp,
+        end: str | pd.Timestamp,
+        *,
+        eids: Iterable[int] | None = None,
+        tokens: Iterable[int] | None = None,
+    ) -> pd.DataFrame:
+        """Query events in an inclusive calendar-date interval."""
+        start_day = np.datetime64(pd.Timestamp(start).date(), "D")
+        end_day = np.datetime64(pd.Timestamp(end).date(), "D")
+        mask = (self.events["event_date"] >= start_day) & (self.events["event_date"] <= end_day)
+        if eids is not None:
+            mask &= np.isin(self.events["eid"], list(eids))
+        if tokens is not None:
+            mask &= np.isin(self.events["token"], list(tokens))
+        return self.to_frame(self.events[mask])
+
+    def anchors_before(self, eid: int, date: str | pd.Timestamp) -> pd.DataFrame:
+        """Return a participant's event history strictly before an exposure anchor."""
+        day = np.datetime64(pd.Timestamp(date).date(), "D")
+        mask = (self.events["eid"] == int(eid)) & (self.events["event_date"] < day)
+        return self.to_frame(self.events[mask])
+
+    def first_event(self, token: int) -> pd.DataFrame:
+        """Return each participant's first date for a requested token."""
+        selected = self.events[self.events["token"] == int(token)]
+        # Arrays produced by prepare_event_dates.py are already deduplicated;
+        # sorting makes this safe for externally produced compatible arrays too.
+        order = np.lexsort((selected["event_date"], selected["eid"]))
+        selected = selected[order]
+        _, first = np.unique(selected["eid"], return_index=True)
+        return self.to_frame(selected[first])
--- a/extra_info_types_all.txt
+++ b/extra_info_types_all.txt
@@ -0,0 +1,268 @@
+# All other-info variables (field_type=1 and field_type=2)
+# Generated from field_ids_enriched.csv using prepare_data.py other-info type ordering.
+# Format: <extra_info_type_id> # <var_name> | <full_name>
+1 # waist_circumference | Waist circumference
+2 # hip_circumference | Hip circumference
+3 # standing_height | Standing height
+4 # fasting_time | Fasting time
+5 # pulse_rate | Pulse rate automated reading
+6 # dbp | Diastolic blood pressure automated reading
+7 # sbp | Systolic blood pressure automated reading
+8 # fev1_best | Forced expiratory volume in 1-second (FEV1) Best measure
+9 # fvc_best | Forced vital capacity (FVC) Best measure
+10 # fev1_fvc_ratio | FEV1/ FVC ratio Z-score
+11 # bmi | Body mass index (BMI)
+12 # WBC | White blood cell (leukocyte) count
+13 # RBC | Red blood cell (erythrocyte) count
+14 # hemoglobin | Haemoglobin concentration
+15 # hematocrit | Haematocrit percentage
+16 # MCV | Mean corpuscular volume
+17 # MCH | Mean corpuscular haemoglobin
+18 # MCHC | Mean corpuscular haemoglobin concentration
+19 # Pc | Platelet count
+20 # MPV | Mean platelet (thrombocyte) volume
+21 # LymC | Lymphocyte count
+22 # MonC | Monocyte count
+23 # NeuC | Neutrophill count
+24 # EosC | Eosinophill count
+25 # BasC | Basophill count
+26 # nRBC | Nucleated red blood cell count
+27 # RC | Reticulocyte count
+28 # MRV | Mean reticulocyte volume
+29 # MSCV | Mean sphered cell volume
+30 # IRF | Immature reticulocyte fraction
+31 # HLSRC | High light scatter reticulocyte count
+32 # MicU | Microalbumin in urine
+33 # CreaU | Creatinine (enzymatic) in urine
+34 # PotU | Potassium in urine
+35 # SodU | Sodium in urine
+36 # Alb | Albumin
+37 # ALP | Alkaline phosphatase
+38 # Alanine | Alanine aminotransferase
+39 # ApoA | Apolipoprotein A
+40 # ApoB | Apolipoprotein B
+41 # AA | Aspartate aminotransferase
+42 # DBil | Direct bilirubin
+43 # Urea | Urea
+44 # Calcium | Calcium
+45 # Cholesterol | Cholesterol
+46 # Creatinine | Creatinine
+47 # CRP | C-reactive protein
+48 # CystatinC | Cystatin C
+49 # GGT | Gamma glutamyltransferase
+50 # Glu | Glucose
+51 # HbA1c | Glycated haemoglobin (HbA1c)
+52 # HDL | HDL cholesterol
+53 # IGF1 | IGF-1
+54 # LDL | LDL direct
+55 # LpA | Lipoprotein A
+56 # Oestradiol | Oestradiol
+57 # Phosphate | Phosphate
+58 # Rheu | Rheumatoid factor
+59 # SHBG | SHBG
+60 # TotalBil | Total bilirubin
+61 # Testosterone | Testosterone
+62 # TotalProtein | Total protein
+63 # Tri | Triglycerides
+64 # Urate | Urate
+65 # VitaminD | Vitamin D
+66 # smoking | Current tobacco smoking
+67 # alcohol | Alcohol intake frequency.
+68 # ipaq_activity_group | IPAQ activity group
+69 # moderate_activity_met_minutes_week | MET minutes per week for moderate activity
+70 # vigorous_activity_met_minutes_week | MET minutes per week for vigorous activity
+71 # walking_met_minutes_week | MET minutes per week for walking
+72 # total_activity_met_minutes_week | Summed MET minutes per week for all activity
+73 # total_activity_days | Summed days activity
+74 # total_activity_minutes | Summed minutes activity
+75 # heavy_diy_duration | Duration of heavy DIY
+76 # light_diy_duration | Duration of light DIY
+77 # moderate_activity_duration | Duration of moderate activity
+78 # other_exercise_duration | Duration of other exercises
+79 # strenuous_sport_duration | Duration of strenuous sports
+80 # vigorous_activity_duration | Duration of vigorous activity
+81 # walking_duration | Duration of walks
+82 # pleasure_walking_duration | Duration walking for pleasure
+83 # heavy_diy_frequency_4_weeks | Frequency of heavy DIY in last 4 weeks
+84 # light_diy_frequency_4_weeks | Frequency of light DIY in last 4 weeks
+85 # other_exercise_frequency_4_weeks | Frequency of other exercises in last 4 weeks
+86 # stair_climbing_frequency_4_weeks | Frequency of stair climbing in last 4 weeks
+87 # strenuous_sport_frequency_4_weeks | Frequency of strenuous sports in last 4 weeks
+88 # pleasure_walking_frequency_4_weeks | Frequency of walking for pleasure in last 4 weeks
+89 # moderate_activity_days_week_10min | Number of days/week of moderate physical activity 10+ minutes
+90 # vigorous_activity_days_week_10min | Number of days/week of vigorous physical activity 10+ minutes
+91 # walking_days_week_10min | Number of days/week walked 10+ minutes
+92 # driving_time | Time spent driving
+93 # computer_use_time | Time spent using computer
+94 # tv_watching_time | Time spent watching television (TV)
+95 # physical_activity_types_4_weeks | Types of physical activity in last 4 weeks
+96 # nonwork_transport_types | Types of transport used (excluding work)
+97 # usual_walking_pace | Usual walking pace
+98 # mobile_phone_use_duration | Length of mobile phone use
+99 # mobile_phone_use_weekly_3_months | Weekly usage of mobile phone in last 3 months
+100 # computer_game_playing | Plays computer games
+101 # sleep_duration | Sleep duration
+102 # chronotype | Morning/evening person (chronotype)
+103 # daytime_napping | Nap during day
+104 # insomnia | Sleeplessness / insomnia
+105 # daytime_dozing | Daytime dozing / sleeping
+106 # ever_smoked | Ever smoked
+107 # smoking_pack_years | Pack years of smoking
+108 # smoking_status | Smoking status
+109 # past_tobacco_smoking | Past tobacco smoking
+110 # lifetime_smoking_100_plus | Light smokers, at least 100 smokes in lifetime
+111 # current_tobacco_type | Type of tobacco currently smoked
+112 # current_cigarettes_per_day | Number of cigarettes currently smoked daily (current cigarette smokers)
+113 # previous_cigarettes_per_day_current_cigar_pipe_smokers | Number of cigarettes previously smoked daily (current cigar/pipe smokers)
+114 # time_to_first_cigarette | Time from waking to first cigarette
+115 # ever_tried_smoking_cessation | Ever tried to stop smoking
+116 # smoking_change_vs_10_years_ago | Smoking compared to 10 years previous
+117 # previous_tobacco_type | Type of tobacco previously smoked
+118 # previous_cigarettes_per_day | Number of cigarettes previously smoked daily
+119 # ever_stopped_smoking_6_months | Ever stopped smoking for 6+ months
+120 # household_smokers | Smoking/smokers in household
+121 # home_secondhand_smoke_exposure | Exposure to tobacco smoke at home
+122 # nonhome_secondhand_smoke_exposure | Exposure to tobacco smoke outside home
+123 # cooked_vegetable_intake | Cooked vegetable intake
+124 # raw_vegetable_intake | Salad / raw vegetable intake
+125 # fresh_fruit_intake | Fresh fruit intake
+126 # dried_fruit_intake | Dried fruit intake
+127 # oily_fish_intake | Oily fish intake
+128 # non_oily_fish_intake | Non-oily fish intake
+129 # processed_meat_intake | Processed meat intake
+130 # poultry_intake | Poultry intake
+131 # beef_intake | Beef intake
+132 # lamb_mutton_intake | Lamb/mutton intake
+133 # pork_intake | Pork intake
+134 # age_last_ate_meat | Age when last ate meat
+135 # food_avoidance_eggs_dairy_wheat_sugar | Never eat eggs, dairy, wheat, sugar
+136 # cheese_intake | Cheese intake
+137 # milk_type | Milk type used
+138 # spread_type | Spread type
+139 # bread_intake | Bread intake
+140 # bread_type | Bread type
+141 # cereal_intake | Cereal intake
+142 # cereal_type | Cereal type
+143 # added_salt | Salt added to food
+144 # tea_intake | Tea intake
+145 # coffee_intake | Coffee intake
+146 # coffee_type | Coffee type
+147 # hot_drink_temperature | Hot drink temperature
+148 # water_intake | Water intake
+149 # diet_variation | Variation in diet
+150 # alcohol_drinker_status | Alcohol drinker status
+151 # former_alcohol_drinker | Former alcohol drinker
+152 # red_wine_intake_monthly | Average monthly red wine intake
+153 # champagne_white_wine_intake_monthly | Average monthly champagne plus white wine intake
+154 # beer_cider_intake_monthly | Average monthly beer plus cider intake
+155 # spirits_intake_monthly | Average monthly spirits intake
+156 # fortified_wine_intake_monthly | Average monthly fortified wine intake
+157 # other_alcohol_intake_monthly | Average monthly intake of other alcoholic drinks
+158 # red_wine_intake_weekly | Average weekly red wine intake
+159 # champagne_white_wine_intake_weekly | Average weekly champagne plus white wine intake
+160 # beer_cider_intake_weekly | Average weekly beer plus cider intake
+161 # spirits_intake_weekly | Average weekly spirits intake
+162 # fortified_wine_intake_weekly | Average weekly fortified wine intake
+163 # other_alcohol_intake_weekly | Average weekly intake of other alcoholic drinks
+164 # alcohol_with_meals | Alcohol usually taken with meals
+165 # country_of_birth_uk_elsewhere | Country of birth (UK/elsewhere)
+166 # breastfed_in_infancy | Breastfed as a baby
+167 # comparative_body_size_age_10 | Comparative body size at age 10
+168 # comparative_height_age_10 | Comparative height size at age 10
+169 # handedness | Handedness (chirality/laterality)
+170 # adopted_as_child | Adopted as a child
+171 # multiple_birth | Part of a multiple birth
+172 # maternal_smoking_around_birth | Maternal smoking around birth
+173 # accommodation_type | Type of accommodation lived in
+174 # housing_tenure | Own or rent accommodation lived in
+175 # gas_solid_fuel_use | Gas or solid-fuel cooking/heating
+176 # home_heating_types | Heating type(s) in home
+177 # household_vehicle_count | Number of vehicles in household
+178 # household_income_before_tax | Average total household income before tax
+179 # current_employment_status | Current employment status
+180 # current_employment_status_corrected | Current employment status - corrected
+181 # home_work_distance | Distance between home and job workplace
+182 # main_job_hours_week | Length of working week for main job
+183 # commuting_frequency | Frequency of travelling from home to job workplace
+184 # commuting_transport_type | Transport type for commuting to job workplace
+185 # job_walking_standing | Job involves mainly walking or standing
+186 # job_heavy_manual_work | Job involves heavy manual or physical work
+187 # job_shift_work | Job involves shift work
+188 # job_night_shift_work | Job involves night shift work
+189 # educational_qualifications | Qualifications
+190 # age_completed_full_time_education | Age completed full time education
+191 # friend_family_visit_frequency | Frequency of friend/family visits
+192 # leisure_social_activities | Leisure/social activities
+193 # ability_to_confide | Able to confide
+194 # bipolar_major_depression_status | Bipolar and major depression status
+195 # neuroticism_score | Neuroticism score
+196 # mood_swings | Mood swings
+197 # miserableness | Miserableness
+198 # irritability | Irritability
+199 # sensitivity_hurt_feelings | Sensitivity / hurt feelings
+200 # fed_up_feelings | Fed-up feelings
+201 # nervous_feelings | Nervous feelings
+202 # worry_anxiety_feelings | Worrier / anxious feelings
+203 # tenseness_highly_strung | Tense / 'highly strung'
+204 # suffering_from_nerves | Suffer from 'nerves'
+205 # loneliness_isolation | Loneliness, isolation
+206 # guilty_feelings | Guilty feelings
+207 # risk_taking | Risk taking
+208 # happiness | Happiness
+209 # job_satisfaction | Work/job satisfaction
+210 # health_satisfaction | Health satisfaction
+211 # family_relationship_satisfaction | Family relationship satisfaction
+212 # friendship_satisfaction | Friendships satisfaction
+213 # financial_situation_satisfaction | Financial situation satisfaction
+214 # depressed_mood_frequency_2_weeks | Frequency of depressed mood in last 2 weeks
+215 # disinterest_frequency_2_weeks | Frequency of unenthusiasm / disinterest in last 2 weeks
+216 # tenseness_restlessness_frequency_2_weeks | Frequency of tenseness / restlessness in last 2 weeks
+217 # tiredness_lethargy_frequency_2_weeks | Frequency of tiredness / lethargy in last 2 weeks
+218 # ever_depressed_full_week | Ever depressed for a whole week
+219 # longest_depression_duration | Longest period of depression
+220 # depression_episode_count | Number of depression episodes
+221 # longest_disinterest_duration | Longest period of unenthusiasm / disinterest
+222 # disinterest_episode_count | Number of unenthusiastic/disinterested episodes
+223 # ever_manic_hyper_2_days | Ever manic/hyper for 2 days
+224 # ever_irritable_argumentative_2_days | Ever highly irritable/argumentative for 2 days
+225 # manic_hyper_symptoms | Manic/hyper symptoms
+226 # longest_manic_irritable_episode_duration | Length of longest manic/irritable episode
+227 # manic_irritable_episode_severity | Severity of manic/irritable episodes
+228 # adverse_life_events_2_years | Illness, injury, bereavement, stress in last 2 years
+229 # outdoor_time_summer | Time spend outdoors in summer
+230 # outdoor_time_winter | Time spent outdoors in winter
+231 # skin_tanning_ease | Ease of skin tanning
+232 # childhood_sunburn_frequency | Childhood sunburn occasions
+233 # sun_uv_protection_use | Use of sun/uv protection
+234 # solarium_sunlamp_frequency | Frequency of solarium/sunlamp use
+235 # proximity_to_major_road | Close to major road
+236 # inverse_distance_nearest_major_road | Inverse distance to the nearest major road
+237 # inverse_distance_nearest_road | Inverse distance to the nearest road
+238 # no2_2005 | Nitrogen dioxide air pollution; 2005
+239 # no2_2006 | Nitrogen dioxide air pollution; 2006
+240 # no2_2007 | Nitrogen dioxide air pollution; 2007
+241 # no2_2010 | Nitrogen dioxide air pollution; 2010
+242 # nox_2010 | Nitrogen oxides air pollution; 2010
+243 # pm10_2007 | Particulate matter air pollution (pm10); 2007
+244 # pm10_2010 | Particulate matter air pollution (pm10); 2010
+245 # pm25_absorbance_2010 | Particulate matter air pollution (pm2.5) absorbance; 2010
+246 # pm25_2010 | Particulate matter air pollution (pm2.5); 2010
+247 # pm25_10_2010 | Particulate matter air pollution 2.5-10um; 2010
+248 # major_road_length_100m | Sum of road length of major roads within 100m
+249 # major_road_traffic_load | Total traffic load on major roads
+250 # nearest_major_road_traffic_intensity | Traffic intensity on the nearest major road
+251 # nearest_road_traffic_intensity | Traffic intensity on the nearest road
+252 # noise_level_16h | Average 16-hour sound level of noise pollution
+253 # noise_level_24h | Average 24-hour sound level of noise pollution
+254 # noise_level_daytime | Average daytime sound level of noise pollution
+255 # noise_level_evening | Average evening sound level of noise pollution
+256 # noise_level_nighttime | Average night-time sound level of noise pollution
+257 # natural_environment_percent_1000m | Natural environment percentage, buffer 1000m
+258 # natural_environment_percent_300m | Natural environment percentage, buffer 300m
+259 # greenspace_percent_1000m | Greenspace percentage, buffer 1000m
+260 # greenspace_percent_300m | Greenspace percentage, buffer 300m
+261 # domestic_garden_percent_1000m | Domestic garden percentage, buffer 1000m
+262 # domestic_garden_percent_300m | Domestic garden percentage, buffer 300m
+263 # water_percent_1000m | Water percentage, buffer 1000m
+264 # water_percent_300m | Water percentage, buffer 300m
+265 # distance_to_coast | Distance (Euclidean) to coast
--- a/extra_info_types_assessment_only.txt
+++ b/extra_info_types_assessment_only.txt
@@ -0,0 +1,68 @@
+# Only assessment/body-measurement variables (field_type=1)
+# Generated from field_ids_enriched.csv using prepare_data.py other-info type ordering.
+# Format: <extra_info_type_id> # <var_name> | <full_name>
+1 # waist_circumference | Waist circumference
+2 # hip_circumference | Hip circumference
+3 # standing_height | Standing height
+4 # fasting_time | Fasting time
+5 # pulse_rate | Pulse rate automated reading
+6 # dbp | Diastolic blood pressure automated reading
+7 # sbp | Systolic blood pressure automated reading
+8 # fev1_best | Forced expiratory volume in 1-second (FEV1) Best measure
+9 # fvc_best | Forced vital capacity (FVC) Best measure
+10 # fev1_fvc_ratio | FEV1/ FVC ratio Z-score
+11 # bmi | Body mass index (BMI)
+12 # WBC | White blood cell (leukocyte) count
+13 # RBC | Red blood cell (erythrocyte) count
+14 # hemoglobin | Haemoglobin concentration
+15 # hematocrit | Haematocrit percentage
+16 # MCV | Mean corpuscular volume
+17 # MCH | Mean corpuscular haemoglobin
+18 # MCHC | Mean corpuscular haemoglobin concentration
+19 # Pc | Platelet count
+20 # MPV | Mean platelet (thrombocyte) volume
+21 # LymC | Lymphocyte count
+22 # MonC | Monocyte count
+23 # NeuC | Neutrophill count
+24 # EosC | Eosinophill count
+25 # BasC | Basophill count
+26 # nRBC | Nucleated red blood cell count
+27 # RC | Reticulocyte count
+28 # MRV | Mean reticulocyte volume
+29 # MSCV | Mean sphered cell volume
+30 # IRF | Immature reticulocyte fraction
+31 # HLSRC | High light scatter reticulocyte count
+32 # MicU | Microalbumin in urine
+33 # CreaU | Creatinine (enzymatic) in urine
+34 # PotU | Potassium in urine
+35 # SodU | Sodium in urine
+36 # Alb | Albumin
+37 # ALP | Alkaline phosphatase
+38 # Alanine | Alanine aminotransferase
+39 # ApoA | Apolipoprotein A
+40 # ApoB | Apolipoprotein B
+41 # AA | Aspartate aminotransferase
+42 # DBil | Direct bilirubin
+43 # Urea | Urea
+44 # Calcium | Calcium
+45 # Cholesterol | Cholesterol
+46 # Creatinine | Creatinine
+47 # CRP | C-reactive protein
+48 # CystatinC | Cystatin C
+49 # GGT | Gamma glutamyltransferase
+50 # Glu | Glucose
+51 # HbA1c | Glycated haemoglobin (HbA1c)
+52 # HDL | HDL cholesterol
+53 # IGF1 | IGF-1
+54 # LDL | LDL direct
+55 # LpA | Lipoprotein A
+56 # Oestradiol | Oestradiol
+57 # Phosphate | Phosphate
+58 # Rheu | Rheumatoid factor
+59 # SHBG | SHBG
+60 # TotalBil | Total bilirubin
+61 # Testosterone | Testosterone
+62 # TotalProtein | Total protein
+63 # Tri | Triglycerides
+64 # Urate | Urate
+65 # VitaminD | Vitamin D
--- a/extra_info_types_exposure_only.txt
+++ b/extra_info_types_exposure_only.txt
@@ -0,0 +1,203 @@
+# Only environment/lifestyle exposure variables (field_type=2)
+# Generated from field_ids_enriched.csv using prepare_data.py other-info type ordering.
+# Format: <extra_info_type_id> # <var_name> | <full_name>
+66 # smoking | Current tobacco smoking
+67 # alcohol | Alcohol intake frequency.
+68 # ipaq_activity_group | IPAQ activity group
+69 # moderate_activity_met_minutes_week | MET minutes per week for moderate activity
+70 # vigorous_activity_met_minutes_week | MET minutes per week for vigorous activity
+71 # walking_met_minutes_week | MET minutes per week for walking
+72 # total_activity_met_minutes_week | Summed MET minutes per week for all activity
+73 # total_activity_days | Summed days activity
+74 # total_activity_minutes | Summed minutes activity
+75 # heavy_diy_duration | Duration of heavy DIY
+76 # light_diy_duration | Duration of light DIY
+77 # moderate_activity_duration | Duration of moderate activity
+78 # other_exercise_duration | Duration of other exercises
+79 # strenuous_sport_duration | Duration of strenuous sports
+80 # vigorous_activity_duration | Duration of vigorous activity
+81 # walking_duration | Duration of walks
+82 # pleasure_walking_duration | Duration walking for pleasure
+83 # heavy_diy_frequency_4_weeks | Frequency of heavy DIY in last 4 weeks
+84 # light_diy_frequency_4_weeks | Frequency of light DIY in last 4 weeks
+85 # other_exercise_frequency_4_weeks | Frequency of other exercises in last 4 weeks
+86 # stair_climbing_frequency_4_weeks | Frequency of stair climbing in last 4 weeks
+87 # strenuous_sport_frequency_4_weeks | Frequency of strenuous sports in last 4 weeks
+88 # pleasure_walking_frequency_4_weeks | Frequency of walking for pleasure in last 4 weeks
+89 # moderate_activity_days_week_10min | Number of days/week of moderate physical activity 10+ minutes
+90 # vigorous_activity_days_week_10min | Number of days/week of vigorous physical activity 10+ minutes
+91 # walking_days_week_10min | Number of days/week walked 10+ minutes
+92 # driving_time | Time spent driving
+93 # computer_use_time | Time spent using computer
+94 # tv_watching_time | Time spent watching television (TV)
+95 # physical_activity_types_4_weeks | Types of physical activity in last 4 weeks
+96 # nonwork_transport_types | Types of transport used (excluding work)
+97 # usual_walking_pace | Usual walking pace
+98 # mobile_phone_use_duration | Length of mobile phone use
+99 # mobile_phone_use_weekly_3_months | Weekly usage of mobile phone in last 3 months
+100 # computer_game_playing | Plays computer games
+101 # sleep_duration | Sleep duration
+102 # chronotype | Morning/evening person (chronotype)
+103 # daytime_napping | Nap during day
+104 # insomnia | Sleeplessness / insomnia
+105 # daytime_dozing | Daytime dozing / sleeping
+106 # ever_smoked | Ever smoked
+107 # smoking_pack_years | Pack years of smoking
+108 # smoking_status | Smoking status
+109 # past_tobacco_smoking | Past tobacco smoking
+110 # lifetime_smoking_100_plus | Light smokers, at least 100 smokes in lifetime
+111 # current_tobacco_type | Type of tobacco currently smoked
+112 # current_cigarettes_per_day | Number of cigarettes currently smoked daily (current cigarette smokers)
+113 # previous_cigarettes_per_day_current_cigar_pipe_smokers | Number of cigarettes previously smoked daily (current cigar/pipe smokers)
+114 # time_to_first_cigarette | Time from waking to first cigarette
+115 # ever_tried_smoking_cessation | Ever tried to stop smoking
+116 # smoking_change_vs_10_years_ago | Smoking compared to 10 years previous
+117 # previous_tobacco_type | Type of tobacco previously smoked
+118 # previous_cigarettes_per_day | Number of cigarettes previously smoked daily
+119 # ever_stopped_smoking_6_months | Ever stopped smoking for 6+ months
+120 # household_smokers | Smoking/smokers in household
+121 # home_secondhand_smoke_exposure | Exposure to tobacco smoke at home
+122 # nonhome_secondhand_smoke_exposure | Exposure to tobacco smoke outside home
+123 # cooked_vegetable_intake | Cooked vegetable intake
+124 # raw_vegetable_intake | Salad / raw vegetable intake
+125 # fresh_fruit_intake | Fresh fruit intake
+126 # dried_fruit_intake | Dried fruit intake
+127 # oily_fish_intake | Oily fish intake
+128 # non_oily_fish_intake | Non-oily fish intake
+129 # processed_meat_intake | Processed meat intake
+130 # poultry_intake | Poultry intake
+131 # beef_intake | Beef intake
+132 # lamb_mutton_intake | Lamb/mutton intake
+133 # pork_intake | Pork intake
+134 # age_last_ate_meat | Age when last ate meat
+135 # food_avoidance_eggs_dairy_wheat_sugar | Never eat eggs, dairy, wheat, sugar
+136 # cheese_intake | Cheese intake
+137 # milk_type | Milk type used
+138 # spread_type | Spread type
+139 # bread_intake | Bread intake
+140 # bread_type | Bread type
+141 # cereal_intake | Cereal intake
+142 # cereal_type | Cereal type
+143 # added_salt | Salt added to food
+144 # tea_intake | Tea intake
+145 # coffee_intake | Coffee intake
+146 # coffee_type | Coffee type
+147 # hot_drink_temperature | Hot drink temperature
+148 # water_intake | Water intake
+149 # diet_variation | Variation in diet
+150 # alcohol_drinker_status | Alcohol drinker status
+151 # former_alcohol_drinker | Former alcohol drinker
+152 # red_wine_intake_monthly | Average monthly red wine intake
+153 # champagne_white_wine_intake_monthly | Average monthly champagne plus white wine intake
+154 # beer_cider_intake_monthly | Average monthly beer plus cider intake
+155 # spirits_intake_monthly | Average monthly spirits intake
+156 # fortified_wine_intake_monthly | Average monthly fortified wine intake
+157 # other_alcohol_intake_monthly | Average monthly intake of other alcoholic drinks
+158 # red_wine_intake_weekly | Average weekly red wine intake
+159 # champagne_white_wine_intake_weekly | Average weekly champagne plus white wine intake
+160 # beer_cider_intake_weekly | Average weekly beer plus cider intake
+161 # spirits_intake_weekly | Average weekly spirits intake
+162 # fortified_wine_intake_weekly | Average weekly fortified wine intake
+163 # other_alcohol_intake_weekly | Average weekly intake of other alcoholic drinks
+164 # alcohol_with_meals | Alcohol usually taken with meals
+165 # country_of_birth_uk_elsewhere | Country of birth (UK/elsewhere)
+166 # breastfed_in_infancy | Breastfed as a baby
+167 # comparative_body_size_age_10 | Comparative body size at age 10
+168 # comparative_height_age_10 | Comparative height size at age 10
+169 # handedness | Handedness (chirality/laterality)
+170 # adopted_as_child | Adopted as a child
+171 # multiple_birth | Part of a multiple birth
+172 # maternal_smoking_around_birth | Maternal smoking around birth
+173 # accommodation_type | Type of accommodation lived in
+174 # housing_tenure | Own or rent accommodation lived in
+175 # gas_solid_fuel_use | Gas or solid-fuel cooking/heating
+176 # home_heating_types | Heating type(s) in home
+177 # household_vehicle_count | Number of vehicles in household
+178 # household_income_before_tax | Average total household income before tax
+179 # current_employment_status | Current employment status
+180 # current_employment_status_corrected | Current employment status - corrected
+181 # home_work_distance | Distance between home and job workplace
+182 # main_job_hours_week | Length of working week for main job
+183 # commuting_frequency | Frequency of travelling from home to job workplace
+184 # commuting_transport_type | Transport type for commuting to job workplace
+185 # job_walking_standing | Job involves mainly walking or standing
+186 # job_heavy_manual_work | Job involves heavy manual or physical work
+187 # job_shift_work | Job involves shift work
+188 # job_night_shift_work | Job involves night shift work
+189 # educational_qualifications | Qualifications
+190 # age_completed_full_time_education | Age completed full time education
+191 # friend_family_visit_frequency | Frequency of friend/family visits
+192 # leisure_social_activities | Leisure/social activities
+193 # ability_to_confide | Able to confide
+194 # bipolar_major_depression_status | Bipolar and major depression status
+195 # neuroticism_score | Neuroticism score
+196 # mood_swings | Mood swings
+197 # miserableness | Miserableness
+198 # irritability | Irritability
+199 # sensitivity_hurt_feelings | Sensitivity / hurt feelings
+200 # fed_up_feelings | Fed-up feelings
+201 # nervous_feelings | Nervous feelings
+202 # worry_anxiety_feelings | Worrier / anxious feelings
+203 # tenseness_highly_strung | Tense / 'highly strung'
+204 # suffering_from_nerves | Suffer from 'nerves'
+205 # loneliness_isolation | Loneliness, isolation
+206 # guilty_feelings | Guilty feelings
+207 # risk_taking | Risk taking
+208 # happiness | Happiness
+209 # job_satisfaction | Work/job satisfaction
+210 # health_satisfaction | Health satisfaction
+211 # family_relationship_satisfaction | Family relationship satisfaction
+212 # friendship_satisfaction | Friendships satisfaction
+213 # financial_situation_satisfaction | Financial situation satisfaction
+214 # depressed_mood_frequency_2_weeks | Frequency of depressed mood in last 2 weeks
+215 # disinterest_frequency_2_weeks | Frequency of unenthusiasm / disinterest in last 2 weeks
+216 # tenseness_restlessness_frequency_2_weeks | Frequency of tenseness / restlessness in last 2 weeks
+217 # tiredness_lethargy_frequency_2_weeks | Frequency of tiredness / lethargy in last 2 weeks
+218 # ever_depressed_full_week | Ever depressed for a whole week
+219 # longest_depression_duration | Longest period of depression
+220 # depression_episode_count | Number of depression episodes
+221 # longest_disinterest_duration | Longest period of unenthusiasm / disinterest
+222 # disinterest_episode_count | Number of unenthusiastic/disinterested episodes
+223 # ever_manic_hyper_2_days | Ever manic/hyper for 2 days
+224 # ever_irritable_argumentative_2_days | Ever highly irritable/argumentative for 2 days
+225 # manic_hyper_symptoms | Manic/hyper symptoms
+226 # longest_manic_irritable_episode_duration | Length of longest manic/irritable episode
+227 # manic_irritable_episode_severity | Severity of manic/irritable episodes
+228 # adverse_life_events_2_years | Illness, injury, bereavement, stress in last 2 years
+229 # outdoor_time_summer | Time spend outdoors in summer
+230 # outdoor_time_winter | Time spent outdoors in winter
+231 # skin_tanning_ease | Ease of skin tanning
+232 # childhood_sunburn_frequency | Childhood sunburn occasions
+233 # sun_uv_protection_use | Use of sun/uv protection
+234 # solarium_sunlamp_frequency | Frequency of solarium/sunlamp use
+235 # proximity_to_major_road | Close to major road
+236 # inverse_distance_nearest_major_road | Inverse distance to the nearest major road
+237 # inverse_distance_nearest_road | Inverse distance to the nearest road
+238 # no2_2005 | Nitrogen dioxide air pollution; 2005
+239 # no2_2006 | Nitrogen dioxide air pollution; 2006
+240 # no2_2007 | Nitrogen dioxide air pollution; 2007
+241 # no2_2010 | Nitrogen dioxide air pollution; 2010
+242 # nox_2010 | Nitrogen oxides air pollution; 2010
+243 # pm10_2007 | Particulate matter air pollution (pm10); 2007
+244 # pm10_2010 | Particulate matter air pollution (pm10); 2010
+245 # pm25_absorbance_2010 | Particulate matter air pollution (pm2.5) absorbance; 2010
+246 # pm25_2010 | Particulate matter air pollution (pm2.5); 2010
+247 # pm25_10_2010 | Particulate matter air pollution 2.5-10um; 2010
+248 # major_road_length_100m | Sum of road length of major roads within 100m
+249 # major_road_traffic_load | Total traffic load on major roads
+250 # nearest_major_road_traffic_intensity | Traffic intensity on the nearest major road
+251 # nearest_road_traffic_intensity | Traffic intensity on the nearest road
+252 # noise_level_16h | Average 16-hour sound level of noise pollution
+253 # noise_level_24h | Average 24-hour sound level of noise pollution
+254 # noise_level_daytime | Average daytime sound level of noise pollution
+255 # noise_level_evening | Average evening sound level of noise pollution
+256 # noise_level_nighttime | Average night-time sound level of noise pollution
+257 # natural_environment_percent_1000m | Natural environment percentage, buffer 1000m
+258 # natural_environment_percent_300m | Natural environment percentage, buffer 300m
+259 # greenspace_percent_1000m | Greenspace percentage, buffer 1000m
+260 # greenspace_percent_300m | Greenspace percentage, buffer 300m
+261 # domestic_garden_percent_1000m | Domestic garden percentage, buffer 1000m
+262 # domestic_garden_percent_300m | Domestic garden percentage, buffer 300m
+263 # water_percent_1000m | Water percentage, buffer 1000m
+264 # water_percent_300m | Water percentage, buffer 300m
+265 # distance_to_coast | Distance (Euclidean) to coast
--- a/extra_info_types_none.txt
+++ b/extra_info_types_none.txt
@@ -0,0 +1,3 @@
+# No extra-info variables.
+# Use this file with --extra_info_types_file to train/evaluate with disease history only.
+# Keep this file free of numeric type ids; the loader parses it as an empty list.
--- a/extra_info_types_smoking_alcohol_bmi.txt
+++ b/extra_info_types_smoking_alcohol_bmi.txt
@@ -0,0 +1,6 @@
+# Only smoking, alcohol, and BMI variables
+# Generated from field_ids_enriched.csv using prepare_data.py other-info type ordering.
+# Format: <extra_info_type_id> # <var_name> | <full_name>
+11 # bmi | Body mass index (BMI)
+66 # smoking | Current tobacco smoking
+67 # alcohol | Alcohol intake frequency.
--- a/field_ids_enriched.csv
+++ b/field_ids_enriched.csv
--- a/future_risk.py
+++ b/future_risk.py
@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+import torch.nn.functional as F
+
+
+def death_token(vocab_size: int) -> int:
+    if int(vocab_size) <= 0:
+        raise ValueError(f"vocab_size must be positive, got {vocab_size}")
+    return int(vocab_size) - 1
+
+
+def probabilities_from_logits(
+    logits: torch.Tensor,
+    tau_years: float | torch.Tensor,
+    *,
+    dist_mode: str = "exponential",
+    rho: torch.Tensor | None = None,
+    death_rho: torch.Tensor | None = None,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    """
+    Convert all-future logits to tau-year event probabilities.
+
+    Death is always treated as token vocab_size - 1. For dist_mode="mixed",
+    non-death tokens use exponential hazards and death uses death_rho.
+    """
+    if logits.ndim != 2:
+        raise ValueError(f"logits must have shape (N, V), got {tuple(logits.shape)}")
+    if float(torch.as_tensor(tau_years).detach().min().cpu()) < 0:
+        raise ValueError("tau_years must be non-negative")
+
+    mode = str(dist_mode).lower()
+    if mode not in {"exponential", "weibull", "mixed"}:
+        raise ValueError("dist_mode must be one of: exponential, weibull, mixed")
+
+    rate = F.softplus(logits) + float(eps)
+    tau = torch.as_tensor(tau_years, dtype=rate.dtype, device=rate.device)
+    if tau.ndim == 0:
+        tau = tau.expand(logits.shape[0])
+    if tau.ndim != 1 or tau.shape[0] != logits.shape[0]:
+        raise ValueError(
+            "tau_years must be a scalar or a 1D tensor with length N, got "
+            f"{tuple(tau.shape)} for N={logits.shape[0]}"
+        )
+
+    if mode == "exponential":
+        exposure = tau[:, None].expand_as(rate)
+    elif mode == "weibull":
+        if rho is None or rho.shape != logits.shape:
+            raise ValueError("rho must have the same shape as logits for dist_mode='weibull'")
+        exposure = torch.pow(tau[:, None].clamp_min(float(eps)), rho.to(rate.dtype))
+    else:
+        exposure = tau[:, None].expand_as(rate).clone()
+        if death_rho is None:
+            raise ValueError("death_rho is required for dist_mode='mixed'")
+        death_idx = death_token(logits.shape[1])
+        death_shape = tuple(death_rho.shape)
+        death_rho = death_rho.to(device=rate.device, dtype=rate.dtype)
+        if death_rho.ndim == 2 and death_rho.shape[1] == 1:
+            death_rho = death_rho.squeeze(1)
+        if death_rho.ndim != 1 or death_rho.shape[0] != logits.shape[0]:
+            raise ValueError(
+                "death_rho must have shape (N,) or (N, 1), got "
+                f"{death_shape} for N={logits.shape[0]}"
+            )
+        exposure[:, death_idx] = torch.pow(tau.clamp_min(float(eps)), death_rho)
+
+    return -torch.expm1(-rate * exposure)
+
+
+def death_risk_from_probabilities(probabilities: torch.Tensor) -> torch.Tensor:
+    """Return p_death(t, tau), with death fixed to token vocab_size - 1."""
+    if probabilities.ndim != 2:
+        raise ValueError(
+            f"probabilities must have shape (N, V), got {tuple(probabilities.shape)}"
+        )
+    return probabilities[:, death_token(probabilities.shape[1])]
+
+
+def new_disease_risk_from_probabilities(
+    probabilities: torch.Tensor,
+    occurred: torch.Tensor,
+    disease_ids: Sequence[int],
+) -> torch.Tensor:
+    """
+    Compute P(at least one selected disease newly occurs within tau years).
+
+    Already occurred diseases are masked out. Death is not included here and
+    should be reported separately with death_risk_from_probabilities.
+    """
+    if probabilities.ndim != 2 or occurred.shape != probabilities.shape:
+        raise ValueError(
+            "probabilities and occurred must both have shape (N, V), got "
+            f"{tuple(probabilities.shape)} and {tuple(occurred.shape)}"
+        )
+    if not disease_ids:
+        return probabilities.new_zeros(probabilities.shape[0])
+
+    death_idx = death_token(probabilities.shape[1])
+    ids = [
+        idx
+        for idx in dict.fromkeys(int(x) for x in disease_ids)
+        if 0 <= idx < probabilities.shape[1] and idx != death_idx
+    ]
+    if not ids:
+        return probabilities.new_zeros(probabilities.shape[0])
+
+    idx_tensor = torch.as_tensor(ids, dtype=torch.long, device=probabilities.device)
+    p = probabilities[:, idx_tensor].clamp(0.0, 1.0 - 1e-7)
+    new_mask = ~occurred[:, idx_tensor].to(dtype=torch.bool)
+    log_no_new = torch.log1p(-p) * new_mask.to(dtype=p.dtype)
+    return -torch.expm1(log_no_new.sum(dim=1))
--- a/labels.csv
+++ b/labels.csv
--- a/losses.py
+++ b/losses.py
@@ -0,0 +1,400 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+PAD_IDX = 0
+CHECKUP_IDX = 1
+NO_EVENT_IDX = 2
+
+
+def _make_ignore_mask(
+    vocab_size: int,
+    ignored_idx: Iterable[int],
+    device: torch.device,
+) -> torch.Tensor:
+    ignore_mask = torch.zeros(vocab_size, dtype=torch.bool, device=device)
+    for idx in ignored_idx:
+        idx = int(idx)
+        if 0 <= idx < vocab_size:
+            ignore_mask[idx] = True
+    return ignore_mask
+
+
+def _valid_vocab_mask(
+    vocab_size: int,
+    ignored_idx: Iterable[int],
+    device: torch.device,
+) -> torch.Tensor:
+    return ~_make_ignore_mask(vocab_size, ignored_idx, device)
+
+
+def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor:
+    return logits.sum() * 0.0
+
+
+class Delphi2MLoss(nn.Module):
+    """Next-token plus exponential time-to-next-token supervision."""
+
+    def __init__(
+        self,
+        t_min: float = 1.0 / 365.25,
+        ignored_tokens: Iterable[int] | None = None,
+        exclude_ignored_from_intensity: bool = True,
+        max_exp_input: float = 60.0,
+        ce_weight: float = 1.0,
+        time_weight: float = 1.0,
+    ):
+        super().__init__()
+        self.t_min = float(t_min)
+        self.ignored_tokens = (
+            [PAD_IDX, CHECKUP_IDX]
+            if ignored_tokens is None
+            else [int(x) for x in ignored_tokens]
+        )
+        self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)
+        self.max_exp_input = float(max_exp_input)
+        self.ce_weight = float(ce_weight)
+        self.time_weight = float(time_weight)
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        target_events: torch.Tensor,
+        target_times: torch.Tensor,
+        current_times: torch.Tensor,
+        padding_mask: torch.Tensor,
+        return_components: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        if logits.dim() != 3:
+            raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
+        bsz, seq_len, vocab_size = logits.shape
+
+        expected = (bsz, seq_len)
+        if target_events.shape != expected:
+            raise ValueError(f"target_events must be {expected}, got {tuple(target_events.shape)}")
+        if target_times.shape != expected:
+            raise ValueError(f"target_times must be {expected}, got {tuple(target_times.shape)}")
+        if current_times.shape != expected:
+            raise ValueError(f"current_times must be {expected}, got {tuple(current_times.shape)}")
+        if padding_mask.shape != expected:
+            raise ValueError(f"padding_mask must be {expected}, got {tuple(padding_mask.shape)}")
+
+        valid_mask = padding_mask.bool()
+        for idx in self.ignored_tokens:
+            valid_mask = valid_mask & (target_events != int(idx))
+        valid_mask = valid_mask & (target_events > PAD_IDX)
+
+        if not valid_mask.any():
+            total_loss = _zero_loss_like(logits)
+            if return_components:
+                return total_loss, {
+                    "ce": total_loss.detach(),
+                    "time": total_loss.detach(),
+                    "total": total_loss.detach(),
+                }
+            return total_loss
+
+        logits_valid = logits[valid_mask]
+        target_events_valid = target_events[valid_mask]
+        target_times_valid = target_times[valid_mask]
+        current_times_valid = current_times[valid_mask]
+
+        logits_safe = torch.nan_to_num(
+            logits_valid,
+            nan=0.0,
+            posinf=self.max_exp_input,
+            neginf=-self.max_exp_input,
+        )
+
+        loss_ce = F.cross_entropy(
+            logits_safe,
+            target_events_valid,
+            reduction="mean",
+        )
+
+        logits_for_lse = logits_safe
+        if self.exclude_ignored_from_intensity:
+            ignore_mask = _make_ignore_mask(vocab_size, self.ignored_tokens, logits.device)
+            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
+
+        log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
+        log_lambda_total = -torch.log(torch.exp(-log_lambda_total) + self.t_min)
+
+        dt = torch.clamp(target_times_valid - current_times_valid, min=self.t_min)
+        log_dt_inv = -torch.log(dt + self.t_min)
+        loss_dt = -(
+            log_lambda_total
+            - torch.exp(
+                torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input)
+            )
+        )
+        loss_dt = loss_dt.mean()
+
+        total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt
+        if return_components:
+            return total_loss, {
+                "ce": loss_ce.detach(),
+                "time": loss_dt.detach(),
+                "total": total_loss.detach(),
+            }
+        return total_loss
+
+
+class UniqueTimeSetExponentialLoss(nn.Module):
+    """Next distinct timestamp event-set supervision with sum reduction."""
+
+    def __init__(
+        self,
+        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+        t_min: float = 1.0 / 365.25,
+        max_exp_input: float = 60.0,
+        exclude_ignored_from_intensity: bool = True,
+    ):
+        super().__init__()
+        self.ignored_idx = [int(x) for x in ignored_idx]
+        self.t_min = float(t_min)
+        self.max_exp_input = float(max_exp_input)
+        self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        target_multi_hot: torch.Tensor,
+        target_dt_unique: torch.Tensor,
+        readout_mask: torch.Tensor,
+        return_components: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        if logits.dim() != 3:
+            raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
+        bsz, seq_len, vocab_size = logits.shape
+
+        if target_multi_hot.shape != (bsz, seq_len, vocab_size):
+            raise ValueError(
+                "target_multi_hot must match logits shape, "
+                f"got {tuple(target_multi_hot.shape)} vs {tuple(logits.shape)}"
+            )
+        if target_dt_unique.shape != (bsz, seq_len):
+            raise ValueError(
+                f"target_dt_unique must be {(bsz, seq_len)}, got {tuple(target_dt_unique.shape)}"
+            )
+        if readout_mask.shape != (bsz, seq_len):
+            raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}")
+
+        ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)
+
+        num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1)
+        valid_mask = readout_mask.bool() & (num_targets > 0)
+
+        if not valid_mask.any():
+            total_loss = _zero_loss_like(logits)
+            if return_components:
+                return total_loss, {
+                    "observed": total_loss.detach(),
+                    "penalty": total_loss.detach(),
+                    "total": total_loss.detach(),
+                }
+            return total_loss
+
+        logits_safe = torch.nan_to_num(
+            logits[valid_mask],
+            nan=0.0,
+            posinf=self.max_exp_input,
+            neginf=-self.max_exp_input,
+        )
+        target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype)
+        target_valid[:, ignore_mask] = 0.0
+
+        observed_term = (logits_safe * target_valid).sum(dim=-1)
+        penalty_scale = target_valid.sum(dim=-1)
+
+        logits_for_lse = logits_safe
+        if self.exclude_ignored_from_intensity:
+            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
+
+        dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min)
+        log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
+        log_penalty = log_lambda_total + dt_clamped.log()
+        penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input))
+
+        observed_loss = -observed_term
+        penalty_loss = penalty_scale * penalty
+        total_loss = (observed_loss + penalty_loss).mean()
+
+        if return_components:
+            return total_loss, {
+                "observed": observed_loss.mean().detach(),
+                "penalty": penalty_loss.mean().detach(),
+                "total": total_loss.detach(),
+            }
+        return total_loss
+
+
+class ExponentialLoss(nn.Module):
+    """Query-conditioned all-future-event exponential point-process loss."""
+
+    def __init__(
+        self,
+        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+        eps: float = 1e-8,
+    ):
+        super().__init__()
+        self.ignored_idx = tuple(int(i) for i in ignored_idx)
+        self.eps = eps
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        targets: torch.Tensor,
+        exposure: torch.Tensor,
+    ) -> torch.Tensor:
+        _, vocab_size = logits.shape
+        rate = F.softplus(logits) + self.eps
+        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
+
+        penalty = exposure.to(rate.dtype) * rate[:, valid_vocab].sum(dim=-1)
+        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
+        for idx in self.ignored_idx:
+            target_valid &= targets != idx
+
+        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
+        observed = rate.log().gather(1, safe_targets)
+        observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1)
+        return (-observed + penalty).mean()
+
+
+class WeibullLoss(nn.Module):
+    """Query-conditioned all-future-event Weibull point-process loss."""
+
+    def __init__(
+        self,
+        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+        eps: float = 1e-8,
+    ):
+        super().__init__()
+        self.ignored_idx = tuple(int(i) for i in ignored_idx)
+        self.eps = eps
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        weibull_rho: torch.Tensor,
+        targets: torch.Tensor,
+        dt: torch.Tensor,
+        exposure: torch.Tensor,
+    ) -> torch.Tensor:
+        _, vocab_size = logits.shape
+        if weibull_rho is None:
+            raise ValueError("weibull_rho is required for WeibullLoss")
+        if weibull_rho.shape != logits.shape:
+            raise ValueError(
+                "weibull_rho must have the same shape as logits. "
+                f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}"
+            )
+
+        dtype = logits.dtype
+        rate = F.softplus(logits) + self.eps
+        rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
+        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
+
+        t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1)
+        penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1)
+
+        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
+        for idx in self.ignored_idx:
+            target_valid &= targets != idx
+
+        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
+        target_rate = rate.gather(1, safe_targets)
+        target_rho = rho.gather(1, safe_targets)
+        target_dt = dt.to(dtype).clamp_min(self.eps)
+        log_intensity = (
+            target_rate.log()
+            + target_rho.log()
+            + (target_rho - 1.0) * target_dt.log()
+        )
+        observed = (log_intensity * target_valid.to(dtype)).sum(dim=-1)
+        return (-observed + penalty).mean()
+
+
+class MixedLoss(nn.Module):
+    """Exponential diseases plus one Weibull death endpoint."""
+
+    def __init__(
+        self,
+        death_idx: int,
+        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+        eps: float = 1e-8,
+    ):
+        super().__init__()
+        self.death_idx = int(death_idx)
+        self.ignored_idx = tuple(int(i) for i in ignored_idx)
+        self.eps = eps
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        death_rho: torch.Tensor,
+        targets: torch.Tensor,
+        dt: torch.Tensor,
+        exposure: torch.Tensor,
+    ) -> torch.Tensor:
+        _, vocab_size = logits.shape
+        dtype = logits.dtype
+        rate = F.softplus(logits) + self.eps
+
+        if death_rho.dim() == 2:
+            death_rho = death_rho.squeeze(-1)
+        death_rho = death_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
+
+        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
+        valid_disease_vocab = valid_vocab.clone()
+        valid_disease_vocab[self.death_idx] = False
+
+        t_exp = exposure.to(dtype).clamp_min(self.eps)
+        disease_penalty = t_exp * rate[:, valid_disease_vocab].sum(dim=-1)
+        death_rate = rate[:, self.death_idx]
+        death_penalty = death_rate * torch.pow(t_exp, death_rho)
+        penalty = disease_penalty + death_penalty
+
+        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
+        for idx in self.ignored_idx:
+            target_valid &= targets != idx
+
+        disease_event_mask = target_valid & (targets != self.death_idx)
+        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
+        disease_log_rate = rate.log().gather(1, safe_targets)
+        observed_disease = (disease_log_rate * disease_event_mask.to(dtype)).sum(dim=-1)
+
+        death_event_mask = target_valid & (targets == self.death_idx)
+        death_observed = death_event_mask.any(dim=1)
+        death_dt = (dt.to(dtype).clamp_min(self.eps) * death_event_mask.to(dtype)).sum(dim=1)
+        death_log_intensity = (
+            death_rate.log()
+            + death_rho.log()
+            + (death_rho - 1.0) * death_dt.clamp_min(self.eps).log()
+        )
+        observed_death = death_log_intensity * death_observed.to(dtype)
+
+        return (-observed_disease - observed_death + penalty).mean()
+
+
+def build_loss(name: str, **kwargs) -> nn.Module:
+    name = name.lower()
+    if name in {"delphi2m", "d2m", "next_token"}:
+        return Delphi2MLoss(**kwargs)
+    if name in {"uts", "unique_time_set", "unique_time_exponential"}:
+        return UniqueTimeSetExponentialLoss(**kwargs)
+    if name in {"exponential", "query_exponential"}:
+        return ExponentialLoss(**kwargs)
+    if name in {"weibull", "query_weibull"}:
+        return WeibullLoss(**kwargs)
+    if name in {"mixed", "query_mixed"}:
+        return MixedLoss(**kwargs)
+    raise ValueError(
+        f"Unknown loss {name!r}. Available: delphi2m, uts, exponential, weibull, mixed."
+    )
--- a/models.py
+++ b/models.py
@@ -0,0 +1,481 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from backbones import (
+    AgeSinusoidalEncoding,
+    GPTBlock,
+    GaussianRBFTimeBasis,
+    TimeRoPE,
+    TokenAutoDiscretization,
+)
+from targets import PAD_IDX
+
+
+@dataclass
+class DeepHealthOutput:
+    hidden: torch.Tensor
+    time_seq: torch.Tensor
+    padding_mask: torch.Tensor
+    event_len: int
+
+
+class OtherInfoTokenizer(nn.Module):
+    PAD_KIND = 0
+    CONT_KIND = 1
+    CATE_KIND = 2
+
+    def __init__(
+        self,
+        n_embd: int,
+        n_types: int,
+        n_cont_types: int,
+        n_categories: int,
+        cont_type_ids: list[int],
+        n_value_kinds: int = 3,
+        n_bins: int = 16,
+    ):
+        super().__init__()
+        if len(cont_type_ids) != n_cont_types:
+            raise ValueError(
+                "cont_type_ids length must match n_cont_types, got "
+                f"{len(cont_type_ids)} vs {n_cont_types}"
+            )
+        if n_types <= 0:
+            raise ValueError(f"n_types must include PAD and be > 0, got {n_types}")
+        if n_categories <= 0:
+            raise ValueError(
+                f"n_categories must include PAD and be > 0, got {n_categories}"
+            )
+        if n_value_kinds <= self.CATE_KIND:
+            raise ValueError(
+                f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
+            )
+
+        self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
+        self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
+        self.cont_value_encoder = (
+            TokenAutoDiscretization(
+                n_cont_types=n_cont_types,
+                n_bins=n_bins,
+                n_embd=n_embd,
+            )
+            if n_cont_types > 0
+            else None
+        )
+        self.cate_value_emb = nn.Embedding(
+            n_categories,
+            n_embd,
+            padding_idx=0,
+        )
+
+        cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
+        for idx, type_id in enumerate(cont_type_ids):
+            if type_id <= 0 or type_id >= n_types:
+                raise ValueError(
+                    f"continuous type id {type_id} must be in [1, {n_types})"
+                )
+            cont_type_index[type_id] = idx
+        self.register_buffer(
+            "cont_type_index",
+            cont_type_index,
+            persistent=False,
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.type_emb.weight[0])
+        nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.kind_emb.weight[0])
+        nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.cate_value_emb.weight[0])
+
+    def forward(
+        self,
+        other_type: torch.LongTensor,
+        other_value: torch.Tensor,
+        other_value_kind: torch.LongTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if other_type.shape != other_value.shape:
+            raise ValueError(
+                "other_type and other_value must have the same shape, got "
+                f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}"
+            )
+        if other_type.shape != other_value_kind.shape:
+            raise ValueError(
+                "other_type and other_value_kind must have the same shape, got "
+                f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}"
+            )
+
+        other_valid = other_type > 0
+        type_emb = self.type_emb(other_type)
+        kind_emb = self.kind_emb(other_value_kind)
+        value_emb = torch.zeros_like(type_emb)
+
+        cont_pos = other_valid & (other_value_kind == self.CONT_KIND)
+        if cont_pos.any():
+            if self.cont_value_encoder is None:
+                raise ValueError("continuous tokens found but n_cont_types is 0")
+            cont_idx = self.cont_type_index[other_type[cont_pos]]
+            if (cont_idx < 0).any():
+                bad_type = other_type[cont_pos][cont_idx < 0][0].item()
+                raise ValueError(
+                    f"type_id={bad_type} is marked continuous but is not in "
+                    "cont_type_ids"
+                )
+            value_emb[cont_pos] = self.cont_value_encoder(
+                cont_type_idx=cont_idx,
+                value=other_value[cont_pos].to(type_emb.dtype),
+            )
+
+        cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
+        if cate_pos.any():
+            cate_id = other_value[cate_pos].long()
+            value_emb[cate_pos] = self.cate_value_emb(cate_id)
+
+        out = type_emb + kind_emb + value_emb
+        out = out * other_valid.unsqueeze(-1).to(out.dtype)
+        return out, other_valid
+
+
+class DeepHealth(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        n_embd: int,
+        n_head: int,
+        n_hist_layer: int,
+        n_tab_layer: int,
+        n_types: int,
+        n_cont_types: int,
+        n_categories: int,
+        cont_type_ids: list[int],
+        n_value_kinds: int = 3,
+        n_bins: int = 16,
+        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"
+        extra_pool_reduce: str = "mean",
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        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'")
+        if extra_pool_reduce not in {"mean", "sum"}:
+            raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
+        self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
+        self.gender_embedding = nn.Embedding(
+            2, n_embd)  # Assuming binary gender
+        self.tokenizer = OtherInfoTokenizer(
+            n_embd=n_embd,
+            n_types=n_types,
+            n_cont_types=n_cont_types,
+            n_categories=n_categories,
+            cont_type_ids=cont_type_ids,
+            n_value_kinds=n_value_kinds,
+            n_bins=n_bins,
+        )
+        self.target_mode = target_mode
+        self.time_mode = time_mode
+        self.dist_mode = dist_mode
+        self.extra_pool_reduce = extra_pool_reduce
+        self.n_embd = n_embd
+        self.vocab_size = vocab_size
+        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.token_embedding.weight[0])
+        nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
+        if dist_mode == "weibull":
+            self.rho_head = nn.Linear(n_embd, vocab_size)
+            nn.init.zeros_(self.rho_head.weight)
+            nn.init.constant_(self.rho_head.bias, 0.5413)
+
+        if dist_mode == "mixed":
+            self.death_idx = vocab_size - 1
+            self.rho_death_head = nn.Linear(n_embd, 1)
+            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)
+        if target_mode == "next_token":
+            self.risk_head.weight = self.token_embedding.weight
+        self.query_token = nn.Parameter(torch.zeros(n_embd))
+        nn.init.normal_(self.query_token, mean=0.0, std=0.02)
+
+    def _make_history_attn_mask(
+        self,
+        padding_mask: torch.Tensor,
+        time_seq: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        valid_key = padding_mask[:, None, :]  # (B, 1, L)
+        visible_by_time = time_seq[:, None, :] <= time_seq[:, :, None]
+        valid = valid_key & visible_by_time
+        return torch.zeros(
+            valid.shape,
+            device=valid.device,
+            dtype=dtype,
+        ).masked_fill(~valid, -1e4)[:, None, :, :]
+
+    def _pool_other_by_time(
+        self,
+        h_other: torch.Tensor,
+        other_time: torch.Tensor,
+        other_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, n_other, n_embd = h_other.shape
+        if n_other == 0:
+            empty_h = h_other.new_zeros(batch_size, 0, n_embd)
+            empty_t = other_time.new_zeros(batch_size, 0)
+            empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device)
+            return empty_h, empty_t, empty_m
+
+        masked_time = other_time.masked_fill(~other_mask, float("inf"))
+        _sorted_time_with_pad, order = masked_time.sort(dim=1)
+        sorted_time = other_time.gather(1, order)
+        sorted_mask = other_mask.gather(1, order)
+        sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd))
+
+        group_start = torch.zeros_like(sorted_mask)
+        group_start[:, 0] = sorted_mask[:, 0]
+        group_start[:, 1:] = sorted_mask[:, 1:] & (
+            sorted_time[:, 1:] != sorted_time[:, :-1]
+        )
+        group_id = group_start.long().cumsum(dim=1) - 1
+        max_groups = int(group_start.sum(dim=1).max().item())
+
+        pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
+        pooled_time = other_time.new_zeros(batch_size, max_groups)
+        pooled_mask = torch.zeros(
+            batch_size,
+            max_groups,
+            dtype=torch.bool,
+            device=h_other.device,
+        )
+        if max_groups == 0:
+            return pooled_h, pooled_time, pooled_mask
+
+        safe_group_id = group_id.clamp_min(0)
+        pooled_h.scatter_add_(
+            1,
+            safe_group_id.unsqueeze(-1).expand_as(sorted_h),
+            sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype),
+        )
+        if self.extra_pool_reduce == "mean":
+            counts = h_other.new_zeros(batch_size, max_groups, 1)
+            counts.scatter_add_(
+                1,
+                safe_group_id.unsqueeze(-1),
+                sorted_mask.unsqueeze(-1).to(h_other.dtype),
+            )
+            pooled_h = pooled_h / counts.clamp_min(1.0)
+
+        pooled_time.scatter_add_(
+            1,
+            safe_group_id,
+            sorted_time * group_start.to(sorted_time.dtype),
+        )
+        group_count = group_start.sum(dim=1)
+        arange_groups = torch.arange(max_groups, device=h_other.device)
+        pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
+        return pooled_h, pooled_time, pooled_mask
+
+    def _forward_shared(
+        self,
+        event_seq: torch.LongTensor,
+        time_seq: torch.FloatTensor,
+        sex: torch.LongTensor,
+        mode: str,
+        padding_mask: torch.Tensor | None = None,
+        t_query: torch.FloatTensor | None = None,
+        other_type: torch.LongTensor | None = None,
+        other_value: torch.Tensor | None = None,
+        other_value_kind: torch.LongTensor | None = None,
+        other_time: torch.FloatTensor | None = None,
+        return_output: bool = False,
+        **unused_kwargs,
+    ) -> torch.Tensor | DeepHealthOutput:
+        if unused_kwargs:
+            unknown = ", ".join(sorted(unused_kwargs))
+            raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
+        if mode not in {"next_token", "all_future"}:
+            raise ValueError("mode must be either 'next_token' or 'all_future'")
+        if mode == "all_future" and t_query is None:
+            raise ValueError("t_query is required when mode='all_future'")
+        if (
+            other_type is None
+            or other_value is None
+            or other_value_kind is None
+            or other_time is None
+        ):
+            raise ValueError(
+                "DeepHealth expects other_type, other_value, "
+                "other_value_kind, and other_time."
+            )
+
+        if padding_mask is None:
+            padding_mask = event_seq > PAD_IDX
+        else:
+            padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
+
+        event_len = event_seq.size(1)
+        h_disease = self.token_embedding(event_seq)
+        t_disease = time_seq
+
+        if other_time.shape != other_type.shape:
+            raise ValueError(
+                "other_time must have the same shape as other_type, got "
+                f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
+            )
+        other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
+        h_other, other_mask = self.tokenizer(
+            other_type=other_type,
+            other_value=other_value,
+            other_value_kind=other_value_kind,
+        )
+        h_other = h_other.to(device=event_seq.device)
+        other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
+
+        h_disease = torch.cat([h_disease, h_other], dim=1)
+        t_disease = torch.cat([t_disease, other_time], dim=1)
+        padding_mask = torch.cat([padding_mask, other_mask], dim=1)
+        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+
+        if mode == "all_future":
+            batch_size = event_seq.size(0)
+            query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
+            h_disease = torch.cat([h_disease, query], dim=1)
+            t_disease = torch.cat([t_disease, t_query[:, None]], dim=1)
+            query_mask = torch.ones(
+                batch_size,
+                1,
+                dtype=torch.bool,
+                device=event_seq.device,
+            )
+            padding_mask = torch.cat([padding_mask, query_mask], dim=1)
+
+        sex_emb = self.gender_embedding(sex)[:, None, :]
+        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,
+            time_seq=t_disease,
+            dtype=h_disease.dtype,
+        )
+        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)
+
+        h_disease = self.final_ln(h_disease)
+        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+
+        if mode == "all_future":
+            hidden = h_disease[:, -1, :]
+            if return_output:
+                return DeepHealthOutput(
+                    hidden=hidden,
+                    time_seq=t_query[:, None],
+                    padding_mask=torch.ones(
+                        hidden.size(0),
+                        1,
+                        dtype=torch.bool,
+                        device=hidden.device,
+                    ),
+                    event_len=event_len,
+                )
+            return hidden
+        if return_output:
+            h_event = h_disease[:, :event_len, :]
+            t_event = t_disease[:, :event_len]
+            event_mask = padding_mask[:, :event_len]
+            h_extra, t_extra, extra_mask = self._pool_other_by_time(
+                h_other=h_disease[:, event_len:, :],
+                other_time=t_disease[:, event_len:],
+                other_mask=padding_mask[:, event_len:],
+            )
+            return DeepHealthOutput(
+                hidden=torch.cat([h_event, h_extra], dim=1),
+                time_seq=torch.cat([t_event, t_extra], dim=1),
+                padding_mask=torch.cat([event_mask, extra_mask], dim=1),
+                event_len=event_len,
+            )
+        return h_disease[:, :event_len, :]
+
+    def forward_next_token(self, **kwargs) -> torch.Tensor:
+        return self._forward_shared(mode="next_token", **kwargs)
+
+    def forward_all_future(self, **kwargs) -> torch.Tensor:
+        return self._forward_shared(mode="all_future", **kwargs)
+
+    def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor:
+        mode = self.target_mode if target_mode is None else target_mode
+        return self._forward_shared(mode=mode, **kwargs)
+
+    def calc_risk(self, x: torch.Tensor) -> torch.Tensor:
+        return self.risk_head(x)
+
+    def calc_weibull_rho(self, x: torch.Tensor) -> torch.Tensor:
+        if self.dist_mode != "weibull":
+            raise RuntimeError(
+                f"calc_weibull_rho called with dist_mode={self.dist_mode!r}"
+            )
+        return F.softplus(self.rho_head(x)) + 1e-6
+
+    def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor:
+        if self.dist_mode != "mixed":
+            raise RuntimeError(
+                f"calc_death_rho called with dist_mode={self.dist_mode!r}"
+            )
+        return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6
--- a/prepare_data.py
+++ b/prepare_data.py
@@ -0,0 +1,385 @@
+"""ETL pipeline for UK Biobank data preparation.
+
+This script converts raw UK Biobank CSV exports into the artefacts consumed by
+DeepHealth:
+
+* ``ukb_event_data.npy``: ``(N, 3)`` uint32 array of ``(eid, days, label)``
+  disease/death/checkup events sorted by patient then time.
+* ``ukb_basic_info.csv``: basic patient table indexed by ``eid`` with ``sex``.
+* ``ukb_other_info.npy``: ``(M, 5)`` float64 array of
+  ``(eid, type, value, value_kind, time)`` rows. ``type=0`` is reserved for
+  padding, ``value_kind=1`` means continuous, and ``value_kind=2`` means
+  categorical.
+* ``cate_types.csv``: categorical-variable metadata with
+  ``type,name,n_categories``. Dataset code computes global category ids after
+  experiment-specific variable selection.
+
+Processing steps
+----------------
+1. Stream the raw CSV in 10 000-row chunks to bound memory usage.
+2. Convert date columns to integer day offsets from date of birth.
+3. Extract ICD-10 date fields and cancer date/type fields into long-form
+   events and map codes to integer labels via ``labels.csv``.
+4. De-duplicate events per ``(eid, label)`` keeping the earliest occurrence.
+5. Convert available non-sex tabular fields into unified other-information
+   tokens timestamped by ``date_of_assessment``.
+6. Write event data, sex, unified other-information tokens, and categorical
+   type metadata.
+
+Usage
+-----
+::
+
+    python prepare_data.py
+
+The script expects these files in the working directory:
+
+* ``ukb_data.csv``: raw UK Biobank CSV export.
+* ``field_ids_enriched.csv``: field-ID to column-name mapping.
+* ``icd10_codes_mod.tsv``: ICD-10 field to date-column mapping.
+* ``labels.csv``: disease-code to integer-label mapping.
+"""
+
+import numpy as np  # Numerical operations
+import pandas as pd  # Pandas for data manipulation
+import tqdm  # Progress bar for chunk processing
+
+
+CONT_VALUE_KIND = 1
+CATE_VALUE_KIND = 2
+
+
+def _sort_values(values):
+    """Sort mixed pandas/numpy scalar values deterministically."""
+    try:
+        return sorted(values)
+    except TypeError:
+        return sorted(values, key=lambda x: str(x))
+
+
+def _build_other_info_tokens(
+    table: pd.DataFrame,
+    feature_fields: list[str],
+    *,
+    time_col: str = "date_of_assessment",
+) -> tuple[np.ndarray, pd.DataFrame]:
+    """Convert wide tabular features into (eid, type, value, kind, time) rows."""
+    rows = []
+    cate_meta = []
+    present_features = [
+        col for col in feature_fields
+        if col in table.columns and col not in {time_col, "sex"}
+    ]
+
+    if time_col not in table.columns:
+        raise ValueError(
+            f"{time_col!r} is required to timestamp other-information tokens"
+        )
+
+    token_times = pd.to_numeric(table[time_col], errors="coerce")
+
+    for type_idx, col in enumerate(present_features, start=1):
+        series = table[col]
+        n_unique = series.dropna().nunique()
+        valid_time = token_times.notna()
+
+        if n_unique > 10:
+            numeric = pd.to_numeric(series, errors="coerce")
+            valid = numeric.notna() & valid_time
+            if not valid.any():
+                continue
+            rows.append(
+                np.column_stack(
+                    (
+                        table.index[valid].to_numpy(dtype=np.float64),
+                        np.full(valid.sum(), type_idx, dtype=np.float64),
+                        numeric[valid].to_numpy(dtype=np.float64),
+                        np.full(valid.sum(), CONT_VALUE_KIND, dtype=np.float64),
+                        token_times[valid].to_numpy(dtype=np.float64),
+                    )
+                )
+            )
+        else:
+            unique_vals = _sort_values(series.dropna().unique())
+            value_map = {val: idx + 1 for idx, val in enumerate(unique_vals)}
+            mapped = series.map(value_map)
+            valid = mapped.notna() & valid_time
+            n_categories = len(unique_vals)
+            cate_meta.append(
+                {
+                    "type": type_idx,
+                    "name": col,
+                    "n_categories": n_categories,
+                }
+            )
+            if not valid.any():
+                continue
+            rows.append(
+                np.column_stack(
+                    (
+                        table.index[valid].to_numpy(dtype=np.float64),
+                        np.full(valid.sum(), type_idx, dtype=np.float64),
+                        mapped[valid].to_numpy(dtype=np.float64),
+                        np.full(valid.sum(), CATE_VALUE_KIND, dtype=np.float64),
+                        token_times[valid].to_numpy(dtype=np.float64),
+                    )
+                )
+            )
+
+    cate_types = pd.DataFrame(
+        cate_meta,
+        columns=["type", "name", "n_categories"],
+    )
+    if not rows:
+        return np.empty((0, 5), dtype=np.float64), cate_types
+    out = np.vstack(rows)
+    return out[np.lexsort((out[:, 3], out[:, 1], out[:, 0]))], cate_types
+
+
+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"
+
+field_df = pd.read_csv(field_map_file, low_memory=False)
+required_cols = {"field_instance", "var_name", "field_type"}
+missing_cols = required_cols - set(field_df.columns)
+if missing_cols:
+    raise ValueError(
+        f"{field_map_file} is missing required columns: {sorted(missing_cols)}"
+    )
+
+field_df = field_df.dropna(subset=["field_instance", "var_name", "field_type"])
+field_df["field_instance"] = field_df["field_instance"].astype(str)
+field_df["var_name"] = field_df["var_name"].astype(str)
+field_df["field_type"] = field_df["field_type"].astype(int)
+
+# Map original field ID -> renamed output variable.
+field_dict = dict(zip(field_df["field_instance"], field_df["var_name"]))
+
+# Build source field groups from field_type.
+basic_info_fields = _unique_preserve_order(
+    field_df.loc[field_df["field_type"] == 0, "var_name"].tolist()
+)
+assessment_fields = _unique_preserve_order(
+    field_df.loc[field_df["field_type"] == 1, "var_name"].tolist()
+)
+exposure_fields = _unique_preserve_order(
+    field_df.loc[field_df["field_type"] == 2, "var_name"].tolist()
+)
+
+# Keep only sex and enrollment time for the basic info table.
+basic_info_fields = [
+    f for f in ["sex", "date_of_assessment"] if f in set(basic_info_fields)
+]
+
+# Fields needed for tabular extraction from raw CSV.
+tabular_fields = _unique_preserve_order(
+    basic_info_fields + assessment_fields + exposure_fields
+)
+
+# TSV mapping field IDs to ICD10-related date columns
+field_to_icd_map = "icd10_codes_mod.tsv"
+# Date-like variables to be converted to offsets
+date_vars = []
+with open(field_to_icd_map, encoding="utf-8") as f:  # Open ICD10 mapping
+    for line in f:  # Iterate each mapping row
+        parts = line.strip().split()  # Split on whitespace for TSV
+        if len(parts) >= 6:  # Guard against malformed lines
+            # Map field ID to the date column name
+            field_dict[parts[0]] = parts[5]
+            date_vars.append(parts[5])  # Track date column names in order
+
+for j in range(17):  # Map up to 17 cancer entry slots (dates and types)
+    # Cancer diagnosis date slot j
+    field_dict[f"40005-{j}.0"] = f"cancer_date_{j}"
+    field_dict[f"40006-{j}.0"] = f"cancer_type_{j}"  # Cancer type/code slot j
+
+# Number of ICD-related date columns before adding extras
+len_icd = len(date_vars)
+date_vars.extend(
+    ["Death", "date_of_assessment"]  # Add outcome date and assessment date
+    +
+    # Add cancer date columns
+    [f"cancer_date_{j}" for j in range(17)]
+)
+
+labels_file = "labels.csv"  # File listing label codes
+label_dict = {}  # Map code string -> integer label id
+with open(labels_file, encoding="utf-8") as f:  # Open labels file
+    for idx, line in enumerate(f):  # Enumerate to assign incremental label IDs
+        parts = line.strip().split(" ")  # Split by space
+        if parts and parts[0]:  # Guard against empty lines
+            # Start labels from 1 to reserve 0 for padding, 1 for checkup
+            label_dict[parts[0]] = idx + 2
+
+# Pre-build lookup: ICD/Death column name -> integer label for fast per-column extraction
+icd_label_lookup = {}
+for col_name in date_vars[:len_icd]:
+    if col_name in label_dict:
+        icd_label_lookup[col_name] = label_dict[col_name]
+if "Death" in label_dict:
+    icd_label_lookup["Death"] = label_dict["Death"]
+
+event_list = []  # Accumulator for event arrays across chunks
+tabular_list = []  # Accumulator for tabular feature DataFrames across chunks
+ukb_iterator = pd.read_csv(  # Stream UK Biobank data in chunks
+    "ukb_data.csv",
+    sep=",",
+    chunksize=10000,  # Stream file in manageable chunks to reduce memory footprint
+    # First column (participant ID) becomes DataFrame index
+    index_col=0,
+    low_memory=False,  # Disable type inference optimization for consistent dtypes
+)
+# Iterate chunks with progress
+for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
+    # Rename columns to friendly names
+    ukb_chunk = ukb_chunk.rename(columns=field_dict)
+    # Require sex to be present
+    ukb_chunk.dropna(subset=["sex"], inplace=True)
+    if ukb_chunk.empty:
+        continue
+
+    # Construct date of birth from year and month (day fixed to 1)
+    dob = pd.to_datetime(
+        pd.DataFrame(
+            {"year": ukb_chunk["year"], "month": ukb_chunk["month"], "day": 1}
+        ),
+        errors="coerce",
+    )
+
+    # 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]
+
+    # Convert date-like columns to day offsets from dob (per-column, avoids .apply overhead)
+    for col in present_date_vars:
+        ukb_chunk[col] = (
+            pd.to_datetime(
+                ukb_chunk[col], format="%Y-%m-%d", errors="coerce") - dob
+        ).dt.days
+
+    # Append tabular features (use only columns that exist)
+    present_tabular_fields = [
+        c for c in tabular_fields if c in ukb_chunk.columns]
+    tabular_list.append(ukb_chunk[present_tabular_fields].copy())
+
+    # Extract ICD10 + Death events directly per column (avoids costly melt)
+    icd10_cols = present_date_vars[: len_icd + 1]
+    for col in icd10_cols:
+        if col not in icd_label_lookup:
+            continue
+        label_id = icd_label_lookup[col]
+        series = ukb_chunk[col]
+        valid_mask = series.notna()
+        if not valid_mask.any():
+            continue
+        event_list.append(
+            np.column_stack(
+                (
+                    ukb_chunk.index[valid_mask].values,
+                    series[valid_mask].values.astype(np.int64),
+                    np.full(valid_mask.sum(), label_id, dtype=np.int64),
+                )
+            )
+        )
+
+    # Optimized cancer processing without wide_to_long
+    cancer_frames = []
+    for j in range(17):
+        d_col = f"cancer_date_{j}"
+        t_col = f"cancer_type_{j}"
+        if d_col in ukb_chunk.columns and t_col in ukb_chunk.columns:
+            # Filter rows where both date and type are present
+            mask = ukb_chunk[d_col].notna() & ukb_chunk[t_col].notna()
+            if mask.any():
+                subset_idx = ukb_chunk.index[mask]
+                subset_days = ukb_chunk.loc[mask, d_col]
+                subset_type = ukb_chunk.loc[mask, t_col]
+
+                # Map cancer type to label
+                # Use first 3 chars
+                cancer_codes = subset_type.str.slice(0, 3)
+                labels = cancer_codes.map(label_dict)
+
+                # Filter valid labels
+                valid_label_mask = labels.notna()
+                if valid_label_mask.any():
+                    # Create array: eid, days, label
+                    # Ensure types are correct for numpy
+                    c_eids = subset_idx[valid_label_mask].values
+                    c_days = subset_days[valid_label_mask].values
+                    c_labels = labels[valid_label_mask].values
+
+                    # Stack
+                    chunk_cancer_data = np.column_stack(
+                        (c_eids, c_days, c_labels))
+                    cancer_frames.append(chunk_cancer_data)
+
+    if cancer_frames:
+        event_list.append(np.vstack(cancer_frames))
+
+    # Add checkup events with label=1 using date_of_assessment (already in days from dob)
+    if "date_of_assessment" in ukb_chunk.columns:
+        doa_series = ukb_chunk["date_of_assessment"].dropna()
+        if not doa_series.empty:
+            checkup_data = np.column_stack(
+                (
+                    doa_series.index.values,
+                    doa_series.values.astype(int),
+                    np.ones(len(doa_series), dtype=int),
+                )
+            )
+            event_list.append(checkup_data)
+
+# Combine tabular chunks
+final_tabular = pd.concat(tabular_list, axis=0, ignore_index=False)
+final_tabular.index.name = "eid"  # Ensure index named consistently
+data = np.vstack(event_list)  # Stack all event arrays into one
+
+# Sort by participant then day
+data = data[np.lexsort((data[:, 1], data[:, 0]))]
+
+# Keep only events with non-negative day offsets
+data = data[data[:, 1] >= 0]
+
+# Remove duplicate (participant_id, label) pairs keeping first occurrence (numpy-based).
+_, first_idx = np.unique(data[:, [0, 2]], axis=0, return_index=True)
+first_idx.sort()  # Preserve original sorted order
+data = data[first_idx]
+
+# Store compactly using unsigned 32-bit integers
+data = data.astype(np.uint32)
+
+# Select eid in both data and tabular
+valid_eids = np.intersect1d(data[:, 0], final_tabular.index)
+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()
+basic_info.to_csv("ukb_basic_info.csv")
+
+# Save unified other-information tokens. Missing values simply produce no token.
+other_info_fields = _unique_preserve_order(
+    basic_info_fields + assessment_fields + exposure_fields
+)
+other_info, cate_types = _build_other_info_tokens(
+    final_tabular,
+    other_info_fields,
+    time_col="date_of_assessment",
+)
+np.save("ukb_other_info.npy", other_info)
+cate_types.to_csv("cate_types.csv", index=False)
+
+# Save event data
+np.save("ukb_event_data.npy", data)
--- a/prepare_event_dates.py
+++ b/prepare_event_dates.py
@@ -0,0 +1,203 @@
+"""Create a compact calendar-dated disease-event index from ``ukb_data.csv``.
+
+Unlike ``prepare_data.py``, this ETL does not create model-ready relative-time
+sequences or other-information tokens.  It writes one structured ``.npy`` file
+with exactly three fields:
+
+    eid         int64
+    event_date  datetime64[D]
+    token       int32
+
+``token`` follows the existing ``labels.csv`` convention used by
+``prepare_data.py``: padding=0, checkup=1 (not emitted here), and the first
+label in ``labels.csv`` receives token 2.  Each ``(eid, token)`` is deduplicated
+to the first known event date.
+
+The output is intended for calendar-indexed temperature and air-pollution
+queries.  It contains no date of birth, sex, covariates, or checkup events.
+
+Usage
+-----
+python prepare_event_dates.py --output ukb_disease_event_dates.npy
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+try:
+    from tqdm import tqdm
+except ImportError:  # Keep the ETL runnable in minimal Python environments.
+    tqdm = None
+
+
+EVENT_DTYPE = np.dtype(
+    [
+        ("eid", "<i8"),
+        ("event_date", "datetime64[D]"),
+        ("token", "<i4"),
+    ]
+)
+
+
+def load_label_tokens(labels_file: str | Path) -> dict[str, int]:
+    """Return the label-code -> token mapping shared with ``prepare_data.py``."""
+    token_map: dict[str, int] = {}
+    with Path(labels_file).open(encoding="utf-8") as handle:
+        for index, line in enumerate(handle):
+            code = line.strip().split(" ", maxsplit=1)[0]
+            if code:
+                token_map[code] = index + 2
+    return token_map
+
+
+def build_raw_column_map(
+    field_map_file: str | Path,
+    icd_map_file: str | Path,
+) -> tuple[dict[str, str], list[str]]:
+    """Build raw-column renames and the calendar-date event columns to inspect."""
+    field_df = pd.read_csv(field_map_file, low_memory=False)
+    required = {"field_instance", "var_name"}
+    missing = required - set(field_df.columns)
+    if missing:
+        raise ValueError(f"{field_map_file} is missing columns: {sorted(missing)}")
+
+    raw_to_name = dict(
+        zip(field_df["field_instance"].astype(str), field_df["var_name"].astype(str))
+    )
+    icd_date_columns: list[str] = []
+    with Path(icd_map_file).open(encoding="utf-8") as handle:
+        for line in handle:
+            parts = line.strip().split()
+            if len(parts) >= 6:
+                raw_to_name[parts[0]] = parts[5]
+                icd_date_columns.append(parts[5])
+
+    for slot in range(17):
+        raw_to_name[f"40005-{slot}.0"] = f"cancer_date_{slot}"
+        raw_to_name[f"40006-{slot}.0"] = f"cancer_type_{slot}"
+    return raw_to_name, icd_date_columns
+
+
+def _records_from_icd_columns(
+    chunk: pd.DataFrame,
+    event_columns: list[str],
+    token_map: dict[str, int],
+) -> list[pd.DataFrame]:
+    frames: list[pd.DataFrame] = []
+    for column in event_columns:
+        token = token_map.get(column)
+        if token is None or column not in chunk.columns:
+            continue
+        event_date = pd.to_datetime(chunk[column], format="%Y-%m-%d", errors="coerce")
+        valid = event_date.notna()
+        if valid.any():
+            frames.append(
+                pd.DataFrame(
+                    {
+                        "eid": chunk.index[valid].astype("int64"),
+                        "event_date": event_date.loc[valid].to_numpy(),
+                        "token": np.full(valid.sum(), token, dtype=np.int32),
+                    }
+                )
+            )
+    return frames
+
+
+def _records_from_cancer_columns(chunk: pd.DataFrame, token_map: dict[str, int]) -> list[pd.DataFrame]:
+    frames: list[pd.DataFrame] = []
+    for slot in range(17):
+        date_column = f"cancer_date_{slot}"
+        type_column = f"cancer_type_{slot}"
+        if date_column not in chunk.columns or type_column not in chunk.columns:
+            continue
+        event_date = pd.to_datetime(chunk[date_column], format="%Y-%m-%d", errors="coerce")
+        code = chunk[type_column].astype("string").str.slice(0, 3)
+        token = code.map(token_map)
+        valid = event_date.notna() & token.notna()
+        if valid.any():
+            frames.append(
+                pd.DataFrame(
+                    {
+                        "eid": chunk.index[valid].astype("int64"),
+                        "event_date": event_date.loc[valid].to_numpy(),
+                        "token": token.loc[valid].astype("int32").to_numpy(),
+                    }
+                )
+            )
+    return frames
+
+
+def prepare_event_dates(
+    *,
+    ukb_data_file: str | Path,
+    field_map_file: str | Path,
+    icd_map_file: str | Path,
+    labels_file: str | Path,
+    output_file: str | Path,
+    chunksize: int = 10_000,
+) -> int:
+    """Stream the raw UKB export, then write a sorted structured event array."""
+    token_map = load_label_tokens(labels_file)
+    raw_to_name, icd_date_columns = build_raw_column_map(field_map_file, icd_map_file)
+    event_columns = [*icd_date_columns, "Death"]
+    frames: list[pd.DataFrame] = []
+
+    reader = pd.read_csv(
+        ukb_data_file,
+        chunksize=chunksize,
+        index_col=0,  # UKB participant ID / eid
+        low_memory=False,
+    )
+    chunk_iterator = tqdm(reader, desc="Extracting calendar-dated disease events") if tqdm else reader
+    for raw_chunk in chunk_iterator:
+        chunk = raw_chunk.rename(columns=raw_to_name)
+        frames.extend(_records_from_icd_columns(chunk, event_columns, token_map))
+        frames.extend(_records_from_cancer_columns(chunk, token_map))
+
+    if not frames:
+        result = np.empty(0, dtype=EVENT_DTYPE)
+    else:
+        events = pd.concat(frames, ignore_index=True)
+        events = events.dropna(subset=["eid", "event_date", "token"])
+        events = events.sort_values(["eid", "token", "event_date"], kind="stable")
+        # Match prepare_data.py: first occurrence of each disease/death token.
+        events = events.drop_duplicates(["eid", "token"], keep="first")
+        events = events.sort_values(["eid", "event_date", "token"], kind="stable")
+
+        result = np.empty(len(events), dtype=EVENT_DTYPE)
+        result["eid"] = events["eid"].to_numpy(dtype=np.int64)
+        result["event_date"] = events["event_date"].to_numpy(dtype="datetime64[D]")
+        result["token"] = events["token"].to_numpy(dtype=np.int32)
+
+    np.save(output_file, result)
+    return len(result)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--ukb-data", default="ukb_data.csv")
+    parser.add_argument("--field-map", default="field_ids_enriched.csv")
+    parser.add_argument("--icd-map", default="icd10_codes_mod.tsv")
+    parser.add_argument("--labels", default="labels.csv")
+    parser.add_argument("--output", default="ukb_disease_event_dates.npy")
+    parser.add_argument("--chunksize", type=int, default=10_000)
+    args = parser.parse_args()
+    if args.chunksize <= 0:
+        raise ValueError("chunksize must be positive")
+    count = prepare_event_dates(
+        ukb_data_file=args.ukb_data,
+        field_map_file=args.field_map,
+        icd_map_file=args.icd_map,
+        labels_file=args.labels,
+        output_file=args.output,
+        chunksize=args.chunksize,
+    )
+    print(f"Wrote {count:,} first disease/death events to {args.output}")
+
+
+if __name__ == "__main__":
+    main()
--- a/readouts.py
+++ b/readouts.py
@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class ReadoutOutput:
+    hidden: torch.Tensor
+    readout_mask: torch.Tensor
+
+
+class TokenReadout(nn.Module):
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        time_seq: torch.Tensor,
+        padding_mask: torch.Tensor,
+        readout_mask: torch.Tensor | None = None,
+    ) -> ReadoutOutput:
+        mask = padding_mask if readout_mask is None else readout_mask
+        return ReadoutOutput(hidden=hidden, readout_mask=mask.bool())
+
+
+class SameTimeGroupEndReadout(nn.Module):
+    def __init__(self, reduce: str = "mean"):
+        super().__init__()
+        if reduce not in {"mean", "sum"}:
+            raise ValueError("reduce must be either 'mean' or 'sum'")
+        self.reduce = reduce
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        time_seq: torch.Tensor,
+        padding_mask: torch.Tensor,
+        readout_mask: torch.Tensor | None = None,
+    ) -> ReadoutOutput:
+        if readout_mask is None:
+            next_is_new_time = torch.ones_like(padding_mask, dtype=torch.bool)
+            next_is_new_time[:, :-1] = time_seq[:, 1:] != time_seq[:, :-1]
+            readout_mask = padding_mask.bool() & next_is_new_time
+        else:
+            readout_mask = readout_mask.bool()
+
+        group_start = torch.ones_like(padding_mask, dtype=torch.bool)
+        group_start[:, 1:] = time_seq[:, 1:] != time_seq[:, :-1]
+        group_start = group_start & padding_mask.bool()
+
+        group_id = group_start.long().cumsum(dim=1) - 1
+        group_id = group_id.clamp_min(0)
+        max_groups = hidden.size(1)
+
+        group_sum = hidden.new_zeros(hidden.size(0), max_groups, hidden.size(2))
+        group_sum.scatter_add_(
+            1,
+            group_id.unsqueeze(-1).expand_as(hidden),
+            hidden * padding_mask.unsqueeze(-1).to(hidden.dtype),
+        )
+
+        if self.reduce == "mean":
+            group_count = hidden.new_zeros(hidden.size(0), max_groups, 1)
+            group_count.scatter_add_(
+                1,
+                group_id.unsqueeze(-1),
+                padding_mask.unsqueeze(-1).to(hidden.dtype),
+            )
+            group_sum = group_sum / group_count.clamp_min(1.0)
+
+        out = hidden.clone()
+        out[readout_mask] = group_sum.gather(
+            1,
+            group_id.unsqueeze(-1).expand_as(hidden),
+        )[readout_mask]
+        return ReadoutOutput(hidden=out, readout_mask=readout_mask)
+
+
+class LastValidReadout(nn.Module):
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        time_seq: torch.Tensor,
+        padding_mask: torch.Tensor,
+        readout_mask: torch.Tensor | None = None,
+    ) -> ReadoutOutput:
+        batch_size, seq_len = padding_mask.shape
+        last_idx = padding_mask.long().sum(dim=1).clamp_min(1) - 1
+        out = hidden[torch.arange(batch_size, device=hidden.device), last_idx]
+        mask = torch.ones(batch_size, dtype=torch.bool, device=hidden.device)
+        return ReadoutOutput(hidden=out, readout_mask=mask)
+
+
+def build_readout(name: str, **kwargs) -> nn.Module:
+    name = name.lower()
+    if name == "token":
+        return TokenReadout()
+    if name in {"same_time_group_end", "same_time"}:
+        return SameTimeGroupEndReadout(**kwargs)
+    if name == "last_valid":
+        return LastValidReadout()
+    raise ValueError(
+        "Unknown readout {!r}. Available: token, same_time_group_end, last_valid.".format(
+            name
+        )
+    )
--- a/requirements.txt
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+numpy
+pandas
+torch
+tqdm
+scikit-learn
--- a/targets.py
+++ b/targets.py
@@ -0,0 +1,394 @@
+# targets.py
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+import numpy as np
+
+
+PAD_IDX = 0
+CHECKUP_IDX = 1
+NO_EVENT_IDX = 2
+DAYS_PER_YEAR = 365.25
+
+
+@dataclass(frozen=True)
+class NextTokenTargets:
+    """
+    Delphi2M-style next-token supervision targets.
+
+    Shapes:
+        input_events:      (L,)
+        input_times_years: (L,)
+        target_events:     (L,)
+        target_times_years:(L,)
+
+    where L = N - 1.
+    """
+    input_events: np.ndarray
+    input_times_years: np.ndarray
+    target_events: np.ndarray
+    target_times_years: np.ndarray
+
+
+@dataclass(frozen=True)
+class UniqueTimeSetTargets:
+    """
+    Unique-time set supervision targets.
+
+    Shapes:
+        readout_mask:      (L,)
+        target_dt_unique:  (L,)
+        target_multi_hot:  (L, vocab_size)
+
+    where L = N - 1.
+
+    Only group-end positions can have readout_mask=True.
+    target_dt_unique is measured in years.
+    """
+    readout_mask: np.ndarray
+    target_dt_unique: np.ndarray
+    target_multi_hot: np.ndarray
+
+
+@dataclass(frozen=True)
+class TargetPack:
+    """
+    Combined target package for one patient sequence.
+
+    Contains both next-token targets and unique-time-set targets.
+    The training pipeline decides which one to use.
+    """
+    next_token: NextTokenTargets
+    unique_time_set: UniqueTimeSetTargets
+
+
+def _as_numpy_1d(
+    x: np.ndarray,
+    name: str,
+    dtype: np.dtype | type | None = None,
+) -> np.ndarray:
+    arr = np.asarray(x)
+    if arr.ndim != 1:
+        raise ValueError(f"{name} must be 1D, got shape {arr.shape}")
+    if dtype is not None:
+        arr = arr.astype(dtype)
+    return arr
+
+
+def validate_event_sequence(
+    labels: np.ndarray,
+    times_days: np.ndarray,
+    *,
+    require_sorted: bool = True,
+) -> None:
+    """
+    Validate one patient's event sequence.
+
+    labels:
+        1D integer label ids.
+
+    times_days:
+        1D event times in days.
+
+    require_sorted:
+        If True, raises when times_days is not non-decreasing.
+    """
+    labels = _as_numpy_1d(labels, "labels")
+    times_days = _as_numpy_1d(times_days, "times_days")
+
+    if len(labels) != len(times_days):
+        raise ValueError(
+            f"labels and times_days must have the same length, "
+            f"got {len(labels)} and {len(times_days)}"
+        )
+
+    if len(labels) == 0:
+        raise ValueError("Empty event sequence is not valid.")
+
+    if np.any(labels < 0):
+        raise ValueError("labels contains negative ids.")
+
+    if not np.all(np.isfinite(times_days)):
+        raise ValueError("times_days contains non-finite values.")
+
+    if require_sorted and len(times_days) > 1:
+        if np.any(np.diff(times_days) < 0):
+            raise ValueError("times_days must be non-decreasing.")
+
+
+def build_next_token_targets(
+    labels: np.ndarray,
+    times_days: np.ndarray,
+    *,
+    require_sorted: bool = True,
+) -> NextTokenTargets:
+    """
+    Build Delphi2M-style autoregressive next-token targets.
+
+    Given full sequence:
+
+        labels:     [x0, x1, x2, ..., xN-1]
+        times_days: [t0, t1, t2, ..., tN-1]
+
+    returns:
+
+        input_events:       [x0, x1, ..., xN-2]
+        input_times_years:  [t0, t1, ..., tN-2] / 365.25
+        target_events:      [x1, x2, ..., xN-1]
+        target_times_years: [t1, t2, ..., tN-1] / 365.25
+
+    This function does not ignore PAD/CHECKUP/NO_EVENT. Ignoring belongs to
+    the loss function because different objectives may use different ignore ids.
+    """
+    labels = _as_numpy_1d(labels, "labels", np.int64)
+    times_days = _as_numpy_1d(times_days, "times_days", np.float32)
+    validate_event_sequence(labels, times_days, require_sorted=require_sorted)
+
+    if len(labels) < 2:
+        raise ValueError(
+            "Need at least two events to build next-token targets."
+        )
+
+    input_events = labels[:-1].astype(np.int64)
+    input_times_years = (times_days[:-1] / DAYS_PER_YEAR).astype(np.float32)
+
+    target_events = labels[1:].astype(np.int64)
+    target_times_years = (times_days[1:] / DAYS_PER_YEAR).astype(np.float32)
+
+    return NextTokenTargets(
+        input_events=input_events,
+        input_times_years=input_times_years,
+        target_events=target_events,
+        target_times_years=target_times_years,
+    )
+
+
+def build_unique_time_set_targets(
+    labels: np.ndarray,
+    times_days: np.ndarray,
+    *,
+    vocab_size: int,
+    ignored_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+    require_sorted: bool = True,
+) -> UniqueTimeSetTargets:
+    """
+    Build next-unique-time set targets.
+
+    This is the target construction used by your UTS / default mode.
+
+    For each input position i:
+        - only if i is the last token of its timestamp group;
+        - find the next distinct timestamp group;
+        - target is the set of valid event labels at that next timestamp.
+
+    Example:
+
+        t=49: X
+        t=50: A, B, C
+        t=51: D, E
+
+    Supervises:
+
+        X@49           -> {A, B, C}@50
+        group_end@50   -> {D, E}@51
+
+    It does NOT supervise:
+
+        A@50 -> B@50
+        B@50 -> C@50
+
+    Parameters
+    ----------
+    labels:
+        Full event sequence labels, shape (N,).
+
+    times_days:
+        Full event sequence times in days, shape (N,).
+
+    vocab_size:
+        Size of output vocabulary.
+
+    ignored_target_ids:
+        Label ids that should not enter target_multi_hot.
+        Usually:
+            no no-event:      {0, 1}
+            with no-event:    {0, 1, 2}
+        For UTS, I recommend ignoring <NO_EVENT> unless explicitly testing it
+        as an event target.
+
+    Returns
+    -------
+    UniqueTimeSetTargets
+    """
+    labels = _as_numpy_1d(labels, "labels", np.int64)
+    times_days = _as_numpy_1d(times_days, "times_days", np.float32)
+    validate_event_sequence(labels, times_days, require_sorted=require_sorted)
+
+    if vocab_size <= 0:
+        raise ValueError(f"vocab_size must be positive, got {vocab_size}")
+
+    if len(labels) < 2:
+        raise ValueError(
+            "Need at least two events to build unique-time-set targets."
+        )
+
+    input_len = len(labels) - 1
+
+    readout_mask = np.zeros(input_len, dtype=bool)
+    target_dt_unique = np.zeros(input_len, dtype=np.float32)
+    target_multi_hot = np.zeros((input_len, vocab_size), dtype=bool)
+
+    ignored = {int(x) for x in ignored_target_ids}
+
+    unique_times = np.unique(times_days)
+    time_to_group_idx = {t: i for i, t in enumerate(unique_times)}
+    group_indices = np.array([time_to_group_idx[t]
+                             for t in times_days], dtype=np.int64)
+
+    for i in range(input_len):
+        current_group = group_indices[i]
+
+        is_last_in_group = (
+            i == input_len - 1
+            or group_indices[i + 1] != current_group
+        )
+        if not is_last_in_group:
+            continue
+
+        next_group_idx = current_group + 1
+        if next_group_idx >= len(unique_times):
+            continue
+
+        next_time = unique_times[next_group_idx]
+        next_labels = labels[group_indices == next_group_idx]
+
+        valid_next_labels: list[int] = []
+        for lab in next_labels:
+            lab_int = int(lab)
+            if lab_int in ignored:
+                continue
+            if lab_int < 0 or lab_int >= vocab_size:
+                continue
+            valid_next_labels.append(lab_int)
+
+        # If next timestamp contains only technical tokens, do not supervise UTS.
+        if len(valid_next_labels) == 0:
+            continue
+
+        readout_mask[i] = True
+        target_dt_unique[i] = float(next_time - times_days[i]) / DAYS_PER_YEAR
+        target_multi_hot[i, valid_next_labels] = True
+
+    return UniqueTimeSetTargets(
+        readout_mask=readout_mask,
+        target_dt_unique=target_dt_unique.astype(np.float32),
+        target_multi_hot=target_multi_hot,
+    )
+
+
+def build_all_targets(
+    labels: np.ndarray,
+    times_days: np.ndarray,
+    *,
+    vocab_size: int,
+    ignored_uts_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
+    require_sorted: bool = True,
+) -> TargetPack:
+    """
+    Build both next-token targets and unique-time-set targets for one patient.
+
+    This is the function dataset.py should usually call during initialization.
+
+    The dataset can then store:
+        event_seq           = target_pack.next_token.input_events
+        time_seq            = target_pack.next_token.input_times_years
+
+        target_event_seq    = target_pack.next_token.target_events
+        target_time_seq     = target_pack.next_token.target_times_years
+
+        readout_mask        = target_pack.unique_time_set.readout_mask
+        target_dt_unique    = target_pack.unique_time_set.target_dt_unique
+        target_multi_hot    = target_pack.unique_time_set.target_multi_hot
+    """
+    next_token = build_next_token_targets(
+        labels=labels,
+        times_days=times_days,
+        require_sorted=require_sorted,
+    )
+
+    unique_time_set = build_unique_time_set_targets(
+        labels=labels,
+        times_days=times_days,
+        vocab_size=vocab_size,
+        ignored_target_ids=ignored_uts_target_ids,
+        require_sorted=require_sorted,
+    )
+
+    return TargetPack(
+        next_token=next_token,
+        unique_time_set=unique_time_set,
+    )
+
+
+def get_group_end_mask_from_times(
+    times_days: np.ndarray,
+    *,
+    input_len: int | None = None,
+) -> np.ndarray:
+    """
+    Convenience utility for debugging.
+
+    Returns a bool mask indicating the last token of each same-time group
+    within the input sequence.
+
+    If input_len is None, uses len(times_days) - 1, matching model input length.
+    """
+    times_days = _as_numpy_1d(times_days, "times_days", np.float32)
+
+    if input_len is None:
+        input_len = len(times_days) - 1
+
+    if input_len < 0 or input_len > len(times_days):
+        raise ValueError(
+            f"Invalid input_len={input_len} for sequence length {len(times_days)}"
+        )
+
+    out = np.zeros(input_len, dtype=bool)
+
+    for i in range(input_len):
+        is_last_in_group = (
+            i == input_len - 1
+            or times_days[i + 1] != times_days[i]
+        )
+        out[i] = is_last_in_group
+
+    return out
+
+
+def summarize_targets(
+    target_pack: TargetPack,
+) -> dict[str, int | float]:
+    """
+    Small debugging helper for logging.
+    """
+    nt = target_pack.next_token
+    uts = target_pack.unique_time_set
+
+    n_tokens = int(len(nt.input_events))
+    n_readout = int(uts.readout_mask.sum())
+    n_positive_labels = int(uts.target_multi_hot.sum())
+
+    mean_set_size = (
+        float(n_positive_labels / n_readout)
+        if n_readout > 0
+        else 0.0
+    )
+
+    return {
+        "n_input_tokens": n_tokens,
+        "n_uts_readouts": n_readout,
+        "n_uts_positive_labels": n_positive_labels,
+        "mean_uts_set_size": mean_set_size,
+    }
--- a/train_next_step.py
+++ b/train_next_step.py
@@ -0,0 +1,667 @@
+"""
+Train DeepHealth with next-token / next-time-point supervision.
+
+The next-step dataset uses observed event histories, including CHECKUP state
+tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
+same-time group ends.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import time
+from pathlib import Path
+from typing import Any, Dict
+
+import numpy as np
+import torch
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, RandomSampler
+from tqdm.auto import tqdm
+
+from dataset import HealthDataset, collate_fn
+from losses import build_loss
+from models import DeepHealth, DeepHealthOutput
+from readouts import build_readout
+from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+from train_util import (
+    configure_torch_for_training,
+    create_unique_run_dir,
+    format_extra_info_types,
+    load_extra_info_types_file,
+    resolve_device,
+    save_checkpoint,
+    save_config,
+    set_optimizer_lr,
+    set_seed,
+    setup_logging,
+    split_dataset,
+    split_dataset_by_eid_files,
+)
+
+
+MODEL_INPUT_KEYS = (
+    "event_seq",
+    "time_seq",
+    "sex",
+    "padding_mask",
+    "other_type",
+    "other_value",
+    "other_value_kind",
+    "other_time",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Train DeepHealth with next-token/point supervision")
+
+    parser.add_argument("--data_prefix", type=str, default="ukb")
+    parser.add_argument("--labels_file", type=str, default="labels.csv")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--extra_info_types_file", type=str, default=None)
+    parser.add_argument("--no_event_interval_years", type=float, default=5.0)
+    parser.add_argument("--include_no_event_in_uts_target", action="store_true")
+
+    parser.add_argument("--train_ratio", type=float, default=0.7)
+    parser.add_argument("--val_ratio", type=float, default=0.15)
+    parser.add_argument("--test_ratio", type=float, default=0.15)
+    parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
+    parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
+    parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
+
+    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("--n_tab_layer", type=int, default=4)
+    parser.add_argument("--n_bins", type=int, default=16)
+    parser.add_argument("--extra_pool_reduce", type=str, default="mean",
+                        choices=["mean", "sum"])
+    parser.add_argument("--time_mode", type=str, default="relative",
+                        choices=["relative", "absolute"])
+    parser.add_argument("--dropout", type=float, default=0.0)
+
+    parser.add_argument("--target_mode", type=str, default="uts",
+                        choices=["delphi2m", "uts"])
+    parser.add_argument("--readout_name", type=str, default=None,
+                        choices=["token", "same_time_group_end", "last_valid"])
+    parser.add_argument("--readout_reduce", type=str, default="mean",
+                        choices=["mean", "sum"])
+    parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
+    parser.add_argument("--max_exp_input", type=float, default=60.0)
+    parser.add_argument("--ce_weight", type=float, default=1.0)
+    parser.add_argument("--time_weight", type=float, default=1.0)
+    parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true")
+
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--base_lr", type=float, default=3e-4)
+    parser.add_argument("--weight_decay", type=float, default=0.1)
+    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--max_epochs", type=int, default=200)
+    parser.add_argument("--warmup_epochs", type=int, default=10)
+    parser.add_argument("--patience", type=int, default=15)
+    parser.add_argument("--min_lr_ratio", type=float, default=0.1)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--progress_interval", type=int, default=20)
+
+    args = parser.parse_args()
+    use_eid_split = all(
+        getattr(args, name)
+        for name in ("train_eid_file", "val_eid_file", "test_eid_file")
+    )
+    if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
+        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
+    if args.target_mode == "uts":
+        args.readout_name = args.readout_name or "same_time_group_end"
+        args.include_no_event_in_uts_target = True
+    else:
+        args.readout_name = args.readout_name or "token"
+    args.extra_info_types = (
+        load_extra_info_types_file(args.extra_info_types_file)
+        if args.extra_info_types_file is not None
+        else None
+    )
+    return args
+
+
+def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
+    if epoch < args.warmup_epochs:
+        return adaptive_lr * (epoch + 1) / args.warmup_epochs
+    progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
+    cosine = 0.5 * (1 + math.cos(math.pi * progress))
+    return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
+
+
+def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
+    non_blocking = device.type == "cuda"
+    return {
+        key: value.to(device, non_blocking=non_blocking)
+        if isinstance(value, torch.Tensor)
+        else value
+        for key, value in batch.items()
+    }
+
+
+def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
+    return DeepHealth(
+        vocab_size=dataset.vocab_size,
+        n_embd=args.n_embd,
+        n_head=args.n_head,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
+        n_types=dataset.n_types,
+        n_cont_types=dataset.n_cont_types,
+        n_categories=dataset.n_categories,
+        cont_type_ids=dataset.cont_type_ids,
+        n_bins=args.n_bins,
+        extra_pool_reduce=args.extra_pool_reduce,
+        target_mode="next_token",
+        time_mode=args.time_mode,
+        dist_mode="exponential",
+        dropout=args.dropout,
+    )
+
+
+def build_next_step_readout(args: argparse.Namespace):
+    if args.readout_name == "same_time_group_end":
+        return build_readout("same_time_group_end", reduce=args.readout_reduce)
+    return build_readout(args.readout_name)
+
+
+def build_next_step_loss(args: argparse.Namespace):
+    if args.target_mode == "delphi2m":
+        ignored_tokens = {PAD_IDX, CHECKUP_IDX}
+        if args.ignore_no_event_in_delphi2m:
+            ignored_tokens.add(NO_EVENT_IDX)
+        return build_loss(
+            "delphi2m",
+            ignored_tokens=ignored_tokens,
+            t_min=args.t_min,
+            max_exp_input=args.max_exp_input,
+            ce_weight=args.ce_weight,
+            time_weight=args.time_weight,
+        )
+    return build_loss(
+        "uts",
+        ignored_idx={PAD_IDX, CHECKUP_IDX},
+        t_min=args.t_min,
+        max_exp_input=args.max_exp_input,
+    )
+
+
+def build_augmented_next_step_targets(
+    batch_cpu: Dict[str, torch.Tensor],
+    model_out: DeepHealthOutput,
+    include_uts_targets: bool,
+) -> Dict[str, torch.Tensor]:
+    hidden_len = model_out.hidden.size(1)
+    event_len = int(model_out.event_len)
+    extra_len = hidden_len - event_len
+    device = model_out.hidden.device
+    non_blocking = device.type == "cuda"
+    if extra_len <= 0:
+        targets = {
+            "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
+            "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
+            "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
+        }
+        if include_uts_targets:
+            targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
+                device, non_blocking=non_blocking
+            )
+            targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
+                device, non_blocking=non_blocking
+            )
+        return targets
+
+    bsz = batch_cpu["target_event_seq"].size(0)
+    vocab_size = (
+        batch_cpu["target_multi_hot"].size(2)
+        if include_uts_targets
+        else None
+    )
+    other_valid = batch_cpu["other_type"] > 0
+    extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
+    extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
+    for b in range(bsz):
+        unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
+        n_time = min(int(unique_time.numel()), extra_len)
+        if n_time > 0:
+            extra_time[b, :n_time] = unique_time[:n_time]
+            extra_mask[b, :n_time] = True
+
+    target_event_seq = torch.cat(
+        [
+            batch_cpu["target_event_seq"],
+            torch.full(
+                (bsz, extra_len),
+                PAD_IDX,
+                dtype=batch_cpu["target_event_seq"].dtype,
+            ),
+        ],
+        dim=1,
+    )
+    target_time_seq = torch.cat(
+        [
+            batch_cpu["target_time_seq"],
+            torch.zeros(
+                bsz,
+                extra_len,
+                dtype=batch_cpu["target_time_seq"].dtype,
+            ),
+        ],
+        dim=1,
+    )
+    readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
+    target_dt_unique = None
+    target_multi_hot = None
+    if include_uts_targets:
+        target_dt_unique = torch.cat(
+            [
+                batch_cpu["target_dt_unique"],
+                torch.zeros(
+                    bsz,
+                    extra_len,
+                    dtype=batch_cpu["target_dt_unique"].dtype,
+                ),
+            ],
+            dim=1,
+        )
+        target_multi_hot = torch.cat(
+            [
+                batch_cpu["target_multi_hot"],
+                torch.zeros(
+                    bsz,
+                    extra_len,
+                    vocab_size,
+                    dtype=batch_cpu["target_multi_hot"].dtype,
+                ),
+            ],
+            dim=1,
+        )
+
+    for b in range(bsz):
+        valid_event = batch_cpu["padding_mask"][b].bool()
+        if not valid_event.any():
+            continue
+        n_event = int(valid_event.sum().item())
+        events = torch.cat(
+            [
+                batch_cpu["event_seq"][b, :n_event],
+                batch_cpu["target_event_seq"][b, n_event - 1:n_event],
+            ]
+        )
+        times = torch.cat(
+            [
+                batch_cpu["time_seq"][b, :n_event],
+                batch_cpu["target_time_seq"][b, n_event - 1:n_event],
+            ]
+        )
+        valid_full = events > PAD_IDX
+        events = events[valid_full]
+        times = times[valid_full]
+        if events.numel() == 0:
+            continue
+
+        for j in range(extra_len):
+            if not bool(extra_mask[b, j]):
+                continue
+            pos = event_len + j
+            t = extra_time[b, j]
+            future = times > t
+            if not future.any():
+                readout_mask[b, pos] = False
+                continue
+
+            first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
+            next_time = times[first_idx]
+            next_event = events[first_idx]
+            target_event_seq[b, pos] = next_event
+            target_time_seq[b, pos] = next_time
+
+            if not include_uts_targets:
+                continue
+
+            same_next_time = times == next_time
+            next_events = events[same_next_time]
+            valid_next_events = next_events[
+                (next_events > PAD_IDX) & (next_events < vocab_size)
+            ].long()
+            if valid_next_events.numel() == 0:
+                readout_mask[b, pos] = False
+                continue
+            target_multi_hot[b, pos, valid_next_events] = True
+            target_dt_unique[b, pos] = next_time - t
+
+    targets = {
+        "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
+        "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
+        "readout_mask": readout_mask.to(device, non_blocking=non_blocking),
+    }
+    if include_uts_targets:
+        targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
+        targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
+    return targets
+
+
+def compute_next_step_loss(
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout,
+    criterion,
+    batch: Dict[str, torch.Tensor],
+    device: torch.device,
+) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    batch_cpu = batch
+    batch = move_batch_to_device(
+        {key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
+        device,
+    )
+    model_out = model(
+        event_seq=batch["event_seq"],
+        time_seq=batch["time_seq"],
+        sex=batch["sex"],
+        padding_mask=batch["padding_mask"],
+        other_type=batch["other_type"],
+        other_value=batch["other_value"],
+        other_value_kind=batch["other_value_kind"],
+        other_time=batch["other_time"],
+        target_mode="next_token",
+        return_output=True,
+    )
+    if not isinstance(model_out, DeepHealthOutput):
+        raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
+    targets = build_augmented_next_step_targets(
+        batch_cpu=batch_cpu,
+        model_out=model_out,
+        include_uts_targets=args.target_mode == "uts",
+    )
+    readout_out = readout(
+        hidden=model_out.hidden,
+        time_seq=model_out.time_seq,
+        padding_mask=model_out.padding_mask,
+        readout_mask=targets["readout_mask"]
+        if args.readout_name == "same_time_group_end"
+        else None,
+    )
+    logits = model.calc_risk(readout_out.hidden)
+
+    if args.target_mode == "delphi2m":
+        loss, parts = criterion(
+            logits=logits,
+            target_events=targets["target_event_seq"],
+            target_times=targets["target_time_seq"],
+            current_times=model_out.time_seq,
+            padding_mask=readout_out.readout_mask,
+            return_components=True,
+        )
+    else:
+        loss, parts = criterion(
+            logits=logits,
+            target_multi_hot=targets["target_multi_hot"],
+            target_dt_unique=targets["target_dt_unique"],
+            readout_mask=readout_out.readout_mask,
+            return_components=True,
+        )
+    if not torch.isfinite(loss):
+        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
+    return loss, parts
+
+
+def run_epoch(
+    logger: logging.Logger,
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout,
+    criterion,
+    loader: DataLoader,
+    optimizer: AdamW | None,
+    device: torch.device,
+    is_train: bool,
+) -> float:
+    model.train(is_train)
+    readout.train(is_train)
+    total = torch.zeros((), device=device)
+    n_batches = 0
+    skipped = 0
+    parts_sum: Dict[str, torch.Tensor] = {}
+    desc = "train" if is_train else "val"
+    progress_interval = max(1, int(args.progress_interval))
+
+    progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
+    for batch_idx, batch in enumerate(progress):
+        try:
+            loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device)
+            if is_train:
+                if optimizer is None:
+                    raise ValueError("optimizer is required for training")
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                if args.grad_clip > 0:
+                    clip_grad_norm_(model.parameters(), args.grad_clip)
+                optimizer.step()
+
+            total = total + loss.detach()
+            n_batches += 1
+            for name, value in parts.items():
+                parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
+            if (batch_idx + 1) % progress_interval == 0:
+                avg = total / max(1, n_batches)
+                postfix = {
+                    "loss": f"{float(loss.detach().cpu()):.4f}",
+                    "avg": f"{float(avg.detach().cpu()):.4f}",
+                    "skipped": skipped,
+                }
+                for name, value in parts_sum.items():
+                    postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}"
+                progress.set_postfix(postfix)
+        except RuntimeError as exc:
+            if "Loss is not finite" not in str(exc):
+                raise
+            skipped += 1
+            logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")
+
+    if skipped:
+        logger.info(f"Skipped {skipped} batches due to non-finite loss")
+    return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
+
+
+def build_metadata(
+    args: argparse.Namespace,
+    dataset: HealthDataset,
+    run_name: str,
+    train_subset,
+    val_subset,
+    test_subset,
+) -> Dict[str, Any]:
+    return {
+        "run_name": run_name,
+        "dataset_class": "NextStepHealthDataset",
+        "collate_fn": "next_step_collate_fn",
+        "model_class": "DeepHealth",
+        "model_target_mode": "next_token",
+        "target_mode": args.target_mode,
+        "dist_mode": "exponential",
+        "extra_info_types_file": (
+            Path(args.extra_info_types_file).name
+            if args.extra_info_types_file is not None
+            else None
+        ),
+        "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        "dataset_metadata": {
+            "vocab_size": int(dataset.vocab_size),
+            "n_types": int(dataset.n_types),
+            "n_cont_types": int(dataset.n_cont_types),
+            "n_categories": int(dataset.n_categories),
+            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
+            "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        },
+        "split_sizes": {
+            "train": int(len(train_subset)),
+            "val": int(len(val_subset)),
+            "test": int(len(test_subset)),
+        },
+        "resolved_readout_name": args.readout_name,
+        "resolved_loss_name": args.target_mode,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+    device = resolve_device(args.device)
+    configure_torch_for_training(device)
+
+    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}"
+        )
+    )
+    logger = setup_logging(run_dir)
+
+    logger.info(f"Starting next-step training run: {run_name}")
+    logger.info(f"Device: {device}")
+    logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
+    logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
+
+    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,
+        extra_info_types=args.extra_info_types,
+    )
+    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(
+            dataset=dataset,
+            train_eid_file=args.train_eid_file,
+            val_eid_file=args.val_eid_file,
+            test_eid_file=args.test_eid_file,
+        )
+        logger.info(
+            "Using eid split files: "
+            f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
+        )
+    else:
+        train_subset, val_subset, test_subset = split_dataset(
+            dataset=dataset,
+            train_ratio=args.train_ratio,
+            val_ratio=args.val_ratio,
+            test_ratio=args.test_ratio,
+            seed=args.seed,
+        )
+        logger.info(
+            f"Using random ratio split: train={args.train_ratio}, "
+            f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
+        )
+    logger.info(
+        f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
+    )
+
+    train_loader = DataLoader(
+        train_subset,
+        batch_size=args.batch_size,
+        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
+        collate_fn=collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+    val_loader = DataLoader(
+        val_subset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+    test_loader = DataLoader(
+        test_subset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+
+    model = build_model(args, dataset).to(device)
+    readout = build_next_step_readout(args).to(device)
+    criterion = build_next_step_loss(args)
+    optimizer = AdamW(
+        model.parameters(),
+        lr=args.base_lr,
+        betas=tuple(args.betas),
+        weight_decay=args.weight_decay,
+    )
+    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
+
+    save_config(
+        args,
+        run_dir / "train_config.json",
+        extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
+    )
+
+    best_val = float("inf")
+    patience = 0
+    history = []
+    best_model_path = run_dir / "best_model.pt"
+    start = time.time()
+
+    for epoch in range(args.max_epochs):
+        lr = get_lr(epoch, args, adaptive_lr)
+        set_optimizer_lr(optimizer, lr)
+
+        train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True)
+        with torch.no_grad():
+            val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False)
+
+        is_best = val_loss < best_val
+        if is_best:
+            best_val = val_loss
+            patience = 0
+            save_checkpoint(model, best_model_path)
+        else:
+            patience += 1
+
+        logger.info(
+            f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
+            f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
+            f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
+            f"elapsed={time.time() - start:.1f}s"
+        )
+        history.append({
+            "epoch": epoch + 1,
+            "lr": lr,
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "best_val_loss": best_val,
+            "is_best": int(is_best),
+        })
+        if patience >= args.patience:
+            logger.info(f"Early stopping triggered at epoch {epoch + 1}")
+            break
+
+    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
+        json.dump(history, f, indent=2)
+
+    logger.info("Evaluating best model on next-step test split...")
+    model.load_state_dict(torch.load(best_model_path, map_location=device))
+    with torch.no_grad():
+        test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False)
+    logger.info(f"Test loss: {test_loss:.6f}")
+    logger.info(f"Best checkpoint: {best_model_path}")
+
+
+if __name__ == "__main__":
+    main()
--- a/train_util.py
+++ b/train_util.py
@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import json
+import logging
+import sys
+import time
+import csv
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Iterable, Tuple
+
+import numpy as np
+import torch
+from torch.optim import AdamW
+from torch.utils.data import Subset
+
+from dataset import AllFutureHealthDataset, HealthDataset
+from models import DeepHealth
+
+
+def create_unique_run_dir(name_fn, runs_root: Path = Path("runs")) -> tuple[Path, str]:
+    while True:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        run_name = name_fn(timestamp)
+        run_dir = runs_root / run_name
+        try:
+            run_dir.mkdir(parents=True, exist_ok=False)
+            return run_dir, run_name
+        except FileExistsError:
+            time.sleep(1.0)
+
+
+def setup_logging(run_dir: Path) -> logging.Logger:
+    run_dir.mkdir(parents=True, exist_ok=True)
+    logger = logging.getLogger("DeepHealth")
+    logger.setLevel(logging.INFO)
+    logger.handlers.clear()
+
+    formatter = logging.Formatter(
+        "%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(logging.INFO)
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+
+    file_handler = logging.FileHandler(run_dir / "train.log", mode="w")
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+def set_seed(seed: int) -> None:
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+
+
+def load_extra_info_types_file(path: str) -> list[int]:
+    file_path = Path(path)
+    if not file_path.is_file():
+        raise FileNotFoundError(f"extra_info_types_file not found: {path}")
+
+    text = file_path.read_text(encoding="utf-8").strip()
+    if not text:
+        return []
+
+    if text.startswith("["):
+        raw_items = json.loads(text)
+        if not isinstance(raw_items, list):
+            raise ValueError("extra_info_types_file JSON must be a list")
+    else:
+        raw_items = []
+        for line in text.splitlines():
+            line = line.split("#", 1)[0].strip()
+            if line:
+                raw_items.extend(line.replace(",", " ").replace(";", " ").split())
+
+    try:
+        return [int(x) for x in raw_items]
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid extra info type id in {path}") from exc
+
+
+def format_extra_info_types(extra_info_types: Iterable[int] | None) -> str:
+    if extra_info_types is None:
+        return "all"
+    values = [int(x) for x in extra_info_types]
+    if not values:
+        return "none"
+    return str(values)
+
+
+def load_eid_file(path: str | Path) -> set[int]:
+    file_path = Path(path)
+    if not file_path.is_file():
+        raise FileNotFoundError(f"eid split file not found: {file_path}")
+    with file_path.open(newline="", encoding="utf-8-sig") as f:
+        reader = csv.DictReader(f)
+        if reader.fieldnames is None or "eid" not in reader.fieldnames:
+            raise ValueError(
+                f"eid split file must contain an 'eid' column: {file_path}"
+            )
+        out: set[int] = set()
+        for row in reader:
+            raw = (row.get("eid") or "").strip()
+            if raw:
+                out.add(int(raw))
+    if not out:
+        raise ValueError(f"eid split file is empty: {file_path}")
+    return out
+
+
+def configure_torch_for_training(device: torch.device) -> None:
+    if device.type != "cuda":
+        return
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    if hasattr(torch, "set_float32_matmul_precision"):
+        torch.set_float32_matmul_precision("high")
+
+
+def resolve_device(device_arg: str) -> torch.device:
+    requested = device_arg.strip().lower()
+    if requested == "cpu":
+        return torch.device("cpu")
+    if requested == "cuda":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if requested.startswith("cuda:"):
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        index = int(requested.split(":", 1)[1])
+        if index < 0 or index >= torch.cuda.device_count():
+            raise ValueError(f"Requested CUDA device is out of range: {device_arg}")
+        return torch.device(f"cuda:{index}")
+    raise ValueError(f"Unsupported device: {device_arg}")
+
+
+def split_dataset(
+    dataset: HealthDataset,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[Subset, Subset, Subset]:
+    total = train_ratio + val_ratio + test_ratio
+    if not np.isclose(total, 1.0, atol=1e-6):
+        raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
+
+    indices = np.random.RandomState(seed).permutation(len(dataset))
+    n_train = int(len(dataset) * train_ratio)
+    n_val = int(len(dataset) * val_ratio)
+    return (
+        Subset(dataset, indices[:n_train]),
+        Subset(dataset, indices[n_train:n_train + n_val]),
+        Subset(dataset, indices[n_train + n_val:]),
+    )
+
+
+def split_dataset_by_eid_files(
+    dataset: HealthDataset,
+    train_eid_file: str | Path,
+    val_eid_file: str | Path,
+    test_eid_file: str | Path,
+) -> Tuple[Subset, Subset, Subset]:
+    split_sets = {
+        "train": load_eid_file(train_eid_file),
+        "val": load_eid_file(val_eid_file),
+        "test": load_eid_file(test_eid_file),
+    }
+    overlaps = (
+        split_sets["train"] & split_sets["val"],
+        split_sets["train"] & split_sets["test"],
+        split_sets["val"] & split_sets["test"],
+    )
+    if any(overlaps):
+        raise ValueError("eid split files must be disjoint")
+
+    split_indices: Dict[str, list[int]] = {"train": [], "val": [], "test": []}
+    for idx, sample in enumerate(dataset.samples):
+        eid = int(sample["eid"])
+        for split_name, eid_set in split_sets.items():
+            if eid in eid_set:
+                split_indices[split_name].append(idx)
+                break
+
+    missing = [name for name, indices in split_indices.items() if not indices]
+    if missing:
+        raise ValueError(f"Empty dataset split(s) after eid filtering: {missing}")
+
+    return (
+        Subset(dataset, np.asarray(split_indices["train"], dtype=np.int64)),
+        Subset(dataset, np.asarray(split_indices["val"], dtype=np.int64)),
+        Subset(dataset, np.asarray(split_indices["test"], dtype=np.int64)),
+    )
+
+
+def split_all_future_datasets(
+    train_dataset: AllFutureHealthDataset,
+    val_dataset: AllFutureHealthDataset,
+    test_dataset: AllFutureHealthDataset,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[Subset, Subset, Subset]:
+    total = train_ratio + val_ratio + test_ratio
+    if not np.isclose(total, 1.0, atol=1e-6):
+        raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
+
+    patient_indices = np.random.RandomState(seed).permutation(len(train_dataset.patients))
+    n_train = int(len(patient_indices) * train_ratio)
+    n_val = int(len(patient_indices) * val_ratio)
+    train_patient_idx = patient_indices[:n_train]
+    val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val])
+    test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:])
+
+    val_query_idx = [
+        i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
+        if int(pidx) in val_patient_set
+    ]
+    test_query_idx = [
+        i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
+        if int(pidx) in test_patient_set
+    ]
+    if not val_query_idx:
+        raise ValueError("All-future validation split has no valid query samples.")
+    if not test_query_idx:
+        raise ValueError("All-future test split has no valid query samples.")
+
+    return (
+        Subset(train_dataset, train_patient_idx),
+        Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)),
+        Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)),
+    )
+
+
+def split_all_future_datasets_by_eid_files(
+    train_dataset: AllFutureHealthDataset,
+    val_dataset: AllFutureHealthDataset,
+    test_dataset: AllFutureHealthDataset,
+    train_eid_file: str | Path,
+    val_eid_file: str | Path,
+    test_eid_file: str | Path,
+) -> Tuple[Subset, Subset, Subset]:
+    split_sets = {
+        "train": load_eid_file(train_eid_file),
+        "val": load_eid_file(val_eid_file),
+        "test": load_eid_file(test_eid_file),
+    }
+    overlaps = (
+        split_sets["train"] & split_sets["val"],
+        split_sets["train"] & split_sets["test"],
+        split_sets["val"] & split_sets["test"],
+    )
+    if any(overlaps):
+        raise ValueError("eid split files must be disjoint")
+
+    train_patient_idx = [
+        idx
+        for idx, patient in enumerate(train_dataset.patients)
+        if int(patient["eid"]) in split_sets["train"]
+    ]
+    val_query_idx = [
+        idx
+        for idx, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
+        if int(val_dataset.patients[int(pidx)]["eid"]) in split_sets["val"]
+    ]
+    test_query_idx = [
+        idx
+        for idx, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
+        if int(test_dataset.patients[int(pidx)]["eid"]) in split_sets["test"]
+    ]
+
+    if not train_patient_idx:
+        raise ValueError("All-future training eid split has no patients.")
+    if not val_query_idx:
+        raise ValueError("All-future validation eid split has no valid query samples.")
+    if not test_query_idx:
+        raise ValueError("All-future test eid split has no valid query samples.")
+
+    return (
+        Subset(train_dataset, np.asarray(train_patient_idx, dtype=np.int64)),
+        Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)),
+        Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)),
+    )
+
+
+def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
+    return AdamW(
+        model.parameters(),
+        lr=args.base_lr,
+        betas=tuple(args.betas),
+        weight_decay=args.weight_decay,
+    )
+
+
+def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+
+def save_checkpoint(model: DeepHealth, checkpoint_path: Path) -> None:
+    torch.save(model.state_dict(), checkpoint_path)
+
+
+def save_config(
+    args: Any,
+    config_path: Path,
+    extra: Dict[str, Any] | None = None,
+) -> None:
+    config: Dict[str, Any] = {}
+    for key, value in vars(args).items():
+        if isinstance(value, tuple):
+            config[key] = list(value)
+        elif isinstance(value, list):
+            config[key] = value
+        elif isinstance(value, (int, float, str, bool, type(None))):
+            config[key] = value
+        else:
+            config[key] = str(value)
+    if extra:
+        config.update(extra)
+    config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")