Add training scripts for all-future and next-step supervision with DeepHealth

- Implement `train_all_future.py` for training with query-conditioned all-future supervision.
- Implement `train_next_step.py` for training with next-token/next-time-point supervision.
- Introduce `train_util.py` for shared utility functions including logging, dataset splitting, and model checkpointing.
- Enhance argument parsing for both training scripts to accommodate new parameters.
- Update loss functions and model configurations to support the new training paradigms.

README.md

@@ -223,7 +223,7 @@ all-future / query-conditioned 监督：
 
 ## 训练
 
-当前 `train.py` 支持 next-token 和 all-future 两类训练入口：
+当前 `train_next_step.py` / `train_all_future.py` 支持 next-token 和 all-future 两类训练入口：
 
 - `--model_target_mode next_token`
   - 使用 `NextStepHealthDataset`
@@ -235,7 +235,7 @@ all-future / query-conditioned 监督：
   - 不使用 readout，直接对 query hidden 计算风险
   - `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
 
-当前 `train.py` 支持所有已有训练目标定义的组合：
+当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合：
 
 | 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout |
 | --- | --- | --- | --- |
@@ -248,7 +248,7 @@ all-future / query-conditioned 监督：
 示例：
 
 ```bash
-python train.py \
+python train_next_step.py \
   --data_prefix ukb \
   --labels_file labels.csv \
   --model_target_mode next_token \
@@ -262,7 +262,7 @@ python train.py \
 all-future 示例：
 
 ```bash
-python train.py \
+python train_next_step.py \
   --data_prefix ukb \
   --labels_file labels.csv \
   --model_target_mode all_future \
@@ -273,10 +273,10 @@ python train.py \
 选择额外信息变量：
 
 ```bash
-python train.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt
+python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt
 ```
 
-`train.py` 只接受 `--extra_info_types_file` 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 `#` 注释；如果不传 `--extra_info_types_file`，默认使用全部 other-info type。
+`train_next_step.py` / `train_all_future.py` 只接受 `--extra_info_types_file` 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 `#` 注释；如果不传 `--extra_info_types_file`，默认使用全部 other-info type。
 
 训练输出的 `train_config.json` 会记录：
 

dataset.py

@@ -93,6 +93,7 @@ def _cache_file_path(
     split: str | None = None,
     min_history_events: int | None = None,
     min_future_events: int | None = None,
+    validation_query_seed: int | None = None,
 ) -> str:
     event_path = f"{data_prefix}_event_data.npy"
     basic_path = f"{data_prefix}_basic_info.csv"
@@ -122,6 +123,7 @@ def _cache_file_path(
         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)),
+        "" if validation_query_seed is None else str(int(validation_query_seed)),
     ]
 
     for path in (event_path, basic_path, other_path, cate_types_path, labels_file):
@@ -305,7 +307,11 @@ class _ExpoBaseDataset(Dataset):
                 "other_time": (rows[:, 4].astype(np.float32) / DAYS_PER_YEAR),
             }
 
-    def _iter_patient_events(self) -> Iterable[tuple[int, np.ndarray, np.ndarray]]:
+    def _iter_patient_events(
+        self,
+        *,
+        impute_no_event_gaps: bool,
+    ) -> Iterable[tuple[int, np.ndarray, np.ndarray]]:
         unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True)
         ends = np.concatenate([starts[1:], [len(self.event_data)]])
         for eid_raw, start, end in zip(unique_eids, starts, ends):
@@ -313,7 +319,19 @@ class _ExpoBaseDataset(Dataset):
             rows = self.event_data[start:end]
             times_days_raw = rows[:, 1].astype(np.float32)
             labels_raw = rows[:, 2].astype(np.int64)
+
+            disease_mask = labels_raw != CHECKUP_IDX
+            times_days_raw = times_days_raw[disease_mask]
+            labels_raw = labels_raw[disease_mask]
+            if len(labels_raw) == 0:
+                yield eid, times_days_raw, labels_raw
+                continue
+
             labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw)
+            if not impute_no_event_gaps:
+                yield eid, times_days_raw, labels_raw
+                continue
+
             times_days, labels = _insert_gap_no_event_tokens(
                 times_days_raw,
                 labels_raw,
@@ -374,7 +392,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot
     """
 
-    CACHE_VERSION = 1
+    CACHE_VERSION = 2
 
     def __init__(
         self,
@@ -406,7 +424,9 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         )
 
         self.samples: List[Dict] = []
-        for eid, times_days, labels in self._iter_patient_events():
+        for eid, times_days, labels in self._iter_patient_events(
+            impute_no_event_gaps=True,
+        ):
             if len(labels) < 2:
                 continue
 
@@ -463,10 +483,12 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
     targets.
 
     Train samples one query time per patient at each __getitem__ call.
-    Valid/test use deterministic pre-event query points.
+    Valid/test use random-but-fixed query points. For each patient with N real
+    disease events, N - 2 query points are sampled from the eligible observed
+    time range, with at least one future event after every query.
     """
 
-    CACHE_VERSION = 1
+    CACHE_VERSION = 4
 
     def __init__(
         self,
@@ -477,6 +499,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         include_no_event_in_uts_target: bool = False,
         min_history_events: int = 1,
         min_future_events: int = 1,
+        validation_query_seed: int = 42,
         extra_info_types: Iterable[int] | None = None,
     ) -> None:
         if split not in {"train", "valid", "test"}:
@@ -492,6 +515,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             split=split,
             min_history_events=min_history_events,
             min_future_events=min_future_events,
+            validation_query_seed=validation_query_seed if split in {"valid", "test"} else None,
         )
         cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
         if cached_state is not None:
@@ -509,10 +533,17 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         self.split = split
         self.min_history_events = int(min_history_events)
         self.min_future_events = int(min_future_events)
+        self.validation_query_seed = int(validation_query_seed)
         self.patients: List[Dict] = []
         self.valid_queries: List[Tuple[int, float]] = []
+        validation_rng = None
+        if split in {"valid", "test"}:
+            split_offset = 0 if split == "valid" else 1_000_003
+            validation_rng = np.random.RandomState(self.validation_query_seed + split_offset)
 
-        for eid, times_days, labels in self._iter_patient_events():
+        for eid, times_days, labels in self._iter_patient_events(
+            impute_no_event_gaps=False,
+        ):
             times_years = (times_days / DAYS_PER_YEAR).astype(np.float32)
             unique_times = np.unique(times_years)
             if len(labels) < 2 or len(unique_times) < 2:
@@ -534,13 +565,18 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             self.patients.append(patient)
 
             if split in {"valid", "test"}:
-                for target_time in unique_times[1:]:
+                if validation_rng is None:
-                    t_query = float(target_time) - ONE_DAY_YEARS
+                    raise RuntimeError("validation_rng was not initialized")
-                    if self._is_valid_query(patient, t_query):
+                self.valid_queries.extend(
-                        self.valid_queries.append((pidx, t_query))
+                    (pidx, t_query)
+                    for t_query in self._sample_fixed_validation_queries(
+                        patient,
+                        validation_rng,
+                    )
+                )
 
         if split in {"valid", "test"} and not self.valid_queries:
-            raise ValueError("No valid deterministic query points were built.")
+            raise ValueError("No random-but-fixed validation query points were built.")
 
         self._save_cache(cache_path, self.CACHE_VERSION)
 
@@ -554,6 +590,45 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             and patient["t_obs"] > t_query
         )
 
+    def _sample_fixed_validation_queries(
+        self,
+        patient: Dict,
+        rng: np.random.RandomState,
+    ) -> List[float]:
+        times = np.asarray(patient["times"], dtype=np.float32)
+        labels = np.asarray(patient["labels"], dtype=np.int64)
+        real_event_mask = ~np.isin(
+            labels,
+            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
+        )
+        real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
+        n_real_events = int(real_times.size)
+        n_queries = max(0, n_real_events - 2)
+        if n_queries == 0:
+            return []
+
+        min_hist = int(self.min_history_events)
+        min_future = int(self.min_future_events)
+        if n_real_events < min_hist + min_future:
+            return []
+
+        left = float(real_times[min_hist - 1])
+        right_event_time = float(real_times[n_real_events - min_future])
+        right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
+        if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
+            return []
+
+        queries: List[float] = []
+        max_attempts = max(100, n_queries * 50)
+        for _ in range(max_attempts):
+            if len(queries) >= n_queries:
+                break
+            t_query = float(rng.uniform(left, float(right)))
+            if self._is_valid_query(patient, t_query):
+                queries.append(t_query)
+
+        return queries
+
     def _sample_train_query(self, patient: Dict) -> float:
         unique_times = np.unique(patient["times"])
         if len(unique_times) < 2:

