refactor: isolate Delphi2M next-token pipeline

README.md

@@ -59,8 +59,8 @@ python prepare_data.py
 `dataset.py` 提供两个 dataset：
 
 - `NextStepHealthDataset`
-  - 用于 next-token / next-time-point 监督
-  - 对应 `Delphi2MLoss` 和 `UniqueTimeSetExponentialLoss`
+  - 仅用于 absolute-time Delphi2M next-token 复现
+  - 对应 `Delphi2MLoss`
 
 - `AllFutureHealthDataset`
   - 用于 query-conditioned all-future 监督
@@ -173,7 +173,7 @@ extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入
 如果需要拿到完整 next-token 输出，可使用结构化返回：
 
 ```python
-out = model(..., target_mode="next_token", return_output=True)
+out = model(..., return_output=True)
 out.hidden        # disease tokens + pooled extra-info readout tokens
 out.time_seq      # 与 hidden 对齐的时间
 out.padding_mask  # 与 hidden 对齐的有效位置
@@ -187,8 +187,7 @@ model = DeepHealth(
     vocab_size=dataset.vocab_size,
     n_embd=120,
     n_head=10,
-    n_hist_layer=12,
-    n_tab_layer=4,  # 兼容旧配置；当前不再创建独立 tabular transformer
+    n_layer=12,
     n_types=dataset.n_types,
     n_cont_types=dataset.n_cont_types,
     n_categories=dataset.n_categories,
@@ -204,14 +203,13 @@ model = DeepHealth(
 next-token 监督：
 
 - `Delphi2MLoss`
-- `UniqueTimeSetExponentialLoss`
 
 next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括：
 
 - 原 disease token readout 位置
 - 同一时间点 extra hidden 池化后的 pooled extra-info readout token
 
-pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的下一个 disease 事件时间；`delphi2m` 使用第一个未来事件作为 next-token target，`uts` 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。
+pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的第一个 disease 事件作为 next-token target。若该时间点之后没有未来 disease target，则该位置不参与 loss。
 
 all-future / query-conditioned 监督：
 
@@ -221,17 +219,15 @@ all-future / query-conditioned 监督：
 
 all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。
 
-`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
-
 ## 训练
 
 当前提供两类训练入口：
 
 - `train_next_step.py`
   - 使用 `NextStepHealthDataset`
-  - `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
-  - `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
-  - 当前 next-token 训练只支持 exponential time loss
+  - 仅用于 Delphi2M 复现，固定 `time_mode=absolute`、`target_mode=delphi2m`
+  - 固定使用 `Delphi2MLoss` + `token` readout
+  - next-token 训练只支持 exponential time loss
   - 展开的 extra-info tokens 进入主序列；读出端 pooled extra-info tokens 会加入 prediction/loss 监督
 - `train_all_future.py`
   - 使用 `AllFutureHealthDataset`
@@ -243,8 +239,7 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in
 
 | 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout |
 | --- | --- | --- | --- |
-| `next_token` | `relative`, `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
-| `next_token` | `relative`, `absolute` | `target_mode=uts`, `dist_mode=exponential` | `UniqueTimeSetExponentialLoss` + `same_time_group_end` |
+| `next_token` | `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
 | `all_future` | `relative`, `absolute` | `dist_mode=exponential` | `ExponentialLoss`，无 readout |
 | `all_future` | `relative`, `absolute` | `dist_mode=weibull` | `WeibullLoss`，无 readout |
 | `all_future` | `relative`, `absolute` | `dist_mode=mixed` | `MixedLoss`，无 readout |
@@ -255,11 +250,9 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in
 python train_next_step.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 \
