350 lines
12 KiB
Python
350 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
|
|
from dataset import AllFutureHealthDataset, HealthDataset
|
|
from model_architectures import resolve_model_architecture
|
|
from models import DeepHealth
|
|
from targets import PAD_IDX
|
|
|
|
|
|
def load_json_config(path: str | Path | None) -> Dict[str, Any]:
|
|
if path is None:
|
|
return {}
|
|
config_path = Path(path)
|
|
if not config_path.exists():
|
|
return {}
|
|
with config_path.open("r", encoding="utf-8") as file:
|
|
return json.load(file)
|
|
|
|
|
|
def cfg_get(
|
|
args: argparse.Namespace | Dict[str, Any] | None,
|
|
cfg: Dict[str, Any],
|
|
name: str,
|
|
default: Any,
|
|
) -> Any:
|
|
if args is not None:
|
|
value = (
|
|
args.get(name)
|
|
if isinstance(args, dict)
|
|
else 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:
|
|
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 validate_training_mode_config(cfg: Dict[str, Any]) -> None:
|
|
model_target_mode = str(
|
|
cfg.get("model_target_mode", "next_token")
|
|
).lower()
|
|
if model_target_mode not in {"next_token", "all_future"}:
|
|
raise ValueError(
|
|
"model_target_mode must be next_token or all_future, got "
|
|
f"{model_target_mode!r}"
|
|
)
|
|
if model_target_mode != "next_token":
|
|
return
|
|
|
|
time_mode = str(cfg.get("time_mode", "")).lower()
|
|
target_mode = str(cfg.get("target_mode", "")).lower()
|
|
if time_mode != "absolute" or target_mode != "delphi2m":
|
|
raise ValueError(
|
|
"next_token is reserved for Delphi2M reproduction and requires "
|
|
"time_mode='absolute' and target_mode='delphi2m'; got "
|
|
f"time_mode={time_mode!r}, target_mode={target_mode!r}"
|
|
)
|
|
|
|
|
|
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}"
|
|
)
|
|
indices = np.random.RandomState(seed).permutation(n)
|
|
n_train = int(n * train_ratio)
|
|
n_val = int(n * val_ratio)
|
|
return (
|
|
indices[:n_train],
|
|
indices[n_train:n_train + n_val],
|
|
indices[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:
|
|
validate_training_mode_config(cfg)
|
|
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(
|
|
"model_target_mode must be next_token or all_future, got "
|
|
f"{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", "absolute")),
|
|
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
|
|
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
|
|
model_architecture=model_architecture,
|
|
)
|
|
|
|
|
|
def validate_dataset_metadata(
|
|
dataset: HealthDataset,
|
|
cfg: Dict[str, Any],
|
|
) -> None:
|
|
metadata = cfg.get("dataset_metadata")
|
|
if not isinstance(metadata, 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={metadata.get(key)!r}, current_dataset={value!r}"
|
|
for key, value in actual.items()
|
|
if key in metadata and metadata.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)
|
|
)
|
|
|
|
|
|
def build_first_occurrence_map(
|
|
dataset: HealthDataset,
|
|
subset_indices: np.ndarray,
|
|
) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
|
|
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
|
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
|
sample = dataset.samples[int(dataset_index)]
|
|
sequence_events = np.asarray(sample["event_seq"], dtype=np.int64)
|
|
sequence_times = np.asarray(sample["time_seq"], dtype=np.float32)
|
|
target_events = np.asarray(
|
|
sample["target_event_seq"], dtype=np.int64
|
|
)
|
|
target_times = np.asarray(
|
|
sample["target_time_seq"], dtype=np.float32
|
|
)
|
|
if sequence_events.size == 0 or target_events.size == 0:
|
|
continue
|
|
|
|
full_events = np.concatenate(
|
|
[sequence_events, target_events[-1:]]
|
|
)
|
|
full_times = np.concatenate([sequence_times, target_times[-1:]])
|
|
unique_tokens, first_indices = np.unique(
|
|
full_events, return_index=True
|
|
)
|
|
for token, first_index in zip(
|
|
unique_tokens.tolist(), first_indices.tolist()
|
|
):
|
|
first_lists.setdefault(int(token), []).append(
|
|
(patient_id, float(full_times[int(first_index)]))
|
|
)
|
|
|
|
return {
|
|
int(token): (
|
|
np.asarray([patient for patient, _ in pairs], dtype=np.int32),
|
|
np.asarray([time for _, time in pairs], dtype=np.float32),
|
|
)
|
|
for token, pairs in first_lists.items()
|
|
if pairs
|
|
}
|
|
|
|
|
|
class AllFutureSequenceEvalDataset:
|
|
"""
|
|
Eval-only sequence view for all-future checkpoints.
|
|
|
|
All-future training uses the observed history, including CHECKUP state
|
|
tokens, without reusing the next-step view that contains imputed
|
|
<NO_EVENT> gap tokens.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
data_prefix: str,
|
|
labels_file: str,
|
|
min_history_events: int = 1,
|
|
min_future_events: int = 1,
|
|
extra_info_types: Iterable[int] | None = None,
|
|
) -> None:
|
|
base = AllFutureHealthDataset(
|
|
data_prefix=data_prefix,
|
|
labels_file=labels_file,
|
|
split="train",
|
|
min_history_events=min_history_events,
|
|
min_future_events=min_future_events,
|
|
extra_info_types=extra_info_types,
|
|
)
|
|
|
|
self.base = base
|
|
self.label_code_to_id = base.label_code_to_id
|
|
self.label_id_to_code = base.label_id_to_code
|
|
self.vocab_size = base.vocab_size
|
|
self.n_types = base.n_types
|
|
self.n_cont_types = base.n_cont_types
|
|
self.n_categories = base.n_categories
|
|
self.cont_type_ids = base.cont_type_ids
|
|
self.extra_info_types = base.extra_info_types
|
|
|
|
self.samples: List[Dict[str, Any]] = []
|
|
for patient in base.patients:
|
|
labels = np.asarray(patient["labels"], dtype=np.int64)
|
|
times = np.asarray(patient["times"], dtype=np.float32)
|
|
if labels.size < 2:
|
|
continue
|
|
self.samples.append(
|
|
{
|
|
"eid": int(patient["eid"]),
|
|
"event_seq": labels[:-1],
|
|
"time_seq": times[:-1],
|
|
"target_event_seq": labels[1:],
|
|
"target_time_seq": times[1:],
|
|
"sex": int(patient["sex"]),
|
|
"other_type": np.asarray(patient["other_type"], dtype=np.int64),
|
|
"other_value": np.asarray(patient["other_value"], dtype=np.float32),
|
|
"other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
|
|
"other_time": np.asarray(patient["other_time"], dtype=np.float32),
|
|
}
|
|
)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
|
s = self.samples[idx]
|
|
return {
|
|
"event_seq": torch.from_numpy(s["event_seq"]).long(),
|
|
"time_seq": torch.from_numpy(s["time_seq"]).float(),
|
|
"target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
|
|
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
|
|
"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(),
|
|
}
|
|
|
|
|
|
def load_sequence_eval_dataset(
|
|
*,
|
|
model_target_mode: str,
|
|
data_prefix: str,
|
|
labels_file: str,
|
|
no_event_interval_years: float,
|
|
min_history_events: int,
|
|
min_future_events: int,
|
|
extra_info_types: Iterable[int] | None,
|
|
):
|
|
mode = str(model_target_mode).lower()
|
|
if mode == "next_token":
|
|
return HealthDataset(
|
|
data_prefix=data_prefix,
|
|
labels_file=labels_file,
|
|
no_event_interval_years=no_event_interval_years,
|
|
extra_info_types=extra_info_types,
|
|
)
|
|
if mode == "all_future":
|
|
return AllFutureSequenceEvalDataset(
|
|
data_prefix=data_prefix,
|
|
labels_file=labels_file,
|
|
min_history_events=min_history_events,
|
|
min_future_events=min_future_events,
|
|
extra_info_types=extra_info_types,
|
|
)
|
|
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
|
|
|
|
|
|
def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
|
event_seq = pad_sequence(
|
|
[s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
|
|
)
|
|
time_seq = pad_sequence(
|
|
[s["time_seq"] for s in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
target_event_seq = pad_sequence(
|
|
[s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX
|
|
)
|
|
target_time_seq = pad_sequence(
|
|
[s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
other_type = pad_sequence(
|
|
[s["other_type"] for s in batch], batch_first=True, padding_value=0
|
|
)
|
|
other_value = pad_sequence(
|
|
[s["other_value"] for s in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
other_value_kind = pad_sequence(
|
|
[s["other_value_kind"] for s in batch], batch_first=True, padding_value=0
|
|
)
|
|
other_time = pad_sequence(
|
|
[s["other_time"] for s in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
|
|
return {
|
|
"event_seq": event_seq,
|
|
"time_seq": time_seq,
|
|
"padding_mask": event_seq > PAD_IDX,
|
|
"target_event_seq": target_event_seq,
|
|
"target_time_seq": target_time_seq,
|
|
"readout_mask": event_seq > PAD_IDX,
|
|
"sex": torch.stack([s["sex"] for s in batch]),
|
|
"other_type": other_type,
|
|
"other_value": other_value,
|
|
"other_value_kind": other_value_kind,
|
|
"other_time": other_time,
|
|
}
|