JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jiarui Li
2026-06-12 11:33:32 +08:00
parent 0fa8bbbb9a
commit 3125b6119f
3 changed files with 378 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

344
evaluate_auc_v2.py
View File

@@ -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 <NO_EVENT> 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 <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(
[
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) == "<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":
raise RuntimeError(
"Weibull checkpoints are not supported by evaluate_auc_v2.py. "
"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 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 <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 = (
"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,