+  --n_layer 12 \
   --extra_pool_reduce mean
 ```
 
@@ -424,11 +417,6 @@ python evaluate_auc_v2.py \
 - `losses.py`
   - next-token 和 all-future losses
 
-- `readouts.py`
-  - token readout
-  - same-time group readout
-  - last-valid readout
-
 - `evaluate_auc.py`
   - next-step/token-level 疾病 AUC 评估
   - 使用 prediction offset、sex、age bracket 分层

dataset.py

@@ -14,7 +14,7 @@ from targets import (
     DAYS_PER_YEAR,
     NO_EVENT_IDX,
     PAD_IDX,
-    build_all_targets,
+    build_next_token_targets,
 )
 
 
@@ -86,13 +86,11 @@ class _ExpoBaseDataset(Dataset):
         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
@@ -138,11 +136,6 @@ class _ExpoBaseDataset(Dataset):
         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():
@@ -289,12 +282,7 @@ class _ExpoBaseDataset(Dataset):
 
 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
+    Delphi2M next-token dataset with unified other-info tokens.
     """
 
     CACHE_VERSION = 3
@@ -304,14 +292,12 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         data_prefix: str = "ukb",
         labels_file: str = "labels.csv",
         no_event_interval_years: float = 5.0,
-        include_no_event_in_uts_target: bool = False,
         extra_info_types: Iterable[int] | None = None,
     ) -> None:
         super().__init__(
             data_prefix=data_prefix,
             labels_file=labels_file,
             no_event_interval_years=no_event_interval_years,
-            include_no_event_in_uts_target=include_no_event_in_uts_target,
             extra_info_types=extra_info_types,
         )
 
@@ -326,23 +312,18 @@ class NextStepHealthDataset(_ExpoBaseDataset):
             if features is None:
                 continue
 
-            target_pack = build_all_targets(
+            targets = build_next_token_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,
+                "event_seq": targets.input_events,
+                "time_seq": targets.input_times_years,
+                "target_event_seq": targets.target_events,
+                "target_time_seq": targets.target_times_years,
                 **features,
             })
 
@@ -361,9 +342,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
             "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(),
         }
 
 
@@ -386,7 +364,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         labels_file: str = "labels.csv",
         split: Literal["train", "valid", "test"] = "train",
         no_event_interval_years: float = 5.0,
-        include_no_event_in_uts_target: bool = False,
         min_history_events: int = 1,
         min_future_events: int = 1,
         validation_query_seed: int = 42,
@@ -399,7 +376,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             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,
         )
 
@@ -605,31 +581,12 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict:
         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

eval_data.py

@@ -1,15 +1,209 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Iterable, List
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import numpy as np
 import torch
 from torch.nn.utils.rnn import pad_sequence
 
 from dataset import AllFutureHealthDataset, HealthDataset
+from model_architectures import resolve_model_architecture
+from models import DeepHealth
 from targets import PAD_IDX
 
 
+def load_json_config(path: str | Path | None) -> Dict[str, Any]:
+    if path is None:
+        return {}
+    config_path = Path(path)
+    if not config_path.exists():
+        return {}
+    with config_path.open("r", encoding="utf-8") as file:
+        return json.load(file)
+
+
+def cfg_get(
+    args: argparse.Namespace | Dict[str, Any] | None,
+    cfg: Dict[str, Any],
+    name: str,
+    default: Any,
+) -> Any:
+    if args is not None:
+        value = (
+            args.get(name)
+            if isinstance(args, dict)
+            else getattr(args, name, None)
+        )
+        if value is not None:
+            return value
+    return cfg.get(name, default)
+
+
+def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
+    device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(device_name)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            f"Requested device {device_name!r}, but CUDA is not available."
+        )
+    return device
+
+
+def validate_training_mode_config(cfg: Dict[str, Any]) -> None:
+    model_target_mode = str(
+        cfg.get("model_target_mode", "next_token")
+    ).lower()
+    if model_target_mode not in {"next_token", "all_future"}:
+        raise ValueError(
+            "model_target_mode must be next_token or all_future, got "
+            f"{model_target_mode!r}"
+        )
+    if model_target_mode != "next_token":
+        return
+
+    time_mode = str(cfg.get("time_mode", "")).lower()
+    target_mode = str(cfg.get("target_mode", "")).lower()
+    if time_mode != "absolute" or target_mode != "delphi2m":
+        raise ValueError(
+            "next_token is reserved for Delphi2M reproduction and requires "
+            "time_mode='absolute' and target_mode='delphi2m'; got "
+            f"time_mode={time_mode!r}, target_mode={target_mode!r}"
+        )
+
+
+def split_indices(
+    n: int,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    total = float(train_ratio) + float(val_ratio) + float(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(n)
+    n_train = int(n * train_ratio)
+    n_val = int(n * val_ratio)
+    return (
+        indices[:n_train],
+        indices[n_train:n_train + n_val],
+        indices[n_train + n_val:],
+    )
+
+
+def build_model_from_dataset(
+    args: argparse.Namespace,
+    cfg: Dict[str, Any],
+    dataset: HealthDataset,
+    state_dict: Optional[Dict[str, Any]] = None,
+) -> DeepHealth:
+    validate_training_mode_config(cfg)
+    model_target_mode = str(
+        cfg_get(args, cfg, "model_target_mode", "next_token")
+    ).lower()
+    if model_target_mode not in {"next_token", "all_future"}:
+        raise ValueError(
+            "model_target_mode must be next_token or all_future, got "
+            f"{model_target_mode!r}"
+        )
+    model_architecture = resolve_model_architecture(cfg, state_dict)
+    return DeepHealth(
+        vocab_size=dataset.vocab_size,
+        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
+        n_head=int(cfg_get(args, cfg, "n_head", 10)),
+        n_layer=int(cfg["n_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=int(cfg_get(args, cfg, "n_bins", 16)),
+        extra_pool_reduce=str(
+            cfg_get(args, cfg, "extra_pool_reduce", "mean")
+        ),
+        target_mode=model_target_mode,
+        time_mode=str(cfg_get(args, cfg, "time_mode", "absolute")),
+        dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
+        dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
+        model_architecture=model_architecture,
+    )
+
+
+def validate_dataset_metadata(
+    dataset: HealthDataset,
+    cfg: Dict[str, Any],
+) -> None:
+    metadata = cfg.get("dataset_metadata")
+    if not isinstance(metadata, dict):
+        return
+
+    actual: Dict[str, Any] = {
+        "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],
+    }
+    mismatches = [
+        f"{key}: train_config={metadata.get(key)!r}, current_dataset={value!r}"
+        for key, value in actual.items()
+        if key in metadata and metadata.get(key) != value
+    ]
+    if mismatches:
+        raise RuntimeError(
+            "Current dataset metadata does not match train_config.json. "
+            "Use the same prepared data and extra_info_types as training. "
+            + "; ".join(mismatches)
+        )
+
+
+def build_first_occurrence_map(
+    dataset: HealthDataset,
+    subset_indices: np.ndarray,
+) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
+    first_lists: Dict[int, List[Tuple[int, float]]] = {}
+    for patient_id, dataset_index in enumerate(subset_indices.tolist()):
+        sample = dataset.samples[int(dataset_index)]
+        sequence_events = np.asarray(sample["event_seq"], dtype=np.int64)
+        sequence_times = np.asarray(sample["time_seq"], dtype=np.float32)
+        target_events = np.asarray(
+            sample["target_event_seq"], dtype=np.int64
+        )
+        target_times = np.asarray(
+            sample["target_time_seq"], dtype=np.float32
+        )
+        if sequence_events.size == 0 or target_events.size == 0:
+            continue
+
+        full_events = np.concatenate(
+            [sequence_events, target_events[-1:]]
+        )
+        full_times = np.concatenate([sequence_times, target_times[-1:]])
+        unique_tokens, first_indices = np.unique(
+            full_events, return_index=True
+        )
+        for token, first_index in zip(
+            unique_tokens.tolist(), first_indices.tolist()
+        ):
+            first_lists.setdefault(int(token), []).append(
+                (patient_id, float(full_times[int(first_index)]))
+            )
+
+    return {
+        int(token): (
+            np.asarray([patient for patient, _ in pairs], dtype=np.int32),
+            np.asarray([time for _, time in pairs], dtype=np.float32),
+        )
+        for token, pairs in first_lists.items()
+        if pairs
+    }
+
+
 class AllFutureSequenceEvalDataset:
     """
     Eval-only sequence view for all-future checkpoints.
@@ -52,7 +246,6 @@ class AllFutureSequenceEvalDataset:
             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"]),
@@ -60,7 +253,6 @@ class AllFutureSequenceEvalDataset:
                     "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),
@@ -79,7 +271,6 @@ class AllFutureSequenceEvalDataset:
             "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(),
@@ -94,7 +285,6 @@ def load_sequence_eval_dataset(
     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,
@@ -105,7 +295,6 @@ def load_sequence_eval_dataset(
             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":
@@ -132,9 +321,6 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
     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
     )
@@ -154,7 +340,7 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
         "padding_mask": event_seq > PAD_IDX,
         "target_event_seq": target_event_seq,
         "target_time_seq": target_time_seq,
-        "readout_mask": readout_mask,
+        "readout_mask": event_seq > PAD_IDX,
         "sex": torch.stack([s["sex"] for s in batch]),
         "other_type": other_type,
         "other_value": other_value,

evaluate_auc.py

@@ -19,7 +19,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_config.json-compatible checkpoints/configs.
+and train_config.json-compatible checkpoints/configs.
 """
 from __future__ import annotations
 
@@ -44,10 +44,20 @@ from delphi2m_auc_report import (
     DEFAULT_DELPHI2M_PERIODS_YEARS,
     build_delphi2m_auc_report,
 )
-from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
+from eval_data import (
+    build_first_occurrence_map,
+    build_model_from_dataset,
+    cfg_get,
+    load_json_config,
+    load_sequence_eval_dataset,
+    resolve_eval_device,
+    sequence_eval_collate_fn,
+    split_indices,
+    validate_training_mode_config,
+    validate_dataset_metadata,
+)
 from model_architectures import resolve_model_architecture
 from models import DeepHealth
-from readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
 
 
@@ -156,55 +166,6 @@ def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> T
 SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
 
 
-def _build_first_occurrence_maps(
-    dataset: HealthDataset,
-    subset_indices: np.ndarray,
-) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
-    patient_count = len(subset_indices)
-    followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
-    death_time = np.full(patient_count, np.inf, dtype=np.float32)
-    sex = np.full(patient_count, -1, dtype=np.int8)
-
-    first_lists: Dict[int, List[Tuple[int, float]]] = {}
-
-    for patient_id, dataset_index in enumerate(subset_indices.tolist()):
-        s = dataset.samples[int(dataset_index)]
-
-        seq_event = np.asarray(s["event_seq"], dtype=np.int64)
-        seq_time = np.asarray(s["time_seq"], dtype=np.float32)
-        tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
-        tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
-
-        if seq_event.size == 0 or tgt_event.size == 0:
-            continue
-
-        full_event = np.concatenate([seq_event, tgt_event[-1:]])
-        full_time = np.concatenate([seq_time, tgt_time[-1:]])
-
-        sex[patient_id] = int(s["sex"])
-        followup_end[patient_id] = np.max(full_time).astype(np.float32)
-
-        uniq_tokens, first_idx = np.unique(full_event, return_index=True)
-
-        for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
-            token = int(token)
-            event_time = float(full_time[int(idx)])
-            if token not in first_lists:
-                first_lists[token] = []
-            first_lists[token].append((patient_id, event_time))
-
-    packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
-    for token, pairs in first_lists.items():
-        if not pairs:
-            continue
-        packed[int(token)] = (
-            np.asarray([p for p, _ in pairs], dtype=np.int32),
-            np.asarray([t for _, t in pairs], dtype=np.float32),
-        )
-
-    return packed, followup_end, death_time, sex
-
-
 def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
     death_ids = [int(dataset.vocab_size) - 1]
     print(f"[INFO] death token ids: {death_ids}")
@@ -264,88 +225,6 @@ def select_disease_tokens(
 # Dataset/split/model helpers
 # ---------------------------------------------------------------------------
 
-def load_json_config(path: Optional[str]) -> Dict[str, Any]:
-    if path is None:
-        return {}
-    p = Path(path)
-    if not p.exists():
-        return {}
-    with p.open("r", encoding="utf-8") as f:
-        return json.load(f)
-
-
-def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any], name: str, default: Any) -> Any:
-    """Get a value from CLI args first, then train_config.json, then default.
-
-    This helper intentionally accepts either an argparse.Namespace or a dict.
-    The earlier version passed cfg as both args and cfg, then tried to access
-    args.eval_split, which fails because dict has no attributes.
-    """
-    val = None
-    if args is not None:
-        if isinstance(args, dict):
-            val = args.get(name, None)
-        else:
-            val = getattr(args, name, None)
-    if val is not None:
-        return val
-    return cfg.get(name, default)
-
-
-def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
-    """Resolve evaluation device without inheriting train_config.json device."""
-    device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
-    device = torch.device(device_name)
-    if device.type == "cuda" and not torch.cuda.is_available():
-        raise RuntimeError(
-            f"Requested device {device_name!r}, but CUDA is not available."
-        )
-    return device
-
-
-def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    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}")
-    rng = np.random.RandomState(seed)
-    idx = rng.permutation(n)
-    n_train = int(n * train_ratio)
-    n_val = int(n * val_ratio)
-    return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
-
-
-def build_model_from_dataset(
-    args: argparse.Namespace,
-    cfg: Dict[str, Any],
-    dataset: HealthDataset,
-    state_dict: Optional[Dict[str, Any]] = None,
-) -> DeepHealth:
-    model_target_mode = str(cfg_get(
-        args, cfg, "model_target_mode", "next_token")).lower()
-    if model_target_mode not in {"next_token", "all_future"}:
-        raise ValueError(
-            f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
-        )
-    model_architecture = resolve_model_architecture(cfg, state_dict)
-    return DeepHealth(
-        vocab_size=dataset.vocab_size,
-        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
-        n_head=int(cfg_get(args, cfg, "n_head", 10)),
-        n_layer=int(cfg["n_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=int(cfg_get(args, cfg, "n_bins", 16)),
-        extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")),
-        target_mode=model_target_mode,
-        time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
-        dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
-        dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
-        model_architecture=model_architecture,
-    )
-
-
 def _extract_state_dict(ckpt: Any) -> Dict[str, Any]:
     if isinstance(ckpt, dict) and "model" in ckpt:
         return ckpt["model"]
@@ -430,32 +309,6 @@ def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str
     return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)
 
 
-def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
-    meta = cfg.get("dataset_metadata")
-    if not isinstance(meta, dict):
-        return
-
-    actual: Dict[str, Any] = {
-        "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],
-    }
-    mismatches = [
-        f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
-        for key, value in actual.items()
-        if key in meta and meta.get(key) != value
-    ]
-    if mismatches:
-        raise RuntimeError(
-            "Current dataset metadata does not match train_config.json. "
-            "Use the same prepared data and extra_info_types as training. "
-            + "; ".join(mismatches)
-        )
-
-
 # ---------------------------------------------------------------------------
 # Batched inference + cached hidden states
 # ---------------------------------------------------------------------------
@@ -480,8 +333,6 @@ def infer_readout_hidden(
     loader: DataLoader,
     device: torch.device,
     model_target_mode: str,
-    readout_name: str,
-    readout_reduce: str,
     use_amp: bool,
     hidden_cache_dtype: str = "float16",
 ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
@@ -492,15 +343,6 @@ def infer_readout_hidden(
             f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
         )
 
-    readout = None
-    if model_target_mode == "next_token" and readout_name == "same_time_group_end":
-        readout = build_readout("same_time_group_end",
-                                reduce=readout_reduce).to(device)
-    elif model_target_mode == "next_token":
-        readout = build_readout(readout_name).to(device)
-    if readout is not None:
-        readout.eval()
-
     hidden_parts: List[np.ndarray] = []
     arrays: Dict[str, List[np.ndarray]] = {
         "event_seq": [],
@@ -553,7 +395,6 @@ def infer_readout_hidden(
                         other_value=batch_dev["other_value"][active],
                         other_value_kind=batch_dev["other_value_kind"][active],
                         other_time=batch_dev["other_time"][active],
-                        target_mode="all_future",
                     )
                     hidden[active, pos, :] = hidden_pos.float()
                 readout_mask_np = batch["padding_mask"].cpu().numpy()
@@ -567,16 +408,9 @@ def infer_readout_hidden(
                     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=time_seq,
-                    padding_mask=padding_mask,
-                    readout_mask=batch_dev["readout_mask"],
-                )
-                hidden = ro.hidden
-                readout_mask_np = ro.readout_mask.detach().cpu().numpy()
+                hidden = hidden_raw
+                readout_mask_np = padding_mask.detach().cpu().numpy()
 
         h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
         hidden_parts.append(h)
@@ -1065,8 +899,6 @@ def evaluate_auc_pipeline(
     offsets: Sequence[float],
     device: torch.device,
     model_target_mode: str,
-    readout_name: str,
-    readout_reduce: str,
     num_workers_auc: int,
     use_amp: bool,
     auc_task_chunk_size: int = 0,
@@ -1125,8 +957,6 @@ def evaluate_auc_pipeline(
         loader=loader,
         device=device,
         model_target_mode=model_target_mode,
-        readout_name=readout_name,
-        readout_reduce=readout_reduce,
         use_amp=use_amp,
         hidden_cache_dtype=hidden_cache_dtype,
     )
@@ -1318,6 +1148,7 @@ def main() -> None:
         raise FileNotFoundError(f"best_model.pt not found in {run_path}")
 
     cfg = load_json_config(str(config_path))
+    validate_training_mode_config(cfg)
 
     if args.output_path is None:
         args.output_path = str(run_path)
@@ -1326,9 +1157,7 @@ def main() -> None:
     data_prefix = cfg.get("data_prefix", "ukb")
     labels_file = cfg.get("labels_file", "labels.csv")
     no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
-    include_no_event = cfg.get("include_no_event_in_uts_target", False)
 
-    target_mode = cfg.get("target_mode", "uts")
     model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
     if model_target_mode not in {"next_token", "all_future"}:
         raise ValueError(
@@ -1336,9 +1165,6 @@ def main() -> None:
             f"got {model_target_mode!r}"
         )
     dist_mode_cfg = cfg.get("dist_mode", "exponential")
-    readout_name = cfg.get(
-        "readout_name", "same_time_group_end" if target_mode == "uts" else "token")
-    readout_reduce = cfg.get("readout_reduce", "mean")
 
     device = resolve_eval_device(args.device)
     if device.type == "cuda":
@@ -1350,7 +1176,6 @@ def main() -> None:
         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)),
@@ -1406,8 +1231,9 @@ def main() -> None:
     if disease_spec is None:
         disease_spec = cfg.get("disease_tokens", None)
     diseases = parse_int_list(disease_spec)
-    first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
-        dataset, subset_indices)
+    first_occurrence_by_token = build_first_occurrence_map(
+        dataset, subset_indices
+    )
     include_death = bool(cfg_get(args, cfg, "include_death", True))
     exclude_death = bool(cfg_get(args, cfg, "exclude_death", False))
     auc_offsets = make_auc_offsets(args, cfg)
@@ -1427,8 +1253,6 @@ def main() -> None:
         offsets=auc_offsets,
         device=device,
         model_target_mode=model_target_mode,
-        readout_name=readout_name,
-        readout_reduce=readout_reduce,
         num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(
             1, (os.cpu_count() or 2) - 1))),
         use_amp=bool(cfg_get(args, cfg, "use_amp", False)),

evaluate_auc_v2.py

@@ -35,39 +35,22 @@ from delphi2m_auc_report import (
     DEFAULT_DELPHI2M_PERIODS_YEARS,
     build_delphi2m_auc_report,
 )
-from eval_data import load_sequence_eval_dataset
+from eval_data import (
+    build_first_occurrence_map,
+    build_model_from_dataset,
+    cfg_get,
+    load_json_config,
+    load_sequence_eval_dataset,
+    resolve_eval_device,
+    split_indices,
+    validate_training_mode_config,
+    validate_dataset_metadata,
+)
 from model_architectures import resolve_model_architecture
 from models import DeepHealth
-from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
 
 SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
-_TARGET_AWARE_MODES = {"target_aware", "delphi2m", "d2m"}
-
-
-def load_json_config(path: Path) -> Dict[str, Any]:
-    if not path.exists():
-        return {}
-    with path.open("r", encoding="utf-8") as f:
-        return json.load(f)
-
-
-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)
-
-
-def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
-    """Resolve evaluation device without inheriting train_config.json device."""
-    device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
-    device = torch.device(device_name)
-    if device.type == "cuda" and not torch.cuda.is_available():
-        raise RuntimeError(
-            f"Requested device {device_name!r}, but CUDA is not available."
-        )
-    return device
 
 
 def parse_int_list(value: Any) -> Optional[List[int]]:
@@ -108,17 +91,6 @@ def parse_float_list(value: Any) -> Optional[List[float]]:
     return [float(x.strip()) for x in text.split(",") if x.strip()]
 
 
-def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    total = float(train_ratio) + float(val_ratio) + float(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}")
-    rng = np.random.RandomState(int(seed))
-    idx = rng.permutation(int(n))
-    n_train = int(n * train_ratio)
-    n_val = int(n * val_ratio)
-    return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
-
-
 def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dict[str, Any]) -> np.ndarray:
     train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
     val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
@@ -185,69 +157,11 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
     return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
 
 
-def build_model_from_dataset(
-    args: argparse.Namespace,
-    cfg: Dict[str, Any],
-    dataset: HealthDataset,
-    state_dict: Optional[Dict[str, Any]] = None,
-) -> DeepHealth:
-    model_target_mode = str(cfg_get(
-        args, cfg, "model_target_mode", "next_token")).lower()
-    if model_target_mode not in {"next_token", "all_future"}:
-        raise ValueError(
-            f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
-        )
-    model_architecture = resolve_model_architecture(cfg, state_dict)
-    return DeepHealth(
-        vocab_size=dataset.vocab_size,
-        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
-        n_head=int(cfg_get(args, cfg, "n_head", 10)),
-        n_layer=int(cfg["n_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=int(cfg_get(args, cfg, "n_bins", 16)),
-        extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")),
-        target_mode=model_target_mode,
-        time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
-        dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
-        dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
-        model_architecture=model_architecture,
-    )
-
-
 def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
     resolve_model_architecture(model.model_architecture, state_dict)
     model.load_state_dict(state_dict, strict=True)
 
 
-def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
-    meta = cfg.get("dataset_metadata")
-    if not isinstance(meta, dict):
-        return
-
-    actual: Dict[str, Any] = {
-        "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],
-    }
-    mismatches = [
-        f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
-        for key, value in actual.items()
-        if key in meta and meta.get(key) != value
-    ]
-    if mismatches:
-        raise RuntimeError(
-            "Current dataset metadata does not match train_config.json. "
-            "Use the same prepared data and extra_info_types as training. "
-            + "; ".join(mismatches)
-        )
-
-
 # ---------------------------------------------------------------------------
 # DeLong AUC utilities
 # ---------------------------------------------------------------------------
@@ -375,54 +289,6 @@ def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFr
     return [int(dataset.vocab_size) - 1]
 
 
-def _build_first_occurrence_maps(
-    dataset: HealthDataset,
-    subset_indices: np.ndarray,
-) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
-    patient_count = len(subset_indices)
-    followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
-    death_time = np.full(patient_count, np.inf, dtype=np.float32)
-    sex = np.full(patient_count, -1, dtype=np.int8)
-
-    first_lists: Dict[int, List[Tuple[int, float]]] = {}
-
-    for patient_id, dataset_index in enumerate(subset_indices.tolist()):
-        s = dataset.samples[int(dataset_index)]
-
-        seq_event = np.asarray(s["event_seq"], dtype=np.int64)
-        seq_time = np.asarray(s["time_seq"], dtype=np.float32)
-        tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
-        tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
-
-        if seq_event.size == 0 or tgt_event.size == 0:
-            continue
-
-        full_event = np.concatenate([seq_event, tgt_event[-1:]])
-        full_time = np.concatenate([seq_time, tgt_time[-1:]])
-
-        sex[patient_id] = int(s["sex"])
-        followup_end[patient_id] = np.max(full_time).astype(np.float32)
-
-        uniq_tokens, first_idx = np.unique(full_event, return_index=True)
-        for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
-            token = int(token)
-            event_time = float(full_time[int(idx)])
-            if token not in first_lists:
-                first_lists[token] = []
-            first_lists[token].append((patient_id, event_time))
-
-    packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
-    for token, pairs in first_lists.items():
-        if not pairs:
-            continue
-        packed[int(token)] = (
-            np.asarray([p for p, _ in pairs], dtype=np.int32),
-            np.asarray([t for _, t in pairs], dtype=np.float32),
-        )
-
-    return packed, followup_end, death_time, sex
-
-
 def select_disease_tokens(
     dataset: HealthDataset,
     labels_meta: Optional[pd.DataFrame],
@@ -468,7 +334,6 @@ class LandmarkDataset(Dataset):
         dataset: HealthDataset,
         subset_indices: np.ndarray,
         landmark_ages: np.ndarray,
-        attn_mask_mode: str,
         model_target_mode: str,
         min_history_events: int,
         first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
@@ -477,7 +342,6 @@ class LandmarkDataset(Dataset):
         self.dataset = dataset
         self.subset_indices = np.asarray(subset_indices, dtype=np.int64)
         self.landmark_ages = np.asarray(landmark_ages, dtype=np.float32)
-        self.attn_mask_mode = str(attn_mask_mode).lower()
         self.model_target_mode = str(model_target_mode).lower()
         if self.model_target_mode not in {"next_token", "all_future"}:
             raise ValueError(
@@ -555,10 +419,11 @@ class LandmarkDataset(Dataset):
                             np.array([np.float32(landmark_age)], dtype=np.float32),
                         ]
                     )
-                    if self.attn_mask_mode in _TARGET_AWARE_MODES:
-                        time_seq_landmark[-1] = np.nextafter(
-                            np.float32(landmark_age), np.float32(np.inf), dtype=np.float32
-                        )
+                    time_seq_landmark[-1] = np.nextafter(
+                        np.float32(landmark_age),
+                        np.float32(np.inf),
+                        dtype=np.float32,
+                    )
                     landmark_pos = int(len(event_seq_landmark) - 1)
                     readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
                     readout_mask[-1] = True
@@ -671,8 +536,6 @@ def infer_landmark_hidden(
     loader: DataLoader,
     device: torch.device,
     model_target_mode: str,
-    readout_name: str,
-    readout_reduce: str,
     use_amp: bool,
     hidden_cache_dtype: str,
 ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
@@ -682,15 +545,6 @@ def infer_landmark_hidden(
             f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
         )
 
-    readout = None
-    if model_target_mode == "next_token" and readout_name == "same_time_group_end":
-        readout = build_readout("same_time_group_end",
-                                reduce=readout_reduce).to(device)
-    elif model_target_mode == "next_token":
-        readout = build_readout(readout_name).to(device)
-    if readout is not None:
-        readout.eval()
-
     hidden_parts: List[np.ndarray] = []
     arrays = {
         "patient_id": [],
@@ -727,7 +581,6 @@ def infer_landmark_hidden(
                     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 = model(
@@ -739,18 +592,11 @@ def infer_landmark_hidden(
                     other_value=batch_dev["other_value"],
                     other_value_kind=batch_dev["other_value_kind"],
                     other_time=batch_dev["other_time"],
-                    target_mode="next_token",
                 )
-                readout_out = readout(
-                    hidden=hidden,
-                    time_seq=batch_dev["time_seq"],
-                    padding_mask=batch_dev["padding_mask"],
-                    readout_mask=batch_dev["readout_mask"],
-                )
-                landmark_hidden = readout_out.hidden.gather(
+                landmark_hidden = hidden.gather(
                     1,
                     batch_dev["landmark_pos"].long()[:, None, None].expand(
-                        -1, 1, readout_out.hidden.shape[-1]
+                        -1, 1, hidden.shape[-1]
                     ),
                 ).squeeze(1)
 
@@ -1085,8 +931,6 @@ def evaluate_landmark_auc(
     horizons: np.ndarray,
     device: torch.device,
     model_target_mode: str,
-    readout_name: str,
-    readout_reduce: str,
     num_workers_auc: int,
     auc_task_chunk_size: int,
     min_cases: int,
@@ -1102,8 +946,6 @@ def evaluate_landmark_auc(
         loader=loader,
         device=device,
         model_target_mode=model_target_mode,
-        readout_name=readout_name,
-        readout_reduce=readout_reduce,
         use_amp=use_amp,
         hidden_cache_dtype=hidden_cache_dtype,
     )
@@ -1282,14 +1124,12 @@ def main() -> None:
         raise FileNotFoundError(f"best_model.pt not found in {run_path}")
 
     cfg = load_json_config(config_path)
+    validate_training_mode_config(cfg)
 
     data_prefix = cfg.get("data_prefix", "ukb")
     labels_file = cfg.get("labels_file", "labels.csv")
     no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
-    include_no_event_in_uts_target = cfg.get(
-        "include_no_event_in_uts_target", False)
 
-    target_mode = cfg.get("target_mode", "uts")
     model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
     if model_target_mode not in {"next_token", "all_future"}:
         raise ValueError(
@@ -1297,12 +1137,6 @@ def main() -> None:
             f"got {model_target_mode!r}"
         )
     dist_mode_cfg = str(cfg.get("dist_mode", "exponential"))
-    attn_mask_mode = str(cfg.get(
-        "attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware"))
-    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"))
-    time_mode = str(cfg.get("time_mode", "relative"))
 
     output_path = Path(
         cfg_get(args, cfg, "output_path", None)
@@ -1325,7 +1159,6 @@ def main() -> None:
         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)),
@@ -1347,8 +1180,9 @@ def main() -> None:
 
     subset_indices = make_eval_indices(dataset, args, cfg)
 
-    first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
-        dataset, subset_indices)
+    first_occurrence_by_token = build_first_occurrence_map(
+        dataset, subset_indices
+    )
 
     disease_requested = parse_int_list(
         cfg_get(args, cfg, "diseases_of_interest", None))
@@ -1444,7 +1278,6 @@ def main() -> None:
         dataset=dataset,
         subset_indices=subset_indices,
         landmark_ages=landmark_ages,
-        attn_mask_mode=attn_mask_mode,
         model_target_mode=model_target_mode,
         min_history_events=min_history_events,
         first_occurrence_by_token=first_occurrence_by_token,
@@ -1516,8 +1349,6 @@ def main() -> None:
         horizons=horizons,
         device=device,
         model_target_mode=model_target_mode,
-        readout_name=readout_name,
-        readout_reduce=readout_reduce,
         num_workers_auc=num_workers_auc,
         auc_task_chunk_size=auc_task_chunk_size,
         min_cases=min_cases,

event_date_utils.py

@@ -1,104 +0,0 @@
-"""Read and query calendar-dated disease-event arrays from prepare_event_dates.py."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Iterable
-
-import numpy as np
-import pandas as pd
-
-
-REQUIRED_FIELDS = {"eid", "event_date", "token"}
-
-
-def load_event_dates(path: str | Path) -> np.ndarray:
-    """Load and validate the structured ``.npy`` event array."""
-    events = np.load(path)
-    if events.dtype.names is None or not REQUIRED_FIELDS.issubset(events.dtype.names):
-        raise ValueError(
-            "Expected a structured .npy with eid, event_date, token fields. "
-            "Create it with prepare_event_dates.py."
-        )
-    return events
-
-
-def load_token_labels(labels_file: str | Path) -> dict[int, str]:
-    """Load token -> human-readable code using the project label convention."""
-    labels = {1: "CHECKUP"}
-    with Path(labels_file).open(encoding="utf-8") as handle:
-        for index, line in enumerate(handle):
-            code = line.strip().split(" ", maxsplit=1)[0]
-            if code:
-                labels[index + 2] = code
-    return labels
-
-
-@dataclass
-class EventDateIndex:
-    """Small in-memory query wrapper for exposure-linkage and cohort scripts."""
-
-    events: np.ndarray
-    token_labels: dict[int, str] | None = None
-
-    @classmethod
-    def from_files(
-        cls,
-        event_file: str | Path,
-        labels_file: str | Path | None = None,
-    ) -> "EventDateIndex":
-        labels = load_token_labels(labels_file) if labels_file is not None else None
-        return cls(load_event_dates(event_file), labels)
-
-    def to_frame(self, events: np.ndarray | None = None) -> pd.DataFrame:
-        """Convert records to a convenient, calendar-dated DataFrame."""
-        data = self.events if events is None else events
-        frame = pd.DataFrame(
-            {
-                "eid": data["eid"].astype("int64"),
-                "event_date": pd.to_datetime(data["event_date"]),
-                "token": data["token"].astype("int32"),
-            }
-        )
-        if self.token_labels is not None:
-            frame["label_code"] = frame["token"].map(self.token_labels).fillna("UNKNOWN")
-        return frame.sort_values(["eid", "event_date", "token"], kind="stable").reset_index(drop=True)
-
-    def for_eid(self, eid: int) -> pd.DataFrame:
-        """Return every stored disease/death event for one participant."""
-        return self.to_frame(self.events[self.events["eid"] == int(eid)])
-
-    def between(
-        self,
-        start: str | pd.Timestamp,
-        end: str | pd.Timestamp,
-        *,
-        eids: Iterable[int] | None = None,
-        tokens: Iterable[int] | None = None,
-    ) -> pd.DataFrame:
-        """Query events in an inclusive calendar-date interval."""
-        start_day = np.datetime64(pd.Timestamp(start).date(), "D")
-        end_day = np.datetime64(pd.Timestamp(end).date(), "D")
-        mask = (self.events["event_date"] >= start_day) & (self.events["event_date"] <= end_day)
-        if eids is not None:
-            mask &= np.isin(self.events["eid"], list(eids))
-        if tokens is not None:
-            mask &= np.isin(self.events["token"], list(tokens))
-        return self.to_frame(self.events[mask])
-
-    def anchors_before(self, eid: int, date: str | pd.Timestamp) -> pd.DataFrame:
-        """Return a participant's event history strictly before an exposure anchor."""
-        day = np.datetime64(pd.Timestamp(date).date(), "D")
-        mask = (self.events["eid"] == int(eid)) & (self.events["event_date"] < day)
-        return self.to_frame(self.events[mask])
-
-    def first_event(self, token: int) -> pd.DataFrame:
-        """Return each participant's first date for a requested token."""
-        selected = self.events[self.events["token"] == int(token)]
-        # Arrays produced by prepare_event_dates.py are already deduplicated;
-        # sorting makes this safe for externally produced compatible arrays too.
-        order = np.lexsort((selected["event_date"], selected["eid"]))
-        selected = selected[order]
-        _, first = np.unique(selected["eid"], return_index=True)
-        return self.to_frame(selected[first])

future_risk.py

@@ -1,115 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Sequence
-
-import torch
-import torch.nn.functional as F
-
-
-def death_token(vocab_size: int) -> int:
-    if int(vocab_size) <= 0:
-        raise ValueError(f"vocab_size must be positive, got {vocab_size}")
-    return int(vocab_size) - 1
-
-
-def probabilities_from_logits(
-    logits: torch.Tensor,
-    tau_years: float | torch.Tensor,
-    *,
-    dist_mode: str = "exponential",
-    rho: torch.Tensor | None = None,
-    death_rho: torch.Tensor | None = None,
-    eps: float = 1e-8,
-) -> torch.Tensor:
-    """
-    Convert all-future logits to tau-year event probabilities.
-
-    Death is always treated as token vocab_size - 1. For dist_mode="mixed",
-    non-death tokens use exponential hazards and death uses death_rho.
-    """
-    if logits.ndim != 2:
-        raise ValueError(f"logits must have shape (N, V), got {tuple(logits.shape)}")
-    if float(torch.as_tensor(tau_years).detach().min().cpu()) < 0:
-        raise ValueError("tau_years must be non-negative")
-
-    mode = str(dist_mode).lower()
-    if mode not in {"exponential", "weibull", "mixed"}:
-        raise ValueError("dist_mode must be one of: exponential, weibull, mixed")
-
-    rate = F.softplus(logits) + float(eps)
-    tau = torch.as_tensor(tau_years, dtype=rate.dtype, device=rate.device)
-    if tau.ndim == 0:
-        tau = tau.expand(logits.shape[0])
-    if tau.ndim != 1 or tau.shape[0] != logits.shape[0]:
-        raise ValueError(
-            "tau_years must be a scalar or a 1D tensor with length N, got "
-            f"{tuple(tau.shape)} for N={logits.shape[0]}"
-        )
-
-    if mode == "exponential":
-        exposure = tau[:, None].expand_as(rate)
-    elif mode == "weibull":
-        if rho is None or rho.shape != logits.shape:
-            raise ValueError("rho must have the same shape as logits for dist_mode='weibull'")
-        exposure = torch.pow(tau[:, None].clamp_min(float(eps)), rho.to(rate.dtype))
-    else:
-        exposure = tau[:, None].expand_as(rate).clone()
-        if death_rho is None:
-            raise ValueError("death_rho is required for dist_mode='mixed'")
-        death_idx = death_token(logits.shape[1])
-        death_shape = tuple(death_rho.shape)
-        death_rho = death_rho.to(device=rate.device, dtype=rate.dtype)
-        if death_rho.ndim == 2 and death_rho.shape[1] == 1:
-            death_rho = death_rho.squeeze(1)
-        if death_rho.ndim != 1 or death_rho.shape[0] != logits.shape[0]:
-            raise ValueError(
-                "death_rho must have shape (N,) or (N, 1), got "
-                f"{death_shape} for N={logits.shape[0]}"
-            )
-        exposure[:, death_idx] = torch.pow(tau.clamp_min(float(eps)), death_rho)
-
-    return -torch.expm1(-rate * exposure)
-
-
-def death_risk_from_probabilities(probabilities: torch.Tensor) -> torch.Tensor:
-    """Return p_death(t, tau), with death fixed to token vocab_size - 1."""
-    if probabilities.ndim != 2:
-        raise ValueError(
-            f"probabilities must have shape (N, V), got {tuple(probabilities.shape)}"
-        )
-    return probabilities[:, death_token(probabilities.shape[1])]
-
-
-def new_disease_risk_from_probabilities(
-    probabilities: torch.Tensor,
-    occurred: torch.Tensor,
-    disease_ids: Sequence[int],
-) -> torch.Tensor:
-    """
-    Compute P(at least one selected disease newly occurs within tau years).
-
-    Already occurred diseases are masked out. Death is not included here and
-    should be reported separately with death_risk_from_probabilities.
-    """
-    if probabilities.ndim != 2 or occurred.shape != probabilities.shape:
-        raise ValueError(
-            "probabilities and occurred must both have shape (N, V), got "
-            f"{tuple(probabilities.shape)} and {tuple(occurred.shape)}"
-        )
-    if not disease_ids:
-        return probabilities.new_zeros(probabilities.shape[0])
-
-    death_idx = death_token(probabilities.shape[1])
-    ids = [
-        idx
-        for idx in dict.fromkeys(int(x) for x in disease_ids)
-        if 0 <= idx < probabilities.shape[1] and idx != death_idx
-    ]
-    if not ids:
-        return probabilities.new_zeros(probabilities.shape[0])
-
-    idx_tensor = torch.as_tensor(ids, dtype=torch.long, device=probabilities.device)
-    p = probabilities[:, idx_tensor].clamp(0.0, 1.0 - 1e-7)
-    new_mask = ~occurred[:, idx_tensor].to(dtype=torch.bool)
-    log_no_new = torch.log1p(-p) * new_mask.to(dtype=p.dtype)
-    return -torch.expm1(log_no_new.sum(dim=1))

losses.py

@@ -145,95 +145,6 @@ class Delphi2MLoss(nn.Module):
         return total_loss
 
 
-class UniqueTimeSetExponentialLoss(nn.Module):
-    """Next distinct timestamp event-set supervision with sum reduction."""
-
-    def __init__(
-        self,
-        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
-        t_min: float = 1.0 / 365.25,
-        max_exp_input: float = 60.0,
-        exclude_ignored_from_intensity: bool = True,
-    ):
-        super().__init__()
-        self.ignored_idx = [int(x) for x in ignored_idx]
-        self.t_min = float(t_min)
-        self.max_exp_input = float(max_exp_input)
-        self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)
-
-    def forward(
-        self,
-        logits: torch.Tensor,
-        target_multi_hot: torch.Tensor,
-        target_dt_unique: torch.Tensor,
-        readout_mask: torch.Tensor,
-        return_components: bool = False,
-    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
-        if logits.dim() != 3:
-            raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
-        bsz, seq_len, vocab_size = logits.shape
-
-        if target_multi_hot.shape != (bsz, seq_len, vocab_size):
-            raise ValueError(
-                "target_multi_hot must match logits shape, "
-                f"got {tuple(target_multi_hot.shape)} vs {tuple(logits.shape)}"
-            )
-        if target_dt_unique.shape != (bsz, seq_len):
-            raise ValueError(
-                f"target_dt_unique must be {(bsz, seq_len)}, got {tuple(target_dt_unique.shape)}"
-            )
-        if readout_mask.shape != (bsz, seq_len):
-            raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}")
-
-        ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)
-
-        num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1)
-        valid_mask = readout_mask.bool() & (num_targets > 0)
-
-        if not valid_mask.any():
-            total_loss = _zero_loss_like(logits)
-            if return_components:
-                return total_loss, {
-                    "observed": total_loss.detach(),
-                    "penalty": total_loss.detach(),
-                    "total": total_loss.detach(),
-                }
-            return total_loss
-
-        logits_safe = torch.nan_to_num(
-            logits[valid_mask],
-            nan=0.0,
-            posinf=self.max_exp_input,
-            neginf=-self.max_exp_input,
-        )
-        target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype)
-        target_valid[:, ignore_mask] = 0.0
-
-        observed_term = (logits_safe * target_valid).sum(dim=-1)
-        penalty_scale = target_valid.sum(dim=-1)
-
-        logits_for_lse = logits_safe
-        if self.exclude_ignored_from_intensity:
-            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
-
-        dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min)
-        log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
-        log_penalty = log_lambda_total + dt_clamped.log()
-        penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input))
-
-        observed_loss = -observed_term
-        penalty_loss = penalty_scale * penalty
-        total_loss = (observed_loss + penalty_loss).mean()
-
-        if return_components:
-            return total_loss, {
-                "observed": observed_loss.mean().detach(),
-                "penalty": penalty_loss.mean().detach(),
-                "total": total_loss.detach(),
-            }
-        return total_loss
-
-
 class ExponentialLoss(nn.Module):
     """Query-conditioned all-future-event exponential point-process loss."""
 
@@ -385,10 +296,8 @@ class MixedLoss(nn.Module):
 
 def build_loss(name: str, **kwargs) -> nn.Module:
     name = name.lower()
-    if name in {"delphi2m", "d2m", "next_token"}:
+    if name == "delphi2m":
         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"}:
@@ -396,5 +305,5 @@ def build_loss(name: str, **kwargs) -> nn.Module:
     if name in {"mixed", "query_mixed"}:
         return MixedLoss(**kwargs)
     raise ValueError(
-        f"Unknown loss {name!r}. Available: delphi2m, uts, exponential, weibull, mixed."
+        f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull, mixed."
     )

models.py

@@ -156,7 +156,7 @@ class DeepHealth(nn.Module):
         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"
+        time_mode: str = "absolute",  # next_token requires absolute
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
@@ -169,6 +169,11 @@ class DeepHealth(nn.Module):
         if time_mode not in ["relative", "absolute"]:
             raise ValueError(
                 "time_mode must be either 'relative' or 'absolute'")
+        if target_mode == "next_token" and time_mode != "absolute":
+            raise ValueError(
+                "next_token is reserved for Delphi2M reproduction and "
+                "requires time_mode='absolute'"
+            )
         if dist_mode not in ["exponential", "weibull", "mixed"]:
             raise ValueError(
                 "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
@@ -461,15 +466,8 @@ class DeepHealth(nn.Module):
             )
         return h_disease[:, :event_len, :]
 
-    def forward_next_token(self, **kwargs) -> torch.Tensor:
-        return self._forward_shared(mode="next_token", **kwargs)
-
-    def forward_all_future(self, **kwargs) -> torch.Tensor:
-        return self._forward_shared(mode="all_future", **kwargs)
-
-    def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor:
-        mode = self.target_mode if target_mode is None else target_mode
-        return self._forward_shared(mode=mode, **kwargs)
+    def forward(self, **kwargs) -> torch.Tensor:
+        return self._forward_shared(mode=self.target_mode, **kwargs)
 
     def calc_risk(self, x: torch.Tensor) -> torch.Tensor:
         return self.risk_head(x)

readouts.py

@@ -1,107 +0,0 @@
-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

@@ -2,8 +2,6 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Iterable
-
 import numpy as np
 
 
@@ -32,38 +30,6 @@ class NextTokenTargets:
     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,
@@ -163,232 +129,3 @@ def build_next_token_targets(
         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,
-    }

test_model_architectures.py

@@ -1,247 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import (
-    SwiGLU,
-    TrajMixer,
-    TrajMixerBlock,
-    TransformerFFNBlock,
-    build_backbone_block,
-)
-from model_architectures import (
-    TRAJ_MIXER_ARCHITECTURE,
-    TRANSFORMER_FFN_ARCHITECTURE,
-    detect_model_architecture_from_state_dict,
-    resolve_model_architecture,
-)
-from models import DeepHealth
-
-
-def _build_block(model_architecture: str):
-    return build_backbone_block(
-        model_architecture,
-        n_embd=12,
-        n_head=3,
-        use_time_rope=False,
-        use_rbf_bias=False,
-        mlp_dropout=0.0,
-    )
-
-
-def _as_model_state_dict(block: torch.nn.Module) -> dict[str, torch.Tensor]:
-    return {
-        f"blocks.0.{name}": value.detach().clone()
-        for name, value in block.state_dict().items()
-    }
-
-
-def _build_model(
-    model_architecture: str | None,
-    *,
-    n_layer: int = 1,
-) -> DeepHealth:
-    return DeepHealth(
-        vocab_size=8,
-        n_embd=12,
-        n_head=3,
-        n_layer=n_layer,
-        n_types=2,
-        n_cont_types=0,
-        n_categories=2,
-        cont_type_ids=[],
-        time_mode="absolute",
-        model_architecture=model_architecture,
-    )
-
-
-class ModelArchitectureFactoryTest(unittest.TestCase):
-    def test_factory_builds_both_architectures_with_expected_topology(self) -> None:
-        ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        self.assertIsInstance(ffn_block, TransformerFFNBlock)
-        self.assertIsInstance(ffn_block.mlp, SwiGLU)
-        self.assertTrue(hasattr(ffn_block, "ln1"))
-        self.assertTrue(hasattr(ffn_block, "ln2"))
-
-        traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
-        self.assertIsInstance(traj_block, TrajMixerBlock)
-        self.assertIsInstance(traj_block.mlp, TrajMixer)
-        self.assertTrue(hasattr(traj_block, "ln1"))
-        self.assertFalse(hasattr(traj_block, "ln2"))
-
-    def test_both_architectures_forward_and_backward(self) -> None:
-        for architecture in (
-            TRANSFORMER_FFN_ARCHITECTURE,
-            TRAJ_MIXER_ARCHITECTURE,
-        ):
-            with self.subTest(architecture=architecture):
-                torch.manual_seed(0)
-                block = _build_block(architecture)
-                x = torch.randn(2, 5, 12, requires_grad=True)
-
-                output = block(x)
-                self.assertEqual(output.shape, x.shape)
-                output.square().mean().backward()
-
-                self.assertIsNotNone(x.grad)
-                self.assertTrue(torch.isfinite(x.grad).all())
-                self.assertGreater(x.grad.abs().sum().item(), 0.0)
-                self.assertIsNotNone(block.attn.qkv.weight.grad)
-                self.assertGreater(
-                    block.attn.qkv.weight.grad.abs().sum().item(),
-                    0.0,
-                )
-
-                if architecture == TRANSFORMER_FFN_ARCHITECTURE:
-                    branch_parameters = (
-                        block.mlp.w1.weight,
-                        block.mlp.w2.weight,
-                        block.mlp.w3.weight,
-                    )
-                else:
-                    branch_parameters = (
-                        block.mlp.intra_gate_proj,
-                        block.mlp.intra_value_proj,
-                        block.mlp.output_proj,
-                    )
-                for parameter in branch_parameters:
-                    self.assertIsNotNone(parameter.grad)
-                    self.assertTrue(torch.isfinite(parameter.grad).all())
-                    self.assertGreater(parameter.grad.abs().sum().item(), 0.0)
-
-    def test_unknown_architecture_is_rejected(self) -> None:
-        with self.assertRaises(ValueError):
-            _build_block("unknown_architecture")
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            _build_model(None)
-
-    def test_deephealth_rejects_fewer_than_one_layer(self) -> None:
-        for n_layer in (0, -1):
-            with self.subTest(n_layer=n_layer):
-                with self.assertRaisesRegex(ValueError, "n_layer must be >= 1"):
-                    _build_model(
-                        TRANSFORMER_FFN_ARCHITECTURE,
-                        n_layer=n_layer,
-                    )
-
-    def test_deephealth_uses_factory_and_strictly_reloads_both_models(self) -> None:
-        for architecture, block_class in (
-            (TRANSFORMER_FFN_ARCHITECTURE, TransformerFFNBlock),
-            (TRAJ_MIXER_ARCHITECTURE, TrajMixerBlock),
-        ):
-            with self.subTest(architecture=architecture):
-                model = _build_model(architecture)
-                self.assertEqual(model.model_architecture, architecture)
-                self.assertIsInstance(model.blocks[0], block_class)
-                self.assertEqual(
-                    detect_model_architecture_from_state_dict(
-                        model.state_dict()
-                    ),
-                    architecture,
-                )
-
-                reloaded = _build_model(architecture)
-                incompatible = reloaded.load_state_dict(
-                    model.state_dict(),
-                    strict=True,
-                )
-                self.assertEqual(incompatible.missing_keys, [])
-                self.assertEqual(incompatible.unexpected_keys, [])
-
-
-class ModelArchitectureResolutionTest(unittest.TestCase):
-    def setUp(self) -> None:
-        self.ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        self.traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
-        self.ffn_state = _as_model_state_dict(self.ffn_block)
-        self.traj_state = _as_model_state_dict(self.traj_block)
-
-    def test_state_dict_detection_recognizes_both_architectures(self) -> None:
-        self.assertEqual(
-            detect_model_architecture_from_state_dict(self.ffn_state),
-            TRANSFORMER_FFN_ARCHITECTURE,
-        )
-        self.assertEqual(
-            detect_model_architecture_from_state_dict(self.traj_state),
-            TRAJ_MIXER_ARCHITECTURE,
-        )
-
-    def test_explicit_markers_resolve_when_checkpoint_matches(self) -> None:
-        for architecture, state_dict in (
-            (TRANSFORMER_FFN_ARCHITECTURE, self.ffn_state),
-            (TRAJ_MIXER_ARCHITECTURE, self.traj_state),
-        ):
-            with self.subTest(architecture=architecture):
-                self.assertEqual(
-                    resolve_model_architecture(
-                        {"model_architecture": architecture},
-                        state_dict,
-                    ),
-                    architecture,
-                )
-
-    def test_architecture_marker_is_required_for_checkpoint_loading(self) -> None:
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            resolve_model_architecture({}, self.ffn_state)
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            resolve_model_architecture(None, self.traj_state)
-
-    def test_explicit_marker_conflicting_with_state_dict_is_rejected(self) -> None:
-        conflicts = (
-            (TRANSFORMER_FFN_ARCHITECTURE, self.traj_state),
-            (TRAJ_MIXER_ARCHITECTURE, self.ffn_state),
-        )
-        for architecture, state_dict in conflicts:
-            with self.subTest(architecture=architecture):
-                with self.assertRaises(ValueError):
-                    resolve_model_architecture(
-                        {"model_architecture": architecture},
-                        state_dict,
-                    )
-
-    def test_unknown_marker_and_ambiguous_state_dict_are_rejected(self) -> None:
-        with self.assertRaises(ValueError):
-            resolve_model_architecture(
-                {"model_architecture": "traj_mixer_v4"}
-            )
-
-        ambiguous_state = dict(self.ffn_state)
-        ambiguous_state.update(self.traj_state)
-        with self.assertRaises(ValueError):
-            detect_model_architecture_from_state_dict(ambiguous_state)
-
-        with self.assertRaises(ValueError):
-            detect_model_architecture_from_state_dict(
-                {"token_embedding.weight": torch.empty(2, 2)}
-            )
-
-    def test_ffn_block_schema_is_stable_and_strictly_loadable(self) -> None:
-        expected_keys = {
-            "attn.time_bias_scale",
-            "attn.qkv.weight",
-            "attn.out_proj.weight",
-            "attn.rbf_proj.weight",
-            "mlp.w1.weight",
-            "mlp.w1.bias",
-            "mlp.w2.weight",
-            "mlp.w2.bias",
-            "mlp.w3.weight",
-            "mlp.w3.bias",
-            "ln1.weight",
-            "ln1.bias",
-            "ln2.weight",
-            "ln2.bias",
-        }
-        state = self.ffn_block.state_dict()
-        self.assertSetEqual(set(state), expected_keys)
-        self.assertEqual(tuple(state["mlp.w1.weight"].shape), (30, 12))
-        self.assertEqual(tuple(state["mlp.w3.weight"].shape), (12, 30))
-
-        reloaded = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        incompatible = reloaded.load_state_dict(state, strict=True)
-        self.assertEqual(incompatible.missing_keys, [])
-        self.assertEqual(incompatible.unexpected_keys, [])
-
-
-if __name__ == "__main__":
-    unittest.main()

test_temporal_attention.py

@@ -1,46 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import TemporalAttention
-
-
-class TemporalAttentionTest(unittest.TestCase):
-    def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
-        torch.manual_seed(0)
-        attention = TemporalAttention(
-            n_embd=12,
-            n_head=3,
-            use_time_rope=False,
-            use_rbf_bias=True,
-        )
-        features = torch.randn(2, 4, 4, 16)
-        target = torch.randn(2, 4, 4, 3)
-
-        initial_bias = (
-            attention.time_bias_scale.tanh()
-            * attention.rbf_proj(features)
-        )
-        torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))
-
-        (initial_bias * target).sum().backward()
-        projection_grad = attention.rbf_proj.weight.grad
-        self.assertIsNotNone(projection_grad)
-        self.assertGreater(projection_grad.abs().sum().item(), 0.0)
-
-        with torch.no_grad():
-            attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
-        attention.zero_grad(set_to_none=True)
-        updated_bias = (
-            attention.time_bias_scale.tanh()
-            * attention.rbf_proj(features)
-        )
-        (updated_bias * target).sum().backward()
-
-        scale_grad = attention.time_bias_scale.grad
-        self.assertIsNotNone(scale_grad)
-        self.assertGreater(scale_grad.abs().item(), 0.0)
-
-
-if __name__ == "__main__":
-    unittest.main()

test_traj_mixer.py

@@ -1,166 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import TrajMixer
-
-
-class TrajMixerTest(unittest.TestCase):
-    def test_default_shape_parameters_and_initialization(self) -> None:
-        mixer = TrajMixer(
-            n_embd=120,
-            n_head=10,
-            dropout=0.0,
-        )
-
-        x = torch.randn(2, 7, 120)
-        self.assertEqual(mixer(x).shape, x.shape)
-        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
-        self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
-        self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
-        torch.testing.assert_close(
-            torch.sigmoid(mixer.intra_gate_logits.detach()),
-            torch.full((10, 12), 0.1),
-        )
-        self.assertEqual(mixer.intra_hidden, 48)
-        self.assertEqual(
-            tuple(mixer.intra_gate_proj.shape),
-            (10, 12, 48),
-        )
-        self.assertEqual(
-            tuple(mixer.intra_value_proj.shape),
-            (10, 12, 48),
-        )
-        self.assertEqual(
-            tuple(mixer.intra_output_proj.shape),
-            (10, 48, 12),
-        )
-        self.assertEqual(mixer.hidden_group, 40)
-        self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
-        self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
-        self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
-
-    def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        with torch.no_grad():
-            mixer.output_proj.zero_()
-        x = torch.randn(2, 5, 12)
-        torch.testing.assert_close(mixer(x), x)
-
-    def test_forward_matches_single_outer_residual_formula(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        mixer.eval()
-        x = torch.randn(2, 5, 12)
-
-        grouped = mixer.norm(x).reshape(2, 5, 3, 4)
-        intra_output = mixer._intra_mix(grouped)
-        static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
-            1, 1, 3, 4
-        )
-        mixed_input = grouped + static_gate * intra_output
-        update = mixer._cross_mix(mixed_input).reshape(2, 5, 12)
-
-        torch.testing.assert_close(mixer(x), x + update)
-
-    def test_intra_stage_is_independent_across_groups(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        mixer.eval()
-
-        grouped = torch.randn(2, 4, 3, 4)
-        changed = grouped.clone()
-        changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :])
-
-        original_out = mixer._intra_mix(grouped)
-        changed_out = mixer._intra_mix(changed)
-        unchanged_groups = torch.tensor([0, 2])
-        torch.testing.assert_close(
-            original_out.index_select(2, unchanged_groups),
-            changed_out.index_select(2, unchanged_groups),
-        )
-
-    def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
-        mixer = TrajMixer(6, n_head=3, dropout=0.0)
-        mixer.eval()
-        with torch.no_grad():
-            mixer.gate_proj.zero_()
-            mixer.value_proj.zero_()
-            mixer.output_proj.zero_()
-
-            # Coordinate 0 reads group 0 through hidden unit 0 and writes it
-            # into group 1. Coordinate 1 must remain independent.
-            mixer.gate_proj[0, 0, 0] = 1.0
-            mixer.value_proj[0, 0, 0] = 1.0
-            mixer.output_proj[0, 0, 1] = 1.0
-
-        grouped = torch.tensor(
-            [[[
-                [-1.0, 4.0],
-                [0.0, 5.0],
-                [1.0, 6.0],
-            ]]]
-        )
-        changed = grouped.clone()
-        changed[0, 0, 0, 0] = 2.0
-
-        original_out = mixer._cross_mix(grouped)
-        changed_out = mixer._cross_mix(changed)
-
-        self.assertNotEqual(
-            original_out[0, 0, 1, 0].item(),
-            changed_out[0, 0, 1, 0].item(),
-        )
-        torch.testing.assert_close(
-            original_out[..., 1],
-            changed_out[..., 1],
-        )
-
-    def test_mixer_does_not_mix_sequence_positions(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        mixer.eval()
-        x = torch.randn(2, 5, 12)
-        changed = x.clone()
-        changed[:, 3, :] += torch.randn_like(changed[:, 3, :])
-
-        original_out = mixer(x)
-        changed_out = mixer(changed)
-        unchanged_positions = torch.tensor([0, 1, 2, 4])
-        torch.testing.assert_close(
-            original_out.index_select(1, unchanged_positions),
-            changed_out.index_select(1, unchanged_positions),
-        )
-
-    def test_gradients_reach_every_projection_family(self) -> None:
-        torch.manual_seed(1)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        x = torch.randn(2, 4, 12, requires_grad=True)
-
-        mixer(x).square().mean().backward()
-
-        self.assertIsNotNone(x.grad)
-        self.assertTrue(torch.isfinite(x.grad).all())
-        self.assertGreater(x.grad.abs().sum().item(), 0.0)
-        projection_names = (
-            "intra_gate_proj",
-            "intra_value_proj",
-            "intra_output_proj",
-            "gate_proj",
-            "value_proj",
-            "output_proj",
-        )
-        for name in projection_names:
-            parameter = getattr(mixer, name)
-            self.assertIsNotNone(parameter.grad, name)
-            self.assertTrue(torch.isfinite(parameter.grad).all(), name)
-            self.assertGreater(parameter.grad.abs().sum().item(), 0.0, name)
-
-    def test_invalid_group_partition_is_rejected(self) -> None:
-        with self.assertRaisesRegex(ValueError, "divisible"):
-            TrajMixer(n_embd=121, n_head=10)
-
-
-if __name__ == "__main__":
-    unittest.main()

train_all_future.py

@@ -37,8 +37,10 @@ from train_util import (
     configure_torch_for_training,
     create_unique_run_dir,
     format_extra_info_types,
+    get_lr,
     get_model_parameter_counts,
     load_extra_info_types_file,
+    move_batch_to_device,
     resolve_device,
     save_checkpoint,
     save_config,
@@ -135,24 +137,6 @@ def parse_args() -> argparse.Namespace:
     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,
@@ -214,7 +198,6 @@ def compute_all_future_loss(
         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)
 

train_next_step.py

@@ -1,10 +1,4 @@
-"""
-Train DeepHealth with next-token / next-time-point supervision.
-
-The next-step dataset uses observed event histories, including CHECKUP state
-tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
-same-time group ends.
-"""
+"""Reproduce Delphi2M with absolute-time next-token supervision."""
 from __future__ import annotations
 
 import argparse
@@ -29,14 +23,15 @@ from model_architectures import (
     SUPPORTED_MODEL_ARCHITECTURES,
 )
 from models import DeepHealth, DeepHealthOutput
-from readouts import build_readout
-from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+from targets import CHECKUP_IDX, PAD_IDX
 from train_util import (
     configure_torch_for_training,
     create_unique_run_dir,
     format_extra_info_types,
+    get_lr,
     get_model_parameter_counts,
     load_extra_info_types_file,
+    move_batch_to_device,
     resolve_device,
     save_checkpoint,
     save_config,
@@ -70,7 +65,6 @@ def parse_args() -> argparse.Namespace:
     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)
@@ -85,8 +79,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
-    parser.add_argument("--time_mode", type=str, default="relative",
-                        choices=["relative", "absolute"])
     parser.add_argument("--dropout", type=float, default=0.0)
     parser.add_argument(
         "--model_architecture",
@@ -95,17 +87,10 @@ def parse_args() -> argparse.Namespace:
         choices=SUPPORTED_MODEL_ARCHITECTURES,
     )
 
-    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)
@@ -127,11 +112,6 @@ def parse_args() -> argparse.Namespace:
     )
     if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
         raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
-    if args.target_mode == "uts":
-        args.readout_name = args.readout_name or "same_time_group_end"
-        args.include_no_event_in_uts_target = True
-    else:
-        args.readout_name = args.readout_name or "token"
     args.extra_info_types = (
         load_extra_info_types_file(args.extra_info_types_file)
         if args.extra_info_types_file is not None
@@ -140,24 +120,6 @@ def parse_args() -> argparse.Namespace:
     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,
@@ -171,44 +133,27 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         n_bins=args.n_bins,
         extra_pool_reduce=args.extra_pool_reduce,
         target_mode="next_token",
-        time_mode=args.time_mode,
+        time_mode="absolute",
         dist_mode="exponential",
         dropout=args.dropout,
         model_architecture=args.model_architecture,
     )
 
 
-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},
+        "delphi2m",
+        ignored_tokens={PAD_IDX, CHECKUP_IDX},
         t_min=args.t_min,
         max_exp_input=args.max_exp_input,
+        ce_weight=args.ce_weight,
+        time_weight=args.time_weight,
     )
 
 
 def build_augmented_next_step_targets(
     batch_cpu: Dict[str, torch.Tensor],
     model_out: DeepHealthOutput,
-    include_uts_targets: bool,
 ) -> Dict[str, torch.Tensor]:
     hidden_len = model_out.hidden.size(1)
     event_len = int(model_out.event_len)
@@ -216,26 +161,12 @@ def build_augmented_next_step_targets(
     device = model_out.hidden.device
     non_blocking = device.type == "cuda"
     if extra_len <= 0:
-        targets = {
+        return {
             "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
             "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
-            "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
         }
-        if include_uts_targets:
-            targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
-                device, non_blocking=non_blocking
-            )
-            targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
-                device, non_blocking=non_blocking
-            )
-        return targets
 
     bsz = batch_cpu["target_event_seq"].size(0)
-    vocab_size = (
-        batch_cpu["target_multi_hot"].size(2)
-        if include_uts_targets
-        else None
-    )
     other_valid = batch_cpu["other_type"] > 0
     extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
     extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
@@ -268,34 +199,6 @@ def build_augmented_next_step_targets(
         ],
         dim=1,
     )
-    readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
-    target_dt_unique = None
-    target_multi_hot = None
-    if include_uts_targets:
-        target_dt_unique = torch.cat(
-            [
-                batch_cpu["target_dt_unique"],
-                torch.zeros(
-                    bsz,
-                    extra_len,
-                    dtype=batch_cpu["target_dt_unique"].dtype,
-                ),
-            ],
-            dim=1,
-        )
-        target_multi_hot = torch.cat(
-            [
-                batch_cpu["target_multi_hot"],
-                torch.zeros(
-                    bsz,
-                    extra_len,
-                    vocab_size,
-                    dtype=batch_cpu["target_multi_hot"].dtype,
-                ),
-            ],
-            dim=1,
-        )
-
     for b in range(bsz):
         valid_event = batch_cpu["padding_mask"][b].bool()
         if not valid_event.any():
@@ -326,7 +229,6 @@ def build_augmented_next_step_targets(
             t = extra_time[b, j]
             future = times > t
             if not future.any():
-                readout_mask[b, pos] = False
                 continue
 
             first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
@@ -335,35 +237,15 @@ def build_augmented_next_step_targets(
             target_event_seq[b, pos] = next_event
             target_time_seq[b, pos] = next_time
 
-            if not include_uts_targets:
-                continue
-
-            same_next_time = times == next_time
-            next_events = events[same_next_time]
-            valid_next_events = next_events[
-                (next_events > PAD_IDX) & (next_events < vocab_size)
-            ].long()
-            if valid_next_events.numel() == 0:
-                readout_mask[b, pos] = False
-                continue
-            target_multi_hot[b, pos, valid_next_events] = True
-            target_dt_unique[b, pos] = next_time - t
-
-    targets = {
+    return {
         "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
         "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
-        "readout_mask": readout_mask.to(device, non_blocking=non_blocking),
     }
-    if include_uts_targets:
-        targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
-        targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
-    return targets
 
 
 def compute_next_step_loss(
     args: argparse.Namespace,
     model: DeepHealth,
-    readout,
     criterion,
     batch: Dict[str, torch.Tensor],
     device: torch.device,
@@ -382,7 +264,6 @@ def compute_next_step_loss(
         other_value=batch["other_value"],
         other_value_kind=batch["other_value_kind"],
         other_time=batch["other_time"],
-        target_mode="next_token",
         return_output=True,
     )
     if not isinstance(model_out, DeepHealthOutput):
@@ -390,35 +271,17 @@ def compute_next_step_loss(
     targets = build_augmented_next_step_targets(
         batch_cpu=batch_cpu,
         model_out=model_out,
-        include_uts_targets=args.target_mode == "uts",
     )
-    readout_out = readout(
-        hidden=model_out.hidden,
-        time_seq=model_out.time_seq,
-        padding_mask=model_out.padding_mask,
-        readout_mask=targets["readout_mask"]
-        if args.readout_name == "same_time_group_end"
-        else None,
-    )
-    logits = model.calc_risk(readout_out.hidden)
+    logits = model.calc_risk(model_out.hidden)
 
-    if args.target_mode == "delphi2m":
-        loss, parts = criterion(
-            logits=logits,
-            target_events=targets["target_event_seq"],
-            target_times=targets["target_time_seq"],
-            current_times=model_out.time_seq,
-            padding_mask=readout_out.readout_mask,
-            return_components=True,
-        )
-    else:
-        loss, parts = criterion(
-            logits=logits,
-            target_multi_hot=targets["target_multi_hot"],
-            target_dt_unique=targets["target_dt_unique"],
-            readout_mask=readout_out.readout_mask,
-            return_components=True,
-        )
+    loss, parts = criterion(
+        logits=logits,
+        target_events=targets["target_event_seq"],
+        target_times=targets["target_time_seq"],
+        current_times=model_out.time_seq,
+        padding_mask=model_out.padding_mask,
+        return_components=True,
+    )
     if not torch.isfinite(loss):
         raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
     return loss, parts
@@ -428,7 +291,6 @@ def run_epoch(
     logger: logging.Logger,
     args: argparse.Namespace,
     model: DeepHealth,
-    readout,
     criterion,
     loader: DataLoader,
     optimizer: AdamW | None,
@@ -436,7 +298,6 @@ def run_epoch(
     is_train: bool,
 ) -> float:
     model.train(is_train)
-    readout.train(is_train)
     total = torch.zeros((), device=device)
     n_batches = 0
     skipped = 0
@@ -447,7 +308,9 @@ def run_epoch(
     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)
+            loss, parts = compute_next_step_loss(
+                args, model, criterion, batch, device
+            )
             if is_train:
                 if optimizer is None:
                     raise ValueError("optimizer is required for training")
@@ -497,7 +360,8 @@ def build_metadata(
         "model_class": "DeepHealth",
         "model_architecture": args.model_architecture,
         "model_target_mode": "next_token",
-        "target_mode": args.target_mode,
+        "target_mode": "delphi2m",
+        "time_mode": "absolute",
         "dist_mode": "exponential",
         "extra_info_types_file": (
             Path(args.extra_info_types_file).name
@@ -518,8 +382,8 @@ def build_metadata(
             "val": int(len(val_subset)),
             "test": int(len(test_subset)),
         },
-        "resolved_readout_name": args.readout_name,
-        "resolved_loss_name": args.target_mode,
+        "resolved_readout_name": "token",
+        "resolved_loss_name": "delphi2m",
     }
 
 
@@ -531,7 +395,7 @@ def main() -> None:
 
     run_dir, run_name = create_unique_run_dir(
         lambda timestamp: (
-            f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
+            "absolute_exponential_next_token_delphi2m_"
             f"gap_{args.no_event_interval_years:g}y_{timestamp}"
         ),
         runs_root=Path(args.runs_root) / args.model_architecture,
@@ -542,13 +406,12 @@ def main() -> None:
     logger.info(f"Device: {device}")
     logger.info(f"Model architecture: {args.model_architecture}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
-    logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
+    logger.info("time_mode=absolute, readout=token, target_mode=delphi2m")
 
     dataset = HealthDataset(
         data_prefix=args.data_prefix,
         labels_file=args.labels_file,
         no_event_interval_years=args.no_event_interval_years,
-        include_no_event_in_uts_target=args.include_no_event_in_uts_target,
         extra_info_types=args.extra_info_types,
     )
     if args.train_eid_file and args.val_eid_file and args.test_eid_file:
@@ -616,7 +479,6 @@ def main() -> None:
         f"total={parameter_counts['model_parameter_count']:,}, "
         f"trainable={parameter_counts['trainable_parameter_count']:,}"
     )
-    readout = build_next_step_readout(args).to(device)
     criterion = build_next_step_loss(args)
     optimizer = AdamW(
         model.parameters(),
@@ -646,9 +508,13 @@ def main() -> None:
         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)
+        train_loss = run_epoch(
+            logger, args, model, criterion, train_loader, optimizer, device, True
+        )
         with torch.no_grad():
-            val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False)
+            val_loss = run_epoch(
+                logger, args, model, criterion, val_loader, None, device, False
+            )
 
         is_best = val_loss < best_val
         if is_best:
@@ -682,7 +548,9 @@ def main() -> None:
     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)
+        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}")
 

train_util.py

@@ -6,6 +6,7 @@ import sys
 import time
 import csv
 from datetime import datetime
+import math
 from pathlib import Path
 from typing import Any, Dict, Iterable, Tuple
 
@@ -141,6 +142,31 @@ def resolve_device(device_arg: str) -> torch.device:
     raise ValueError(f"Unsupported device: {device_arg}")
 
 
+def get_lr(epoch: int, args: Any, 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 split_dataset(
     dataset: HealthDataset,
     train_ratio: float,
@@ -291,15 +317,6 @@ def split_all_future_datasets_by_eid_files(
     )
 
 
-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 get_model_parameter_counts(model: torch.nn.Module) -> Dict[str, int]:
     """Return stable total and trainable parameter counts."""
     return {