Add target construction and training script for DeepHealth model

- Implemented target construction in `targets.py` for next-token and unique-time set supervision.
- Added validation functions and utility methods for target building.
- Created a comprehensive training script in `train.py` that includes data loading, model building, optimizer setup, and training loop with early stopping and logging.
- Integrated loss functions and readout mechanisms based on target modes.
- Established dataset splitting and DataLoader configurations for training, validation, and testing.

README.md

@@ -0,0 +1,280 @@
+# DeepHealthNew
+
+这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：
+
+```text
+疾病序列 stream + 统一的额外信息 token stream
+```
+
+疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。
+
+## 数据准备
+
+运行：
+
+```bash
+python prepare_data.py
+```
+
+输入文件：
+
+- `ukb_data.csv`
+- `field_ids_enriched.csv`
+- `icd10_codes_mod.tsv`
+- `labels.csv`
+
+输出文件：
+
+- `ukb_event_data.npy`
+  - 形状为 `(N, 3)`
+  - 每行是 `(eid, days, label)`
+  - 包含疾病、死亡、checkup 事件
+
+- `ukb_basic_info.csv`
+  - index 为 `eid`
+  - 当前只保留 `sex`
+
+- `ukb_other_info.npy`
+  - 形状为 `(M, 5)`
+  - 每行是：
+    ```text
+    (eid, type, value, value_kind, time)
+    ```
+  - `type=0` 预留给 padding
+  - `value_kind=1` 表示连续变量
+  - `value_kind=2` 表示分类变量
+  - 缺失值不会生成 token
+  - 当前 UKB 额外信息只有一个时间点，所以 `time` 暂时都是 `date_of_assessment`
+
+- `cate_types.csv`
+  - 分类变量元信息
+  - 字段：
+    ```text
+    type,name,n_categories
+    ```
+  - `ukb_other_info.npy` 里的分类 value 是变量内部的局部 id；global category id 在 dataset 中根据实验选择的变量动态计算。
+
+## Dataset
+
+`dataset.py` 提供两个 dataset：
+
+- `NextStepHealthDataset`
+  - 用于 next-token / next-time-point 监督
+  - 对应 `Delphi2MLoss` 和 `UniqueTimeSetExponentialLoss`
+
+- `AllFutureHealthDataset`
+  - 用于 query-conditioned all-future 监督
+  - 对应 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
+
+为了兼容旧训练入口：
+
+```python
+HealthDataset = NextStepHealthDataset
+collate_fn = next_step_collate_fn
+```
+
+dataset 会输出：
+
+```python
+event_seq
+time_seq
+padding_mask
+sex
+other_type
+other_value
+other_value_kind
+other_time
+```
+
+其中 `other_type=0` 表示 padding，不额外传 other-token mask。
+
+可以通过 `extra_info_types` 选择纳入哪些额外信息变量：
+
+```python
+dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])
+```
+
+如果不传，则使用全部可用 other-info type。
+
+dataset 会暴露模型初始化需要的元信息：
+
+```python
+dataset.n_types
+dataset.n_cont_types
+dataset.n_categories
+dataset.cont_type_ids
+dataset.vocab_size
+```
+
+## 模型
+
+模型主体定义在 `models.py`，通用网络模块定义在 `backbones.py`。
+
+### BaselineEncoder
+
+`BaselineEncoder` 编码统一的 other-info token：
+
+```python
+other_type        # (B, K)
+other_value       # (B, K)
+other_value_kind  # (B, K)
+```
+
+它暂时不直接使用 `other_time`。时间信息保留给后续 `CrossAttention`，用于建模疾病/query 与 other-info token 的相对时间关系。
+
+连续值使用 `TokenAutoDiscretization`：
+
+```text
+type_id -> continuous type index -> soft bin embedding
+```
+
+分类值使用 dataset 动态计算后的 global category id：
+
+```text
+selected type offsets + local category id
+```
+
+### CrossAttention
+
+`CrossAttention` 让 disease-side hidden state 注意到 other-info token：
+
+```python
+h_disease  # (B, L, D)
+t_disease  # (B, L)
+h_token    # (B, K, D)
+t_token    # (B, K)
+```
+
+时间信息通过两种方式进入注意力：
+
+- `TimeRoPE`
+  - 使用 query time 和 key time 旋转 q/k
+  - 让 q-k 相似度带有时间位置信息
+
+- `GaussianRBFTimeBasis`
+  - 对 `t_disease - t_token` 做 RBF 编码
+  - 投影成每个 attention head 的时间 bias
+
+注意力是时间因果的：
+
+```text
+other_info_time <= disease_or_query_time
+```
+
+如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。
+
+### DeepHealth
+
+`DeepHealth` 的统一路径是：
+
+```text
+disease-side sequence
+-> disease temporal backbone
+-> CrossAttention 到 other-info tokens
+-> risk head
+```
+
+两种目标模式共用同一套语义：
+
+- `next_token`
+  - `h_disease` 长度为 `L`
+  - 输出 `(B, L, D)`
+
+- `all_future`
+  - 在 disease-side 序列末尾拼接一个 query token
+  - query token 的时间是 `t_query`
+  - `h_disease` 长度为 `L + 1`
+  - 输出最后一个 query hidden，形状为 `(B, D)`
+
+模型初始化示例：
+
+```python
+model = DeepHealth(
+    vocab_size=dataset.vocab_size,
+    n_embd=120,
+    n_head=10,
+    n_hist_layer=12,
+    n_tab_layer=4,
+    n_types=dataset.n_types,
+    n_cont_types=dataset.n_cont_types,
+    n_categories=dataset.n_categories,
+    cont_type_ids=dataset.cont_type_ids,
+)
+```
+
+## Loss
+
+`losses.py` 中保留：
+
+next-token 监督：
+
+- `Delphi2MLoss`
+- `UniqueTimeSetExponentialLoss`
+
+all-future / query-conditioned 监督：
+
+- `ExponentialLoss`
+- `WeibullLoss`
+- `MixedLoss`
+
+`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
+
+## 训练
+
+当前 `train.py` 是 next-step 训练入口，使用：
+
+```python
+HealthDataset = NextStepHealthDataset
+```
+
+示例：
+
+```bash
+python train.py \
+  --data_prefix ukb \
+  --labels_file labels.csv \
+  --target_mode uts \
+  --n_embd 120 \
+  --n_head 10 \
+  --n_hist_layer 12 \
+  --n_tab_layer 4
+```
+
+选择额外信息变量：
+
+```bash
+python train.py --extra_info_types 1 3 7
+```
+
+如果不传 `--extra_info_types`，默认使用全部 other-info type。
+
+## 主要文件
+
+- `prepare_data.py`
+  - UKB 原始数据到模型输入文件的 ETL
+
+- `dataset.py`
+  - next-step 和 all-future dataset
+  - 动态选择 other-info type
+  - 动态计算 categorical global id
+
+- `models.py`
+  - `DeepHealth`
+
+- `backbones.py`
+  - `TimeRoPE`
+  - `GaussianRBFTimeBasis`
+  - `TemporalAttention`
+  - `GPTBlock`
+  - `TokenAutoDiscretization`
+  - `BaselineEncoder`
+  - `CrossAttention`
+  - `AgeSinusoidalEncoding`
+
+- `losses.py`
+  - next-token 和 all-future losses
+
+- `readouts.py`
+  - token readout
+  - same-time group readout
+  - last-valid readout

backbones.py

@@ -0,0 +1,630 @@
+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: (dim // 2,) — 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 BaselineEncoder(nn.Module):
+    PAD_KIND = 0
+    CONT_KIND = 1
+    CATE_KIND = 2
+
+    def __init__(
+        self,
+        n_embd: int,
+        n_head: 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,
+        n_tab_layer: int = 2,
+        dropout: float = 0.0,
+    ):
+        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.n_embd = n_embd
+        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.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_tab_layer)
+        ])
+        self.ln = nn.LayerNorm(n_embd)
+        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 _make_attn_mask(self, mask: torch.Tensor, dtype: torch.dtype):
+        return torch.zeros(
+            mask.size(0),
+            1,
+            1,
+            mask.size(1),
+            device=mask.device,
+            dtype=dtype,
+        ).masked_fill(~mask[:, None, None, :], -1e4)
+
+    def forward(
+        self,
+        other_type: torch.LongTensor,        # (B, K), 0 = padding
+        other_value: torch.Tensor,           # (B, K), cate stores global id
+        other_value_kind: torch.LongTensor,  # (B, K), 0=PAD, 1=CONT, 2=CATE
+    ):
+        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)
+
+        f = type_emb + kind_emb + value_emb
+        f = f * other_valid.unsqueeze(-1).to(f.dtype)
+
+        if f.size(1) == 0:
+            return f, other_valid
+
+        attn_mask = self._make_attn_mask(other_valid, f.dtype)
+        for block in self.blocks:
+            f = block(f, attn_mask=attn_mask)
+            f = f * other_valid.unsqueeze(-1).to(f.dtype)
+
+        h = self.ln(f)
+        h = h * other_valid.unsqueeze(-1).to(h.dtype)
+        return h, other_valid
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        n_embd: int,
+        n_head: int,
+        dropout: float = 0.0,
+        n_rbf_bases: int = 16,
+        max_time_diff: float = 40.0,
+    ):
+        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.dropout = dropout
+        self.mask_value = -1e4
+
+        self.q_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.k_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.v_proj = nn.Linear(n_embd, n_embd, bias=False)
+        self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
+
+        self.time_rope = TimeRoPE(self.d_head)
+        self.rbf_time_basis = GaussianRBFTimeBasis(
+            n_bases=n_rbf_bases,
+            max_time_diff=max_time_diff,
+        )
+        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.ln = nn.LayerNorm(n_embd)
+        self.out_ln = nn.LayerNorm(n_embd)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.rbf_proj.weight, mean=0.0, std=0.02)
+
+    def _make_attn_mask(
+        self,
+        token_mask: torch.Tensor,  # (B, K), True = valid
+        t_disease: torch.Tensor,     # (B, L)
+        t_token: torch.Tensor,       # (B, K)
+        dtype: torch.dtype,
+    ):
+        valid_token = token_mask[:, None, :]  # (B, 1, K)
+        visible_by_time = t_token[:, None, :] <= t_disease[:, :, None]
+        valid = visible_by_time & valid_token  # (B, L, K)
+
+        has_any_valid = valid.any(dim=-1, keepdim=True)
+        safe_valid = valid.clone()
+        safe_valid[..., 0:1] = torch.where(
+            has_any_valid,
+            safe_valid[..., 0:1],
+            torch.ones_like(safe_valid[..., 0:1]),
+        )
+
+        attn_mask = torch.zeros(
+            safe_valid.shape,
+            device=safe_valid.device,
+            dtype=dtype,
+        ).masked_fill(~safe_valid, self.mask_value)
+        return attn_mask[:, None, :, :], valid
+
+    def _cross_rbf_cache(
+        self,
+        t_disease: torch.Tensor,
+        t_token: torch.Tensor,
+    ) -> torch.Tensor:
+        tau = torch.cat([t_disease, t_token], dim=1)
+        rbf_cache = self.rbf_time_basis.precompute_cache(tau)
+        n_disease = t_disease.size(1)
+        return rbf_cache[:, :n_disease, n_disease:, :]
+
+    def forward(
+        self,
+        h_disease: torch.Tensor,          # (B, L, D)
+        t_disease: torch.Tensor,          # (B, L)
+        h_token: torch.Tensor,            # (B, K, D)
+        t_token: torch.Tensor,            # (B, K)
+        token_mask: torch.Tensor,         # (B, K), True = valid
+        need_weights: bool = False,
+    ):
+        B, L, _ = h_disease.shape
+        K = h_token.size(1)
+        H, Dh = self.n_head, self.d_head
+        if K == 0:
+            empty_weights = h_disease.new_zeros(B, L, 0)
+            if need_weights:
+                return h_disease, empty_weights
+            return h_disease
+
+        attn_mask, valid = self._make_attn_mask(
+            token_mask=token_mask.to(device=h_disease.device, dtype=torch.bool),
+            t_disease=t_disease,
+            t_token=t_token,
+            dtype=h_disease.dtype,
+        )
+
+        q = self.q_proj(h_disease).reshape(B, L, H, Dh).transpose(1, 2)
+        k = self.k_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2)
+        v = self.v_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2)
+
+        q_cos, q_sin = self.time_rope.precompute_cache(t_disease)
+        k_cos, k_sin = self.time_rope.precompute_cache(t_token)
+        q = q * q_cos + TimeRoPE._rotate_half(q) * q_sin
+        k = k * k_cos + TimeRoPE._rotate_half(k) * k_sin
+
+        rbf_cache = self._cross_rbf_cache(t_disease, t_token)  # (B, L, K, R)
+        time_bias = self.rbf_proj(rbf_cache).permute(0, 3, 1, 2)
+        time_bias = self.time_bias_scale.tanh() * time_bias
+
+        attn_bias = attn_mask.to(time_bias.dtype) + time_bias
+        attn_out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_bias,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=False,
+            scale=self.scale,
+        )
+        attn_out = attn_out.transpose(1, 2).reshape(B, L, H * Dh)
+        attn_out = self.resid_drop(self.out_proj(attn_out))
+
+        has_any_valid = valid.any(dim=-1)  # (B, L)
+        attn_out = attn_out * has_any_valid.unsqueeze(-1).to(attn_out.dtype)
+
+        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        attn_scores = attn_scores + attn_bias
+        attn_weights = torch.softmax(attn_scores, dim=-1).mean(dim=1)
+        attn_weights = attn_weights * valid.to(attn_weights.dtype)
+        weight_sum = attn_weights.sum(dim=-1, keepdim=True).clamp_min(1e-12)
+        norm_weights = attn_weights / weight_sum
+        norm_weights = norm_weights * has_any_valid.unsqueeze(-1).to(
+            norm_weights.dtype
+        )
+        out = self.out_ln(h_disease + attn_out)
+        out = torch.where(has_any_valid.unsqueeze(-1), out, h_disease)
+
+        if need_weights:
+            return out, norm_weights
+        return out
+
+
+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

dataset.py

@@ -0,0 +1,723 @@
+# dataset.py
+from __future__ import annotations
+
+import hashlib
+import os
+import pickle
+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]
+
+
+def _cache_file_path(
+    data_prefix: str,
+    labels_file: str,
+    no_event_interval_years: float,
+    include_no_event_in_uts_target: bool,
+    dataset_kind: str,
+    extra_info_types: Iterable[int] | None = None,
+    split: str | None = None,
+    min_history_events: int | None = None,
+    min_future_events: int | None = None,
+) -> str:
+    event_path = f"{data_prefix}_event_data.npy"
+    basic_path = f"{data_prefix}_basic_info.csv"
+    other_path = f"{data_prefix}_other_info.npy"
+    cate_types_path = "cate_types.csv"
+    selected_types = ""
+    if extra_info_types is not None:
+        seen_types: set[int] = set()
+        selected = []
+        for raw_type in extra_info_types:
+            type_id = int(raw_type)
+            if type_id not in seen_types:
+                seen_types.add(type_id)
+                selected.append(type_id)
+        selected_types = ",".join(str(t) for t in selected)
+    signature_parts = [
+        "deephealthnew_dataset_cache_v2",
+        dataset_kind,
+        split or "",
+        event_path,
+        basic_path,
+        other_path,
+        cate_types_path,
+        selected_types,
+        labels_file,
+        f"{no_event_interval_years:.8f}",
+        str(int(include_no_event_in_uts_target)),
+        "" if min_history_events is None else str(int(min_history_events)),
+        "" if min_future_events is None else str(int(min_future_events)),
+    ]
+
+    for path in (event_path, basic_path, other_path, cate_types_path, labels_file):
+        try:
+            stat = os.stat(path)
+            signature_parts.append(f"{path}:{stat.st_mtime_ns}:{stat.st_size}")
+        except OSError:
+            signature_parts.append(f"{path}:missing")
+
+    digest = hashlib.sha1("|".join(signature_parts).encode("utf-8")).hexdigest()
+    cache_dir = os.path.dirname(event_path) or "."
+    return os.path.join(cache_dir, f"{data_prefix}_{dataset_kind}_cache_{digest}.pkl")
+
+
+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) -> 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)
+            labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw)
+            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,
+        }
+
+    @staticmethod
+    def _load_cache(cache_path: str, cache_version: int) -> Optional[Dict]:
+        try:
+            with open(cache_path, "rb") as f:
+                payload = pickle.load(f)
+        except OSError:
+            return None
+        except Exception:
+            return None
+
+        if not isinstance(payload, dict):
+            return None
+        if payload.get("_cache_version") != cache_version:
+            return None
+        state = payload.get("state")
+        if not isinstance(state, dict):
+            return None
+        return state
+
+    def _save_cache(self, cache_path: str, cache_version: int) -> None:
+        payload = {
+            "_cache_version": cache_version,
+            "state": {key: value for key, value in self.__dict__.items()},
+        }
+        try:
+            cache_dir = os.path.dirname(cache_path)
+            if cache_dir:
+                os.makedirs(cache_dir, exist_ok=True)
+            with open(cache_path, "wb") as f:
+                pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
+        except OSError:
+            return
+
+
+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 = 1
+
+    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:
+        cache_path = _cache_file_path(
+            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,
+            dataset_kind="next_step",
+            extra_info_types=extra_info_types,
+        )
+        cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
+        if cached_state is not None:
+            self.__dict__.update(cached_state)
+            return
+
+        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():
+            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,
+            })
+
+        self._save_cache(cache_path, self.CACHE_VERSION)
+
+    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 deterministic pre-event query points.
+    """
+
+    CACHE_VERSION = 1
+
+    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,
+        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}")
+
+        cache_path = _cache_file_path(
+            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,
+            dataset_kind="all_future",
+            extra_info_types=extra_info_types,
+            split=split,
+            min_history_events=min_history_events,
+            min_future_events=min_future_events,
+        )
+        cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
+        if cached_state is not None:
+            self.__dict__.update(cached_state)
+            return
+
+        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.patients: List[Dict] = []
+        self.valid_queries: List[Tuple[int, float]] = []
+
+        for eid, times_days, labels in self._iter_patient_events():
+            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"}:
+                for target_time in unique_times[1:]:
+                    t_query = float(target_time) - ONE_DAY_YEARS
+                    if self._is_valid_query(patient, t_query):
+                        self.valid_queries.append((pidx, t_query))
+
+        if split in {"valid", "test"} and not self.valid_queries:
+            raise ValueError("No valid deterministic query points were built.")
+
+        self._save_cache(cache_path, self.CACHE_VERSION)
+
+    def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
+        times = patient["times"]
+        n_hist = int((times <= t_query).sum())
+        n_future = int((times > t_query).sum())
+        return (
+            n_hist >= self.min_history_events
+            and n_future >= self.min_future_events
+            and patient["t_obs"] > t_query
+        )
+
+    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

losses.py

@@ -0,0 +1,403 @@
+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_safe = torch.nan_to_num(
+            logits,
+            nan=0.0,
+            posinf=self.max_exp_input,
+            neginf=-self.max_exp_input,
+        )
+
+        loss_ce = F.cross_entropy(
+            logits_safe.reshape(-1, vocab_size)[valid_mask.reshape(-1)],
+            target_events.reshape(-1)[valid_mask.reshape(-1)],
+            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.clone()
+            logits_for_lse[:, :, ignore_mask] = 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 - current_times, 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[valid_mask].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)}")
+
+        logits_safe = torch.nan_to_num(
+            logits,
+            nan=0.0,
+            posinf=self.max_exp_input,
+            neginf=-self.max_exp_input,
+        )
+        ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)
+
+        target_valid = target_multi_hot.float().clone()
+        target_valid[:, :, ignore_mask] = 0.0
+        num_targets = target_valid.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
+
+        log_lambda_targets = logits_safe * target_valid
+        log_lambda_targets = torch.where(
+            target_valid.bool(),
+            log_lambda_targets,
+            torch.zeros_like(log_lambda_targets),
+        )
+
+        observed_term = log_lambda_targets.sum(dim=-1)
+        penalty_scale = num_targets
+
+        logits_for_lse = logits_safe
+        if self.exclude_ignored_from_intensity:
+            logits_for_lse = logits_safe.clone()
+            logits_for_lse[:, :, ignore_mask] = float("-inf")
+
+        dt_clamped = torch.clamp(target_dt_unique, 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)[valid_mask].mean()
+
+        if return_components:
+            return total_loss, {
+                "observed": observed_loss[valid_mask].mean().detach(),
+                "penalty": penalty_loss[valid_mask].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."
+    )

models.py

@@ -0,0 +1,272 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from backbones import (
+    AgeSinusoidalEncoding,
+    BaselineEncoder,
+    CrossAttention,
+    GPTBlock,
+    GaussianRBFTimeBasis,
+    TimeRoPE,
+)
+
+
+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"
+        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'")
+        self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
+        self.gender_embedding = nn.Embedding(
+            2, n_embd)  # Assuming binary gender
+        self.token_encoder = BaselineEncoder(
+            n_embd=n_embd,
+            n_head=n_head,
+            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,
+            n_tab_layer=n_tab_layer,
+            dropout=dropout,
+        )
+        self.cross_attention = CrossAttention(
+            n_embd=n_embd,
+            n_head=n_head,
+            dropout=dropout,
+            n_rbf_bases=16,
+        )
+        self.target_mode = target_mode
+        self.time_mode = time_mode
+        self.dist_mode = dist_mode
+        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
+        if 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)
+        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 _encode_other_tokens(
+        self,
+        other_type: torch.LongTensor,
+        other_value: torch.Tensor,
+        other_value_kind: torch.LongTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.token_encoder(
+            other_type=other_type,
+            other_value=other_value,
+            other_value_kind=other_value_kind,
+        )
+
+    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,
+        **unused_kwargs,
+    ) -> torch.Tensor:
+        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 > 0
+        else:
+            padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
+        h_disease = self.token_embedding(event_seq)
+        t_disease = time_seq
+
+        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([time_seq, 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)
+
+        h_token, token_mask = self._encode_other_tokens(
+            other_type=other_type,
+            other_value=other_value,
+            other_value_kind=other_value_kind,
+        )
+        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)}"
+            )
+        token_time = other_time.to(device=h_token.device, dtype=t_disease.dtype)
+
+        h_disease = self.cross_attention(
+            h_disease=h_disease,
+            t_disease=t_disease,
+            h_token=h_token,
+            t_token=token_time,
+            token_mask=token_mask,
+        )
+        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+
+        if mode == "all_future":
+            return h_disease[:, -1, :]
+        return h_disease
+
+    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

prepare_data.R

@@ -0,0 +1,26 @@
+library(data.table)
+setDTthreads(40)
+library(readr)
+
+field_id <- read.csv("field_ids_enriched.csv", header = TRUE, stringsAsFactors = FALSE)
+uid <- unique(c("eid", field_id$field_instance))
+
+big_path   <- "/mnt/storage/shared_data/UKBB/20230518-from-zhourong/HHdata_221103_0512.csv"
+header_dt  <- fread(big_path, nrows = 0)     # Read 0 rows => only column names
+all_names  <- names(header_dt)
+keep_names <- intersect(all_names,uid)
+ukb_disease <- fread(big_path,
+                 select     = keep_names,
+                 showProgress = TRUE)
+
+big_path <- "/mnt/storage/shared_data/UKBB/20230518-from-zhourong/HH_data_220812_0512.csv"
+header_dt  <- fread(big_path, nrows = 0)     # Read 0 rows => only column names
+all_names  <- names(header_dt)
+keep_names <- intersect(all_names,uid)
+ukb_others <- fread(big_path,
+                 select     = keep_names,
+                 showProgress = TRUE)
+
+# merge disease and other data by "eid"
+ukb_data <- merge(ukb_disease, ukb_others, by = "eid", all = TRUE)
+fwrite(ukb_data, "ukb_data.csv")

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)

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
+        )
+    )

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,
+    }

train.py

@@ -0,0 +1,996 @@
+"""
+Training script for DeepHealth model.
+
+Implements the complete training pipeline:
+1. Load data and split train/val/test
+2. Build model, optimizer, readout, and loss
+3. Train with adaptive learning rate (warmup + cosine annealing)
+4. Early stopping based on validation loss
+5. Save checkpoints and metrics
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import os
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Iterable, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, RandomSampler, Subset
+from torch.optim import AdamW
+from torch.nn.utils import clip_grad_norm_
+from tqdm.auto import tqdm
+
+from dataset import HealthDataset, collate_fn
+from models import DeepHealth
+from readouts import build_readout
+from losses import build_loss
+from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
+
+
+# ---------------------------------------------------------------------------
+# Setup Logging
+# ---------------------------------------------------------------------------
+
+def setup_logging(run_dir: Path) -> logging.Logger:
+    """Configure logging to both console and file."""
+    run_dir.mkdir(parents=True, exist_ok=True)
+    log_file = run_dir / "train.log"
+
+    logger = logging.getLogger("DeepHealth")
+    logger.setLevel(logging.INFO)
+    logger.handlers.clear()
+
+    # Console handler
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(logging.INFO)
+    console_formatter = logging.Formatter(
+        "%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S"
+    )
+    console_handler.setFormatter(console_formatter)
+    logger.addHandler(console_handler)
+
+    # File handler
+    file_handler = logging.FileHandler(log_file, mode="w")
+    file_handler.setLevel(logging.INFO)
+    file_formatter = logging.Formatter(
+        "%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S"
+    )
+    file_handler.setFormatter(file_formatter)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+# ---------------------------------------------------------------------------
+# Utilities
+# ---------------------------------------------------------------------------
+
+def set_seed(seed: int) -> None:
+    """Set random seed for reproducibility."""
+    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]:
+    """
+    Load other-information type ids from a text or JSON file.
+
+    Text files may use whitespace, commas, or semicolons as separators. Lines may
+    contain comments after '#'. JSON files should contain a top-level list of ids.
+    """
+    file_path = Path(path)
+    if not file_path.exists():
+        raise FileNotFoundError(f"extra_info_types_file not found: {path}")
+    if not file_path.is_file():
+        raise ValueError(f"extra_info_types_file is not a file: {path}")
+
+    text = file_path.read_text(encoding="utf-8").strip()
+    if not text:
+        return []
+
+    if text.startswith("["):
+        try:
+            raw_items = json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise ValueError(
+                f"Invalid JSON in extra_info_types_file: {path}"
+            ) from exc
+        if not isinstance(raw_items, list):
+            raise ValueError(
+                f"extra_info_types_file JSON must be a list, got {type(raw_items).__name__}"
+            )
+    else:
+        tokens: list[str] = []
+        for line in text.splitlines():
+            line = line.split("#", 1)[0].strip()
+            if not line:
+                continue
+            tokens.extend(line.replace(",", " ").replace(";", " ").split())
+        raw_items = tokens
+
+    parsed: list[int] = []
+    for item in raw_items:
+        try:
+            type_id = int(item)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"Invalid extra info type id {item!r} in {path}; expected integers."
+            ) from exc
+        parsed.append(type_id)
+
+    return parsed
+
+
+def merge_extra_info_types(*sources: Iterable[int] | None) -> list[int] | None:
+    """Merge optional type-id lists while preserving first-seen order."""
+    merged: list[int] = []
+    seen: set[int] = set()
+    for source in sources:
+        if source is None:
+            continue
+        for raw_type in source:
+            type_id = int(raw_type)
+            if type_id in seen:
+                continue
+            seen.add(type_id)
+            merged.append(type_id)
+    return merged or None
+
+
+def configure_torch_for_training(device: torch.device) -> None:
+    """Enable backend settings that can improve training throughput on CUDA."""
+    if device.type == "cuda":
+        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:
+    """Resolve and validate requested device string."""
+    requested = device_arg.strip().lower()
+
+    if requested == "cpu":
+        return torch.device("cpu")
+
+    if requested == "cuda":
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        return torch.device("cuda")
+
+    if requested.startswith("cuda:"):
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        try:
+            index = int(requested.split(":", 1)[1])
+        except ValueError as exc:
+            raise ValueError(
+                f"Invalid CUDA device format '{device_arg}'. Use 'cuda' or 'cuda:<index>'."
+            ) from exc
+
+        n_cuda = torch.cuda.device_count()
+        if index < 0 or index >= n_cuda:
+            raise ValueError(
+                f"Requested device '{device_arg}' is out of range. "
+                f"Available CUDA devices: 0..{max(0, n_cuda - 1)}"
+            )
+        return torch.device(f"cuda:{index}")
+
+    raise ValueError(
+        f"Unsupported device '{device_arg}'. Use 'cpu', 'cuda', or 'cuda:<index>'."
+    )
+
+
+def split_dataset(
+    dataset: HealthDataset,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[Subset, Subset, Subset]:
+    """
+    Split dataset into train/val/test subsets.
+
+    Parameters
+    ----------
+    train_ratio, val_ratio, test_ratio : float
+        Ratios must sum to 1.0 (within 1e-6 tolerance).
+    seed : int
+        Random seed for splitting.
+
+    Returns
+    -------
+    train_subset, val_subset, test_subset
+    """
+    total = train_ratio + val_ratio + test_ratio
+    if not np.isclose(total, 1.0, atol=1e-6):
+        raise ValueError(
+            f"train_ratio + val_ratio + test_ratio must equal 1.0, "
+            f"got {total}"
+        )
+
+    n = len(dataset)
+    rng = np.random.RandomState(seed)
+    indices = rng.permutation(n)
+
+    n_train = int(n * train_ratio)
+    n_val = int(n * val_ratio)
+
+    train_indices = indices[:n_train]
+    val_indices = indices[n_train: n_train + n_val]
+    test_indices = indices[n_train + n_val:]
+
+    return (
+        Subset(dataset, train_indices),
+        Subset(dataset, val_indices),
+        Subset(dataset, test_indices),
+    )
+
+
+def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
+    """
+    Build DeepHealth model using metadata from dataset.
+
+    Uses unified other-information metadata computed during dataset initialization.
+    """
+    model = 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,
+        target_mode="next_token",
+        time_mode=args.time_mode,
+        dist_mode="exponential",
+        dropout=args.dropout,
+    )
+
+    return model
+
+
+def build_optimizer(args: argparse.Namespace, model: nn.Module) -> AdamW:
+    """Build AdamW optimizer."""
+    return AdamW(
+        model.parameters(),
+        lr=args.base_lr,
+        betas=args.betas,
+        weight_decay=args.weight_decay,
+    )
+
+
+def get_lr(
+    epoch: int,
+    args: argparse.Namespace,
+    adaptive_lr: float,
+    total_epochs: int,
+) -> float:
+    """
+    Calculate learning rate for the given epoch.
+
+    Warmup: linear from 0 to adaptive_lr over warmup_epochs
+    After: cosine annealing from adaptive_lr to adaptive_lr * min_lr_ratio
+    """
+    if epoch < args.warmup_epochs:
+        # Linear warmup
+        return adaptive_lr * (epoch + 1) / args.warmup_epochs
+    else:
+        # Cosine annealing
+        progress = (epoch - args.warmup_epochs) / \
+            (total_epochs - args.warmup_epochs)
+        return adaptive_lr * (args.min_lr_ratio + 0.5 * (1 + math.cos(math.pi * progress)) * (1 - args.min_lr_ratio))
+
+
+def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
+    """Set learning rate for all parameter groups."""
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+
+def move_batch_to_device(batch: Dict, device: torch.device) -> Dict:
+    """Move all tensors in batch to specified device."""
+    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 compute_loss(
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout: nn.Module,
+    criterion: nn.Module,
+    batch: Dict[str, torch.Tensor],
+    device: torch.device,
+) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    """
+    Compute loss for one batch.
+
+    Flow: forward model, apply readout, compute risk logits, compute loss.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Training configuration with target_mode and loss options.
+    model : DeepHealth
+        The model.
+    readout : nn.Module
+        Readout module.
+    criterion : nn.Module
+        Loss criterion.
+    batch : Dict
+        Batch data from DataLoader.
+    device : torch.device
+        Device to compute on.
+
+    Returns
+    -------
+    loss : torch.Tensor
+        Scalar loss tensor.
+    """
+    # Move batch to device
+    batch = move_batch_to_device(batch, device)
+
+    event_seq = batch["event_seq"]           # (B, L)
+    time_seq = batch["time_seq"]             # (B, L)
+    padding_mask = batch["padding_mask"]     # (B, L)
+    sex = batch["sex"]                       # (B,)
+
+    hidden = model(
+        event_seq=event_seq,
+        time_seq=time_seq,
+        sex=sex,
+        padding_mask=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",
+    )
+
+    # Apply readout
+    readout_mask = (
+        batch["readout_mask"]
+        if args.readout_name == "same_time_group_end"
+        else None
+    )
+    readout_out = readout(
+        hidden=hidden,
+        time_seq=time_seq,
+        padding_mask=padding_mask,
+        readout_mask=readout_mask,
+    )
+
+    # Compute risk logits
+    logits = model.calc_risk(readout_out.hidden)
+
+    # Compute loss based on target_mode
+    if args.target_mode == "delphi2m":
+        loss_out = criterion(
+            logits=logits,
+            target_events=batch["target_event_seq"],
+            target_times=batch["target_time_seq"],
+            current_times=batch["time_seq"],
+            padding_mask=readout_out.readout_mask,
+            return_components=True,
+        )
+    elif args.target_mode == "uts":
+        loss_out = criterion(
+            logits=logits,
+            target_multi_hot=batch["target_multi_hot"],
+            target_dt_unique=batch["target_dt_unique"],
+            readout_mask=readout_out.readout_mask,
+            return_components=True,
+        )
+    else:
+        raise ValueError(f"Unknown target_mode: {args.target_mode}")
+
+    loss, loss_parts = loss_out
+
+    # Check for NaN/Inf
+    if not torch.isfinite(loss):
+        raise RuntimeError(
+            f"Loss is not finite: {loss.item()}. "
+            f"batch_idx info: event_seq shape {event_seq.shape}, "
+            f"logits shape {logits.shape}, logits range [{logits.min():.4f}, {logits.max():.4f}]"
+        )
+
+    return loss, loss_parts
+
+
+def run_one_epoch(
+    logger: logging.Logger,
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout: nn.Module,
+    criterion: nn.Module,
+    train_loader: DataLoader,
+    optimizer: AdamW,
+    device: torch.device,
+    is_train: bool = True,
+) -> float:
+    """
+    Run one epoch of training or validation.
+
+    Parameters
+    ----------
+    logger : logging.Logger
+    args : argparse.Namespace
+    model : DeepHealth
+    readout : nn.Module
+    criterion : nn.Module
+    train_loader : DataLoader
+    optimizer : AdamW
+        Unused if is_train=False.
+    device : torch.device
+    is_train : bool
+        If True, perform training updates. Otherwise just evaluate.
+
+    Returns
+    -------
+    avg_loss : float
+        Average loss over the epoch.
+    """
+    if is_train:
+        model.train()
+    else:
+        model.eval()
+
+    total_loss = 0.0
+    n_batches = 0
+    n_skipped = 0
+    component_sums: Dict[str, float] = {}
+
+    epoch_desc = "train" if is_train else "val"
+    progress = tqdm(
+        train_loader,
+        desc=epoch_desc,
+        total=len(train_loader),
+        leave=False,
+        dynamic_ncols=True,
+    )
+
+    for batch_idx, batch in enumerate(progress):
+        try:
+            loss, loss_parts = compute_loss(
+                args=args,
+                model=model,
+                readout=readout,
+                criterion=criterion,
+                batch=batch,
+                device=device,
+            )
+
+            if is_train:
+                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_loss += loss.item()
+            n_batches += 1
+
+            for name, value in loss_parts.items():
+                component_sums[name] = component_sums.get(name, 0.0) + float(
+                    value.detach().item()
+                )
+
+            current_loss = loss.item()
+            avg_loss = total_loss / max(1, n_batches)
+            postfix = {
+                "loss": f"{current_loss:.4f}",
+                "avg": f"{avg_loss:.4f}",
+                "skipped": n_skipped,
+            }
+            for name, sum_value in component_sums.items():
+                postfix[name] = f"{sum_value / max(1, n_batches):.4f}"
+            progress.set_postfix(postfix)
+
+        except RuntimeError as e:
+            if "Loss is not finite" in str(e):
+                n_skipped += 1
+                logger.warning(f"Batch {batch_idx} skipped: {str(e)[:100]}")
+                progress.set_postfix(
+                    loss="nan",
+                    avg=f"{total_loss / max(1, n_batches):.4f}" if n_batches > 0 else "inf",
+                    skipped=n_skipped,
+                )
+                continue
+            else:
+                raise
+
+    if n_skipped > 0:
+        logger.info(f"Skipped {n_skipped} batches due to non-finite loss")
+
+    if component_sums and n_batches > 0:
+        component_summary = ", ".join(
+            f"{name}={sum_value / n_batches:.4f}"
+            for name, sum_value in sorted(component_sums.items())
+        )
+        logger.info(f"Epoch loss breakdown: {component_summary}")
+
+    avg_loss = total_loss / max(1, n_batches) if n_batches > 0 else float("inf")
+    return avg_loss
+
+
+def evaluate(
+    logger: logging.Logger,
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout: nn.Module,
+    criterion: nn.Module,
+    val_loader: DataLoader,
+    device: torch.device,
+) -> float:
+    """
+    Evaluate model on validation set.
+
+    Returns
+    -------
+    val_loss : float
+    """
+    with torch.no_grad():
+        val_loss = run_one_epoch(
+            logger=logger,
+            args=args,
+            model=model,
+            readout=readout,
+            criterion=criterion,
+            train_loader=val_loader,
+            optimizer=None,
+            device=device,
+            is_train=False,
+        )
+    return val_loss
+
+
+def save_checkpoint(
+    model: DeepHealth,
+    checkpoint_path: Path,
+) -> None:
+    """Save model state_dict."""
+    torch.save(model.state_dict(), checkpoint_path)
+
+
+def save_config(args: argparse.Namespace, config_path: Path) -> None:
+    """Save training config as JSON."""
+    config_dict = vars(args)
+    # Convert non-serializable types
+    config_dict = {
+        k: str(v) if not isinstance(
+            v, (int, float, str, bool, type(None))) else v
+        for k, v in config_dict.items()
+    }
+    with open(config_path, "w") as f:
+        json.dump(config_dict, f, indent=2)
+
+
+def normalize_training_config(args: argparse.Namespace) -> None:
+    """Fill in and validate training options that depend on other flags."""
+    if args.target_mode not in {"delphi2m", "uts"}:
+        raise ValueError(f"Unknown target_mode: {args.target_mode}")
+
+    # gap_5y is always enabled, so preserve NO_EVENT target behavior.
+    args.ignore_no_event_in_delphi2m = False
+    if args.target_mode == "uts":
+        args.include_no_event_in_uts_target = True
+
+
+def normalize_loss_and_distribution_config(args: argparse.Namespace) -> None:
+    """Validate and resolve loss/distribution options after auto-selection."""
+    if args.loss_name not in {"delphi2m", "uts"}:
+        raise ValueError(
+            "Unknown loss_name. Supported values: delphi2m, uts."
+        )
+
+    if args.loss_name == "delphi2m" and args.target_mode != "delphi2m":
+        raise ValueError(
+            "loss_name=delphi2m requires target_mode=delphi2m."
+        )
+
+    if args.loss_name == "uts" and args.target_mode != "uts":
+        raise ValueError(
+            "loss_name=uts requires target_mode=uts."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Main Training Loop
+# ---------------------------------------------------------------------------
+
+def main():
+    """Main training function."""
+    parser = argparse.ArgumentParser(description="DeepHealth Training")
+
+    # ---- Data & Output ----
+    parser.add_argument("--data_prefix", type=str, default="ukb",
+                        help="Prefix for data files")
+    parser.add_argument("--labels_file", type=str, default="labels.csv",
+                        help="Path to labels file")
+    parser.add_argument("--seed", type=int, default=42,
+                        help="Random seed")
+
+    # ---- Dataset ----
+    parser.add_argument("--no_event_interval_years", type=float, default=5.0,
+                        help="Interval in years for no-event insertion")
+    parser.add_argument("--include_no_event_in_uts_target", action="store_true",
+                        help="Include NO_EVENT in UTS target multi-hot")
+
+    # ---- Split ----
+    parser.add_argument("--train_ratio", type=float, default=0.7,
+                        help="Training set ratio")
+    parser.add_argument("--val_ratio", type=float, default=0.15,
+                        help="Validation set ratio")
+    parser.add_argument("--test_ratio", type=float, default=0.15,
+                        help="Test set ratio")
+
+    # ---- Model ----
+    parser.add_argument("--n_embd", type=int, default=120,
+                        help="Embedding dimension")
+    parser.add_argument("--n_head", type=int, default=10,
+                        help="Number of attention heads")
+    parser.add_argument("--n_hist_layer", type=int, default=12,
+                        help="Number of history encoder layers")
+    parser.add_argument("--n_tab_layer", type=int, default=4,
+                        help="Number of self-attention layers for other-token encoder")
+    parser.add_argument("--n_bins", type=int, default=16,
+                        help="Number of bins for continuous other-token values")
+    parser.add_argument("--time_mode", type=str, default="relative",
+                        choices=["relative", "absolute"],
+                        help="Time encoding mode for disease history")
+    parser.add_argument("--dropout", type=float, default=0.0,
+                        help="Dropout rate")
+    parser.add_argument("--extra_info_types", type=int, nargs="*", default=None,
+                        help="Optional list of other-information type ids to include")
+    parser.add_argument("--extra_info_types_file", type=str, default=None,
+                        help="Optional file containing other-information type ids to include")
+
+    # ---- Training Protocol ----
+    parser.add_argument("--target_mode", type=str, default="uts",
+                        choices=["delphi2m", "uts"],
+                        help="Target supervision mode")
+    parser.add_argument("--readout_name", type=str, default=None,
+                        help="Readout name (auto-selected if None)")
+    parser.add_argument("--readout_reduce", type=str, default="mean",
+                        choices=["mean", "sum"],
+                        help="Readout reduction for SameTimeGroupEndReadout")
+
+    # ---- Loss ----
+    parser.add_argument("--loss_name", type=str, default=None,
+                        help="Loss name (auto-selected if None): delphi2m, uts")
+    parser.add_argument("--t_min", type=float, default=0.0027378507871321013,
+                        help="Minimum time for loss (1/365.25)")
+    parser.add_argument("--max_exp_input", type=float, default=60.0,
+                        help="Max exponent input for loss")
+    parser.add_argument("--ce_weight", type=float, default=1.0,
+                        help="Cross-entropy weight in delphi2m loss")
+    parser.add_argument("--time_weight", type=float, default=1.0,
+                        help="Time loss weight in delphi2m loss")
+    parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true",
+                        help="Ignore NO_EVENT in delphi2m loss")
+
+    # ---- Optimization ----
+    parser.add_argument("--batch_size", type=int, default=128,
+                        help="Batch size")
+    parser.add_argument("--base_lr", type=float, default=3e-4,
+                        help="Base learning rate")
+    parser.add_argument("--weight_decay", type=float, default=0.1,
+                        help="Weight decay (L2 regularization)")
+    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99),
+                        help="AdamW betas")
+    parser.add_argument("--grad_clip", type=float, default=1.0,
+                        help="Gradient clipping norm")
+    parser.add_argument("--max_epochs", type=int, default=200,
+                        help="Maximum number of epochs")
+    parser.add_argument("--warmup_epochs", type=int, default=10,
+                        help="Number of warmup epochs")
+    parser.add_argument("--patience", type=int, default=15,
+                        help="Early stopping patience")
+    parser.add_argument("--min_lr_ratio", type=float, default=0.1,
+                        help="Minimum LR as ratio of adaptive_lr")
+    parser.add_argument("--num_workers", type=int, default=4,
+                        help="Number of DataLoader workers")
+    parser.add_argument("--device", type=str, default="cuda",
+                        help="Device to use for training: cpu, cuda, or cuda:<index>")
+
+    args = parser.parse_args()
+    file_extra_info_types = (
+        load_extra_info_types_file(args.extra_info_types_file)
+        if args.extra_info_types_file is not None
+        else None
+    )
+    args.extra_info_types = merge_extra_info_types(
+        args.extra_info_types,
+        file_extra_info_types,
+    )
+
+    # ---- Setup ----
+    set_seed(args.seed)
+    device = resolve_device(args.device)
+    configure_torch_for_training(device)
+
+    normalize_training_config(args)
+    runs_root = Path("runs")
+    while True:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        run_name = (
+            f"{args.time_mode}_exponential_{args.target_mode}_"
+            f"other_tokens_gap_5y_{timestamp}"
+        )
+        run_dir = runs_root / run_name
+        if not run_dir.exists():
+            break
+        time.sleep(1)
+
+    run_dir.mkdir(parents=True, exist_ok=False)
+
+    logger = setup_logging(run_dir)
+    logger.info(f"Starting training run: {run_name}")
+    logger.info(f"Device: {device}")
+    logger.info(f"extra_info_types: {args.extra_info_types or 'all'}")
+    logger.info(
+        f"Resolved no_event config: ignore_no_event_in_delphi2m={args.ignore_no_event_in_delphi2m}, "
+        f"include_no_event_in_uts_target={args.include_no_event_in_uts_target}"
+    )
+
+    # ---- Validate Arguments ----
+    total_ratio = args.train_ratio + args.val_ratio + args.test_ratio
+    if not np.isclose(total_ratio, 1.0, atol=1e-6):
+        raise ValueError(
+            f"train_ratio + val_ratio + test_ratio must equal 1.0, got {total_ratio}"
+        )
+
+    # Auto-select readout if not specified
+    if args.readout_name is None:
+        args.readout_name = (
+            "token" if args.target_mode == "delphi2m"
+            else "same_time_group_end"
+        )
+
+    # Auto-select loss if not specified
+    if args.loss_name is None:
+        args.loss_name = (
+            "delphi2m" if args.target_mode == "delphi2m"
+            else "uts"
+        )
+
+    normalize_loss_and_distribution_config(args)
+
+    logger.info(f"Auto-selected readout: {args.readout_name}")
+    logger.info(f"Auto-selected loss: {args.loss_name}")
+
+    # ---- Load Dataset ----
+    logger.info("Loading dataset...")
+    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,
+    )
+    logger.info(
+        f"Dataset loaded: {len(dataset)} samples, vocab_size={dataset.vocab_size}")
+
+    # ---- Split Dataset ----
+    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"Dataset split: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
+    )
+
+    # ---- Build DataLoaders ----
+    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,
+    )
+
+    # ---- Build Model ----
+    logger.info("Building model...")
+    model = build_model(args, dataset)
+    model.to(device)
+    n_params = sum(p.numel() for p in model.parameters())
+    logger.info(f"Model built: {n_params:,} parameters")
+
+    # ---- Build Optimizer ----
+    optimizer = build_optimizer(args, model)
+    logger.info(
+        f"Optimizer: AdamW, base_lr={args.base_lr}, weight_decay={args.weight_decay}")
+
+    # ---- Compute Adaptive LR ----
+    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
+    logger.info(
+        f"Adaptive LR: {adaptive_lr:.6f} (base_lr * sqrt(batch_size/128))")
+
+    # ---- Build Readout ----
+    if args.readout_name == "token":
+        readout = build_readout("token")
+    elif args.readout_name == "same_time_group_end":
+        readout = build_readout("same_time_group_end",
+                                reduce=args.readout_reduce)
+    elif args.readout_name == "last_valid":
+        readout = build_readout("last_valid")
+    else:
+        raise ValueError(f"Unknown readout: {args.readout_name}")
+    logger.info(f"Readout: {args.readout_name}")
+
+    # ---- Build Loss ----
+    if args.loss_name == "delphi2m":
+        ignored_tokens = {PAD_IDX, CHECKUP_IDX}
+        if args.ignore_no_event_in_delphi2m:
+            ignored_tokens.add(NO_EVENT_IDX)
+        criterion = 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,
+        )
+        logger.info(f"Loss: delphi2m, ignored_tokens={ignored_tokens}")
+    elif args.loss_name == "uts":
+        ignored_idx = {PAD_IDX, CHECKUP_IDX}
+        criterion = build_loss(
+            "uts",
+            ignored_idx=ignored_idx,
+            t_min=args.t_min,
+            max_exp_input=args.max_exp_input,
+        )
+        logger.info(f"Loss: uts, ignored_idx={ignored_idx}")
+    else:
+        raise ValueError(f"Unknown loss: {args.loss_name}")
+
+    # ---- Save Config ----
+    save_config(args, run_dir / "train_config.json")
+    logger.info(f"Config saved to {run_dir / 'train_config.json'}")
+
+    # ---- Training Loop ----
+    logger.info("Starting training...")
+    best_val_loss = float("inf")
+    patience_counter = 0
+    metrics = []
+
+    best_model_path = run_dir / "best_model.pt"
+    history_path = run_dir / "history.json"
+
+    start_time = time.time()
+
+    for epoch in range(args.max_epochs):
+        lr = get_lr(epoch, args, adaptive_lr, args.max_epochs)
+        set_optimizer_lr(optimizer, lr)
+
+        # Train
+        train_loss = run_one_epoch(
+            logger=logger,
+            args=args,
+            model=model,
+            readout=readout,
+            criterion=criterion,
+            train_loader=train_loader,
+            optimizer=optimizer,
+            device=device,
+            is_train=True,
+        )
+
+        # Validate
+        val_loss = evaluate(
+            logger=logger,
+            args=args,
+            model=model,
+            readout=readout,
+            criterion=criterion,
+            val_loader=val_loader,
+            device=device,
+        )
+
+        # Early stopping
+        is_best = False
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            patience_counter = 0
+            is_best = True
+            save_checkpoint(model, best_model_path)
+        else:
+            patience_counter += 1
+
+        elapsed = time.time() - start_time
+        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_loss:.6f} | patience={patience_counter}/{args.patience} | "
+            f"elapsed={elapsed:.1f}s"
+        )
+
+        metrics.append({
+            "epoch": epoch + 1,
+            "lr": lr,
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "best_val_loss": best_val_loss,
+            "is_best": int(is_best),
+        })
+
+        if patience_counter >= args.patience:
+            logger.info(f"Early stopping triggered at epoch {epoch+1}")
+            break
+
+    # ---- Save Training History ----
+    with open(history_path, "w") as f:
+        json.dump(metrics, f, indent=2)
+    logger.info(f"History saved to {history_path}")
+
+    # ---- Test on Best Model ----
+    logger.info("Evaluating best model on test set...")
+    best_state_dict = torch.load(best_model_path, map_location=device)
+    model.load_state_dict(best_state_dict)
+
+    test_loss = evaluate(
+        logger=logger,
+        args=args,
+        model=model,
+        readout=readout,
+        criterion=criterion,
+        val_loader=test_loader,
+        device=device,
+    )
+    logger.info(f"Test loss: {test_loss:.6f}")
+
+    total_time = time.time() - start_time
+    logger.info(f"Training completed in {total_time:.1f}s")
+    logger.info(f"Best checkpoint: {best_model_path}")
+
+
+if __name__ == "__main__":
+    main()