eval_data.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, List
+
+import numpy as np
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from dataset import AllFutureHealthDataset, HealthDataset
+from targets import PAD_IDX
+
+
+class AllFutureSequenceEvalDataset:
+    """
+    Eval-only sequence view for all-future checkpoints.
+
+    All-future training uses pure disease histories, so token-level and landmark
+    evaluation should not reuse the next-step dataset view that contains
+    imputed <NO_EVENT> gap tokens.
+    """
+
+    def __init__(
+        self,
+        data_prefix: str,
+        labels_file: str,
+        min_history_events: int = 1,
+        min_future_events: int = 1,
+        extra_info_types: Iterable[int] | None = None,
+    ) -> None:
+        base = AllFutureHealthDataset(
+            data_prefix=data_prefix,
+            labels_file=labels_file,
+            split="train",
+            min_history_events=min_history_events,
+            min_future_events=min_future_events,
+            extra_info_types=extra_info_types,
+        )
+
+        self.base = base
+        self.label_code_to_id = base.label_code_to_id
+        self.label_id_to_code = base.label_id_to_code
+        self.vocab_size = base.vocab_size
+        self.n_types = base.n_types
+        self.n_cont_types = base.n_cont_types
+        self.n_categories = base.n_categories
+        self.cont_type_ids = base.cont_type_ids
+        self.extra_info_types = base.extra_info_types
+
+        self.samples: List[Dict[str, Any]] = []
+        for patient in base.patients:
+            labels = np.asarray(patient["labels"], dtype=np.int64)
+            times = np.asarray(patient["times"], dtype=np.float32)
+            if labels.size < 2:
+                continue
+            input_len = int(labels.size - 1)
+            self.samples.append(
+                {
+                    "eid": int(patient["eid"]),
+                    "event_seq": labels[:-1],
+                    "time_seq": times[:-1],
+                    "target_event_seq": labels[1:],
+                    "target_time_seq": times[1:],
+                    "readout_mask": np.ones(input_len, dtype=bool),
+                    "sex": int(patient["sex"]),
+                    "other_type": np.asarray(patient["other_type"], dtype=np.int64),
+                    "other_value": np.asarray(patient["other_value"], dtype=np.float32),
+                    "other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
+                    "other_time": np.asarray(patient["other_time"], dtype=np.float32),
+                }
+            )
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        s = self.samples[idx]
+        return {
+            "event_seq": torch.from_numpy(s["event_seq"]).long(),
+            "time_seq": torch.from_numpy(s["time_seq"]).float(),
+            "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
+            "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
+            "readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
+            "sex": torch.tensor(s["sex"], dtype=torch.long),
+            "other_type": torch.from_numpy(s["other_type"]).long(),
+            "other_value": torch.from_numpy(s["other_value"]).float(),
+            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
+            "other_time": torch.from_numpy(s["other_time"]).float(),
+        }
+
+
+def load_sequence_eval_dataset(
+    *,
+    model_target_mode: str,
+    data_prefix: str,
+    labels_file: str,
+    no_event_interval_years: float,
+    include_no_event_in_uts_target: bool,
+    min_history_events: int,
+    min_future_events: int,
+    extra_info_types: Iterable[int] | None,
+):
+    mode = str(model_target_mode).lower()
+    if mode == "next_token":
+        return HealthDataset(
+            data_prefix=data_prefix,
+            labels_file=labels_file,
+            no_event_interval_years=no_event_interval_years,
+            include_no_event_in_uts_target=include_no_event_in_uts_target,
+            extra_info_types=extra_info_types,
+        )
+    if mode == "all_future":
+        return AllFutureSequenceEvalDataset(
+            data_prefix=data_prefix,
+            labels_file=labels_file,
+            min_history_events=min_history_events,
+            min_future_events=min_future_events,
+            extra_info_types=extra_info_types,
+        )
+    raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
+
+
+def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    event_seq = pad_sequence(
+        [s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    time_seq = pad_sequence(
+        [s["time_seq"] for s in batch], batch_first=True, padding_value=0.0
+    )
+    target_event_seq = pad_sequence(
+        [s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    target_time_seq = pad_sequence(
+        [s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0
+    )
+    readout_mask = pad_sequence(
+        [s["readout_mask"] for s in batch], batch_first=True, padding_value=False
+    )
+    other_type = pad_sequence(
+        [s["other_type"] for s in batch], batch_first=True, padding_value=0
+    )
+    other_value = pad_sequence(
+        [s["other_value"] for s in batch], batch_first=True, padding_value=0.0
+    )
+    other_value_kind = pad_sequence(
+        [s["other_value_kind"] for s in batch], batch_first=True, padding_value=0
+    )
+    other_time = pad_sequence(
+        [s["other_time"] for s in batch], batch_first=True, padding_value=0.0
+    )
+
+    return {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "target_event_seq": target_event_seq,
+        "target_time_seq": target_time_seq,
+        "readout_mask": readout_mask,
+        "sex": torch.stack([s["sex"] for s in batch]),
+        "other_type": other_type,
+        "other_value": other_value,
+        "other_value_kind": other_value_kind,
+        "other_time": other_time,
+    }

evaluate_auc.py

@@ -18,7 +18,7 @@ Efficiency notes:
     avoiding repeated pickling of arrays for every disease.
 
 Run from the DeepHealth code directory containing dataset.py, models.py,
-readouts.py, and train.py-compatible checkpoints/configs.
+readouts.py, and train_config.json-compatible checkpoints/configs.
 """
 from __future__ import annotations
 
@@ -38,7 +38,8 @@ import torch
 from torch.utils.data import DataLoader, Subset
 from tqdm.auto import tqdm
 
-from dataset import HealthDataset, collate_fn
+from dataset import HealthDataset
+from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
 from models import DeepHealth
 from readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
@@ -1330,11 +1331,14 @@ def main() -> None:
         torch.backends.cudnn.benchmark = True
 
     print("Loading dataset...")
-    dataset = HealthDataset(
+    dataset = load_sequence_eval_dataset(
+        model_target_mode=model_target_mode,
         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,
+        min_history_events=int(cfg.get("all_future_min_history_events", 1)),
+        min_future_events=int(cfg.get("all_future_min_future_events", 1)),
         extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
     )
     validate_dataset_metadata(dataset, cfg)
@@ -1346,7 +1350,7 @@ def main() -> None:
         subset,
         batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
         shuffle=False,
-        collate_fn=collate_fn,
+        collate_fn=sequence_eval_collate_fn,
         num_workers=int(cfg_get(args, cfg, "num_workers", 4)),
         pin_memory=device.type == "cuda",
         persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0,

evaluate_auc_v2.py

@@ -28,6 +28,7 @@ from torch.utils.data import DataLoader, Dataset
 from tqdm.auto import tqdm
 
 from dataset import HealthDataset
+from eval_data import load_sequence_eval_dataset
 from models import DeepHealth
 from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -1391,11 +1392,14 @@ def main() -> None:
         labels_meta = pd.read_csv(str(labels_meta_path))
 
     print("Loading dataset...")
-    dataset = HealthDataset(
+    dataset = load_sequence_eval_dataset(
+        model_target_mode=model_target_mode,
         data_prefix=data_prefix,
         labels_file=labels_file,
         no_event_interval_years=float(no_event_interval_years),
         include_no_event_in_uts_target=bool(include_no_event_in_uts_target),
+        min_history_events=int(cfg.get("all_future_min_history_events", 1)),
+        min_future_events=int(cfg.get("all_future_min_future_events", 1)),
         extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
     )
     validate_dataset_metadata(dataset, cfg)

evaluate_doa_auc.py

@@ -0,0 +1,547 @@
+"""Evaluate disease AUC at date of assessment (DOA).
+
+Cases are patients whose first occurrence of a disease is after DOA and within
+the requested horizon. Controls are patients who never have that disease in the
+full observed record. Patients prevalent at/before DOA or incident after the
+horizon are not used for that disease-horizon AUC.
+
+The script adapts automatically to checkpoint target mode:
+  - next_token: use the DOA token position, inserting <NO_EVENT> at DOA when no
+    real disease token exists at DOA;
+  - all_future: query the model directly with t_query=DOA, allowing empty
+    disease history because other-info tokens still describe the DOA state.
+"""
+from __future__ import annotations
+
+import argparse
+import contextlib
+import json
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset
+from tqdm.auto import tqdm
+
+from dataset import _ExpoBaseDataset
+from evaluate_auc_v2 import (
+    build_metadata_for_merge,
+    build_model_from_dataset,
+    get_auc_delong_var,
+    load_checkpoint_state_dict,
+    load_json_config,
+    load_model_state,
+    parse_float_list,
+    parse_int_list,
+    project_distribution_chunk,
+    resolve_dist_mode_for_checkpoint,
+    select_disease_tokens,
+    validate_dataset_metadata,
+    _score_to_probability,
+)
+from readouts import build_readout
+from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX, DAYS_PER_YEAR
+
+
+SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
+
+
+def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any:
+    value = getattr(args, name, None)
+    if value is not None:
+        return value
+    return cfg.get(name, default)
+
+
+class DOAStatusDataset(_ExpoBaseDataset):
+    def __init__(
+        self,
+        data_prefix: str,
+        labels_file: str,
+        model_target_mode: str,
+        extra_info_types: Iterable[int] | None = None,
+    ) -> None:
+        super().__init__(
+            data_prefix=data_prefix,
+            labels_file=labels_file,
+            no_event_interval_years=5.0,
+            include_no_event_in_uts_target=False,
+            extra_info_types=extra_info_types,
+        )
+        self.model_target_mode = str(model_target_mode).lower()
+        if self.model_target_mode not in {"next_token", "all_future"}:
+            raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
+
+        self.records: List[Dict[str, Any]] = []
+        self.first_occurrence_by_token: Dict[int, Tuple[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)]])
+        first_lists: Dict[int, List[Tuple[int, float]]] = {}
+
+        for eid_raw, start, end in zip(unique_eids, starts, ends):
+            eid = int(eid_raw)
+            rows = self.event_data[start:end]
+            checkup_rows = rows[rows[:, 2].astype(np.int64) == CHECKUP_IDX]
+            if len(checkup_rows) == 0:
+                continue
+
+            features = self._split_features(eid)
+            if features is None:
+                continue
+
+            doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32)))
+            doa_years = np.float32(doa_days / DAYS_PER_YEAR)
+
+            disease_rows = rows[rows[:, 2].astype(np.int64) != CHECKUP_IDX]
+            disease_times = disease_rows[:, 1].astype(np.float32) / DAYS_PER_YEAR
+            disease_labels_raw = disease_rows[:, 2].astype(np.int64)
+            disease_labels = np.where(
+                disease_labels_raw >= NO_EVENT_IDX,
+                disease_labels_raw + 1,
+                disease_labels_raw,
+            ).astype(np.int64)
+            order = np.lexsort((disease_labels, disease_times))
+            disease_times = disease_times[order].astype(np.float32)
+            disease_labels = disease_labels[order].astype(np.int64)
+
+            patient_id = len(self.records)
+            for token in np.unique(disease_labels).tolist():
+                token = int(token)
+                if token in SPECIAL_TOKENS:
+                    continue
+                hit = np.where(disease_labels == token)[0]
+                if hit.size:
+                    first_lists.setdefault(token, []).append(
+                        (patient_id, float(disease_times[int(hit[0])]))
+                    )
+
+            hist = disease_times <= doa_years
+            hist_events = disease_labels[hist]
+            hist_times = disease_times[hist]
+
+            if self.model_target_mode == "next_token":
+                at_doa = np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6)
+                if hist_events.size == 0 or not np.any(at_doa):
+                    event_seq = np.concatenate([
+                        hist_events,
+                        np.array([NO_EVENT_IDX], dtype=np.int64),
+                    ])
+                    time_seq = np.concatenate([
+                        hist_times,
+                        np.array([doa_years], dtype=np.float32),
+                    ])
+                else:
+                    event_seq = hist_events
+                    time_seq = hist_times
+                readout_pos = int(len(event_seq) - 1)
+            else:
+                event_seq = hist_events
+                time_seq = hist_times
+                readout_pos = -1
+
+            self.records.append(
+                {
+                    "patient_id": patient_id,
+                    "eid": eid,
+                    "doa": doa_years,
+                    "event_seq": event_seq.astype(np.int64),
+                    "time_seq": time_seq.astype(np.float32),
+                    "readout_pos": readout_pos,
+                    "full_events": disease_labels,
+                    "full_times": disease_times,
+                    "sex": int(features["sex"]),
+                    "other_type": features["other_type"],
+                    "other_value": features["other_value"],
+                    "other_value_kind": features["other_value_kind"],
+                    "other_time": features["other_time"],
+                }
+            )
+
+        for token, pairs in first_lists.items():
+            self.first_occurrence_by_token[int(token)] = (
+                np.asarray([p for p, _ in pairs], dtype=np.int32),
+                np.asarray([t for _, t in pairs], dtype=np.float32),
+            )
+
+        if not self.records:
+            raise RuntimeError("No DOA records were built from checkup events.")
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        s = self.records[idx]
+        return {
+            "event_seq": torch.from_numpy(s["event_seq"]).long(),
+            "time_seq": torch.from_numpy(s["time_seq"]).float(),
+            "readout_pos": torch.tensor(s["readout_pos"], dtype=torch.long),
+            "t_query": torch.tensor(float(s["doa"]), dtype=torch.float32),
+            "patient_id": torch.tensor(s["patient_id"], dtype=torch.long),
+            "sex": torch.tensor(s["sex"], dtype=torch.long),
+            "other_type": torch.from_numpy(s["other_type"]).long(),
+            "other_value": torch.from_numpy(s["other_value"]).float(),
+            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
+            "other_time": torch.from_numpy(s["other_time"]).float(),
+        }
+
+
+def collate_doa_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    event_seq = pad_sequence(
+        [x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    time_seq = pad_sequence(
+        [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    other_type = pad_sequence(
+        [x["other_type"] for x in batch], batch_first=True, padding_value=0
+    )
+    other_value = pad_sequence(
+        [x["other_value"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    other_value_kind = pad_sequence(
+        [x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
+    )
+    other_time = pad_sequence(
+        [x["other_time"] for x in batch], batch_first=True, padding_value=0.0
+    )
+
+    readout_mask = torch.zeros_like(event_seq, dtype=torch.bool)
+    readout_pos = torch.stack([x["readout_pos"] for x in batch])
+    for i, pos in enumerate(readout_pos.tolist()):
+        if pos >= 0:
+            readout_mask[i, int(pos)] = True
+
+    return {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "readout_mask": readout_mask,
+        "readout_pos": readout_pos,
+        "t_query": torch.stack([x["t_query"] for x in batch]),
+        "patient_id": torch.stack([x["patient_id"] for x in batch]),
+        "sex": torch.stack([x["sex"] for x in batch]),
+        "other_type": other_type,
+        "other_value": other_value,
+        "other_value_kind": other_value_kind,
+        "other_time": other_time,
+    }
+
+
+@torch.inference_mode()
+def infer_doa_hidden(
+    model,
+    loader: DataLoader,
+    device: torch.device,
+    model_target_mode: str,
+    readout_name: str,
+    readout_reduce: str,
+    use_amp: bool,
+) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
+    model_target_mode = str(model_target_mode).lower()
+    readout = None
+    if model_target_mode == "next_token":
+        if readout_name == "same_time_group_end":
+            readout = build_readout("same_time_group_end", reduce=readout_reduce).to(device)
+        else:
+            readout = build_readout(readout_name).to(device)
+        readout.eval()
+
+    hidden_parts: List[np.ndarray] = []
+    patient_parts: List[np.ndarray] = []
+    sex_parts: List[np.ndarray] = []
+    autocast_enabled = bool(use_amp and device.type == "cuda")
+
+    for batch in tqdm(loader, desc="DOA inference", leave=False, dynamic_ncols=True):
+        batch_dev = {
+            k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
+            for k, v in batch.items()
+        }
+        amp_context = (
+            torch.autocast(device_type=device.type, dtype=torch.float16)
+            if autocast_enabled
+            else contextlib.nullcontext()
+        )
+        with amp_context:
+            if model_target_mode == "all_future":
+                hidden = model(
+                    event_seq=batch_dev["event_seq"],
+                    time_seq=batch_dev["time_seq"],
+                    sex=batch_dev["sex"],
+                    padding_mask=batch_dev["padding_mask"],
+                    t_query=batch_dev["t_query"],
+                    other_type=batch_dev["other_type"],
+                    other_value=batch_dev["other_value"],
+                    other_value_kind=batch_dev["other_value_kind"],
+                    other_time=batch_dev["other_time"],
+                    target_mode="all_future",
+                )
+            else:
+                hidden_raw = model(
+                    event_seq=batch_dev["event_seq"],
+                    time_seq=batch_dev["time_seq"],
+                    sex=batch_dev["sex"],
+                    padding_mask=batch_dev["padding_mask"],
+                    other_type=batch_dev["other_type"],
+                    other_value=batch_dev["other_value"],
+                    other_value_kind=batch_dev["other_value_kind"],
+                    other_time=batch_dev["other_time"],
+                    target_mode="next_token",
+                )
+                ro = readout(
+                    hidden=hidden_raw,
+                    time_seq=batch_dev["time_seq"],
+                    padding_mask=batch_dev["padding_mask"],
+                    readout_mask=batch_dev["readout_mask"],
+                )
+                if ro.hidden.dim() == 2:
+                    hidden = ro.hidden
+                else:
+                    hidden = ro.hidden[batch_dev["readout_mask"]]
+
+        hidden_parts.append(hidden.detach().float().cpu().numpy().astype(np.float32, copy=False))
+        patient_parts.append(batch["patient_id"].cpu().numpy().astype(np.int32, copy=False))
+        sex_parts.append(batch["sex"].cpu().numpy().astype(np.int8, copy=False))
+
+    return (
+        np.concatenate(hidden_parts, axis=0),
+        {
+            "patient_id": np.concatenate(patient_parts, axis=0),
+            "sex": np.concatenate(sex_parts, axis=0),
+        },
+    )
+
+
+def first_time_array(
+    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
+    token: int,
+    patient_count: int,
+) -> np.ndarray:
+    out = np.full(patient_count, np.inf, dtype=np.float32)
+    pairs = first_occurrence_by_token.get(int(token))
+    if pairs is not None:
+        p, t = pairs
+        out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
+    return out
+
+
+def evaluate_doa_auc(
+    dataset: DOAStatusDataset,
+    hidden_all: np.ndarray,
+    row_arrays: Dict[str, np.ndarray],
+    model,
+    disease_ids: Sequence[int],
+    horizons: np.ndarray,
+    dist_mode: str,
+    score_mode: str,
+    min_cases: int,
+    device: torch.device,
+    logit_batch_size: int,
+    use_amp: bool,
+) -> pd.DataFrame:
+    logits_all, rho_all = project_distribution_chunk(
+        model=model,
+        hidden_all=hidden_all,
+        disease_ids=disease_ids,
+        dist_mode=dist_mode,
+        device=device,
+        logit_batch_size=logit_batch_size,
+        use_amp=use_amp,
+    )
+    patient_ids = row_arrays["patient_id"].astype(np.int32)
+    sex = row_arrays["sex"].astype(np.int8)
+    doa = np.asarray([r["doa"] for r in dataset.records], dtype=np.float32)[patient_ids]
+    patient_count = len(dataset.records)
+    death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
+
+    rows: List[Dict[str, Any]] = []
+    for col, token in enumerate([int(x) for x in disease_ids]):
+        first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids]
+        never = np.isinf(first_time)
+        incident_after_doa = first_time > doa
+
+        for horizon in horizons.tolist():
+            horizon = float(horizon)
+            case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon))
+            control_mask = never
+            if int(case_mask.sum()) < min_cases or int(control_mask.sum()) < min_cases:
+                continue
+
+            rho_col = None if rho_all is None else rho_all[:, col]
+            scores = _score_to_probability(
+                logits=logits_all[:, col],
+                rho=rho_col,
+                score_mode=score_mode,
+                horizon=horizon,
+                dist_mode=dist_mode,
+                token=token,
+                death_idx=death_idx,
+            )
+
+            for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]:
+                if sex_value == -1:
+                    sex_mask = np.ones_like(case_mask, dtype=bool)
+                else:
+                    sex_mask = sex == sex_value
+                cm = case_mask & sex_mask
+                nm = control_mask & sex_mask
+                if int(cm.sum()) < min_cases or int(nm.sum()) < min_cases:
+                    continue
+                auc, var = get_auc_delong_var(scores[cm], scores[nm])
+                rows.append(
+                    {
+                        "token": token,
+                        "horizon": horizon,
+                        "sex": sex_name,
+                        "n_case": int(cm.sum()),
+                        "n_control": int(nm.sum()),
+                        "auc": auc,
+                        "auc_var": var,
+                        "auc_se": float(np.sqrt(max(var, 0.0))) if np.isfinite(var) else np.nan,
+                    }
+                )
+    return pd.DataFrame(rows)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate DOA fixed-horizon disease AUC")
+    parser.add_argument("--run_path", type=str, required=True)
+    parser.add_argument("--output_path", type=str, default=None)
+    parser.add_argument("--batch_size", type=int, default=None)
+    parser.add_argument("--num_workers", type=int, default=None)
+    parser.add_argument("--logit_batch_size", type=int, default=None)
+    parser.add_argument("--horizons", type=str, default=None)
+    parser.add_argument("--score_mode", type=str, choices=["risk", "eta"], default=None)
+    parser.add_argument("--filter_min_total", type=int, default=None)
+    parser.add_argument("--min_cases", type=int, default=None)
+    parser.add_argument("--labels_meta_path", type=str, default=None)
+    parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None)
+    args = parser.parse_args()
+
+    run_path = Path(args.run_path)
+    cfg = load_json_config(run_path / "train_config.json")
+    ckpt_path = run_path / "best_model.pt"
+    if not ckpt_path.exists():
+        raise FileNotFoundError(f"best_model.pt not found in {run_path}")
+
+    output_path = Path(args.output_path or run_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
+    if model_target_mode not in {"next_token", "all_future"}:
+        raise ValueError(f"Unsupported model_target_mode={model_target_mode!r}")
+
+    labels_meta_path = cfg_get(args, cfg, "labels_meta_path", None)
+    if labels_meta_path is None:
+        labels_meta_path = cfg.get("labels_meta_path", "delphi_labels_chapters_colours_icd.csv")
+    labels_meta = pd.read_csv(labels_meta_path) if labels_meta_path and Path(labels_meta_path).exists() else None
+
+    dataset = DOAStatusDataset(
+        data_prefix=cfg.get("data_prefix", "ukb"),
+        labels_file=cfg.get("labels_file", "labels.csv"),
+        model_target_mode=model_target_mode,
+        extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
+    )
+    validate_dataset_metadata(dataset, cfg)
+
+    disease_requested = parse_int_list(cfg_get(args, cfg, "diseases_of_interest", None))
+    disease_ids = select_disease_tokens(
+        dataset=dataset,
+        labels_meta=labels_meta,
+        requested_tokens=disease_requested,
+        filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
+        first_occurrence_by_token=dataset.first_occurrence_by_token,
+    )
+    if not disease_ids:
+        raise RuntimeError("No disease tokens selected after filtering.")
+
+    horizons = np.asarray(
+        parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0],
+        dtype=np.float32,
+    )
+    score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower()
+    min_cases = int(cfg_get(args, cfg, "min_cases", 2))
+
+    state_dict = load_checkpoint_state_dict(ckpt_path, map_location="cpu")
+    dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
+    cfg_model = dict(cfg)
+    cfg_model["dist_mode"] = dist_mode
+
+    device = torch.device(cfg.get("device", "cuda") if torch.cuda.is_available() else "cpu")
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
+    load_model_state(model, state_dict)
+    model.eval()
+
+    if model_target_mode == "next_token" and (
+        model.token_embedding.num_embeddings <= NO_EVENT_IDX
+        or model.risk_head.out_features <= NO_EVENT_IDX
+    ):
+        raise RuntimeError("Next-token DOA evaluation requires <NO_EVENT> in the model vocabulary.")
+
+    loader = DataLoader(
+        dataset,
+        batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
+        shuffle=False,
+        collate_fn=collate_doa_fn,
+        num_workers=int(cfg_get(args, cfg, "num_workers", 4)),
+        pin_memory=device.type == "cuda",
+        persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0,
+        prefetch_factor=2 if int(cfg_get(args, cfg, "num_workers", 4)) > 0 else None,
+    )
+
+    target_mode = cfg.get("target_mode", "uts")
+    readout_name = str(cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
+    readout_reduce = str(cfg.get("readout_reduce", "mean"))
+
+    print(f"DOA records: {len(dataset)}")
+    print(f"Model target mode: {model_target_mode}")
+    print(f"Dist mode: {dist_mode}")
+    print(f"Score mode: {score_mode}")
+    print(f"Horizons: {horizons.tolist()}")
+    print(f"Disease tokens: {len(disease_ids)}")
+
+    hidden_all, row_arrays = infer_doa_hidden(
+        model=model,
+        loader=loader,
+        device=device,
+        model_target_mode=model_target_mode,
+        readout_name=readout_name,
+        readout_reduce=readout_reduce,
+        use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
+    )
+    result = evaluate_doa_auc(
+        dataset=dataset,
+        hidden_all=hidden_all,
+        row_arrays=row_arrays,
+        model=model,
+        disease_ids=disease_ids,
+        horizons=horizons,
+        dist_mode=dist_mode,
+        score_mode=score_mode,
+        min_cases=min_cases,
+        device=device,
+        logit_batch_size=int(cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
+        use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
+    )
+    if result.empty:
+        raise RuntimeError("No DOA AUC rows produced. Check disease selection and min_cases.")
+
+    meta = build_metadata_for_merge(dataset, labels_meta)
+    result = result.merge(meta, on="token", how="left")
+    out_file = output_path / "doa_auc.csv"
+    result.to_csv(out_file, index=False)
+
+    summary = result.groupby(["token", "label_code", "horizon"], dropna=False, as_index=False).agg(
+        auc_mean=("auc", "mean"),
+        n_case=("n_case", "sum"),
+        n_control=("n_control", "sum"),
+    )
+    summary.to_csv(output_path / "doa_auc_summary.csv", index=False)
+    print(f"Wrote {out_file}")
+
+
+if __name__ == "__main__":
+    main()

evaluate_landmark_auc.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from evaluate_auc_v2 import main
+
+
+if __name__ == "__main__":
+    main()

evaluate_token_auc.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from evaluate_auc import main
+
+
+if __name__ == "__main__":
+    main()

train_all_future.py

@@ -0,0 +1,448 @@
+"""
+Train DeepHealth with query-conditioned all-future supervision.
+
+Training samples are patient-level. For each patient and each __getitem__ call,
+AllFutureHealthDataset randomly samples a query time t_query, uses events at or
+before t_query as history, and uses events after t_query as the future target set.
+
+Validation/test samples are deterministic query points built from future event
+times, then split by patient.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict
+
+import numpy as np
+import torch
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, RandomSampler
+from tqdm.auto import tqdm
+
+from dataset import AllFutureHealthDataset, all_future_collate_fn
+from losses import build_loss
+from models import DeepHealth
+from targets import CHECKUP_IDX, PAD_IDX
+from train_util import (
+    configure_torch_for_training,
+    load_extra_info_types_file,
+    resolve_device,
+    save_checkpoint,
+    save_config,
+    set_optimizer_lr,
+    set_seed,
+    setup_logging,
+    split_all_future_datasets,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Train DeepHealth with all-future supervision")
+
+    parser.add_argument("--data_prefix", type=str, default="ukb")
+    parser.add_argument("--labels_file", type=str, default="labels.csv")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--extra_info_types_file", type=str, default=None)
+
+    parser.add_argument("--train_ratio", type=float, default=0.7)
+    parser.add_argument("--val_ratio", type=float, default=0.15)
+    parser.add_argument("--test_ratio", type=float, default=0.15)
+    parser.add_argument("--min_history_events", type=int, default=1)
+    parser.add_argument("--min_future_events", type=int, default=1)
+    parser.add_argument("--validation_query_seed", type=int, default=None)
+
+    parser.add_argument("--n_embd", type=int, default=120)
+    parser.add_argument("--n_head", type=int, default=10)
+    parser.add_argument("--n_hist_layer", type=int, default=12)
+    parser.add_argument("--n_tab_layer", type=int, default=4)
+    parser.add_argument("--n_bins", type=int, default=16)
+    parser.add_argument("--time_mode", type=str, default="relative",
+                        choices=["relative", "absolute"])
+    parser.add_argument("--dist_mode", type=str, default="exponential",
+                        choices=["exponential", "weibull", "mixed"])
+    parser.add_argument("--dropout", type=float, default=0.0)
+
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--base_lr", type=float, default=3e-4)
+    parser.add_argument("--weight_decay", type=float, default=0.1)
+    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--max_epochs", type=int, default=200)
+    parser.add_argument("--warmup_epochs", type=int, default=10)
+    parser.add_argument("--patience", type=int, default=15)
+    parser.add_argument("--min_lr_ratio", type=float, default=0.1)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--device", type=str, default="cuda")
+
+    args = parser.parse_args()
+    if args.min_history_events < 1:
+        raise ValueError("min_history_events must be >= 1")
+    if args.min_future_events < 1:
+        raise ValueError("min_future_events must be >= 1")
+    if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
+        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
+    if args.validation_query_seed is None:
+        args.validation_query_seed = int(args.seed)
+    args.extra_info_types = (
+        load_extra_info_types_file(args.extra_info_types_file)
+        if args.extra_info_types_file is not None
+        else None
+    )
+    return args
+
+
+def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
+    if epoch < args.warmup_epochs:
+        return adaptive_lr * (epoch + 1) / args.warmup_epochs
+    progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
+    cosine = 0.5 * (1 + math.cos(math.pi * progress))
+    return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
+
+
+def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
+    non_blocking = device.type == "cuda"
+    return {
+        key: value.to(device, non_blocking=non_blocking)
+        if isinstance(value, torch.Tensor)
+        else value
+        for key, value in batch.items()
+    }
+
+
+def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
+    return DeepHealth(
+        vocab_size=dataset.vocab_size,
+        n_embd=args.n_embd,
+        n_head=args.n_head,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
+        n_types=dataset.n_types,
+        n_cont_types=dataset.n_cont_types,
+        n_categories=dataset.n_categories,
+        cont_type_ids=dataset.cont_type_ids,
+        n_bins=args.n_bins,
+        target_mode="all_future",
+        time_mode=args.time_mode,
+        dist_mode=args.dist_mode,
+        dropout=args.dropout,
+    )
+
+
+def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset):
+    ignored_idx = {PAD_IDX, CHECKUP_IDX}
+    if args.dist_mode == "exponential":
+        return build_loss("exponential", ignored_idx=ignored_idx)
+    if args.dist_mode == "weibull":
+        return build_loss("weibull", ignored_idx=ignored_idx)
+    if args.dist_mode == "mixed":
+        return build_loss(
+            "mixed",
+            death_idx=dataset.vocab_size - 1,
+            ignored_idx=ignored_idx,
+        )
+    raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
+
+
+def compute_all_future_loss(
+    args: argparse.Namespace,
+    model: DeepHealth,
+    criterion,
+    batch: Dict[str, torch.Tensor],
+    device: torch.device,
+) -> torch.Tensor:
+    batch = move_batch_to_device(batch, device)
+
+    hidden = model(
+        event_seq=batch["event_seq"],
+        time_seq=batch["time_seq"],
+        sex=batch["sex"],
+        padding_mask=batch["padding_mask"],
+        t_query=batch["t_query"],
+        other_type=batch["other_type"],
+        other_value=batch["other_value"],
+        other_value_kind=batch["other_value_kind"],
+        other_time=batch["other_time"],
+        target_mode="all_future",
+    )
+    logits = model.calc_risk(hidden)
+
+    if args.dist_mode == "exponential":
+        loss = criterion(
+            logits=logits,
+            targets=batch["future_targets"],
+            exposure=batch["exposure"],
+        )
+    elif args.dist_mode == "weibull":
+        loss = criterion(
+            logits=logits,
+            weibull_rho=model.calc_weibull_rho(hidden),
+            targets=batch["future_targets"],
+            dt=batch["future_dt"],
+            exposure=batch["exposure"],
+        )
+    else:
+        loss = criterion(
+            logits=logits,
+            death_rho=model.calc_death_rho(hidden),
+            targets=batch["future_targets"],
+            dt=batch["future_dt"],
+            exposure=batch["exposure"],
+        )
+
+    if not torch.isfinite(loss):
+        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
+    return loss
+
+
+def run_epoch(
+    logger: logging.Logger,
+    args: argparse.Namespace,
+    model: DeepHealth,
+    criterion,
+    loader: DataLoader,
+    optimizer: AdamW | None,
+    device: torch.device,
+    is_train: bool,
+) -> float:
+    model.train(is_train)
+    total = 0.0
+    n_batches = 0
+    skipped = 0
+    desc = "train" if is_train else "val"
+
+    progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
+    for batch_idx, batch in enumerate(progress):
+        try:
+            loss = compute_all_future_loss(args, model, criterion, batch, device)
+            if is_train:
+                if optimizer is None:
+                    raise ValueError("optimizer is required for training")
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                if args.grad_clip > 0:
+                    clip_grad_norm_(model.parameters(), args.grad_clip)
+                optimizer.step()
+
+            total += float(loss.detach().cpu())
+            n_batches += 1
+            avg = total / max(1, n_batches)
+            progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped)
+        except RuntimeError as exc:
+            if "Loss is not finite" not in str(exc):
+                raise
+            skipped += 1
+            logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")
+
+    if skipped:
+        logger.info(f"Skipped {skipped} batches due to non-finite loss")
+    return total / max(1, n_batches) if n_batches else float("inf")
+
+
+def build_metadata(
+    args: argparse.Namespace,
+    dataset: AllFutureHealthDataset,
+    run_name: str,
+    train_subset,
+    val_subset,
+    test_subset,
+) -> Dict[str, Any]:
+    return {
+        "run_name": run_name,
+        "dataset_class": "AllFutureHealthDataset",
+        "collate_fn": "all_future_collate_fn",
+        "model_class": "DeepHealth",
+        "model_target_mode": "all_future",
+        "target_mode": "all_future",
+        "dist_mode": args.dist_mode,
+        "all_future_min_history_events": int(args.min_history_events),
+        "all_future_min_future_events": int(args.min_future_events),
+        "all_future_validation_query_seed": int(args.validation_query_seed),
+        "extra_info_types_file": (
+            Path(args.extra_info_types_file).name
+            if args.extra_info_types_file is not None
+            else None
+        ),
+        "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        "dataset_metadata": {
+            "vocab_size": int(dataset.vocab_size),
+            "n_types": int(dataset.n_types),
+            "n_cont_types": int(dataset.n_cont_types),
+            "n_categories": int(dataset.n_categories),
+            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
+            "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        },
+        "split_sizes": {
+            "train": int(len(train_subset)),
+            "val": int(len(val_subset)),
+            "test": int(len(test_subset)),
+        },
+        "resolved_readout_name": "none",
+        "resolved_loss_name": args.dist_mode,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+    device = resolve_device(args.device)
+    configure_torch_for_training(device)
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    run_name = f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
+    run_dir = Path("runs") / run_name
+    run_dir.mkdir(parents=True, exist_ok=False)
+    logger = setup_logging(run_dir)
+
+    logger.info(f"Starting all-future training run: {run_name}")
+    logger.info(f"Device: {device}")
+    logger.info(f"extra_info_types: {args.extra_info_types or 'all'}")
+
+    logger.info("Loading all-future datasets...")
+    train_dataset = AllFutureHealthDataset(
+        data_prefix=args.data_prefix,
+        labels_file=args.labels_file,
+        split="train",
+        min_history_events=args.min_history_events,
+        min_future_events=args.min_future_events,
+        validation_query_seed=args.validation_query_seed,
+        extra_info_types=args.extra_info_types,
+    )
+    val_dataset = AllFutureHealthDataset(
+        data_prefix=args.data_prefix,
+        labels_file=args.labels_file,
+        split="valid",
+        min_history_events=args.min_history_events,
+        min_future_events=args.min_future_events,
+        validation_query_seed=args.validation_query_seed,
+        extra_info_types=args.extra_info_types,
+    )
+    test_dataset = AllFutureHealthDataset(
+        data_prefix=args.data_prefix,
+        labels_file=args.labels_file,
+        split="test",
+        min_history_events=args.min_history_events,
+        min_future_events=args.min_future_events,
+        validation_query_seed=args.validation_query_seed,
+        extra_info_types=args.extra_info_types,
+    )
+    train_subset, val_subset, test_subset = split_all_future_datasets(
+        train_dataset=train_dataset,
+        val_dataset=val_dataset,
+        test_dataset=test_dataset,
+        train_ratio=args.train_ratio,
+        val_ratio=args.val_ratio,
+        test_ratio=args.test_ratio,
+        seed=args.seed,
+    )
+    logger.info(
+        f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
+    )
+
+    train_loader = DataLoader(
+        train_subset,
+        batch_size=args.batch_size,
+        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
+        collate_fn=all_future_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=all_future_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=all_future_collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+
+    model = build_model(args, train_dataset).to(device)
+    optimizer = AdamW(
+        model.parameters(),
+        lr=args.base_lr,
+        betas=tuple(args.betas),
+        weight_decay=args.weight_decay,
+    )
+    criterion = build_criterion(args, train_dataset)
+    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
+
+    save_config(
+        args,
+        run_dir / "train_config.json",
+        extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset),
+    )
+
+    best_val = float("inf")
+    patience = 0
+    history = []
+    best_model_path = run_dir / "best_model.pt"
+    start = time.time()
+
+    for epoch in range(args.max_epochs):
+        lr = get_lr(epoch, args, adaptive_lr)
+        set_optimizer_lr(optimizer, lr)
+
+        train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True)
+        with torch.no_grad():
+            val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False)
+
+        is_best = val_loss < best_val
+        if is_best:
+            best_val = val_loss
+            patience = 0
+            save_checkpoint(model, best_model_path)
+        else:
+            patience += 1
+
+        logger.info(
+            f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
+            f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
+            f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
+            f"elapsed={time.time() - start:.1f}s"
+        )
+        history.append({
+            "epoch": epoch + 1,
+            "lr": lr,
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "best_val_loss": best_val,
+            "is_best": int(is_best),
+        })
+        if patience >= args.patience:
+            logger.info(f"Early stopping triggered at epoch {epoch + 1}")
+            break
+
+    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
+        json.dump(history, f, indent=2)
+
+    logger.info("Evaluating best model on all-future test queries...")
+    model.load_state_dict(torch.load(best_model_path, map_location=device))
+    with torch.no_grad():
+        test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False)
+    logger.info(f"Test loss: {test_loss:.6f}")
+    logger.info(f"Best checkpoint: {best_model_path}")
+
+
+if __name__ == "__main__":
+    main()

train_next_step.py

@@ -0,0 +1,454 @@
+"""
+Train DeepHealth with next-token / next-time-point supervision.
+
+The dataset remains the current next-step construction: pure disease events plus
+optional gap <NO_EVENT> imputation are shifted into autoregressive inputs and
+targets. UTS training reads out only same-time group ends.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict
+
+import numpy as np
+import torch
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, RandomSampler
+from tqdm.auto import tqdm
+
+from dataset import HealthDataset, collate_fn
+from losses import build_loss
+from models import DeepHealth
+from readouts import build_readout
+from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+from train_util import (
+    configure_torch_for_training,
+    load_extra_info_types_file,
+    resolve_device,
+    save_checkpoint,
+    save_config,
+    set_optimizer_lr,
+    set_seed,
+    setup_logging,
+    split_dataset,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Train DeepHealth with next-token/point supervision")
+
+    parser.add_argument("--data_prefix", type=str, default="ukb")
+    parser.add_argument("--labels_file", type=str, default="labels.csv")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--extra_info_types_file", type=str, default=None)
+    parser.add_argument("--no_event_interval_years", type=float, default=5.0)
+    parser.add_argument("--include_no_event_in_uts_target", action="store_true")
+
+    parser.add_argument("--train_ratio", type=float, default=0.7)
+    parser.add_argument("--val_ratio", type=float, default=0.15)
+    parser.add_argument("--test_ratio", type=float, default=0.15)
+
+    parser.add_argument("--n_embd", type=int, default=120)
+    parser.add_argument("--n_head", type=int, default=10)
+    parser.add_argument("--n_hist_layer", type=int, default=12)
+    parser.add_argument("--n_tab_layer", type=int, default=4)
+    parser.add_argument("--n_bins", type=int, default=16)
+    parser.add_argument("--time_mode", type=str, default="relative",
+                        choices=["relative", "absolute"])
+    parser.add_argument("--dropout", type=float, default=0.0)
+
+    parser.add_argument("--target_mode", type=str, default="uts",
+                        choices=["delphi2m", "uts"])
+    parser.add_argument("--readout_name", type=str, default=None,
+                        choices=["token", "same_time_group_end", "last_valid"])
+    parser.add_argument("--readout_reduce", type=str, default="mean",
+                        choices=["mean", "sum"])
+    parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
+    parser.add_argument("--max_exp_input", type=float, default=60.0)
+    parser.add_argument("--ce_weight", type=float, default=1.0)
+    parser.add_argument("--time_weight", type=float, default=1.0)
+    parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true")
+
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--base_lr", type=float, default=3e-4)
+    parser.add_argument("--weight_decay", type=float, default=0.1)
+    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--max_epochs", type=int, default=200)
+    parser.add_argument("--warmup_epochs", type=int, default=10)
+    parser.add_argument("--patience", type=int, default=15)
+    parser.add_argument("--min_lr_ratio", type=float, default=0.1)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--device", type=str, default="cuda")
+
+    args = parser.parse_args()
+    if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
+        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
+    if args.target_mode == "uts":
+        args.readout_name = args.readout_name or "same_time_group_end"
+        args.include_no_event_in_uts_target = True
+    else:
+        args.readout_name = args.readout_name or "token"
+    args.extra_info_types = (
+        load_extra_info_types_file(args.extra_info_types_file)
+        if args.extra_info_types_file is not None
+        else None
+    )
+    return args
+
+
+def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
+    if epoch < args.warmup_epochs:
+        return adaptive_lr * (epoch + 1) / args.warmup_epochs
+    progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
+    cosine = 0.5 * (1 + math.cos(math.pi * progress))
+    return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
+
+
+def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
+    non_blocking = device.type == "cuda"
+    return {
+        key: value.to(device, non_blocking=non_blocking)
+        if isinstance(value, torch.Tensor)
+        else value
+        for key, value in batch.items()
+    }
+
+
+def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
+    return DeepHealth(
+        vocab_size=dataset.vocab_size,
+        n_embd=args.n_embd,
+        n_head=args.n_head,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
+        n_types=dataset.n_types,
+        n_cont_types=dataset.n_cont_types,
+        n_categories=dataset.n_categories,
+        cont_type_ids=dataset.cont_type_ids,
+        n_bins=args.n_bins,
+        target_mode="next_token",
+        time_mode=args.time_mode,
+        dist_mode="exponential",
+        dropout=args.dropout,
+    )
+
+
+def build_next_step_readout(args: argparse.Namespace):
+    if args.readout_name == "same_time_group_end":
+        return build_readout("same_time_group_end", reduce=args.readout_reduce)
+    return build_readout(args.readout_name)
+
+
+def build_next_step_loss(args: argparse.Namespace):
+    if args.target_mode == "delphi2m":
+        ignored_tokens = {PAD_IDX, CHECKUP_IDX}
+        if args.ignore_no_event_in_delphi2m:
+            ignored_tokens.add(NO_EVENT_IDX)
+        return build_loss(
+            "delphi2m",
+            ignored_tokens=ignored_tokens,
+            t_min=args.t_min,
+            max_exp_input=args.max_exp_input,
+            ce_weight=args.ce_weight,
+            time_weight=args.time_weight,
+        )
+    return build_loss(
+        "uts",
+        ignored_idx={PAD_IDX, CHECKUP_IDX},
+        t_min=args.t_min,
+        max_exp_input=args.max_exp_input,
+    )
+
+
+def compute_next_step_loss(
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout,
+    criterion,
+    batch: Dict[str, torch.Tensor],
+    device: torch.device,
+) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    batch = move_batch_to_device(batch, device)
+    hidden = model(
+        event_seq=batch["event_seq"],
+        time_seq=batch["time_seq"],
+        sex=batch["sex"],
+        padding_mask=batch["padding_mask"],
+        other_type=batch["other_type"],
+        other_value=batch["other_value"],
+        other_value_kind=batch["other_value_kind"],
+        other_time=batch["other_time"],
+        target_mode="next_token",
+    )
+    readout_out = readout(
+        hidden=hidden,
+        time_seq=batch["time_seq"],
+        padding_mask=batch["padding_mask"],
+        readout_mask=batch["readout_mask"]
+        if args.readout_name == "same_time_group_end"
+        else None,
+    )
+    logits = model.calc_risk(readout_out.hidden)
+
+    if args.target_mode == "delphi2m":
+        loss, parts = criterion(
+            logits=logits,
+            target_events=batch["target_event_seq"],
+            target_times=batch["target_time_seq"],
+            current_times=batch["time_seq"],
+            padding_mask=readout_out.readout_mask,
+            return_components=True,
+        )
+    else:
+        loss, parts = 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,
+        )
+    if not torch.isfinite(loss):
+        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
+    return loss, parts
+
+
+def run_epoch(
+    logger: logging.Logger,
+    args: argparse.Namespace,
+    model: DeepHealth,
+    readout,
+    criterion,
+    loader: DataLoader,
+    optimizer: AdamW | None,
+    device: torch.device,
+    is_train: bool,
+) -> float:
+    model.train(is_train)
+    readout.train(is_train)
+    total = 0.0
+    n_batches = 0
+    skipped = 0
+    parts_sum: Dict[str, float] = {}
+    desc = "train" if is_train else "val"
+
+    progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
+    for batch_idx, batch in enumerate(progress):
+        try:
+            loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device)
+            if is_train:
+                if optimizer is None:
+                    raise ValueError("optimizer is required for training")
+                optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                if args.grad_clip > 0:
+                    clip_grad_norm_(model.parameters(), args.grad_clip)
+                optimizer.step()
+
+            total += float(loss.detach().cpu())
+            n_batches += 1
+            for name, value in parts.items():
+                parts_sum[name] = parts_sum.get(name, 0.0) + float(value.detach().cpu())
+            avg = total / max(1, n_batches)
+            postfix = {"loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{avg:.4f}", "skipped": skipped}
+            for name, value in parts_sum.items():
+                postfix[name] = f"{value / max(1, n_batches):.4f}"
+            progress.set_postfix(postfix)
+        except RuntimeError as exc:
+            if "Loss is not finite" not in str(exc):
+                raise
+            skipped += 1
+            logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")
+
+    if skipped:
+        logger.info(f"Skipped {skipped} batches due to non-finite loss")
+    return total / max(1, n_batches) if n_batches else float("inf")
+
+
+def build_metadata(
+    args: argparse.Namespace,
+    dataset: HealthDataset,
+    run_name: str,
+    train_subset,
+    val_subset,
+    test_subset,
+) -> Dict[str, Any]:
+    return {
+        "run_name": run_name,
+        "dataset_class": "NextStepHealthDataset",
+        "collate_fn": "next_step_collate_fn",
+        "model_class": "DeepHealth",
+        "model_target_mode": "next_token",
+        "target_mode": args.target_mode,
+        "dist_mode": "exponential",
+        "extra_info_types_file": (
+            Path(args.extra_info_types_file).name
+            if args.extra_info_types_file is not None
+            else None
+        ),
+        "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        "dataset_metadata": {
+            "vocab_size": int(dataset.vocab_size),
+            "n_types": int(dataset.n_types),
+            "n_cont_types": int(dataset.n_cont_types),
+            "n_categories": int(dataset.n_categories),
+            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
+            "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        },
+        "split_sizes": {
+            "train": int(len(train_subset)),
+            "val": int(len(val_subset)),
+            "test": int(len(test_subset)),
+        },
+        "resolved_readout_name": args.readout_name,
+        "resolved_loss_name": args.target_mode,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+    device = resolve_device(args.device)
+    configure_torch_for_training(device)
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    run_name = (
+        f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
+        f"gap_{args.no_event_interval_years:g}y_{timestamp}"
+    )
+    run_dir = Path("runs") / run_name
+    run_dir.mkdir(parents=True, exist_ok=False)
+    logger = setup_logging(run_dir)
+
+    logger.info(f"Starting next-step training run: {run_name}")
+    logger.info(f"Device: {device}")
+    logger.info(f"extra_info_types: {args.extra_info_types or 'all'}")
+    logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
+
+    dataset = HealthDataset(
+        data_prefix=args.data_prefix,
+        labels_file=args.labels_file,
+        no_event_interval_years=args.no_event_interval_years,
+        include_no_event_in_uts_target=args.include_no_event_in_uts_target,
+        extra_info_types=args.extra_info_types,
+    )
+    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"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
+    )
+
+    train_loader = DataLoader(
+        train_subset,
+        batch_size=args.batch_size,
+        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
+        collate_fn=collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+    val_loader = DataLoader(
+        val_subset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+    test_loader = DataLoader(
+        test_subset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collate_fn,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        prefetch_factor=2 if args.num_workers > 0 else None,
+    )
+
+    model = build_model(args, dataset).to(device)
+    readout = build_next_step_readout(args).to(device)
+    criterion = build_next_step_loss(args)
+    optimizer = AdamW(
+        model.parameters(),
+        lr=args.base_lr,
+        betas=tuple(args.betas),
+        weight_decay=args.weight_decay,
+    )
+    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
+
+    save_config(
+        args,
+        run_dir / "train_config.json",
+        extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
+    )
+
+    best_val = float("inf")
+    patience = 0
+    history = []
+    best_model_path = run_dir / "best_model.pt"
+    start = time.time()
+
+    for epoch in range(args.max_epochs):
+        lr = get_lr(epoch, args, adaptive_lr)
+        set_optimizer_lr(optimizer, lr)
+
+        train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True)
+        with torch.no_grad():
+            val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False)
+
+        is_best = val_loss < best_val
+        if is_best:
+            best_val = val_loss
+            patience = 0
+            save_checkpoint(model, best_model_path)
+        else:
+            patience += 1
+
+        logger.info(
+            f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
+            f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
+            f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
+            f"elapsed={time.time() - start:.1f}s"
+        )
+        history.append({
+            "epoch": epoch + 1,
+            "lr": lr,
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "best_val_loss": best_val,
+            "is_best": int(is_best),
+        })
+        if patience >= args.patience:
+            logger.info(f"Early stopping triggered at epoch {epoch + 1}")
+            break
+
+    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
+        json.dump(history, f, indent=2)
+
+    logger.info("Evaluating best model on next-step test split...")
+    model.load_state_dict(torch.load(best_model_path, map_location=device))
+    with torch.no_grad():
+        test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False)
+    logger.info(f"Test loss: {test_loss:.6f}")
+    logger.info(f"Best checkpoint: {best_model_path}")
+
+
+if __name__ == "__main__":
+    main()

train_util.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import json
+import logging
+import sys
+from pathlib import Path
+from typing import Any, Dict, Iterable, Tuple
+
+import numpy as np
+import torch
+from torch.optim import AdamW
+from torch.utils.data import Subset
+
+from dataset import AllFutureHealthDataset, HealthDataset
+from models import DeepHealth
+
+
+def setup_logging(run_dir: Path) -> logging.Logger:
+    run_dir.mkdir(parents=True, exist_ok=True)
+    logger = logging.getLogger("DeepHealth")
+    logger.setLevel(logging.INFO)
+    logger.handlers.clear()
+
+    formatter = logging.Formatter(
+        "%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(logging.INFO)
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+
+    file_handler = logging.FileHandler(run_dir / "train.log", mode="w")
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+def set_seed(seed: int) -> None:
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+
+
+def load_extra_info_types_file(path: str) -> list[int]:
+    file_path = Path(path)
+    if not file_path.is_file():
+        raise FileNotFoundError(f"extra_info_types_file not found: {path}")
+
+    text = file_path.read_text(encoding="utf-8").strip()
+    if not text:
+        return []
+
+    if text.startswith("["):
+        raw_items = json.loads(text)
+        if not isinstance(raw_items, list):
+            raise ValueError("extra_info_types_file JSON must be a list")
+    else:
+        raw_items = []
+        for line in text.splitlines():
+            line = line.split("#", 1)[0].strip()
+            if line:
+                raw_items.extend(line.replace(",", " ").replace(";", " ").split())
+
+    try:
+        return [int(x) for x in raw_items]
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid extra info type id in {path}") from exc
+
+
+def configure_torch_for_training(device: torch.device) -> None:
+    if device.type != "cuda":
+        return
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    if hasattr(torch, "set_float32_matmul_precision"):
+        torch.set_float32_matmul_precision("high")
+
+
+def resolve_device(device_arg: str) -> torch.device:
+    requested = device_arg.strip().lower()
+    if requested == "cpu":
+        return torch.device("cpu")
+    if requested == "cuda":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if requested.startswith("cuda:"):
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        index = int(requested.split(":", 1)[1])
+        if index < 0 or index >= torch.cuda.device_count():
+            raise ValueError(f"Requested CUDA device is out of range: {device_arg}")
+        return torch.device(f"cuda:{index}")
+    raise ValueError(f"Unsupported device: {device_arg}")
+
+
+def split_dataset(
+    dataset: HealthDataset,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[Subset, Subset, Subset]:
+    total = train_ratio + val_ratio + test_ratio
+    if not np.isclose(total, 1.0, atol=1e-6):
+        raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
+
+    indices = np.random.RandomState(seed).permutation(len(dataset))
+    n_train = int(len(dataset) * train_ratio)
+    n_val = int(len(dataset) * val_ratio)
+    return (
+        Subset(dataset, indices[:n_train]),
+        Subset(dataset, indices[n_train:n_train + n_val]),
+        Subset(dataset, indices[n_train + n_val:]),
+    )
+
+
+def split_all_future_datasets(
+    train_dataset: AllFutureHealthDataset,
+    val_dataset: AllFutureHealthDataset,
+    test_dataset: AllFutureHealthDataset,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[Subset, Subset, Subset]:
+    total = train_ratio + val_ratio + test_ratio
+    if not np.isclose(total, 1.0, atol=1e-6):
+        raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
+
+    patient_indices = np.random.RandomState(seed).permutation(len(train_dataset.patients))
+    n_train = int(len(patient_indices) * train_ratio)
+    n_val = int(len(patient_indices) * val_ratio)
+    train_patient_idx = patient_indices[:n_train]
+    val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val])
+    test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:])
+
+    val_query_idx = [
+        i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
+        if int(pidx) in val_patient_set
+    ]
+    test_query_idx = [
+        i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
+        if int(pidx) in test_patient_set
+    ]
+    if not val_query_idx:
+        raise ValueError("All-future validation split has no valid query samples.")
+    if not test_query_idx:
+        raise ValueError("All-future test split has no valid query samples.")
+
+    return (
+        Subset(train_dataset, train_patient_idx),
+        Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)),
+        Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)),
+    )
+
+
+def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
+    return AdamW(
+        model.parameters(),
+        lr=args.base_lr,
+        betas=tuple(args.betas),
+        weight_decay=args.weight_decay,
+    )
+
+
+def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+
+def save_checkpoint(model: DeepHealth, checkpoint_path: Path) -> None:
+    torch.save(model.state_dict(), checkpoint_path)
+
+
+def save_config(
+    args: Any,
+    config_path: Path,
+    extra: Dict[str, Any] | None = None,
+) -> None:
+    config: Dict[str, Any] = {}
+    for key, value in vars(args).items():
+        if isinstance(value, tuple):
+            config[key] = list(value)
+        elif isinstance(value, list):
+            config[key] = value
+        elif isinstance(value, (int, float, str, bool, type(None))):
+            config[key] = value
+        else:
+            config[key] = str(value)
+    if extra:
+        config.update(extra)
+    config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")