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

refactor: isolate Delphi2M next-token pipeline

Browse Source
This commit is contained in:
Jiarui Li
2026-07-25 14:22:36 +08:00
parent 15ace878f4
commit 315f552301
17 changed files with 330 additions and 1817 deletions
Download Patch File Download Diff File Expand all files Collapse all files

213
evaluate_auc_v2.py
View File

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