Refactor AUC evaluation scripts to support model target modes and improve distribution handling

README.md

@@ -261,11 +261,17 @@ python train.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt
 
 `evaluate_auc.py` 评估的是 **next-step / token-level 预测位置上的疾病 AUC**。
 
+查询方式由 `train_config.json` 中的 `model_target_mode` 决定：
+
+- `model_target_mode="next_token"`：使用训练 readout 对应的历史 token hidden 作为预测点。
+- `model_target_mode="all_future"`：不使用 readout token，直接把每个预测点年龄作为 `t_query` 传入模型，取 query hidden 作为预测点。
+
 核心流程：
 
 - 按训练配置重新构建 `HealthDataset` 和 `DeepHealth`。
-- 对评估 split 中的患者做一次模型推理，缓存每个 disease-token readout hidden。
+- 对评估 split 中的患者做模型推理，缓存每个预测点的 hidden。
 - 对疾病 token 分块投影到 `risk_head`，避免一次性保存全词表 logits。
+- AUC score 使用疾病对应的 eta/logit 排序分数；`dist_mode` 只用于正确构建模型，不会把分数转换成 horizon-specific risk probability。
 - 对每个疾病、性别、年龄段、prediction offset 分别计算 AUC。
 - 输出未池化分层结果和按疾病汇总后的结果。
 
@@ -294,13 +300,20 @@ python evaluate_auc.py \
 
 `evaluate_auc_v2.py` 评估的是 **landmark fixed-horizon incident disease AUC**。
 
-它不是使用已有序列中的普通 readout 位置，而是在指定 landmark age 人工插入一个 `<NO_EVENT>` query token，然后评估该 landmark 后固定 horizon 内是否发生 incident disease。
+它不是使用已有序列中的普通 readout 位置，而是在指定 landmark age 构造一个 landmark query，然后评估该 landmark 后固定 horizon 内是否发生 incident disease。
+
+查询方式由 `train_config.json` 中的 `model_target_mode` 决定：
+
+- `model_target_mode="next_token"`：在 landmark age 人工插入一个 `<NO_EVENT>` token，取该 token 的 hidden 做风险分数。
+- `model_target_mode="all_future"`：不插入 `<NO_EVENT>`，直接把 landmark age 作为 `t_query` 传入模型，取 query hidden 做风险分数。
 
 核心流程：
 
 - 为每个患者和 landmark age 构造 landmark query 样本。
-- 在 landmark age 插入 `<NO_EVENT>` token，取该位置 hidden。
+- 根据模型模式插入 `<NO_EVENT>` token 或直接传 `t_query`，取 landmark/query hidden。
-- 对疾病 token 分块投影到 `risk_head`。
+- 对疾病 token 分块投影到 `risk_head`；`score_mode="risk"` 时会根据 `dist_mode` 把线性输出转换为固定 horizon 风险概率。
+- 分布转换规则与 all-future 训练损失一致：`exponential` 使用 `1 - exp(-rate * horizon)`；`weibull` 使用 `1 - exp(-rate * horizon ** rho)`；`mixed` 中普通疾病使用 exponential，死亡 endpoint 使用 Weibull death rho。
+- `score_mode="eta"` 是诊断用排序分数，不使用 `rho`，因此不区分不同分布的风险曲线。
 - 按疾病、性别、landmark age、horizon 计算 incident disease AUC。
 - 可选择排除 horizon 内先于目标疾病发生的死亡竞争风险。
 
@@ -333,8 +346,9 @@ python evaluate_auc_v2.py \
 | 项目 | `evaluate_auc.py` | `evaluate_auc_v2.py` |
 | --- | --- | --- |
 | 评估口径 | next-step/token-level 预测点 | landmark fixed-horizon incident risk |
-| 查询位置 | 原始序列中满足 offset 条件的最新 readout token | 人工插入的 `<NO_EVENT>` landmark token |
+| 查询位置 | next-token 用满足 offset 条件的最新 readout token；all-future 直接用该预测点年龄作为 `t_query` | next-token 用人工插入的 `<NO_EVENT>` landmark token；all-future 直接用 `t_query` |
 | 时间参数 | `offsets`：预测点至少早于目标事件多少年 | `landmark_*` 和 `horizons`：固定年龄点与未来窗口 |
+| score 与分布 | 使用 eta/logit 排序分数；不按 `dist_mode` 转换风险概率 | `score_mode="risk"` 按 `dist_mode` 区分 exponential / Weibull / mixed；`score_mode="eta"` 不区分分布 |
 | 病例定义 | target table 中出现目标疾病的患者/事件 | landmark 后 horizon 内首次发生目标疾病 |
 | 对照定义 | 从未出现该疾病的患者的 eligible target occurrence | landmark 时未患病，且 horizon 内未发病并有足够随访 |
 | 分层 | sex + age bracket + offset | sex + landmark age + horizon |
