From 3125b6119f803d92144d6685c8fd4735e4e2d1e7 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 12 Jun 2026 11:33:32 +0800 Subject: [PATCH] Refactor AUC evaluation scripts to support model target modes and improve distribution handling --- README.md | 30 +++- evaluate_auc.py | 123 ++++++++++++---- evaluate_auc_v2.py | 344 ++++++++++++++++++++++++++++++++++----------- 3 files changed, 378 insertions(+), 119 deletions(-) diff --git a/README.md b/README.md index eda11e8..279d164 100644 --- a/README.md +++ b/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 人工插入一个 `` 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 人工插入一个 `` token，取该 token 的 hidden 做风险分数。 +- `model_target_mode="all_future"`：不插入 ``，直接把 landmark age 作为 `t_query` 传入模型，取 query hidden 做风险分数。 核心流程： - 为每个患者和 landmark age 构造 landmark query 样本。 -- 在 landmark age 插入 `` token，取该位置 hidden。 -- 对疾病 token 分块投影到 `risk_head`。 +- 根据模型模式插入 `` token 或直接传 `t_query`，取 landmark/query hidden。 +- 对疾病 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 | 人工插入的 `` landmark token | +| 查询位置 | next-token 用满足 offset 条件的最新 readout token；all-future 直接用该预测点年龄作为 `t_query` | next-token 用人工插入的 `` 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 评估 - - 通过插入 `` landmark token 查询固定年龄点风险 + - next-token 模型通过插入 `` landmark token 查询固定年龄点风险 + - all-future 模型直接通过 `t_query` 查询固定年龄点风险 + - `score_mode="risk"` 按 exponential / Weibull / mixed 分布计算固定 horizon 风险 diff --git a/evaluate_auc.py b/evaluate_auc.py index 1449059..568d129 100644 --- a/evaluate_auc.py +++ b/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]]: - """ - Run the expensive transformer/readout path exactly once and cache hidden states. + """Cache per-position hidden states used by the unchanged AUC logic.""" + 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 - though only the final tied risk-head columns changed. That is usually the - 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 = 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: Dict[str, List[np.ndarray]] = { @@ -501,25 +519,55 @@ def infer_readout_hidden( if autocast_enabled else contextlib.nullcontext() ) with amp_context: - hidden = 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, - time_seq=time_seq, - padding_mask=padding_mask, - readout_mask=batch_dev["readout_mask"], - ) + if model_target_mode == "all_future": + batch_size, seq_len = event_seq.shape + hidden = torch.zeros( + batch_size, + seq_len, + model.n_embd, + device=event_seq.device, + dtype=torch.float32, + ) + for pos in range(seq_len): + active = padding_mask[:, pos].bool() + if not active.any(): + continue + hidden_pos = model( + event_seq=event_seq[active], + 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( diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index e09b1d2..acdddec 100644 --- a/evaluate_auc_v2.py +++ b/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: - DeepHealth + exponential distribution + no_event imputation. +This script supports DeepHealth fixed-horizon risk scores for exponential, +Weibull, and mixed all-future distributions. -Landmark querying is implemented by inserting a token at landmark age. -No t_query interface is used. +Landmark querying depends on the model target mode saved in train_config.json: + - next_token: insert a 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( - [ - 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 + 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), + ] ) - - readout_mask = np.zeros(len(event_seq_landmark), dtype=bool) - readout_mask[-1] = True + 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 + ) + 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( - event_seq=batch_dev["event_seq"], - time_seq=batch_dev["time_seq"], - sex=batch_dev["sex"], - padding_mask=batch_dev["padding_mask"], - other_type=batch_dev["other_type"], - other_value=batch_dev["other_value"], - other_value_kind=batch_dev["other_value_kind"], - other_time=batch_dev["other_time"], - target_mode="next_token", - ) - 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) + if model_target_mode == "all_future": + landmark_hidden = model( + event_seq=batch_dev["event_seq"], + time_seq=batch_dev["time_seq"], + sex=batch_dev["sex"], + padding_mask=batch_dev["padding_mask"], + t_query=batch_dev["t_query"], + other_type=batch_dev["other_type"], + other_value=batch_dev["other_value"], + other_value_kind=batch_dev["other_value_kind"], + other_time=batch_dev["other_time"], + target_mode="all_future", + ) + else: + hidden = model( + event_seq=batch_dev["event_seq"], + time_seq=batch_dev["time_seq"], + sex=batch_dev["sex"], + padding_mask=batch_dev["padding_mask"], + other_type=batch_dev["other_type"], + other_value=batch_dev["other_value"], + other_value_kind=batch_dev["other_value_kind"], + other_time=batch_dev["other_time"], + target_mode="next_token", + ) + 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 - return np.concatenate(out_parts, axis=0) + if rho is not None: + 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) == "" 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 imputation. " "Landmark AUC requires inserting a 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": - raise RuntimeError( - "Weibull checkpoints are not supported by evaluate_auc_v2.py. " - "This landmark evaluation requires 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 dist_mode not in {"exponential", "weibull", "mixed"}: + raise ValueError( + f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed." ) 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 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 = ( - "insert_no_event_landmark_exponential_risk" - if score_mode == "risk" - else "insert_no_event_landmark_eta_diagnostic" + landmark_query_mode = ( + "insert_no_event_token" + if model_target_mode == "next_token" + 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("Landmark query mode: insert_no_event_token") - print("Landmark token mode: no_event") + print(f"No-event support: {bool(has_no_event)}") + print(f"Model target mode: {model_target_mode}") + 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_token_mode": "no_event", + "landmark_query_mode": landmark_query_mode, + "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,