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

214
evaluate_auc.py
View File

@@ -19,7 +19,7 @@ Efficiency notes:
avoiding repeated pickling of arrays for every disease.
Run from the DeepHealth code directory containing dataset.py, models.py,
readouts.py, and train_config.json-compatible checkpoints/configs.
and train_config.json-compatible checkpoints/configs.
"""
from __future__ import annotations
@@ -44,10 +44,20 @@ from delphi2m_auc_report import (
DEFAULT_DELPHI2M_PERIODS_YEARS,
build_delphi2m_auc_report,
)
from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
from eval_data import (
build_first_occurrence_map,
build_model_from_dataset,
cfg_get,
load_json_config,
load_sequence_eval_dataset,
resolve_eval_device,
sequence_eval_collate_fn,
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 PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
@@ -156,55 +166,6 @@ def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> T
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
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 _get_death_token_ids(dataset: HealthDataset) -> List[int]:
death_ids = [int(dataset.vocab_size) - 1]
print(f"[INFO] death token ids: {death_ids}")
@@ -264,88 +225,6 @@ def select_disease_tokens(
# Dataset/split/model helpers
# ---------------------------------------------------------------------------
def load_json_config(path: Optional[str]) -> Dict[str, Any]:
if path is None:
return {}
p = Path(path)
if not p.exists():
return {}
with p.open("r", encoding="utf-8") as f:
return json.load(f)
def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any], name: str, default: Any) -> Any:
"""Get a value from CLI args first, then train_config.json, then default.
This helper intentionally accepts either an argparse.Namespace or a dict.
The earlier version passed cfg as both args and cfg, then tried to access
args.eval_split, which fails because dict has no attributes.
"""
val = None
if args is not None:
if isinstance(args, dict):
val = args.get(name, None)
else:
val = getattr(args, name, None)
if val is not None:
return val
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 split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
total = train_ratio + val_ratio + 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(seed)
idx = rng.permutation(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 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 _extract_state_dict(ckpt: Any) -> Dict[str, Any]:
if isinstance(ckpt, dict) and "model" in ckpt:
return ckpt["model"]
@@ -430,32 +309,6 @@ def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str
return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)
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)
)
# ---------------------------------------------------------------------------
# Batched inference + cached hidden states
# ---------------------------------------------------------------------------
@@ -480,8 +333,6 @@ def infer_readout_hidden(
loader: DataLoader,
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
use_amp: bool,
hidden_cache_dtype: str = "float16",
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
@@ -492,15 +343,6 @@ def infer_readout_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: Dict[str, List[np.ndarray]] = {
"event_seq": [],
@@ -553,7 +395,6 @@ def infer_readout_hidden(
other_value=batch_dev["other_value"][active],
other_value_kind=batch_dev["other_value_kind"][active],
other_time=batch_dev["other_time"][active],
target_mode="all_future",
)
hidden[active, pos, :] = hidden_pos.float()
readout_mask_np = batch["padding_mask"].cpu().numpy()
@@ -567,16 +408,9 @@ def infer_readout_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",
)
ro = readout(
hidden=hidden_raw,
time_seq=time_seq,
padding_mask=padding_mask,
readout_mask=batch_dev["readout_mask"],
)
hidden = ro.hidden
readout_mask_np = ro.readout_mask.detach().cpu().numpy()
hidden = hidden_raw
readout_mask_np = padding_mask.detach().cpu().numpy()
h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
hidden_parts.append(h)
@@ -1065,8 +899,6 @@ def evaluate_auc_pipeline(
offsets: Sequence[float],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
num_workers_auc: int,
use_amp: bool,
auc_task_chunk_size: int = 0,
@@ -1125,8 +957,6 @@ def evaluate_auc_pipeline(
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,
)
@@ -1318,6 +1148,7 @@ def main() -> None:
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
cfg = load_json_config(str(config_path))
validate_training_mode_config(cfg)
if args.output_path is None:
args.output_path = str(run_path)
@@ -1326,9 +1157,7 @@ def main() -> None:
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 = 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(
@@ -1336,9 +1165,6 @@ def main() -> None:
f"got {model_target_mode!r}"
)
dist_mode_cfg = cfg.get("dist_mode", "exponential")
readout_name = cfg.get(
"readout_name", "same_time_group_end" if target_mode == "uts" else "token")
readout_reduce = cfg.get("readout_reduce", "mean")
device = resolve_eval_device(args.device)
if device.type == "cuda":
@@ -1350,7 +1176,6 @@ def main() -> None:
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event,
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)),
@@ -1406,8 +1231,9 @@ def main() -> None:
if disease_spec is None:
disease_spec = cfg.get("disease_tokens", None)
diseases = parse_int_list(disease_spec)
first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
dataset, subset_indices)
first_occurrence_by_token = build_first_occurrence_map(
dataset, subset_indices
)
include_death = bool(cfg_get(args, cfg, "include_death", True))
exclude_death = bool(cfg_get(args, cfg, "exclude_death", False))
auc_offsets = make_auc_offsets(args, cfg)
@@ -1427,8 +1253,6 @@ def main() -> None:
offsets=auc_offsets,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(
1, (os.cpu_count() or 2) - 1))),
use_amp=bool(cfg_get(args, cfg, "use_amp", False)),