@@ -380,7 +394,11 @@ python evaluate_auc_v2.py \
 - `evaluate_auc.py`
   - next-step/token-level 疾病 AUC 评估
   - 使用 prediction offset、sex、age bracket 分层
+  - next-token 模型使用 readout hidden；all-future 模型使用 `t_query` hidden
+  - score 是疾病 eta/logit，不按分布转换为固定 horizon 风险概率
 
 - `evaluate_auc_v2.py`
   - landmark fixed-horizon incident disease AUC 评估
-  - 通过插入 `<NO_EVENT>` landmark token 查询固定年龄点风险
+  - next-token 模型通过插入 `<NO_EVENT>` landmark token 查询固定年龄点风险
+  - all-future 模型直接通过 `t_query` 查询固定年龄点风险
+  - `score_mode="risk"` 按 exponential / Weibull / mixed 分布计算固定 horizon 风险

evaluate_auc.py

@@ -309,6 +309,12 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> 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}"
+        )
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
@@ -320,7 +326,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         n_categories=dataset.n_categories,
         cont_type_ids=dataset.cont_type_ids,
         n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
-        target_mode="next_token",
+        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)),
@@ -348,15 +354,25 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
     mode = str(cfg_dist_mode).lower()
     has_rho_head = any(str(k).startswith("rho_head.")
                        for k in state_dict.keys())
