1628 lines
64 KiB
Python
1628 lines
64 KiB
Python
"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth.
|
|
|
|
This script supports DeepHealth fixed-horizon risk scores for exponential,
|
|
Weibull, and mixed all-future distributions.
|
|
|
|
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
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import math
|
|
import multiprocessing as mp
|
|
import os
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
from pathlib import Path
|
|
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
|
|
|
|
from dataset import HealthDataset
|
|
from eval_data import load_sequence_eval_dataset
|
|
from models import DeepHealth
|
|
from readouts import build_readout
|
|
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
|
|
|
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
|
_TARGET_AWARE_MODES = {"target_aware", "delphi2m", "d2m"}
|
|
|
|
|
|
def load_json_config(path: Path) -> Dict[str, Any]:
|
|
if not path.exists():
|
|
return {}
|
|
with path.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any:
|
|
value = getattr(args, name, None)
|
|
if value is not None:
|
|
return value
|
|
return cfg.get(name, default)
|
|
|
|
|
|
def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
|
|
"""Resolve evaluation device without inheriting train_config.json device."""
|
|
device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
|
|
device = torch.device(device_name)
|
|
if device.type == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError(
|
|
f"Requested device {device_name!r}, but CUDA is not available."
|
|
)
|
|
return device
|
|
|
|
|
|
def parse_int_list(value: Any) -> Optional[List[int]]:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, (list, tuple, np.ndarray)):
|
|
return [int(x) for x in value]
|
|
text = str(value).strip()
|
|
if text == "":
|
|
return None
|
|
if text.startswith("["):
|
|
try:
|
|
values = json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"Invalid integer list: {text!r}") from exc
|
|
if not isinstance(values, list):
|
|
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
|
|
return [int(x) for x in values]
|
|
return [int(x.strip()) for x in text.split(",") if x.strip()]
|
|
|
|
|
|
def parse_float_list(value: Any) -> Optional[List[float]]:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, (list, tuple, np.ndarray)):
|
|
return [float(x) for x in value]
|
|
text = str(value).strip()
|
|
if text == "":
|
|
return None
|
|
if text.startswith("["):
|
|
try:
|
|
values = json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"Invalid float list: {text!r}") from exc
|
|
if not isinstance(values, list):
|
|
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
|
|
return [float(x) for x in values]
|
|
return [float(x.strip()) for x in text.split(",") if x.strip()]
|
|
|
|
|
|
def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
total = float(train_ratio) + float(val_ratio) + float(test_ratio)
|
|
if not np.isclose(total, 1.0, atol=1e-6):
|
|
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
|
rng = np.random.RandomState(int(seed))
|
|
idx = rng.permutation(int(n))
|
|
n_train = int(n * train_ratio)
|
|
n_val = int(n * val_ratio)
|
|
return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
|
|
|
|
|
|
def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dict[str, Any]) -> np.ndarray:
|
|
train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
|
|
val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
|
|
test_ratio = float(cfg_get(args, cfg, "test_ratio", 0.15))
|
|
seed = int(cfg_get(args, cfg, "seed", 42))
|
|
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
|
if eval_split in {"valid", "validation"}:
|
|
eval_split = "val"
|
|
|
|
train_idx, val_idx, test_idx = split_indices(
|
|
len(dataset), train_ratio, val_ratio, test_ratio, seed
|
|
)
|
|
split_map = {
|
|
"train": train_idx,
|
|
"val": val_idx,
|
|
"test": test_idx,
|
|
"all": np.arange(len(dataset), dtype=np.int64),
|
|
}
|
|
if eval_split not in split_map:
|
|
raise ValueError(f"Unsupported eval_split={eval_split!r}")
|
|
indices = split_map[eval_split]
|
|
|
|
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
|
if subset_size is not None and int(subset_size) > 0:
|
|
indices = indices[: int(subset_size)]
|
|
return np.asarray(indices, dtype=np.int64)
|
|
|
|
|
|
def load_checkpoint_state_dict(checkpoint_path: Path, map_location: str | torch.device = "cpu") -> Dict[str, Any]:
|
|
payload = torch.load(str(checkpoint_path), map_location=map_location)
|
|
if isinstance(payload, dict) and "model" in payload:
|
|
payload = payload["model"]
|
|
elif isinstance(payload, dict) and "state_dict" in payload:
|
|
payload = payload["state_dict"]
|
|
if not isinstance(payload, dict):
|
|
raise TypeError(f"Unsupported checkpoint payload type: {type(payload)}")
|
|
return payload
|
|
|
|
|
|
def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
|
|
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"
|
|
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)),
|
|
n_head=int(cfg_get(args, cfg, "n_head", 10)),
|
|
n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
|
|
n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
|
|
n_types=dataset.n_types,
|
|
n_cont_types=dataset.n_cont_types,
|
|
n_categories=dataset.n_categories,
|
|
cont_type_ids=dataset.cont_type_ids,
|
|
n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
|
|
extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")),
|
|
target_mode=model_target_mode,
|
|
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
|
|
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
|
|
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
|
|
)
|
|
|
|
|
|
def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
|
|
model.load_state_dict(state_dict, strict=True)
|
|
|
|
|
|
def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
|
|
meta = cfg.get("dataset_metadata")
|
|
if not isinstance(meta, dict):
|
|
return
|
|
|
|
actual: Dict[str, Any] = {
|
|
"vocab_size": int(dataset.vocab_size),
|
|
"n_types": int(dataset.n_types),
|
|
"n_cont_types": int(dataset.n_cont_types),
|
|
"n_categories": int(dataset.n_categories),
|
|
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
|
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
|
}
|
|
mismatches = [
|
|
f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
|
|
for key, value in actual.items()
|
|
if key in meta and meta.get(key) != value
|
|
]
|
|
if mismatches:
|
|
raise RuntimeError(
|
|
"Current dataset metadata does not match train_config.json. "
|
|
"Use the same prepared data and extra_info_types as training. "
|
|
+ "; ".join(mismatches)
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeLong AUC utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def compute_midrank(x: np.ndarray) -> np.ndarray:
|
|
x = np.asarray(x, dtype=np.float64)
|
|
order = np.argsort(x)
|
|
sorted_x = x[order]
|
|
n = len(x)
|
|
ranks = np.zeros(n, dtype=np.float64)
|
|
i = 0
|
|
while i < n:
|
|
j = i
|
|
while j < n and sorted_x[j] == sorted_x[i]:
|
|
j += 1
|
|
ranks[i:j] = 0.5 * (i + j - 1)
|
|
i = j
|
|
out = np.empty(n, dtype=np.float64)
|
|
out[order] = ranks + 1.0
|
|
return out
|
|
|
|
|
|
def fast_delong(predictions_sorted_transposed: np.ndarray, label_1_count: int) -> Tuple[np.ndarray, np.ndarray]:
|
|
predictions_sorted_transposed = np.asarray(
|
|
predictions_sorted_transposed, dtype=np.float64)
|
|
m = int(label_1_count)
|
|
n = int(predictions_sorted_transposed.shape[1] - m)
|
|
if m <= 0 or n <= 0:
|
|
return np.array([np.nan], dtype=np.float64), np.array([[np.nan]], dtype=np.float64)
|
|
|
|
positive_examples = predictions_sorted_transposed[:, :m]
|
|
negative_examples = predictions_sorted_transposed[:, m:]
|
|
k = int(predictions_sorted_transposed.shape[0])
|
|
|
|
tx = np.empty((k, m), dtype=np.float64)
|
|
ty = np.empty((k, n), dtype=np.float64)
|
|
tz = np.empty((k, m + n), dtype=np.float64)
|
|
for r in range(k):
|
|
tx[r] = compute_midrank(positive_examples[r])
|
|
ty[r] = compute_midrank(negative_examples[r])
|
|
tz[r] = compute_midrank(predictions_sorted_transposed[r])
|
|
|
|
aucs = tz[:, :m].sum(axis=1) / m / n - float(m + 1.0) / 2.0 / n
|
|
|
|
v01 = (tz[:, :m] - tx) / n
|
|
v10 = 1.0 - (tz[:, m:] - ty) / m
|
|
|
|
if k == 1:
|
|
sx = np.var(v01[0], ddof=1) if m > 1 else 0.0
|
|
sy = np.var(v10[0], ddof=1) if n > 1 else 0.0
|
|
delong_cov = np.array([[sx / m + sy / n]], dtype=np.float64)
|
|
else:
|
|
sx = np.cov(v01, rowvar=True) if m > 1 else np.zeros(
|
|
(k, k), dtype=np.float64)
|
|
sy = np.cov(v10, rowvar=True) if n > 1 else np.zeros(
|
|
(k, k), dtype=np.float64)
|
|
delong_cov = np.atleast_2d(sx) / m + np.atleast_2d(sy) / n
|
|
|
|
return aucs, delong_cov
|
|
|
|
|
|
def get_auc_delong_var(case_scores: np.ndarray, control_scores: np.ndarray) -> Tuple[float, float]:
|
|
case_scores = np.asarray(case_scores, dtype=np.float64)
|
|
control_scores = np.asarray(control_scores, dtype=np.float64)
|
|
if case_scores.size == 0 or control_scores.size == 0:
|
|
return np.nan, np.nan
|
|
|
|
y_true = np.array([1] * len(case_scores) + [0] *
|
|
len(control_scores), dtype=np.int8)
|
|
y_score = np.concatenate([case_scores, control_scores]).astype(
|
|
np.float64, copy=False)
|
|
order = (-y_true).argsort()
|
|
|
|
aucs, cov = fast_delong(y_score[np.newaxis, order], int(y_true.sum()))
|
|
var = float(np.asarray(cov).reshape(-1)[0])
|
|
return float(aucs[0]), var
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Metadata and token selection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _first_existing_column(df: pd.DataFrame, candidates: Sequence[str]) -> Optional[str]:
|
|
for col in candidates:
|
|
if col in df.columns:
|
|
return col
|
|
return None
|
|
|
|
|
|
def build_metadata_for_merge(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> pd.DataFrame:
|
|
base_rows = []
|
|
for token, code in dataset.label_id_to_code.items():
|
|
token = int(token)
|
|
code_text = str(code)
|
|
if token in SPECIAL_TOKENS or code_text.startswith("<"):
|
|
continue
|
|
base_rows.append({"token": token, "label_code": code_text})
|
|
base = pd.DataFrame(base_rows)
|
|
if labels_meta is None or labels_meta.empty:
|
|
return base
|
|
|
|
meta = labels_meta.copy()
|
|
code_col = _first_existing_column(
|
|
meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
|
|
if code_col is not None:
|
|
meta["_label_code"] = meta[code_col].astype(
|
|
str).map(lambda s: s.split()[0].strip())
|
|
merged = base.merge(meta, left_on="label_code",
|
|
right_on="_label_code", how="left")
|
|
return merged.drop(columns=["_label_code"], errors="ignore")
|
|
|
|
if "index" in meta.columns:
|
|
idx = pd.to_numeric(meta["index"], errors="coerce")
|
|
has_no_event = (
|
|
NO_EVENT_IDX in dataset.label_id_to_code
|
|
and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
|
|
)
|
|
if has_no_event:
|
|
idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
|
|
meta["_index_int"] = idx.astype("Int64")
|
|
merged = base.merge(meta, left_on="token",
|
|
right_on="_index_int", how="left")
|
|
return merged.drop(columns=["_index_int"], errors="ignore")
|
|
|
|
return base
|
|
|
|
|
|
def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> Dict[int, float]:
|
|
if labels_meta is None or labels_meta.empty or "count" not in labels_meta.columns:
|
|
return {}
|
|
|
|
out: Dict[int, float] = {}
|
|
meta = labels_meta.copy()
|
|
count_series = pd.to_numeric(meta["count"], errors="coerce")
|
|
|
|
code_col = _first_existing_column(
|
|
meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
|
|
if code_col is not None:
|
|
for code_text, count in zip(meta[code_col].astype(str).tolist(), count_series.tolist()):
|
|
code = code_text.split()[0].strip()
|
|
if code in dataset.label_code_to_id and pd.notna(count):
|
|
out[int(dataset.label_code_to_id[code])] = float(count)
|
|
if out:
|
|
return out
|
|
|
|
if "index" in meta.columns:
|
|
idx = pd.to_numeric(meta["index"], errors="coerce")
|
|
has_no_event = (
|
|
NO_EVENT_IDX in dataset.label_id_to_code
|
|
and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
|
|
)
|
|
if has_no_event:
|
|
idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
|
|
for token, count in zip(idx.tolist(), count_series.tolist()):
|
|
if pd.notna(token) and pd.notna(count):
|
|
out[int(token)] = float(count)
|
|
return out
|
|
|
|
|
|
def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]:
|
|
ids: List[int] = []
|
|
|
|
if labels_meta is not None and not labels_meta.empty:
|
|
meta = labels_meta.copy()
|
|
if "ICD-10 Chapter (short)" in meta.columns:
|
|
death_rows = meta[meta["ICD-10 Chapter (short)"].astype(
|
|
str) == "Death"]
|
|
code_col = _first_existing_column(
|
|
death_rows, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
|
|
if code_col is not None:
|
|
for raw in death_rows[code_col].astype(str).tolist():
|
|
code = raw.split()[0].strip()
|
|
if code in dataset.label_code_to_id:
|
|
ids.append(int(dataset.label_code_to_id[code]))
|
|
elif "index" in death_rows.columns:
|
|
idx = pd.to_numeric(death_rows["index"], errors="coerce")
|
|
has_no_event = (
|
|
NO_EVENT_IDX in dataset.label_id_to_code
|
|
and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
|
|
)
|
|
if has_no_event:
|
|
idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
|
|
ids.extend(int(x) for x in idx.dropna().astype(int).tolist())
|
|
|
|
exact_codes = {"death", "<death>", "dth", "deceased", "mortality"}
|
|
for token, code in dataset.label_id_to_code.items():
|
|
token = int(token)
|
|
if token in SPECIAL_TOKENS:
|
|
continue
|
|
text = str(code).strip().lower()
|
|
if text in exact_codes or ("death" in text) or ("mortality" in text):
|
|
ids.append(token)
|
|
|
|
return sorted(set(int(x) for x in ids if int(x) not in SPECIAL_TOKENS))
|
|
|
|
|
|
def _build_first_occurrence_maps(
|
|
dataset: HealthDataset,
|
|
subset_indices: np.ndarray,
|
|
) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
|
|
patient_count = len(subset_indices)
|
|
followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
|
|
death_time = np.full(patient_count, np.inf, dtype=np.float32)
|
|
sex = np.full(patient_count, -1, dtype=np.int8)
|
|
|
|
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
|
|
|
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
|
s = dataset.samples[int(dataset_index)]
|
|
|
|
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
|
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
|
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
|
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
|
|
|
if seq_event.size == 0 or tgt_event.size == 0:
|
|
continue
|
|
|
|
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
|
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
|
|
|
sex[patient_id] = int(s["sex"])
|
|
followup_end[patient_id] = np.max(full_time).astype(np.float32)
|
|
|
|
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
|
|
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
|
|
token = int(token)
|
|
event_time = float(full_time[int(idx)])
|
|
if token not in first_lists:
|
|
first_lists[token] = []
|
|
first_lists[token].append((patient_id, event_time))
|
|
|
|
packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
|
for token, pairs in first_lists.items():
|
|
if not pairs:
|
|
continue
|
|
packed[int(token)] = (
|
|
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
|
np.asarray([t for _, t in pairs], dtype=np.float32),
|
|
)
|
|
|
|
return packed, followup_end, death_time, sex
|
|
|
|
|
|
def select_disease_tokens(
|
|
dataset: HealthDataset,
|
|
labels_meta: Optional[pd.DataFrame],
|
|
requested_tokens: Optional[Sequence[int]],
|
|
filter_min_total: int,
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
) -> List[int]:
|
|
base = [
|
|
int(token)
|
|
for token, code in dataset.label_id_to_code.items()
|
|
if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<")
|
|
]
|
|
base_set = set(base)
|
|
|
|
if requested_tokens is not None:
|
|
return sorted(set(int(x) for x in requested_tokens if int(x) in base_set))
|
|
|
|
disease_ids = sorted(base)
|
|
if int(filter_min_total) <= 0:
|
|
return disease_ids
|
|
|
|
meta_counts = _metadata_count_map(dataset, labels_meta)
|
|
if meta_counts:
|
|
return [token for token in disease_ids if float(meta_counts.get(token, 0.0)) > float(filter_min_total)]
|
|
|
|
split_counts = {}
|
|
for token, pairs in first_occurrence_by_token.items():
|
|
token = int(token)
|
|
if token not in base_set:
|
|
continue
|
|
split_counts[token] = len(np.unique(pairs[0]))
|
|
return [token for token in disease_ids if int(split_counts.get(token, 0)) > int(filter_min_total)]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Landmark dataset
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class LandmarkDataset(Dataset):
|
|
def __init__(
|
|
self,
|
|
dataset: HealthDataset,
|
|
subset_indices: np.ndarray,
|
|
landmark_ages: np.ndarray,
|
|
attn_mask_mode: str,
|
|
model_target_mode: str,
|
|
min_history_events: int,
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
death_token_ids: Sequence[int],
|
|
) -> None:
|
|
self.dataset = dataset
|
|
self.subset_indices = np.asarray(subset_indices, dtype=np.int64)
|
|
self.landmark_ages = np.asarray(landmark_ages, dtype=np.float32)
|
|
self.attn_mask_mode = str(attn_mask_mode).lower()
|
|
self.model_target_mode = str(model_target_mode).lower()
|
|
if self.model_target_mode not in {"next_token", "all_future"}:
|
|
raise ValueError(
|
|
"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
|
|
self.death_token_ids = sorted(set(int(x) for x in death_token_ids))
|
|
|
|
rows: List[Dict[str, Any]] = []
|
|
|
|
self.patient_followup_end = np.full(
|
|
len(self.subset_indices), -np.inf, dtype=np.float32)
|
|
self.patient_death_time = np.full(
|
|
len(self.subset_indices), np.inf, dtype=np.float32)
|
|
self.patient_sex = np.full(len(self.subset_indices), -1, dtype=np.int8)
|
|
|
|
for patient_id, dataset_index in enumerate(self.subset_indices.tolist()):
|
|
s = dataset.samples[int(dataset_index)]
|
|
|
|
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
|
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
|
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
|
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
|
|
|
if seq_event.size == 0 or tgt_event.size == 0:
|
|
continue
|
|
|
|
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
|
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
|
|
|
self.patient_sex[patient_id] = int(s["sex"])
|
|
followup_end = float(np.max(full_time))
|
|
self.patient_followup_end[patient_id] = np.float32(followup_end)
|
|
|
|
d_time = np.inf
|
|
for death_token in self.death_token_ids:
|
|
hit = np.where(full_event == int(death_token))[0]
|
|
if hit.size > 0:
|
|
d_time = min(d_time, float(full_time[int(hit[0])]))
|
|
self.patient_death_time[patient_id] = np.float32(d_time)
|
|
|
|
for landmark_age in self.landmark_ages.tolist():
|
|
landmark_age = float(landmark_age)
|
|
|
|
if not (followup_end > landmark_age):
|
|
continue
|
|
if not (float(self.patient_death_time[patient_id]) > landmark_age):
|
|
continue
|
|
|
|
prefix_mask = full_time <= landmark_age
|
|
if not np.any(prefix_mask):
|
|
continue
|
|
|
|
prefix_events = full_event[prefix_mask]
|
|
prefix_times = full_time[prefix_mask]
|
|
|
|
valid_history_mask = ~np.isin(prefix_events, np.array(
|
|
[PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64))
|
|
if valid_history_mask.sum() < self.min_history_events:
|
|
continue
|
|
|
|
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),
|
|
]
|
|
)
|
|
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(
|
|
{
|
|
"patient_id": int(patient_id),
|
|
"dataset_index": int(dataset_index),
|
|
"sex": int(s["sex"]),
|
|
"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": landmark_pos,
|
|
"t_query": np.float32(landmark_age),
|
|
"event_seq": event_seq_landmark,
|
|
"time_seq": time_seq_landmark,
|
|
"readout_mask": readout_mask,
|
|
"other_type": np.asarray(s["other_type"], dtype=np.int64),
|
|
"other_value": np.asarray(s["other_value"], dtype=np.float32),
|
|
"other_value_kind": np.asarray(s["other_value_kind"], dtype=np.int64),
|
|
"other_time": np.asarray(s["other_time"], dtype=np.float32),
|
|
}
|
|
)
|
|
|
|
if not rows:
|
|
raise RuntimeError(
|
|
"No eligible landmark query samples were produced. Check eval split, landmark ages, and min_history_events."
|
|
)
|
|
|
|
self.rows = rows
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.rows)
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
|
s = self.rows[idx]
|
|
return {
|
|
"event_seq": torch.from_numpy(s["event_seq"]).long(),
|
|
"time_seq": torch.from_numpy(s["time_seq"]).float(),
|
|
"readout_mask": torch.from_numpy(s["readout_mask"]),
|
|
"sex": torch.tensor(s["sex"], dtype=torch.long),
|
|
"other_type": torch.from_numpy(s["other_type"]).long(),
|
|
"other_value": torch.from_numpy(s["other_value"]).float(),
|
|
"other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
|
|
"other_time": torch.from_numpy(s["other_time"]).float(),
|
|
"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),
|
|
"death_time": torch.tensor(float(s["death_time"]), dtype=torch.float32),
|
|
}
|
|
|
|
|
|
def collate_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
|
event_seq = pad_sequence(
|
|
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX)
|
|
time_seq = pad_sequence([x["time_seq"]
|
|
for x in batch], batch_first=True, padding_value=0.0)
|
|
readout_mask = pad_sequence(
|
|
[x["readout_mask"] for x in batch], batch_first=True, padding_value=False)
|
|
other_type = pad_sequence(
|
|
[x["other_type"] for x in batch], batch_first=True, padding_value=0)
|
|
other_value = pad_sequence(
|
|
[x["other_value"] for x in batch], batch_first=True, padding_value=0.0)
|
|
other_value_kind = pad_sequence(
|
|
[x["other_value_kind"] for x in batch], batch_first=True, padding_value=0)
|
|
other_time = pad_sequence(
|
|
[x["other_time"] for x in batch], batch_first=True, padding_value=0.0)
|
|
|
|
return {
|
|
"event_seq": event_seq,
|
|
"time_seq": time_seq,
|
|
"padding_mask": event_seq > PAD_IDX,
|
|
"readout_mask": readout_mask,
|
|
"sex": torch.stack([x["sex"] for x in batch]),
|
|
"other_type": other_type,
|
|
"other_value": other_value,
|
|
"other_value_kind": other_value_kind,
|
|
"other_time": other_time,
|
|
"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]),
|
|
"death_time": torch.stack([x["death_time"] for x in batch]),
|
|
}
|
|
|
|
|
|
def _numpy_hidden_dtype(name: str) -> np.dtype:
|
|
key = str(name).lower()
|
|
if key in {"float16", "fp16", "half", "bfloat16", "bf16"}:
|
|
return np.float16
|
|
if key in {"float32", "fp32", "single"}:
|
|
return np.float32
|
|
raise ValueError(
|
|
f"hidden_cache_dtype must be float16 or float32, got {name!r}")
|
|
|
|
|
|
@torch.inference_mode()
|
|
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]]:
|
|
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)
|
|
elif model_target_mode == "next_token":
|
|
readout = build_readout(readout_name).to(device)
|
|
if readout is not None:
|
|
readout.eval()
|
|
|
|
hidden_parts: List[np.ndarray] = []
|
|
arrays = {
|
|
"patient_id": [],
|
|
"sex": [],
|
|
"landmark_age": [],
|
|
"followup_end_time": [],
|
|
"death_time": [],
|
|
}
|
|
|
|
out_dtype = _numpy_hidden_dtype(hidden_cache_dtype)
|
|
amp_enabled = bool(use_amp and device.type == "cuda")
|
|
|
|
for batch in tqdm(loader, desc="Landmark inference", leave=False, dynamic_ncols=True):
|
|
batch_dev = {
|
|
k: (v.to(device, non_blocking=True)
|
|
if isinstance(v, torch.Tensor) else v)
|
|
for k, v in batch.items()
|
|
}
|
|
|
|
amp_ctx = (
|
|
torch.autocast(device_type=device.type, dtype=torch.float16)
|
|
if amp_enabled
|
|
else contextlib.nullcontext()
|
|
)
|
|
with amp_ctx:
|
|
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))
|
|
arrays["patient_id"].append(
|
|
batch["patient_id"].cpu().numpy().astype(np.int32, copy=False))
|
|
arrays["sex"].append(
|
|
batch["sex"].cpu().numpy().astype(np.int8, copy=False))
|
|
arrays["landmark_age"].append(
|
|
batch["landmark_age"].cpu().numpy().astype(np.float32, copy=False))
|
|
arrays["followup_end_time"].append(
|
|
batch["followup_end_time"].cpu().numpy().astype(np.float32, copy=False))
|
|
arrays["death_time"].append(
|
|
batch["death_time"].cpu().numpy().astype(np.float32, copy=False))
|
|
|
|
hidden_all = np.concatenate(hidden_parts, axis=0)
|
|
row_arrays = {
|
|
k: np.concatenate(v, axis=0) for k, v in arrays.items()
|
|
}
|
|
return hidden_all, row_arrays
|
|
|
|
|
|
@torch.inference_mode()
|
|
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,
|
|
) -> 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 = None
|
|
if model.risk_head.bias is not None:
|
|
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())
|
|
if bias is not None:
|
|
logits = logits + 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))
|
|
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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Parallel AUC workers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_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,
|
|
row_followup_end: np.ndarray,
|
|
row_death_time: np.ndarray,
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
patient_count: int,
|
|
horizons: np.ndarray,
|
|
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")
|
|
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
|
|
os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
|
|
|
|
_WORKER.clear()
|
|
_WORKER.update(
|
|
{
|
|
"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),
|
|
"row_followup_end": np.asarray(row_followup_end, dtype=np.float32),
|
|
"row_death_time": np.asarray(row_death_time, dtype=np.float32),
|
|
"first_occurrence_by_token": first_occurrence_by_token,
|
|
"patient_count": int(patient_count),
|
|
"horizons": np.asarray(horizons, dtype=np.float32),
|
|
"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": {},
|
|
}
|
|
)
|
|
|
|
|
|
def _first_time_by_patient(token: int) -> np.ndarray:
|
|
cache = _WORKER["first_time_cache"]
|
|
if int(token) in cache:
|
|
return cache[int(token)]
|
|
|
|
arr = np.full(int(_WORKER["patient_count"]), np.inf, dtype=np.float32)
|
|
pairs = _WORKER["first_occurrence_by_token"].get(int(token))
|
|
if pairs is not None:
|
|
p, t = pairs
|
|
arr[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
|
|
cache[int(token)] = arr
|
|
return arr
|
|
|
|
|
|
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)
|
|
|
|
|
|
def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
|
|
j, token, score_mode = task
|
|
|
|
token = int(token)
|
|
row_patient_id = _WORKER["row_patient_id"]
|
|
row_sex = _WORKER["row_sex"]
|
|
row_landmark_age = _WORKER["row_landmark_age"]
|
|
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"]
|
|
horizons = _WORKER["horizons"]
|
|
|
|
out_rows: List[Dict[str, Any]] = []
|
|
|
|
for sex_value, sex_name in [(0, "female"), (1, "male")]:
|
|
sex_mask_rows = row_sex == sex_value
|
|
if not np.any(sex_mask_rows):
|
|
continue
|
|
|
|
lm_values = np.unique(row_landmark_age[sex_mask_rows])
|
|
for a in lm_values.tolist():
|
|
a = float(a)
|
|
stratum_mask = sex_mask_rows & np.isclose(
|
|
row_landmark_age, np.float32(a))
|
|
if not np.any(stratum_mask):
|
|
continue
|
|
|
|
idx = np.flatnonzero(stratum_mask)
|
|
pid = row_patient_id[idx]
|
|
followup = row_followup_end[idx]
|
|
death_time = row_death_time[idx]
|
|
first_time = first_time_patient[pid]
|
|
|
|
prevalent = first_time <= np.float32(a)
|
|
if np.all(prevalent):
|
|
continue
|
|
|
|
for horizon in horizons.tolist():
|
|
horizon = float(horizon)
|
|
h_end = np.float32(a + horizon)
|
|
|
|
cases = (first_time > np.float32(a)) & (first_time <= h_end)
|
|
controls = (~prevalent) & (
|
|
((~np.isinf(first_time)) & (first_time > h_end))
|
|
| (np.isinf(first_time) & (followup >= h_end))
|
|
)
|
|
|
|
if bool(_WORKER["exclude_death_competing"]) and (not is_death_target):
|
|
death_in_window = (death_time > np.float32(a)) & (
|
|
death_time <= h_end)
|
|
death_before_disease = death_time < first_time
|
|
competing_death = death_in_window & death_before_disease
|
|
controls = controls & (~competing_death)
|
|
|
|
case_idx = np.flatnonzero(cases)
|
|
control_idx = np.flatnonzero(controls)
|
|
if case_idx.size < int(_WORKER["min_cases"]) or control_idx.size == 0:
|
|
continue
|
|
|
|
case_scores = _score_to_probability(
|
|
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]],
|
|
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):
|
|
continue
|
|
|
|
out_rows.append(
|
|
{
|
|
"token": token,
|
|
"sex": sex_name,
|
|
"landmark_age": float(a),
|
|
"horizon": float(horizon),
|
|
"auc": float(auc),
|
|
"auc_delong": float(auc),
|
|
"auc_variance_delong": float(auc_var),
|
|
"n_diseased": int(case_scores.size),
|
|
"n_healthy": int(control_scores.size),
|
|
}
|
|
)
|
|
|
|
return out_rows
|
|
|
|
|
|
def _token_task_block(tasks: Sequence[Tuple[int, int, str]]) -> List[Dict[str, Any]]:
|
|
out: List[Dict[str, Any]] = []
|
|
for t in tasks:
|
|
out.extend(_eval_token(t))
|
|
return out
|
|
|
|
|
|
def _split_tasks(tasks: Sequence[Tuple[int, int, str]], workers: int, chunk_size: int) -> List[List[Tuple[int, int, str]]]:
|
|
if not tasks:
|
|
return []
|
|
if int(chunk_size) <= 0:
|
|
chunk_size = max(1, math.ceil(len(tasks) / max(1, workers * 4)))
|
|
return [list(tasks[i:i + chunk_size]) for i in range(0, len(tasks), chunk_size)]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def evaluate_landmark_auc(
|
|
model: DeepHealth,
|
|
loader: DataLoader,
|
|
landmark_dataset: LandmarkDataset,
|
|
output_path: Path,
|
|
labels_meta: Optional[pd.DataFrame],
|
|
disease_ids: Sequence[int],
|
|
disease_chunk_size: int,
|
|
score_mode: str,
|
|
horizons: np.ndarray,
|
|
device: torch.device,
|
|
model_target_mode: str,
|
|
readout_name: str,
|
|
readout_reduce: str,
|
|
num_workers_auc: int,
|
|
auc_task_chunk_size: int,
|
|
min_cases: int,
|
|
exclude_death_competing: bool,
|
|
use_amp: bool,
|
|
hidden_cache_dtype: str,
|
|
logit_batch_size: int,
|
|
meta_info: Dict[str, Any],
|
|
) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
|
model.eval().to(device)
|
|
|
|
hidden_all, row_arrays = infer_landmark_hidden(
|
|
model=model,
|
|
loader=loader,
|
|
device=device,
|
|
model_target_mode=model_target_mode,
|
|
readout_name=readout_name,
|
|
readout_reduce=readout_reduce,
|
|
use_amp=use_amp,
|
|
hidden_cache_dtype=hidden_cache_dtype,
|
|
)
|
|
print(
|
|
f"Cached landmark hidden: shape={hidden_all.shape}, dtype={hidden_all.dtype}")
|
|
|
|
disease_ids = [int(x) for x in disease_ids]
|
|
if disease_chunk_size <= 0:
|
|
disease_chunk_size = len(disease_ids)
|
|
chunks = np.array_split(np.asarray(disease_ids, dtype=np.int64), math.ceil(
|
|
len(disease_ids) / disease_chunk_size))
|
|
|
|
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, 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,
|
|
)
|
|
|
|
tasks = [(j, int(token), score_mode)
|
|
for j, token in enumerate(chunk_ids)]
|
|
workers = max(1, min(int(num_workers_auc), len(tasks)))
|
|
|
|
if workers <= 1:
|
|
_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"],
|
|
row_followup_end=row_arrays["followup_end_time"],
|
|
row_death_time=row_arrays["death_time"],
|
|
first_occurrence_by_token=landmark_dataset.first_occurrence_by_token,
|
|
patient_count=len(landmark_dataset.subset_indices),
|
|
horizons=horizons,
|
|
min_cases=min_cases,
|
|
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)]
|
|
else:
|
|
ctx = mp.get_context("spawn")
|
|
blocks = _split_tasks(tasks, workers, auc_task_chunk_size)
|
|
with ProcessPoolExecutor(
|
|
max_workers=workers,
|
|
mp_context=ctx,
|
|
initializer=_init_worker,
|
|
initargs=(
|
|
np.asarray(chunk_ids, dtype=np.int64),
|
|
logits_chunk,
|
|
rho_chunk,
|
|
row_arrays["patient_id"],
|
|
row_arrays["sex"],
|
|
row_arrays["landmark_age"],
|
|
row_arrays["followup_end_time"],
|
|
row_arrays["death_time"],
|
|
landmark_dataset.first_occurrence_by_token,
|
|
len(landmark_dataset.subset_indices),
|
|
horizons,
|
|
min_cases,
|
|
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(
|
|
tqdm(
|
|
ex.map(_token_task_block, blocks),
|
|
total=len(blocks),
|
|
desc=f"AUC chunk {chunk_idx}",
|
|
leave=False,
|
|
dynamic_ncols=True,
|
|
)
|
|
)
|
|
|
|
for rows in nested:
|
|
for r in rows:
|
|
r["disease_chunk_idx"] = int(chunk_idx)
|
|
all_rows.append(r)
|
|
|
|
del logits_chunk, rho_chunk
|
|
|
|
if not all_rows:
|
|
raise RuntimeError(
|
|
"No AUC rows produced. Check landmark ages, horizons, min_cases, follow-up, or disease token selection."
|
|
)
|
|
|
|
df_unpooled = pd.DataFrame(all_rows)
|
|
df_unpooled["label_code"] = df_unpooled["token"].map(
|
|
landmark_dataset.dataset.label_id_to_code)
|
|
|
|
for k, v in meta_info.items():
|
|
df_unpooled[k] = v
|
|
|
|
meta_table = build_metadata_for_merge(landmark_dataset.dataset, labels_meta)
|
|
df_unpooled = df_unpooled.merge(
|
|
meta_table, on=["token", "label_code"], how="left")
|
|
|
|
grouped = df_unpooled.groupby(
|
|
["token", "label_code", "horizon"], dropna=False, as_index=False)
|
|
df_merged = grouped.agg(
|
|
auc=("auc_delong", "mean"),
|
|
n_strata=("auc_delong", "size"),
|
|
n_diseased=("n_diseased", "sum"),
|
|
n_healthy=("n_healthy", "sum"),
|
|
auc_variance_sum=("auc_variance_delong", "sum"),
|
|
)
|
|
df_merged["auc_variance_delong"] = (
|
|
df_merged["auc_variance_sum"]
|
|
/ (df_merged["n_strata"].clip(lower=1).astype(np.float64) ** 2)
|
|
)
|
|
df_merged = df_merged.drop(columns=["auc_variance_sum"])
|
|
|
|
keep_meta = [
|
|
c for c in [
|
|
"model_ckpt_path",
|
|
"config_path",
|
|
"target_mode",
|
|
"model_target_mode",
|
|
"dist_mode",
|
|
"time_mode",
|
|
"attn_mask_mode",
|
|
"readout_name",
|
|
"landmark_query_mode",
|
|
"landmark_token_mode",
|
|
"score_mode",
|
|
"eval_split",
|
|
]
|
|
if c in df_unpooled.columns
|
|
]
|
|
for col in keep_meta:
|
|
df_merged[col] = meta_info[col]
|
|
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
df_unpooled.to_csv(
|
|
output_path / "df_auc_landmark_unpooled.csv", index=False)
|
|
df_merged.to_csv(output_path / "df_auc_landmark.csv", index=False)
|
|
|
|
return df_unpooled, df_merged
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Evaluate DeepHealth landmark fixed-horizon incident disease AUC")
|
|
|
|
parser.add_argument("--run_path", type=str, required=True)
|
|
parser.add_argument("--output_path", type=str, default=None)
|
|
parser.add_argument("--eval_split", type=str, default="test",
|
|
choices=["train", "val", "valid", "validation", "test", "all"])
|
|
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
|
|
|
parser.add_argument("--batch_size", type=int, default=None)
|
|
parser.add_argument("--num_workers", type=int, default=None)
|
|
parser.add_argument("--device", type=str, default=None,
|
|
help="Evaluation device, e.g. cpu, cuda, cuda:1. Defaults to cuda if available, else cpu.")
|
|
parser.add_argument("--num_workers_auc", type=int, default=None)
|
|
parser.add_argument("--auc_task_chunk_size", type=int, default=None)
|
|
parser.add_argument("--disease_chunk_size", type=int, default=None)
|
|
parser.add_argument("--filter_min_total", type=int, default=None)
|
|
parser.add_argument("--labels_meta_path", type=str, default=None)
|
|
|
|
parser.add_argument("--landmark_start", type=float, default=None)
|
|
parser.add_argument("--landmark_stop", type=float, default=None)
|
|
parser.add_argument("--landmark_step", type=float, default=None)
|
|
parser.add_argument("--horizons", type=str, default=None)
|
|
|
|
parser.add_argument("--min_cases", type=int, default=None)
|
|
parser.add_argument("--min_history_events", type=int, default=None)
|
|
|
|
parser.add_argument("--score_mode", type=str,
|
|
default=None, choices=["risk", "eta"])
|
|
parser.add_argument("--exclude_death_in_window_without_disease",
|
|
action=argparse.BooleanOptionalAction, default=None)
|
|
parser.add_argument(
|
|
"--use_amp", action=argparse.BooleanOptionalAction, default=None)
|
|
parser.add_argument("--logit_batch_size", type=int, default=None)
|
|
parser.add_argument("--hidden_cache_dtype", type=str,
|
|
default=None, choices=["float16", "float32"])
|
|
|
|
args = parser.parse_args()
|
|
|
|
run_path = Path(args.run_path)
|
|
config_path = run_path / "train_config.json"
|
|
model_ckpt_path = run_path / "best_model.pt"
|
|
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"train_config.json not found in {run_path}")
|
|
if not model_ckpt_path.exists():
|
|
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
|
|
|
cfg = load_json_config(config_path)
|
|
|
|
data_prefix = cfg.get("data_prefix", "ukb")
|
|
labels_file = cfg.get("labels_file", "labels.csv")
|
|
no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
|
|
include_no_event_in_uts_target = cfg.get(
|
|
"include_no_event_in_uts_target", False)
|
|
|
|
target_mode = cfg.get("target_mode", "uts")
|
|
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
|
if model_target_mode not in {"next_token", "all_future"}:
|
|
raise ValueError(
|
|
"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"))
|
|
readout_name = str(cfg.get(
|
|
"readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
|
|
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
|
time_mode = str(cfg.get("time_mode", "relative"))
|
|
|
|
output_path = Path(
|
|
cfg_get(args, cfg, "output_path", None)
|
|
or cfg.get("output_path", None)
|
|
or str(run_path)
|
|
)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
labels_meta_path = cfg_get(args, cfg, "labels_meta_path", None)
|
|
if labels_meta_path is None:
|
|
labels_meta_path = cfg.get(
|
|
"labels_meta_path", "delphi_labels_chapters_colours_icd.csv")
|
|
labels_meta = None
|
|
if labels_meta_path is not None and Path(str(labels_meta_path)).exists():
|
|
labels_meta = pd.read_csv(str(labels_meta_path))
|
|
|
|
print("Loading dataset...")
|
|
dataset = load_sequence_eval_dataset(
|
|
model_target_mode=model_target_mode,
|
|
data_prefix=data_prefix,
|
|
labels_file=labels_file,
|
|
no_event_interval_years=float(no_event_interval_years),
|
|
include_no_event_in_uts_target=bool(include_no_event_in_uts_target),
|
|
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
|
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
|
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
|
|
)
|
|
validate_dataset_metadata(dataset, cfg)
|
|
|
|
has_no_event = (
|
|
NO_EVENT_IDX in dataset.label_id_to_code
|
|
and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
|
|
and dataset.vocab_size > NO_EVENT_IDX
|
|
)
|
|
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. "
|
|
"Please evaluate only runs trained with the current no-event vocabulary."
|
|
)
|
|
return
|
|
|
|
subset_indices = make_eval_indices(dataset, args, cfg)
|
|
|
|
first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
|
|
dataset, subset_indices)
|
|
|
|
disease_requested = parse_int_list(
|
|
cfg_get(args, cfg, "diseases_of_interest", None))
|
|
disease_ids = select_disease_tokens(
|
|
dataset=dataset,
|
|
labels_meta=labels_meta,
|
|
requested_tokens=disease_requested,
|
|
filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
|
|
first_occurrence_by_token=first_occurrence_by_token,
|
|
)
|
|
|
|
if not disease_ids:
|
|
raise RuntimeError("No disease tokens selected after filtering.")
|
|
|
|
landmark_start = float(cfg_get(args, cfg, "landmark_start", 40.0))
|
|
landmark_stop = float(cfg_get(args, cfg, "landmark_stop", 80.0))
|
|
landmark_step = float(cfg_get(args, cfg, "landmark_step", 5.0))
|
|
if landmark_step <= 0:
|
|
raise ValueError("landmark_step must be > 0")
|
|
landmark_ages = np.arange(
|
|
landmark_start, landmark_stop, landmark_step, dtype=np.float32)
|
|
if landmark_ages.size == 0:
|
|
raise ValueError(
|
|
"Landmark ages are empty. Check landmark_start/landmark_stop/landmark_step.")
|
|
|
|
horizons = np.asarray(
|
|
parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [
|
|
1.0, 5.0, 10.0],
|
|
dtype=np.float32,
|
|
)
|
|
if horizons.size == 0:
|
|
raise ValueError("horizons must contain at least one value")
|
|
|
|
score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower()
|
|
if score_mode not in {"risk", "eta"}:
|
|
raise ValueError("score_mode must be 'risk' or 'eta'")
|
|
|
|
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 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 and does not use dist_mode-specific rho parameters.")
|
|
|
|
cfg_model = dict(cfg)
|
|
cfg_model["dist_mode"] = dist_mode
|
|
|
|
device = resolve_eval_device(args.device)
|
|
if device.type == "cuda":
|
|
torch.backends.cudnn.benchmark = True
|
|
|
|
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
|
|
|
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."
|
|
)
|
|
|
|
try:
|
|
load_model_state(model, state_dict)
|
|
except RuntimeError as exc:
|
|
raise RuntimeError(
|
|
"Checkpoint vocabulary shape is incompatible with the dataset/model setup. "
|
|
"Please ensure this run was trained with the same special-token vocabulary and labels file."
|
|
) from exc
|
|
|
|
if model.token_embedding.num_embeddings != dataset.vocab_size or model.risk_head.out_features != dataset.vocab_size:
|
|
raise RuntimeError(
|
|
"Checkpoint/model vocabulary shape mismatch with dataset vocabulary. "
|
|
"Please use a checkpoint trained with the same no-event vocabulary configuration."
|
|
)
|
|
|
|
death_token_ids = _get_death_token_ids(dataset, labels_meta)
|
|
min_history_events = int(cfg_get(args, cfg, "min_history_events", 1))
|
|
landmark_dataset = LandmarkDataset(
|
|
dataset=dataset,
|
|
subset_indices=subset_indices,
|
|
landmark_ages=landmark_ages,
|
|
attn_mask_mode=attn_mask_mode,
|
|
model_target_mode=model_target_mode,
|
|
min_history_events=min_history_events,
|
|
first_occurrence_by_token=first_occurrence_by_token,
|
|
death_token_ids=death_token_ids,
|
|
)
|
|
|
|
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
|
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
|
loader = DataLoader(
|
|
landmark_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
collate_fn=collate_landmark_fn,
|
|
num_workers=num_workers,
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=num_workers > 0,
|
|
prefetch_factor=2 if num_workers > 0 else None,
|
|
)
|
|
|
|
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
|
if eval_split in {"valid", "validation"}:
|
|
eval_split = "val"
|
|
|
|
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)))
|
|
auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0))
|
|
disease_chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 0))
|
|
min_cases = int(cfg_get(args, cfg, "min_cases", 2))
|
|
exclude_death_competing = bool(
|
|
cfg_get(args, cfg, "exclude_death_in_window_without_disease", True))
|
|
use_amp = bool(cfg_get(args, cfg, "use_amp", False))
|
|
hidden_cache_dtype = str(
|
|
cfg_get(args, cfg, "hidden_cache_dtype", "float16"))
|
|
logit_batch_size = int(cfg_get(args, cfg, "logit_batch_size", batch_size))
|
|
|
|
print(f"Eval split: {eval_split}")
|
|
print(f"Number of selected patients: {len(subset_indices)}")
|
|
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()}")
|
|
print(f"Horizons: {horizons.tolist()}")
|
|
print(f"Number of landmark query samples: {len(landmark_dataset)}")
|
|
print(f"Number of disease tokens: {len(disease_ids)}")
|
|
print(f"AUC workers: {num_workers_auc}")
|
|
print(f"Output path: {output_path}")
|
|
|
|
meta_info = {
|
|
"score_mode": score_mode_out,
|
|
"eval_split": eval_split,
|
|
"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": landmark_query_mode,
|
|
"landmark_token_mode": "no_event" if model_target_mode == "next_token" else "none",
|
|
}
|
|
|
|
evaluate_landmark_auc(
|
|
model=model,
|
|
loader=loader,
|
|
landmark_dataset=landmark_dataset,
|
|
output_path=output_path,
|
|
labels_meta=labels_meta,
|
|
disease_ids=disease_ids,
|
|
disease_chunk_size=disease_chunk_size,
|
|
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,
|
|
auc_task_chunk_size=auc_task_chunk_size,
|
|
min_cases=min_cases,
|
|
exclude_death_competing=exclude_death_competing,
|
|
use_amp=use_amp,
|
|
hidden_cache_dtype=hidden_cache_dtype,
|
|
logit_batch_size=logit_batch_size,
|
|
meta_info=meta_info,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|