+    has_rho_death_head = any(str(k).startswith("rho_death_head.")
+                             for k in state_dict.keys())
 
     if has_rho_head and mode != "weibull":
         print(
             "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
         return "weibull"
+    if has_rho_death_head and mode != "mixed":
+        print(
+            "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
+        return "mixed"
     if (not has_rho_head) and mode == "weibull":
         print(
             "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
         return "exponential"
+    if (not has_rho_death_head) and mode == "mixed":
+        print(
+            "[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
+        return "exponential"
     return mode
 
 
@@ -452,25 +468,27 @@ def infer_readout_hidden(
     model: DeepHealth,
     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]]:
-    """
+    """Cache per-position hidden states used by the unchanged AUC logic."""
-    Run the expensive transformer/readout path exactly once and cache hidden states.
+    model_target_mode = str(model_target_mode).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}"
+        )
 
-    The older implementation re-ran the whole model once per disease chunk even
+    readout = None
-    though only the final tied risk-head columns changed. That is usually the
+    if model_target_mode == "next_token" and readout_name == "same_time_group_end":
-    dominant bottleneck. This function computes readout hidden states once; later
-    chunks only perform a cheap selected-vocabulary linear projection.
-    """
-    if readout_name == "same_time_group_end":
         readout = build_readout("same_time_group_end",
                                 reduce=readout_reduce).to(device)
-    else:
+    elif model_target_mode == "next_token":
         readout = build_readout(readout_name).to(device)
-    readout.eval()
+    if readout is not None:
+        readout.eval()
 
     hidden_parts: List[np.ndarray] = []
     arrays: Dict[str, List[np.ndarray]] = {
@@ -501,25 +519,55 @@ def infer_readout_hidden(
             if autocast_enabled else contextlib.nullcontext()
         )
         with amp_context:
-            hidden = model(
+            if model_target_mode == "all_future":
-                event_seq=event_seq,
+                batch_size, seq_len = event_seq.shape
-                time_seq=time_seq,
+                hidden = torch.zeros(
-                sex=batch_dev["sex"],
+                    batch_size,
-                padding_mask=padding_mask,
+                    seq_len,
-                other_type=batch_dev["other_type"],
+                    model.n_embd,
-                other_value=batch_dev["other_value"],
+                    device=event_seq.device,
-                other_value_kind=batch_dev["other_value_kind"],
+                    dtype=torch.float32,
-                other_time=batch_dev["other_time"],
+                )
-                target_mode="next_token",
+                for pos in range(seq_len):
-            )
+                    active = padding_mask[:, pos].bool()
-            ro = readout(
+                    if not active.any():
-                hidden=hidden,
+                        continue
-                time_seq=time_seq,
+                    hidden_pos = model(
-                padding_mask=padding_mask,
+                        event_seq=event_seq[active],
-                readout_mask=batch_dev["readout_mask"],
+                        time_seq=time_seq[active],
-            )
+                        sex=batch_dev["sex"][active],
+                        padding_mask=padding_mask[active],
+                        t_query=time_seq[active, pos],
+                        other_type=batch_dev["other_type"][active],
+                        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()
+            else:
+                hidden_raw = model(
+                    event_seq=event_seq,
+                    time_seq=time_seq,
+                    sex=batch_dev["sex"],
+                    padding_mask=padding_mask,
+                    other_type=batch_dev["other_type"],
+                    other_value=batch_dev["other_value"],
+                    other_value_kind=batch_dev["other_value_kind"],
+                    other_time=batch_dev["other_time"],
+                    target_mode="next_token",
+                )
+                ro = readout(
+                    hidden=hidden_raw,
+                    time_seq=time_seq,
+                    padding_mask=padding_mask,
+                    readout_mask=batch_dev["readout_mask"],
+                )
+                hidden = ro.hidden
+                readout_mask_np = ro.readout_mask.detach().cpu().numpy()
 
-        h = ro.hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
+        h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
         hidden_parts.append(h)
         max_len = max(max_len, h.shape[1])
 
@@ -529,7 +577,7 @@ def infer_readout_hidden(
                     batch[k].cpu().numpy().astype(np.int8, copy=False))
             else:
                 arrays[k].append(batch[k].cpu().numpy())
-        arrays["readout_mask"][-1] = ro.readout_mask.detach().cpu().numpy()
+        arrays["readout_mask"][-1] = readout_mask_np
 
     def pad_3d(parts: List[np.ndarray], fill: float = 0.0) -> np.ndarray:
         out = np.full(
@@ -999,6 +1047,7 @@ def evaluate_auc_pipeline(
     age_groups: np.ndarray,
     offsets: Sequence[float],
     device: torch.device,
+    model_target_mode: str,
     readout_name: str,
     readout_reduce: str,
     num_workers_auc: int,
@@ -1058,6 +1107,7 @@ def evaluate_auc_pipeline(
         model=model,
         loader=loader,
         device=device,
+        model_target_mode=model_target_mode,
         readout_name=readout_name,
         readout_reduce=readout_reduce,
         use_amp=use_amp,
@@ -1257,6 +1307,12 @@ def main() -> None:
     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(
+            "train_config.json model_target_mode must be next_token or all_future, "
+            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")
@@ -1298,7 +1354,13 @@ def main() -> None:
     dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
     cfg = dict(cfg)
     cfg["dist_mode"] = dist_mode
+    cfg["model_target_mode"] = model_target_mode
     print(f"Resolved dist_mode for evaluation: {dist_mode}")
+    print(f"Model target mode for AUC: {model_target_mode}")
+    print(
+        "AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
+        "dist_mode affects model loading but is not converted to horizon-specific risk probability."
+    )
 
     model = build_model_from_dataset(args, cfg, dataset).to(device)
     load_model_state(model, str(model_ckpt_path),
@@ -1335,6 +1397,7 @@ def main() -> None:
         age_groups=age_groups,
         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(

evaluate_auc_v2.py

@@ -1,10 +1,11 @@
 """Evaluate landmark fixed-horizon incident disease AUC for DeepHealth.
 
-This script is intentionally strict and supports only:
+This script supports DeepHealth fixed-horizon risk scores for exponential,
-    DeepHealth + exponential distribution + no_event imputation.
+Weibull, and mixed all-future distributions.
 
-Landmark querying is implemented by inserting a <NO_EVENT> token at landmark age.
+Landmark querying depends on the model target mode saved in train_config.json:
-No t_query interface is used.
+    - next_token: insert a <NO_EVENT> token at landmark age and read it out;
+    - all_future: pass landmark age directly as t_query.
 """
 from __future__ import annotations
 
@@ -21,6 +22,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple
 import numpy as np
 import pandas as pd
 import torch
+import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import DataLoader, Dataset
 from tqdm.auto import tqdm
@@ -140,12 +142,36 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
     mode = str(cfg_dist_mode).lower()
     has_rho_head = any(str(k).startswith("rho_head.")
                        for k in state_dict.keys())
+    has_rho_death_head = any(str(k).startswith("rho_death_head.")
+                             for k in state_dict.keys())
     if has_rho_head:
+        if mode != "weibull":
+            print(
+                "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
         return "weibull"
-    return mode if mode in {"exponential", "weibull"} else "exponential"
+    if has_rho_death_head:
+        if mode != "mixed":
+            print(
+                "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
+        return "mixed"
+    if mode == "weibull":
+        print(
+            "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
+        return "exponential"
+    if mode == "mixed":
+        print(
+            "[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
+        return "exponential"
+    return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> 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}"
+        )
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
@@ -157,7 +183,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         n_categories=dataset.n_categories,
         cont_type_ids=dataset.cont_type_ids,
         n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
-        target_mode="next_token",
+        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)),
@@ -489,6 +515,7 @@ class LandmarkDataset(Dataset):
         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]],
         death_token_ids: Sequence[int],
@@ -497,6 +524,12 @@ class LandmarkDataset(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(
+                "model_target_mode must be next_token or all_future, got "
+                f"{self.model_target_mode!r}"
+            )
         self.min_history_events = int(min_history_events)
 
         self.first_occurrence_by_token = first_occurrence_by_token
@@ -555,25 +588,33 @@ class LandmarkDataset(Dataset):
                 if valid_history_mask.sum() < self.min_history_events:
                     continue
 
-                event_seq_landmark = np.concatenate(
+                if self.model_target_mode == "next_token":
-                    [
+                    event_seq_landmark = np.concatenate(
-                        prefix_events.astype(np.int64, copy=False),
+                        [
-                        np.array([NO_EVENT_IDX], dtype=np.int64),
+                            prefix_events.astype(np.int64, copy=False),
-                    ]
+                            np.array([NO_EVENT_IDX], dtype=np.int64),
-                )
+                        ]
-                time_seq_landmark = np.concatenate(
-                    [
-                        prefix_times.astype(np.float32, copy=False),
-                        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 = np.concatenate(
-                readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
+                        [
-                readout_mask[-1] = True
+                            prefix_times.astype(np.float32, copy=False),
+                            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
+                        )
+                    landmark_pos = int(len(event_seq_landmark) - 1)
+                    readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
+                    readout_mask[-1] = True
+                else:
+                    event_seq_landmark = prefix_events.astype(
+                        np.int64, copy=False)
+                    time_seq_landmark = prefix_times.astype(
+                        np.float32, copy=False)
+                    landmark_pos = int(len(event_seq_landmark) - 1)
+                    readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
 
                 rows.append(
                     {
@@ -583,7 +624,8 @@ class LandmarkDataset(Dataset):
                         "landmark_age": np.float32(landmark_age),
                         "followup_end_time": np.float32(followup_end),
                         "death_time": np.float32(self.patient_death_time[patient_id]),
-                        "landmark_pos": int(len(event_seq_landmark) - 1),
+                        "landmark_pos": landmark_pos,
+                        "t_query": np.float32(landmark_age),
                         "event_seq": event_seq_landmark,
                         "time_seq": time_seq_landmark,
                         "readout_mask": readout_mask,
@@ -616,6 +658,7 @@ class LandmarkDataset(Dataset):
             "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
             "other_time": torch.from_numpy(s["other_time"]).float(),
             "landmark_pos": torch.tensor(s["landmark_pos"], dtype=torch.long),
+            "t_query": torch.tensor(float(s["t_query"]), dtype=torch.float32),
             "patient_id": torch.tensor(s["patient_id"], dtype=torch.long),
             "landmark_age": torch.tensor(float(s["landmark_age"]), dtype=torch.float32),
             "followup_end_time": torch.tensor(float(s["followup_end_time"]), dtype=torch.float32),
@@ -650,6 +693,7 @@ def collate_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch
         "other_value_kind": other_value_kind,
         "other_time": other_time,
         "landmark_pos": torch.stack([x["landmark_pos"] for x in batch]),
+        "t_query": torch.stack([x["t_query"] for x in batch]),
         "patient_id": torch.stack([x["patient_id"] for x in batch]),
         "landmark_age": torch.stack([x["landmark_age"] for x in batch]),
         "followup_end_time": torch.stack([x["followup_end_time"] for x in batch]),
@@ -672,17 +716,26 @@ def infer_landmark_hidden(
     model: DeepHealth,
     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]]:
-    if readout_name == "same_time_group_end":
+    model_target_mode = str(model_target_mode).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}"
+        )
+
+    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)
-    else:
+    elif model_target_mode == "next_token":
         readout = build_readout(readout_name).to(device)
-    readout.eval()
+    if readout is not None:
+        readout.eval()
 
     hidden_parts: List[np.ndarray] = []
     arrays = {
@@ -709,29 +762,43 @@ def infer_landmark_hidden(
             else contextlib.nullcontext()
         )
         with amp_ctx:
-            hidden = model(
+            if model_target_mode == "all_future":
-                event_seq=batch_dev["event_seq"],
+                landmark_hidden = model(
-                time_seq=batch_dev["time_seq"],
+                    event_seq=batch_dev["event_seq"],
-                sex=batch_dev["sex"],
+                    time_seq=batch_dev["time_seq"],
-                padding_mask=batch_dev["padding_mask"],
+                    sex=batch_dev["sex"],
-                other_type=batch_dev["other_type"],
+                    padding_mask=batch_dev["padding_mask"],
-                other_value=batch_dev["other_value"],
+                    t_query=batch_dev["t_query"],
-                other_value_kind=batch_dev["other_value_kind"],
+                    other_type=batch_dev["other_type"],
-                other_time=batch_dev["other_time"],
+                    other_value=batch_dev["other_value"],
-                target_mode="next_token",
+                    other_value_kind=batch_dev["other_value_kind"],
-            )
+                    other_time=batch_dev["other_time"],
-            readout_out = readout(
+                    target_mode="all_future",
-                hidden=hidden,
+                )
-                time_seq=batch_dev["time_seq"],
+            else:
-                padding_mask=batch_dev["padding_mask"],
+                hidden = model(
-                readout_mask=batch_dev["readout_mask"],
+                    event_seq=batch_dev["event_seq"],
-            )
+                    time_seq=batch_dev["time_seq"],
-
+                    sex=batch_dev["sex"],
-        landmark_hidden = readout_out.hidden.gather(
+                    padding_mask=batch_dev["padding_mask"],
-            1,
+                    other_type=batch_dev["other_type"],
-            batch_dev["landmark_pos"].long()[:, None, None].expand(-1,
+                    other_value=batch_dev["other_value"],
-                                                                   1, readout_out.hidden.shape[-1]),
+                    other_value_kind=batch_dev["other_value_kind"],
-        ).squeeze(1)
+                    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(
+                    1,
+                    batch_dev["landmark_pos"].long()[:, None, None].expand(
+                        -1, 1, readout_out.hidden.shape[-1]
+                    ),
+                ).squeeze(1)
 
         hidden_parts.append(landmark_hidden.detach(
         ).cpu().numpy().astype(out_dtype, copy=False))
@@ -754,33 +821,76 @@ def infer_landmark_hidden(
 
 
 @torch.inference_mode()
-def project_logits_chunk(
+def project_distribution_chunk(
     model: DeepHealth,
     hidden_all: np.ndarray,
     disease_ids: Sequence[int],
+    dist_mode: str,
     device: torch.device,
     logit_batch_size: int,
     use_amp: bool,
-) -> np.ndarray:
+) -> Tuple[np.ndarray, Optional[np.ndarray]]:
     n = int(hidden_all.shape[0])
     logit_batch_size = max(1, int(logit_batch_size))
     disease_ids = [int(x) for x in disease_ids]
+    dist_mode = str(dist_mode).lower()
 
     compute_dtype = torch.float16 if (
         device.type == "cuda" and use_amp) else torch.float32
     weight = model.risk_head.weight[disease_ids].detach().to(
         device=device, dtype=compute_dtype)
+    bias = model.risk_head.bias[disease_ids].detach().to(
+        device=device, dtype=compute_dtype)
+    rho_weight = None
+    rho_bias = None
+    death_rho_weight = None
+    death_rho_bias = None
+    mixed_death_cols: List[int] = []
+    death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
+
+    if dist_mode == "weibull":
+        rho_weight = model.rho_head.weight[disease_ids].detach().to(
+            device=device, dtype=compute_dtype)
+        rho_bias = model.rho_head.bias[disease_ids].detach().to(
+            device=device, dtype=compute_dtype)
+    elif dist_mode == "mixed":
+        mixed_death_cols = [j for j, token in enumerate(disease_ids)
+                            if int(token) == death_idx]
+        if mixed_death_cols:
+            death_rho_weight = model.rho_death_head.weight.detach().to(
+                device=device, dtype=compute_dtype)
+            death_rho_bias = model.rho_death_head.bias.detach().to(
+                device=device, dtype=compute_dtype)
 
     out_parts: List[np.ndarray] = []
+    rho_parts: List[np.ndarray] = []
     for start in tqdm(range(0, n, logit_batch_size), desc="Risk projection", leave=False, dynamic_ncols=True):
         end = min(start + logit_batch_size, n)
         h = torch.from_numpy(hidden_all[start:end]).to(
             device=device, dtype=compute_dtype, non_blocking=True)
-        logits = torch.matmul(h, weight.t())
+        logits = torch.matmul(h, weight.t()) + bias
+        rho = None
+        if dist_mode == "weibull":
+            assert rho_weight is not None and rho_bias is not None
+            rho = F.softplus(torch.matmul(h, rho_weight.t()) + rho_bias) + 1e-6
+        elif dist_mode == "mixed" and mixed_death_cols:
+            assert death_rho_weight is not None and death_rho_bias is not None
+            rho = torch.ones_like(logits)
+            death_rho = F.softplus(
+                torch.matmul(h, death_rho_weight.t()).squeeze(-1) + death_rho_bias.squeeze(0)
+            ) + 1e-6
+            for col in mixed_death_cols:
+                rho[:, int(col)] = death_rho
+
         out_parts.append(logits.float().cpu(
         ).numpy().astype(np.float32, copy=False))
-        del h, logits
+        if rho is not None:
-    return np.concatenate(out_parts, axis=0)
+            rho_parts.append(rho.float().cpu(
+            ).numpy().astype(np.float32, copy=False))
+        del h, logits, rho
+    logits_all = np.concatenate(out_parts, axis=0)
+    rho_all = np.concatenate(rho_parts, axis=0) if rho_parts else None
+    return logits_all, rho_all
 
 
 # ---------------------------------------------------------------------------
@@ -794,6 +904,7 @@ _WORKER: Dict[str, Any] = {}
 def _init_worker(
     disease_ids: np.ndarray,
     score_chunk: np.ndarray,
+    rho_chunk: Optional[np.ndarray],
     row_patient_id: np.ndarray,
     row_sex: np.ndarray,
     row_landmark_age: np.ndarray,
@@ -805,6 +916,8 @@ def _init_worker(
     min_cases: int,
     exclude_death_competing: bool,
     death_token_ids: np.ndarray,
+    dist_mode: str,
+    model_death_idx: int,
 ) -> None:
     os.environ.setdefault("OMP_NUM_THREADS", "1")
     os.environ.setdefault("MKL_NUM_THREADS", "1")
@@ -816,6 +929,7 @@ def _init_worker(
         {
             "disease_ids": np.asarray(disease_ids, dtype=np.int64),
             "score_chunk": np.asarray(score_chunk, dtype=np.float32),
+            "rho_chunk": None if rho_chunk is None else np.asarray(rho_chunk, dtype=np.float32),
             "row_patient_id": np.asarray(row_patient_id, dtype=np.int32),
             "row_sex": np.asarray(row_sex, dtype=np.int8),
             "row_landmark_age": np.asarray(row_landmark_age, dtype=np.float32),
@@ -827,6 +941,8 @@ def _init_worker(
             "min_cases": int(min_cases),
             "exclude_death_competing": bool(exclude_death_competing),
             "death_token_ids": set(int(x) for x in np.asarray(death_token_ids, dtype=np.int64).tolist()),
+            "dist_mode": str(dist_mode).lower(),
+            "model_death_idx": int(model_death_idx),
             "first_time_cache": {},
         }
     )
@@ -846,11 +962,30 @@ def _first_time_by_patient(token: int) -> np.ndarray:
     return arr
 
 
-def _score_to_probability(logits: np.ndarray, score_mode: str, horizon: float) -> np.ndarray:
+def _score_to_probability(
+    logits: np.ndarray,
+    rho: Optional[np.ndarray],
+    score_mode: str,
+    horizon: float,
+    dist_mode: str,
+    token: int,
+    death_idx: int,
+) -> np.ndarray:
     if score_mode == "eta":
         return logits.astype(np.float64, copy=False)
     rate = np.log1p(np.exp(-np.abs(logits))) + np.maximum(logits, 0.0)
     rate = rate + np.float32(1e-8)
+    dist_mode = str(dist_mode).lower()
+    if dist_mode == "weibull":
+        if rho is None:
+            raise RuntimeError("Weibull risk scoring requires rho parameters.")
+        exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
+        return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
+    if dist_mode == "mixed" and int(token) == int(death_idx):
+        if rho is None:
+            raise RuntimeError("Mixed death risk scoring requires death rho parameters.")
+        exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
+        return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
     return (-np.expm1(-rate * np.float32(horizon))).astype(np.float64, copy=False)
 
 
@@ -864,6 +999,10 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
     row_followup_end = _WORKER["row_followup_end"]
     row_death_time = _WORKER["row_death_time"]
     logits_token = _WORKER["score_chunk"][:, int(j)]
+    rho_chunk = _WORKER["rho_chunk"]
+    rho_token = None if rho_chunk is None else rho_chunk[:, int(j)]
+    dist_mode = _WORKER["dist_mode"]
+    model_death_idx = int(_WORKER["model_death_idx"])
 
     first_time_patient = _first_time_by_patient(token)
     is_death_target = token in _WORKER["death_token_ids"]
@@ -917,9 +1056,23 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
                     continue
 
                 case_scores = _score_to_probability(
-                    logits_token[idx[case_idx]], score_mode=score_mode, horizon=horizon)
+                    logits_token[idx[case_idx]],
+                    None if rho_token is None else rho_token[idx[case_idx]],
+                    score_mode=score_mode,
+                    horizon=horizon,
+                    dist_mode=dist_mode,
+                    token=token,
+                    death_idx=model_death_idx,
+                )
                 control_scores = _score_to_probability(
-                    logits_token[idx[control_idx]], score_mode=score_mode, horizon=horizon)
+                    logits_token[idx[control_idx]],
+                    None if rho_token is None else rho_token[idx[control_idx]],
+                    score_mode=score_mode,
+                    horizon=horizon,
+                    dist_mode=dist_mode,
+                    token=token,
+                    death_idx=model_death_idx,
+                )
 
                 auc, auc_var = get_auc_delong_var(case_scores, control_scores)
                 if np.isnan(auc) or np.isnan(auc_var):
@@ -973,6 +1126,7 @@ def evaluate_landmark_auc(
     score_mode: str,
     horizons: np.ndarray,
     device: torch.device,
+    model_target_mode: str,
     readout_name: str,
     readout_reduce: str,
     num_workers_auc: int,
@@ -990,6 +1144,7 @@ def evaluate_landmark_auc(
         model=model,
         loader=loader,
         device=device,
+        model_target_mode=model_target_mode,
         readout_name=readout_name,
         readout_reduce=readout_reduce,
         use_amp=use_amp,
@@ -1007,10 +1162,11 @@ def evaluate_landmark_auc(
     all_rows: List[Dict[str, Any]] = []
     for chunk_idx, chunk in enumerate(tqdm(chunks, desc="Disease chunks", dynamic_ncols=True)):
         chunk_ids = chunk.tolist()
-        logits_chunk = project_logits_chunk(
+        logits_chunk, rho_chunk = project_distribution_chunk(
             model=model,
             hidden_all=hidden_all,
             disease_ids=chunk_ids,
+            dist_mode=dist_mode,
             device=device,
             logit_batch_size=logit_batch_size,
             use_amp=use_amp,
@@ -1024,6 +1180,7 @@ def evaluate_landmark_auc(
             _init_worker(
                 disease_ids=np.asarray(chunk_ids, dtype=np.int64),
                 score_chunk=logits_chunk,
+                rho_chunk=rho_chunk,
                 row_patient_id=row_arrays["patient_id"],
                 row_sex=row_arrays["sex"],
                 row_landmark_age=row_arrays["landmark_age"],
@@ -1036,6 +1193,8 @@ def evaluate_landmark_auc(
                 exclude_death_competing=exclude_death_competing,
                 death_token_ids=np.asarray(
                     landmark_dataset.death_token_ids, dtype=np.int64),
+                dist_mode=dist_mode,
+                model_death_idx=int(getattr(model, "death_idx", dataset.vocab_size - 1)),
             )
             nested = [_eval_token(t) for t in tqdm(
                 tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
@@ -1049,6 +1208,7 @@ def evaluate_landmark_auc(
                 initargs=(
                     np.asarray(chunk_ids, dtype=np.int64),
                     logits_chunk,
+                    rho_chunk,
                     row_arrays["patient_id"],
                     row_arrays["sex"],
                     row_arrays["landmark_age"],
@@ -1061,6 +1221,8 @@ def evaluate_landmark_auc(
                     exclude_death_competing,
                     np.asarray(landmark_dataset.death_token_ids,
                                dtype=np.int64),
+                    dist_mode,
+                    int(getattr(model, "death_idx", dataset.vocab_size - 1)),
                 ),
             ) as ex:
                 nested = list(
@@ -1078,7 +1240,7 @@ def evaluate_landmark_auc(
                 r["disease_chunk_idx"] = int(chunk_idx)
                 all_rows.append(r)
 
-        del logits_chunk
+        del logits_chunk, rho_chunk
 
     if not all_rows:
         raise RuntimeError(
@@ -1116,6 +1278,7 @@ def evaluate_landmark_auc(
             "model_ckpt_path",
             "config_path",
             "target_mode",
+            "model_target_mode",
             "dist_mode",
             "time_mode",
             "attn_mask_mode",
@@ -1194,6 +1357,12 @@ def main() -> None:
         "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(
+            "train_config.json model_target_mode must be next_token or all_future, "
+            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"))
@@ -1232,7 +1401,7 @@ def main() -> None:
         and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
         and dataset.vocab_size > NO_EVENT_IDX
     )
-    if not has_no_event:
+    if model_target_mode == "next_token" and not has_no_event:
         print(
             "[SKIP] This checkpoint/run does not support <NO_EVENT> imputation. "
             "Landmark AUC requires inserting a <NO_EVENT> query token. "
@@ -1283,20 +1452,14 @@ def main() -> None:
 
     state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
     dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
-    if dist_mode == "weibull":
+    if dist_mode not in {"exponential", "weibull", "mixed"}:
-        raise RuntimeError(
+        raise ValueError(
-            "Weibull checkpoints are not supported by evaluate_auc_v2.py. "
+            f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed."
-            "This landmark evaluation requires <NO_EVENT> imputation, and Weibull runs "
-            "in this project are treated as unsupported for no-event landmark AUC. "
-            "Please use an exponential no-event checkpoint."
-        )
-    if score_mode == "risk" and dist_mode != "exponential":
-        raise RuntimeError(
-            "score_mode='risk' requires dist_mode='exponential' for evaluate_auc_v2.py."
         )
 
     if score_mode == "eta":
-        print("WARNING: eta diagnostic score is not horizon-specific risk.")
+        print(
+            "WARNING: eta diagnostic score is not horizon-specific risk and does not use dist_mode-specific rho parameters.")
 
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
@@ -1308,7 +1471,13 @@ def main() -> None:
 
     model = build_model_from_dataset(args, cfg_model, dataset).to(device)
 
-    if model.token_embedding.num_embeddings <= NO_EVENT_IDX or model.risk_head.out_features <= NO_EVENT_IDX:
+    if (
+        model_target_mode == "next_token"
+        and (
+            model.token_embedding.num_embeddings <= NO_EVENT_IDX
+            or model.risk_head.out_features <= NO_EVENT_IDX
+        )
+    ):
         raise RuntimeError(
             "Model vocabulary does not include <NO_EVENT> token index. "
             "Checkpoint/model shape is incompatible with no-event landmark querying."
@@ -1335,6 +1504,7 @@ def main() -> None:
         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,
         death_token_ids=death_token_ids,
@@ -1357,11 +1527,12 @@ def main() -> None:
     if eval_split in {"valid", "validation"}:
         eval_split = "val"
 
-    score_mode_out = (
+    landmark_query_mode = (
-        "insert_no_event_landmark_exponential_risk"
+        "insert_no_event_token"
-        if score_mode == "risk"
+        if model_target_mode == "next_token"
-        else "insert_no_event_landmark_eta_diagnostic"
+        else "direct_t_query"
     )
+    score_mode_out = f"{landmark_query_mode}_{score_mode}"
 
     num_workers_auc = int(
         cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1)))
@@ -1377,9 +1548,14 @@ def main() -> None:
 
     print(f"Eval split: {eval_split}")
     print(f"Number of selected patients: {len(subset_indices)}")
-    print("No-event support: true")
+    print(f"No-event support: {bool(has_no_event)}")
-    print("Landmark query mode: insert_no_event_token")
+    print(f"Model target mode: {model_target_mode}")
-    print("Landmark token mode: no_event")
+    print(f"Landmark query mode: {landmark_query_mode}")
+    print(
+        "Landmark token mode: no_event"
+        if model_target_mode == "next_token"
+        else "Landmark token mode: none"
+    )
     print(f"Dist mode: {dist_mode}")
     print(f"Score mode: {score_mode}")
     print(f"Landmark ages: {landmark_ages.tolist()}")
@@ -1395,12 +1571,13 @@ def main() -> None:
         "model_ckpt_path": str(model_ckpt_path),
         "config_path": str(config_path),
         "target_mode": str(target_mode),
+        "model_target_mode": str(model_target_mode),
         "dist_mode": str(dist_mode),
         "time_mode": str(time_mode),
         "attn_mask_mode": str(attn_mask_mode),
         "readout_name": str(readout_name),
-        "landmark_query_mode": "insert_no_event_token",
+        "landmark_query_mode": landmark_query_mode,
-        "landmark_token_mode": "no_event",
+        "landmark_token_mode": "no_event" if model_target_mode == "next_token" else "none",
     }
 
     evaluate_landmark_auc(
@@ -1414,6 +1591,7 @@ def main() -> None:
         score_mode=score_mode,
         horizons=horizons,
         device=device,
+        model_target_mode=model_target_mode,
         readout_name=readout_name,
         readout_reduce=readout_reduce,
         num_workers_auc=num_workers_auc,