diff --git a/README.md b/README.md index 17081f6..a0ec3ad 100644 --- a/README.md +++ b/README.md @@ -59,8 +59,8 @@ python prepare_data.py `dataset.py` 提供两个 dataset： - `NextStepHealthDataset` - - 用于 next-token / next-time-point 监督 - - 对应 `Delphi2MLoss` 和 `UniqueTimeSetExponentialLoss` + - 仅用于 absolute-time Delphi2M next-token 复现 + - 对应 `Delphi2MLoss` - `AllFutureHealthDataset` - 用于 query-conditioned all-future 监督 @@ -173,7 +173,7 @@ extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入 如果需要拿到完整 next-token 输出，可使用结构化返回： ```python -out = model(..., target_mode="next_token", return_output=True) +out = model(..., return_output=True) out.hidden # disease tokens + pooled extra-info readout tokens out.time_seq # 与 hidden 对齐的时间 out.padding_mask # 与 hidden 对齐的有效位置 @@ -187,8 +187,7 @@ model = DeepHealth( vocab_size=dataset.vocab_size, n_embd=120, n_head=10, - n_hist_layer=12, - n_tab_layer=4, # 兼容旧配置；当前不再创建独立 tabular transformer + n_layer=12, n_types=dataset.n_types, n_cont_types=dataset.n_cont_types, n_categories=dataset.n_categories, @@ -204,14 +203,13 @@ model = DeepHealth( next-token 监督： - `Delphi2MLoss` -- `UniqueTimeSetExponentialLoss` next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括： - 原 disease token readout 位置 - 同一时间点 extra hidden 池化后的 pooled extra-info readout token -pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的下一个 disease 事件时间；`delphi2m` 使用第一个未来事件作为 next-token target，`uts` 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。 +pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的第一个 disease 事件作为 next-token target。若该时间点之后没有未来 disease target，则该位置不参与 loss。 all-future / query-conditioned 监督： @@ -221,17 +219,15 @@ all-future / query-conditioned 监督： all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。 -`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。 - ## 训练 当前提供两类训练入口： - `train_next_step.py` - 使用 `NextStepHealthDataset` - - `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout - - `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout - - 当前 next-token 训练只支持 exponential time loss + - 仅用于 Delphi2M 复现，固定 `time_mode=absolute`、`target_mode=delphi2m` + - 固定使用 `Delphi2MLoss` + `token` readout + - next-token 训练只支持 exponential time loss - 展开的 extra-info tokens 进入主序列；读出端 pooled extra-info tokens 会加入 prediction/loss 监督 - `train_all_future.py` - 使用 `AllFutureHealthDataset` @@ -243,8 +239,7 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in | 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout | | --- | --- | --- | --- | -| `next_token` | `relative`, `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` | -| `next_token` | `relative`, `absolute` | `target_mode=uts`, `dist_mode=exponential` | `UniqueTimeSetExponentialLoss` + `same_time_group_end` | +| `next_token` | `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` | | `all_future` | `relative`, `absolute` | `dist_mode=exponential` | `ExponentialLoss`，无 readout | | `all_future` | `relative`, `absolute` | `dist_mode=weibull` | `WeibullLoss`，无 readout | | `all_future` | `relative`, `absolute` | `dist_mode=mixed` | `MixedLoss`，无 readout | @@ -255,11 +250,9 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in python train_next_step.py \ --data_prefix ukb \ --labels_file labels.csv \ - --target_mode uts \ --n_embd 120 \ --n_head 10 \ - --n_hist_layer 12 \ - --n_tab_layer 4 \ + --n_layer 12 \ --extra_pool_reduce mean ``` @@ -424,11 +417,6 @@ python evaluate_auc_v2.py \ - `losses.py` - next-token 和 all-future losses -- `readouts.py` - - token readout - - same-time group readout - - last-valid readout - - `evaluate_auc.py` - next-step/token-level 疾病 AUC 评估 - 使用 prediction offset、sex、age bracket 分层 diff --git a/dataset.py b/dataset.py index 2743c73..bfe4e3f 100644 --- a/dataset.py +++ b/dataset.py @@ -14,7 +14,7 @@ from targets import ( DAYS_PER_YEAR, NO_EVENT_IDX, PAD_IDX, - build_all_targets, + build_next_token_targets, ) @@ -86,13 +86,11 @@ class _ExpoBaseDataset(Dataset): data_prefix: str = "ukb", labels_file: str = "labels.csv", no_event_interval_years: float = 5.0, - include_no_event_in_uts_target: bool = False, extra_info_types: Iterable[int] | None = None, ) -> None: self.data_prefix = data_prefix self.labels_file = labels_file self.no_event_interval_years = float(no_event_interval_years) - self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target) self.requested_extra_info_types = ( None if extra_info_types is None @@ -138,11 +136,6 @@ class _ExpoBaseDataset(Dataset): max_id_in_data += 1 self.vocab_size = max(max_id_in_vocab, max_id_in_data) + 1 - if not self.include_no_event_in_uts_target: - self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} - else: - self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX} - def _prepare_sex(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None: sex_values = pd.to_numeric(basic_table["sex"], errors="coerce").to_numpy() if np.isnan(sex_values).any(): @@ -289,12 +282,7 @@ class _ExpoBaseDataset(Dataset): class NextStepHealthDataset(_ExpoBaseDataset): """ - Dataset for next-token and next-time-point losses with unified other-info - tokens. - - Returned targets cover both: - - Delphi2MLoss: target_event_seq, target_time_seq - - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot + Delphi2M next-token dataset with unified other-info tokens. """ CACHE_VERSION = 3 @@ -304,14 +292,12 @@ class NextStepHealthDataset(_ExpoBaseDataset): data_prefix: str = "ukb", labels_file: str = "labels.csv", no_event_interval_years: float = 5.0, - include_no_event_in_uts_target: bool = False, extra_info_types: Iterable[int] | None = None, ) -> None: super().__init__( 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_in_uts_target, extra_info_types=extra_info_types, ) @@ -326,23 +312,18 @@ class NextStepHealthDataset(_ExpoBaseDataset): if features is None: continue - target_pack = build_all_targets( + targets = build_next_token_targets( labels=labels, times_days=times_days, - vocab_size=self.vocab_size, - ignored_uts_target_ids=self.ignored_uts_target_ids, require_sorted=True, ) self.samples.append({ "eid": eid, - "event_seq": target_pack.next_token.input_events, - "time_seq": target_pack.next_token.input_times_years, - "target_event_seq": target_pack.next_token.target_events, - "target_time_seq": target_pack.next_token.target_times_years, - "readout_mask": target_pack.unique_time_set.readout_mask, - "target_dt_unique": target_pack.unique_time_set.target_dt_unique, - "target_multi_hot": target_pack.unique_time_set.target_multi_hot, + "event_seq": targets.input_events, + "time_seq": targets.input_times_years, + "target_event_seq": targets.target_events, + "target_time_seq": targets.target_times_years, **features, }) @@ -361,9 +342,6 @@ class NextStepHealthDataset(_ExpoBaseDataset): "other_time": torch.from_numpy(s["other_time"]).float(), "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(), "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(), - "readout_mask": torch.from_numpy(s["readout_mask"]).bool(), - "target_dt_unique": torch.from_numpy(s["target_dt_unique"]).float(), - "target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(), } @@ -386,7 +364,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset): labels_file: str = "labels.csv", split: Literal["train", "valid", "test"] = "train", no_event_interval_years: float = 5.0, - include_no_event_in_uts_target: bool = False, min_history_events: int = 1, min_future_events: int = 1, validation_query_seed: int = 42, @@ -399,7 +376,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset): 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_in_uts_target, extra_info_types=extra_info_types, ) @@ -605,31 +581,12 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict: batch_first=True, padding_value=0.0, ) - readout_mask = pad_sequence( - [s["readout_mask"] for s in batch], - batch_first=True, - padding_value=False, - ) - target_dt_unique = pad_sequence( - [s["target_dt_unique"] for s in batch], - batch_first=True, - padding_value=0.0, - ) - target_multi_hot = pad_sequence( - [s["target_multi_hot"] for s in batch], - batch_first=True, - padding_value=False, - ) - out = { "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": readout_mask, - "target_dt_unique": target_dt_unique, - "target_multi_hot": target_multi_hot, } out.update(_collate_common_static(batch)) return out diff --git a/eval_data.py b/eval_data.py index f139b3c..5e70b9f 100644 --- a/eval_data.py +++ b/eval_data.py @@ -1,15 +1,209 @@ from __future__ import annotations -from typing import Any, Dict, Iterable, List +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. @@ -52,7 +246,6 @@ class AllFutureSequenceEvalDataset: times = np.asarray(patient["times"], dtype=np.float32) if labels.size < 2: continue - input_len = int(labels.size - 1) self.samples.append( { "eid": int(patient["eid"]), @@ -60,7 +253,6 @@ class AllFutureSequenceEvalDataset: "time_seq": times[:-1], "target_event_seq": labels[1:], "target_time_seq": times[1:], - "readout_mask": np.ones(input_len, dtype=bool), "sex": int(patient["sex"]), "other_type": np.asarray(patient["other_type"], dtype=np.int64), "other_value": np.asarray(patient["other_value"], dtype=np.float32), @@ -79,7 +271,6 @@ class AllFutureSequenceEvalDataset: "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(), - "readout_mask": torch.from_numpy(s["readout_mask"]).bool(), "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(), @@ -94,7 +285,6 @@ def load_sequence_eval_dataset( data_prefix: str, labels_file: str, no_event_interval_years: float, - include_no_event_in_uts_target: bool, min_history_events: int, min_future_events: int, extra_info_types: Iterable[int] | None, @@ -105,7 +295,6 @@ def load_sequence_eval_dataset( 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_in_uts_target, extra_info_types=extra_info_types, ) if mode == "all_future": @@ -132,9 +321,6 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, target_time_seq = pad_sequence( [s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0 ) - readout_mask = pad_sequence( - [s["readout_mask"] for s in batch], batch_first=True, padding_value=False - ) other_type = pad_sequence( [s["other_type"] for s in batch], batch_first=True, padding_value=0 ) @@ -154,7 +340,7 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, "padding_mask": event_seq > PAD_IDX, "target_event_seq": target_event_seq, "target_time_seq": target_time_seq, - "readout_mask": readout_mask, + "readout_mask": event_seq > PAD_IDX, "sex": torch.stack([s["sex"] for s in batch]), "other_type": other_type, "other_value": other_value, diff --git a/evaluate_auc.py b/evaluate_auc.py index ee0ee4d..b2a8047 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -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)), diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 4c32a3c..fdc29de 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -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, diff --git a/event_date_utils.py b/event_date_utils.py deleted file mode 100644 index 24e602e..0000000 --- a/event_date_utils.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Read and query calendar-dated disease-event arrays from prepare_event_dates.py.""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Iterable - -import numpy as np -import pandas as pd - - -REQUIRED_FIELDS = {"eid", "event_date", "token"} - - -def load_event_dates(path: str | Path) -> np.ndarray: - """Load and validate the structured ``.npy`` event array.""" - events = np.load(path) - if events.dtype.names is None or not REQUIRED_FIELDS.issubset(events.dtype.names): - raise ValueError( - "Expected a structured .npy with eid, event_date, token fields. " - "Create it with prepare_event_dates.py." - ) - return events - - -def load_token_labels(labels_file: str | Path) -> dict[int, str]: - """Load token -> human-readable code using the project label convention.""" - labels = {1: "CHECKUP"} - with Path(labels_file).open(encoding="utf-8") as handle: - for index, line in enumerate(handle): - code = line.strip().split(" ", maxsplit=1)[0] - if code: - labels[index + 2] = code - return labels - - -@dataclass -class EventDateIndex: - """Small in-memory query wrapper for exposure-linkage and cohort scripts.""" - - events: np.ndarray - token_labels: dict[int, str] | None = None - - @classmethod - def from_files( - cls, - event_file: str | Path, - labels_file: str | Path | None = None, - ) -> "EventDateIndex": - labels = load_token_labels(labels_file) if labels_file is not None else None - return cls(load_event_dates(event_file), labels) - - def to_frame(self, events: np.ndarray | None = None) -> pd.DataFrame: - """Convert records to a convenient, calendar-dated DataFrame.""" - data = self.events if events is None else events - frame = pd.DataFrame( - { - "eid": data["eid"].astype("int64"), - "event_date": pd.to_datetime(data["event_date"]), - "token": data["token"].astype("int32"), - } - ) - if self.token_labels is not None: - frame["label_code"] = frame["token"].map(self.token_labels).fillna("UNKNOWN") - return frame.sort_values(["eid", "event_date", "token"], kind="stable").reset_index(drop=True) - - def for_eid(self, eid: int) -> pd.DataFrame: - """Return every stored disease/death event for one participant.""" - return self.to_frame(self.events[self.events["eid"] == int(eid)]) - - def between( - self, - start: str | pd.Timestamp, - end: str | pd.Timestamp, - *, - eids: Iterable[int] | None = None, - tokens: Iterable[int] | None = None, - ) -> pd.DataFrame: - """Query events in an inclusive calendar-date interval.""" - start_day = np.datetime64(pd.Timestamp(start).date(), "D") - end_day = np.datetime64(pd.Timestamp(end).date(), "D") - mask = (self.events["event_date"] >= start_day) & (self.events["event_date"] <= end_day) - if eids is not None: - mask &= np.isin(self.events["eid"], list(eids)) - if tokens is not None: - mask &= np.isin(self.events["token"], list(tokens)) - return self.to_frame(self.events[mask]) - - def anchors_before(self, eid: int, date: str | pd.Timestamp) -> pd.DataFrame: - """Return a participant's event history strictly before an exposure anchor.""" - day = np.datetime64(pd.Timestamp(date).date(), "D") - mask = (self.events["eid"] == int(eid)) & (self.events["event_date"] < day) - return self.to_frame(self.events[mask]) - - def first_event(self, token: int) -> pd.DataFrame: - """Return each participant's first date for a requested token.""" - selected = self.events[self.events["token"] == int(token)] - # Arrays produced by prepare_event_dates.py are already deduplicated; - # sorting makes this safe for externally produced compatible arrays too. - order = np.lexsort((selected["event_date"], selected["eid"])) - selected = selected[order] - _, first = np.unique(selected["eid"], return_index=True) - return self.to_frame(selected[first]) diff --git a/future_risk.py b/future_risk.py deleted file mode 100644 index 8480b1a..0000000 --- a/future_risk.py +++ /dev/null @@ -1,115 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence - -import torch -import torch.nn.functional as F - - -def death_token(vocab_size: int) -> int: - if int(vocab_size) <= 0: - raise ValueError(f"vocab_size must be positive, got {vocab_size}") - return int(vocab_size) - 1 - - -def probabilities_from_logits( - logits: torch.Tensor, - tau_years: float | torch.Tensor, - *, - dist_mode: str = "exponential", - rho: torch.Tensor | None = None, - death_rho: torch.Tensor | None = None, - eps: float = 1e-8, -) -> torch.Tensor: - """ - Convert all-future logits to tau-year event probabilities. - - Death is always treated as token vocab_size - 1. For dist_mode="mixed", - non-death tokens use exponential hazards and death uses death_rho. - """ - if logits.ndim != 2: - raise ValueError(f"logits must have shape (N, V), got {tuple(logits.shape)}") - if float(torch.as_tensor(tau_years).detach().min().cpu()) < 0: - raise ValueError("tau_years must be non-negative") - - mode = str(dist_mode).lower() - if mode not in {"exponential", "weibull", "mixed"}: - raise ValueError("dist_mode must be one of: exponential, weibull, mixed") - - rate = F.softplus(logits) + float(eps) - tau = torch.as_tensor(tau_years, dtype=rate.dtype, device=rate.device) - if tau.ndim == 0: - tau = tau.expand(logits.shape[0]) - if tau.ndim != 1 or tau.shape[0] != logits.shape[0]: - raise ValueError( - "tau_years must be a scalar or a 1D tensor with length N, got " - f"{tuple(tau.shape)} for N={logits.shape[0]}" - ) - - if mode == "exponential": - exposure = tau[:, None].expand_as(rate) - elif mode == "weibull": - if rho is None or rho.shape != logits.shape: - raise ValueError("rho must have the same shape as logits for dist_mode='weibull'") - exposure = torch.pow(tau[:, None].clamp_min(float(eps)), rho.to(rate.dtype)) - else: - exposure = tau[:, None].expand_as(rate).clone() - if death_rho is None: - raise ValueError("death_rho is required for dist_mode='mixed'") - death_idx = death_token(logits.shape[1]) - death_shape = tuple(death_rho.shape) - death_rho = death_rho.to(device=rate.device, dtype=rate.dtype) - if death_rho.ndim == 2 and death_rho.shape[1] == 1: - death_rho = death_rho.squeeze(1) - if death_rho.ndim != 1 or death_rho.shape[0] != logits.shape[0]: - raise ValueError( - "death_rho must have shape (N,) or (N, 1), got " - f"{death_shape} for N={logits.shape[0]}" - ) - exposure[:, death_idx] = torch.pow(tau.clamp_min(float(eps)), death_rho) - - return -torch.expm1(-rate * exposure) - - -def death_risk_from_probabilities(probabilities: torch.Tensor) -> torch.Tensor: - """Return p_death(t, tau), with death fixed to token vocab_size - 1.""" - if probabilities.ndim != 2: - raise ValueError( - f"probabilities must have shape (N, V), got {tuple(probabilities.shape)}" - ) - return probabilities[:, death_token(probabilities.shape[1])] - - -def new_disease_risk_from_probabilities( - probabilities: torch.Tensor, - occurred: torch.Tensor, - disease_ids: Sequence[int], -) -> torch.Tensor: - """ - Compute P(at least one selected disease newly occurs within tau years). - - Already occurred diseases are masked out. Death is not included here and - should be reported separately with death_risk_from_probabilities. - """ - if probabilities.ndim != 2 or occurred.shape != probabilities.shape: - raise ValueError( - "probabilities and occurred must both have shape (N, V), got " - f"{tuple(probabilities.shape)} and {tuple(occurred.shape)}" - ) - if not disease_ids: - return probabilities.new_zeros(probabilities.shape[0]) - - death_idx = death_token(probabilities.shape[1]) - ids = [ - idx - for idx in dict.fromkeys(int(x) for x in disease_ids) - if 0 <= idx < probabilities.shape[1] and idx != death_idx - ] - if not ids: - return probabilities.new_zeros(probabilities.shape[0]) - - idx_tensor = torch.as_tensor(ids, dtype=torch.long, device=probabilities.device) - p = probabilities[:, idx_tensor].clamp(0.0, 1.0 - 1e-7) - new_mask = ~occurred[:, idx_tensor].to(dtype=torch.bool) - log_no_new = torch.log1p(-p) * new_mask.to(dtype=p.dtype) - return -torch.expm1(log_no_new.sum(dim=1)) diff --git a/losses.py b/losses.py index 51883a4..ee44c60 100644 --- a/losses.py +++ b/losses.py @@ -145,95 +145,6 @@ class Delphi2MLoss(nn.Module): return total_loss -class UniqueTimeSetExponentialLoss(nn.Module): - """Next distinct timestamp event-set supervision with sum reduction.""" - - def __init__( - self, - ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX), - t_min: float = 1.0 / 365.25, - max_exp_input: float = 60.0, - exclude_ignored_from_intensity: bool = True, - ): - super().__init__() - self.ignored_idx = [int(x) for x in ignored_idx] - self.t_min = float(t_min) - self.max_exp_input = float(max_exp_input) - self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity) - - def forward( - self, - logits: torch.Tensor, - target_multi_hot: torch.Tensor, - target_dt_unique: torch.Tensor, - readout_mask: torch.Tensor, - return_components: bool = False, - ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: - if logits.dim() != 3: - raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}") - bsz, seq_len, vocab_size = logits.shape - - if target_multi_hot.shape != (bsz, seq_len, vocab_size): - raise ValueError( - "target_multi_hot must match logits shape, " - f"got {tuple(target_multi_hot.shape)} vs {tuple(logits.shape)}" - ) - if target_dt_unique.shape != (bsz, seq_len): - raise ValueError( - f"target_dt_unique must be {(bsz, seq_len)}, got {tuple(target_dt_unique.shape)}" - ) - if readout_mask.shape != (bsz, seq_len): - raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}") - - ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device) - - num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1) - valid_mask = readout_mask.bool() & (num_targets > 0) - - if not valid_mask.any(): - total_loss = _zero_loss_like(logits) - if return_components: - return total_loss, { - "observed": total_loss.detach(), - "penalty": total_loss.detach(), - "total": total_loss.detach(), - } - return total_loss - - logits_safe = torch.nan_to_num( - logits[valid_mask], - nan=0.0, - posinf=self.max_exp_input, - neginf=-self.max_exp_input, - ) - target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype) - target_valid[:, ignore_mask] = 0.0 - - observed_term = (logits_safe * target_valid).sum(dim=-1) - penalty_scale = target_valid.sum(dim=-1) - - logits_for_lse = logits_safe - if self.exclude_ignored_from_intensity: - logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf")) - - dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min) - log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1) - log_penalty = log_lambda_total + dt_clamped.log() - penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input)) - - observed_loss = -observed_term - penalty_loss = penalty_scale * penalty - total_loss = (observed_loss + penalty_loss).mean() - - if return_components: - return total_loss, { - "observed": observed_loss.mean().detach(), - "penalty": penalty_loss.mean().detach(), - "total": total_loss.detach(), - } - return total_loss - - class ExponentialLoss(nn.Module): """Query-conditioned all-future-event exponential point-process loss.""" @@ -385,10 +296,8 @@ class MixedLoss(nn.Module): def build_loss(name: str, **kwargs) -> nn.Module: name = name.lower() - if name in {"delphi2m", "d2m", "next_token"}: + if name == "delphi2m": return Delphi2MLoss(**kwargs) - if name in {"uts", "unique_time_set", "unique_time_exponential"}: - return UniqueTimeSetExponentialLoss(**kwargs) if name in {"exponential", "query_exponential"}: return ExponentialLoss(**kwargs) if name in {"weibull", "query_weibull"}: @@ -396,5 +305,5 @@ def build_loss(name: str, **kwargs) -> nn.Module: if name in {"mixed", "query_mixed"}: return MixedLoss(**kwargs) raise ValueError( - f"Unknown loss {name!r}. Available: delphi2m, uts, exponential, weibull, mixed." + f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull, mixed." ) diff --git a/models.py b/models.py index d6104dc..dd5dae6 100644 --- a/models.py +++ b/models.py @@ -156,7 +156,7 @@ class DeepHealth(nn.Module): n_value_kinds: int = 3, n_bins: int = 16, target_mode: str = "next_token", # "next_token" or "all_future" - time_mode: str = "relative", # "relative" or "absolute" + time_mode: str = "absolute", # next_token requires absolute dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" extra_pool_reduce: str = "mean", dropout: float = 0.0, @@ -169,6 +169,11 @@ class DeepHealth(nn.Module): if time_mode not in ["relative", "absolute"]: raise ValueError( "time_mode must be either 'relative' or 'absolute'") + if target_mode == "next_token" and time_mode != "absolute": + raise ValueError( + "next_token is reserved for Delphi2M reproduction and " + "requires time_mode='absolute'" + ) if dist_mode not in ["exponential", "weibull", "mixed"]: raise ValueError( "dist_mode must be either 'exponential', 'weibull' or 'mixed'") @@ -461,15 +466,8 @@ class DeepHealth(nn.Module): ) return h_disease[:, :event_len, :] - def forward_next_token(self, **kwargs) -> torch.Tensor: - return self._forward_shared(mode="next_token", **kwargs) - - def forward_all_future(self, **kwargs) -> torch.Tensor: - return self._forward_shared(mode="all_future", **kwargs) - - def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor: - mode = self.target_mode if target_mode is None else target_mode - return self._forward_shared(mode=mode, **kwargs) + def forward(self, **kwargs) -> torch.Tensor: + return self._forward_shared(mode=self.target_mode, **kwargs) def calc_risk(self, x: torch.Tensor) -> torch.Tensor: return self.risk_head(x) diff --git a/readouts.py b/readouts.py deleted file mode 100644 index b36285c..0000000 --- a/readouts.py +++ /dev/null @@ -1,107 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -import torch -import torch.nn as nn - - -@dataclass -class ReadoutOutput: - hidden: torch.Tensor - readout_mask: torch.Tensor - - -class TokenReadout(nn.Module): - def forward( - self, - hidden: torch.Tensor, - time_seq: torch.Tensor, - padding_mask: torch.Tensor, - readout_mask: torch.Tensor | None = None, - ) -> ReadoutOutput: - mask = padding_mask if readout_mask is None else readout_mask - return ReadoutOutput(hidden=hidden, readout_mask=mask.bool()) - - -class SameTimeGroupEndReadout(nn.Module): - def __init__(self, reduce: str = "mean"): - super().__init__() - if reduce not in {"mean", "sum"}: - raise ValueError("reduce must be either 'mean' or 'sum'") - self.reduce = reduce - - def forward( - self, - hidden: torch.Tensor, - time_seq: torch.Tensor, - padding_mask: torch.Tensor, - readout_mask: torch.Tensor | None = None, - ) -> ReadoutOutput: - if readout_mask is None: - next_is_new_time = torch.ones_like(padding_mask, dtype=torch.bool) - next_is_new_time[:, :-1] = time_seq[:, 1:] != time_seq[:, :-1] - readout_mask = padding_mask.bool() & next_is_new_time - else: - readout_mask = readout_mask.bool() - - group_start = torch.ones_like(padding_mask, dtype=torch.bool) - group_start[:, 1:] = time_seq[:, 1:] != time_seq[:, :-1] - group_start = group_start & padding_mask.bool() - - group_id = group_start.long().cumsum(dim=1) - 1 - group_id = group_id.clamp_min(0) - max_groups = hidden.size(1) - - group_sum = hidden.new_zeros(hidden.size(0), max_groups, hidden.size(2)) - group_sum.scatter_add_( - 1, - group_id.unsqueeze(-1).expand_as(hidden), - hidden * padding_mask.unsqueeze(-1).to(hidden.dtype), - ) - - if self.reduce == "mean": - group_count = hidden.new_zeros(hidden.size(0), max_groups, 1) - group_count.scatter_add_( - 1, - group_id.unsqueeze(-1), - padding_mask.unsqueeze(-1).to(hidden.dtype), - ) - group_sum = group_sum / group_count.clamp_min(1.0) - - out = hidden.clone() - out[readout_mask] = group_sum.gather( - 1, - group_id.unsqueeze(-1).expand_as(hidden), - )[readout_mask] - return ReadoutOutput(hidden=out, readout_mask=readout_mask) - - -class LastValidReadout(nn.Module): - def forward( - self, - hidden: torch.Tensor, - time_seq: torch.Tensor, - padding_mask: torch.Tensor, - readout_mask: torch.Tensor | None = None, - ) -> ReadoutOutput: - batch_size, seq_len = padding_mask.shape - last_idx = padding_mask.long().sum(dim=1).clamp_min(1) - 1 - out = hidden[torch.arange(batch_size, device=hidden.device), last_idx] - mask = torch.ones(batch_size, dtype=torch.bool, device=hidden.device) - return ReadoutOutput(hidden=out, readout_mask=mask) - - -def build_readout(name: str, **kwargs) -> nn.Module: - name = name.lower() - if name == "token": - return TokenReadout() - if name in {"same_time_group_end", "same_time"}: - return SameTimeGroupEndReadout(**kwargs) - if name == "last_valid": - return LastValidReadout() - raise ValueError( - "Unknown readout {!r}. Available: token, same_time_group_end, last_valid.".format( - name - ) - ) diff --git a/targets.py b/targets.py index f64de6a..818c4e2 100644 --- a/targets.py +++ b/targets.py @@ -2,8 +2,6 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Iterable - import numpy as np @@ -32,38 +30,6 @@ class NextTokenTargets: target_times_years: np.ndarray -@dataclass(frozen=True) -class UniqueTimeSetTargets: - """ - Unique-time set supervision targets. - - Shapes: - readout_mask: (L,) - target_dt_unique: (L,) - target_multi_hot: (L, vocab_size) - - where L = N - 1. - - Only group-end positions can have readout_mask=True. - target_dt_unique is measured in years. - """ - readout_mask: np.ndarray - target_dt_unique: np.ndarray - target_multi_hot: np.ndarray - - -@dataclass(frozen=True) -class TargetPack: - """ - Combined target package for one patient sequence. - - Contains both next-token targets and unique-time-set targets. - The training pipeline decides which one to use. - """ - next_token: NextTokenTargets - unique_time_set: UniqueTimeSetTargets - - def _as_numpy_1d( x: np.ndarray, name: str, @@ -163,232 +129,3 @@ def build_next_token_targets( target_events=target_events, target_times_years=target_times_years, ) - - -def build_unique_time_set_targets( - labels: np.ndarray, - times_days: np.ndarray, - *, - vocab_size: int, - ignored_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX), - require_sorted: bool = True, -) -> UniqueTimeSetTargets: - """ - Build next-unique-time set targets. - - This is the target construction used by your UTS / default mode. - - For each input position i: - - only if i is the last token of its timestamp group; - - find the next distinct timestamp group; - - target is the set of valid event labels at that next timestamp. - - Example: - - t=49: X - t=50: A, B, C - t=51: D, E - - Supervises: - - X@49 -> {A, B, C}@50 - group_end@50 -> {D, E}@51 - - It does NOT supervise: - - A@50 -> B@50 - B@50 -> C@50 - - Parameters - ---------- - labels: - Full event sequence labels, shape (N,). - - times_days: - Full event sequence times in days, shape (N,). - - vocab_size: - Size of output vocabulary. - - ignored_target_ids: - Label ids that should not enter target_multi_hot. - Usually: - no no-event: {0, 1} - with no-event: {0, 1, 2} - For UTS, I recommend ignoring unless explicitly testing it - as an event target. - - Returns - ------- - UniqueTimeSetTargets - """ - labels = _as_numpy_1d(labels, "labels", np.int64) - times_days = _as_numpy_1d(times_days, "times_days", np.float32) - validate_event_sequence(labels, times_days, require_sorted=require_sorted) - - if vocab_size <= 0: - raise ValueError(f"vocab_size must be positive, got {vocab_size}") - - if len(labels) < 2: - raise ValueError( - "Need at least two events to build unique-time-set targets." - ) - - input_len = len(labels) - 1 - - readout_mask = np.zeros(input_len, dtype=bool) - target_dt_unique = np.zeros(input_len, dtype=np.float32) - target_multi_hot = np.zeros((input_len, vocab_size), dtype=bool) - - ignored = {int(x) for x in ignored_target_ids} - - unique_times = np.unique(times_days) - time_to_group_idx = {t: i for i, t in enumerate(unique_times)} - group_indices = np.array([time_to_group_idx[t] - for t in times_days], dtype=np.int64) - - for i in range(input_len): - current_group = group_indices[i] - - is_last_in_group = ( - i == input_len - 1 - or group_indices[i + 1] != current_group - ) - if not is_last_in_group: - continue - - next_group_idx = current_group + 1 - if next_group_idx >= len(unique_times): - continue - - next_time = unique_times[next_group_idx] - next_labels = labels[group_indices == next_group_idx] - - valid_next_labels: list[int] = [] - for lab in next_labels: - lab_int = int(lab) - if lab_int in ignored: - continue - if lab_int < 0 or lab_int >= vocab_size: - continue - valid_next_labels.append(lab_int) - - # If next timestamp contains only technical tokens, do not supervise UTS. - if len(valid_next_labels) == 0: - continue - - readout_mask[i] = True - target_dt_unique[i] = float(next_time - times_days[i]) / DAYS_PER_YEAR - target_multi_hot[i, valid_next_labels] = True - - return UniqueTimeSetTargets( - readout_mask=readout_mask, - target_dt_unique=target_dt_unique.astype(np.float32), - target_multi_hot=target_multi_hot, - ) - - -def build_all_targets( - labels: np.ndarray, - times_days: np.ndarray, - *, - vocab_size: int, - ignored_uts_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX), - require_sorted: bool = True, -) -> TargetPack: - """ - Build both next-token targets and unique-time-set targets for one patient. - - This is the function dataset.py should usually call during initialization. - - The dataset can then store: - event_seq = target_pack.next_token.input_events - time_seq = target_pack.next_token.input_times_years - - target_event_seq = target_pack.next_token.target_events - target_time_seq = target_pack.next_token.target_times_years - - readout_mask = target_pack.unique_time_set.readout_mask - target_dt_unique = target_pack.unique_time_set.target_dt_unique - target_multi_hot = target_pack.unique_time_set.target_multi_hot - """ - next_token = build_next_token_targets( - labels=labels, - times_days=times_days, - require_sorted=require_sorted, - ) - - unique_time_set = build_unique_time_set_targets( - labels=labels, - times_days=times_days, - vocab_size=vocab_size, - ignored_target_ids=ignored_uts_target_ids, - require_sorted=require_sorted, - ) - - return TargetPack( - next_token=next_token, - unique_time_set=unique_time_set, - ) - - -def get_group_end_mask_from_times( - times_days: np.ndarray, - *, - input_len: int | None = None, -) -> np.ndarray: - """ - Convenience utility for debugging. - - Returns a bool mask indicating the last token of each same-time group - within the input sequence. - - If input_len is None, uses len(times_days) - 1, matching model input length. - """ - times_days = _as_numpy_1d(times_days, "times_days", np.float32) - - if input_len is None: - input_len = len(times_days) - 1 - - if input_len < 0 or input_len > len(times_days): - raise ValueError( - f"Invalid input_len={input_len} for sequence length {len(times_days)}" - ) - - out = np.zeros(input_len, dtype=bool) - - for i in range(input_len): - is_last_in_group = ( - i == input_len - 1 - or times_days[i + 1] != times_days[i] - ) - out[i] = is_last_in_group - - return out - - -def summarize_targets( - target_pack: TargetPack, -) -> dict[str, int | float]: - """ - Small debugging helper for logging. - """ - nt = target_pack.next_token - uts = target_pack.unique_time_set - - n_tokens = int(len(nt.input_events)) - n_readout = int(uts.readout_mask.sum()) - n_positive_labels = int(uts.target_multi_hot.sum()) - - mean_set_size = ( - float(n_positive_labels / n_readout) - if n_readout > 0 - else 0.0 - ) - - return { - "n_input_tokens": n_tokens, - "n_uts_readouts": n_readout, - "n_uts_positive_labels": n_positive_labels, - "mean_uts_set_size": mean_set_size, - } diff --git a/test_model_architectures.py b/test_model_architectures.py deleted file mode 100644 index 729c528..0000000 --- a/test_model_architectures.py +++ /dev/null @@ -1,247 +0,0 @@ -import unittest - -import torch - -from backbones import ( - SwiGLU, - TrajMixer, - TrajMixerBlock, - TransformerFFNBlock, - build_backbone_block, -) -from model_architectures import ( - TRAJ_MIXER_ARCHITECTURE, - TRANSFORMER_FFN_ARCHITECTURE, - detect_model_architecture_from_state_dict, - resolve_model_architecture, -) -from models import DeepHealth - - -def _build_block(model_architecture: str): - return build_backbone_block( - model_architecture, - n_embd=12, - n_head=3, - use_time_rope=False, - use_rbf_bias=False, - mlp_dropout=0.0, - ) - - -def _as_model_state_dict(block: torch.nn.Module) -> dict[str, torch.Tensor]: - return { - f"blocks.0.{name}": value.detach().clone() - for name, value in block.state_dict().items() - } - - -def _build_model( - model_architecture: str | None, - *, - n_layer: int = 1, -) -> DeepHealth: - return DeepHealth( - vocab_size=8, - n_embd=12, - n_head=3, - n_layer=n_layer, - n_types=2, - n_cont_types=0, - n_categories=2, - cont_type_ids=[], - time_mode="absolute", - model_architecture=model_architecture, - ) - - -class ModelArchitectureFactoryTest(unittest.TestCase): - def test_factory_builds_both_architectures_with_expected_topology(self) -> None: - ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE) - self.assertIsInstance(ffn_block, TransformerFFNBlock) - self.assertIsInstance(ffn_block.mlp, SwiGLU) - self.assertTrue(hasattr(ffn_block, "ln1")) - self.assertTrue(hasattr(ffn_block, "ln2")) - - traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE) - self.assertIsInstance(traj_block, TrajMixerBlock) - self.assertIsInstance(traj_block.mlp, TrajMixer) - self.assertTrue(hasattr(traj_block, "ln1")) - self.assertFalse(hasattr(traj_block, "ln2")) - - def test_both_architectures_forward_and_backward(self) -> None: - for architecture in ( - TRANSFORMER_FFN_ARCHITECTURE, - TRAJ_MIXER_ARCHITECTURE, - ): - with self.subTest(architecture=architecture): - torch.manual_seed(0) - block = _build_block(architecture) - x = torch.randn(2, 5, 12, requires_grad=True) - - output = block(x) - self.assertEqual(output.shape, x.shape) - output.square().mean().backward() - - self.assertIsNotNone(x.grad) - self.assertTrue(torch.isfinite(x.grad).all()) - self.assertGreater(x.grad.abs().sum().item(), 0.0) - self.assertIsNotNone(block.attn.qkv.weight.grad) - self.assertGreater( - block.attn.qkv.weight.grad.abs().sum().item(), - 0.0, - ) - - if architecture == TRANSFORMER_FFN_ARCHITECTURE: - branch_parameters = ( - block.mlp.w1.weight, - block.mlp.w2.weight, - block.mlp.w3.weight, - ) - else: - branch_parameters = ( - block.mlp.intra_gate_proj, - block.mlp.intra_value_proj, - block.mlp.output_proj, - ) - for parameter in branch_parameters: - self.assertIsNotNone(parameter.grad) - self.assertTrue(torch.isfinite(parameter.grad).all()) - self.assertGreater(parameter.grad.abs().sum().item(), 0.0) - - def test_unknown_architecture_is_rejected(self) -> None: - with self.assertRaises(ValueError): - _build_block("unknown_architecture") - with self.assertRaisesRegex(ValueError, "model_architecture is required"): - _build_model(None) - - def test_deephealth_rejects_fewer_than_one_layer(self) -> None: - for n_layer in (0, -1): - with self.subTest(n_layer=n_layer): - with self.assertRaisesRegex(ValueError, "n_layer must be >= 1"): - _build_model( - TRANSFORMER_FFN_ARCHITECTURE, - n_layer=n_layer, - ) - - def test_deephealth_uses_factory_and_strictly_reloads_both_models(self) -> None: - for architecture, block_class in ( - (TRANSFORMER_FFN_ARCHITECTURE, TransformerFFNBlock), - (TRAJ_MIXER_ARCHITECTURE, TrajMixerBlock), - ): - with self.subTest(architecture=architecture): - model = _build_model(architecture) - self.assertEqual(model.model_architecture, architecture) - self.assertIsInstance(model.blocks[0], block_class) - self.assertEqual( - detect_model_architecture_from_state_dict( - model.state_dict() - ), - architecture, - ) - - reloaded = _build_model(architecture) - incompatible = reloaded.load_state_dict( - model.state_dict(), - strict=True, - ) - self.assertEqual(incompatible.missing_keys, []) - self.assertEqual(incompatible.unexpected_keys, []) - - -class ModelArchitectureResolutionTest(unittest.TestCase): - def setUp(self) -> None: - self.ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE) - self.traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE) - self.ffn_state = _as_model_state_dict(self.ffn_block) - self.traj_state = _as_model_state_dict(self.traj_block) - - def test_state_dict_detection_recognizes_both_architectures(self) -> None: - self.assertEqual( - detect_model_architecture_from_state_dict(self.ffn_state), - TRANSFORMER_FFN_ARCHITECTURE, - ) - self.assertEqual( - detect_model_architecture_from_state_dict(self.traj_state), - TRAJ_MIXER_ARCHITECTURE, - ) - - def test_explicit_markers_resolve_when_checkpoint_matches(self) -> None: - for architecture, state_dict in ( - (TRANSFORMER_FFN_ARCHITECTURE, self.ffn_state), - (TRAJ_MIXER_ARCHITECTURE, self.traj_state), - ): - with self.subTest(architecture=architecture): - self.assertEqual( - resolve_model_architecture( - {"model_architecture": architecture}, - state_dict, - ), - architecture, - ) - - def test_architecture_marker_is_required_for_checkpoint_loading(self) -> None: - with self.assertRaisesRegex(ValueError, "model_architecture is required"): - resolve_model_architecture({}, self.ffn_state) - with self.assertRaisesRegex(ValueError, "model_architecture is required"): - resolve_model_architecture(None, self.traj_state) - - def test_explicit_marker_conflicting_with_state_dict_is_rejected(self) -> None: - conflicts = ( - (TRANSFORMER_FFN_ARCHITECTURE, self.traj_state), - (TRAJ_MIXER_ARCHITECTURE, self.ffn_state), - ) - for architecture, state_dict in conflicts: - with self.subTest(architecture=architecture): - with self.assertRaises(ValueError): - resolve_model_architecture( - {"model_architecture": architecture}, - state_dict, - ) - - def test_unknown_marker_and_ambiguous_state_dict_are_rejected(self) -> None: - with self.assertRaises(ValueError): - resolve_model_architecture( - {"model_architecture": "traj_mixer_v4"} - ) - - ambiguous_state = dict(self.ffn_state) - ambiguous_state.update(self.traj_state) - with self.assertRaises(ValueError): - detect_model_architecture_from_state_dict(ambiguous_state) - - with self.assertRaises(ValueError): - detect_model_architecture_from_state_dict( - {"token_embedding.weight": torch.empty(2, 2)} - ) - - def test_ffn_block_schema_is_stable_and_strictly_loadable(self) -> None: - expected_keys = { - "attn.time_bias_scale", - "attn.qkv.weight", - "attn.out_proj.weight", - "attn.rbf_proj.weight", - "mlp.w1.weight", - "mlp.w1.bias", - "mlp.w2.weight", - "mlp.w2.bias", - "mlp.w3.weight", - "mlp.w3.bias", - "ln1.weight", - "ln1.bias", - "ln2.weight", - "ln2.bias", - } - state = self.ffn_block.state_dict() - self.assertSetEqual(set(state), expected_keys) - self.assertEqual(tuple(state["mlp.w1.weight"].shape), (30, 12)) - self.assertEqual(tuple(state["mlp.w3.weight"].shape), (12, 30)) - - reloaded = _build_block(TRANSFORMER_FFN_ARCHITECTURE) - incompatible = reloaded.load_state_dict(state, strict=True) - self.assertEqual(incompatible.missing_keys, []) - self.assertEqual(incompatible.unexpected_keys, []) - - -if __name__ == "__main__": - unittest.main() diff --git a/test_temporal_attention.py b/test_temporal_attention.py deleted file mode 100644 index a1a3475..0000000 --- a/test_temporal_attention.py +++ /dev/null @@ -1,46 +0,0 @@ -import unittest - -import torch - -from backbones import TemporalAttention - - -class TemporalAttentionTest(unittest.TestCase): - def test_zero_rbf_bias_has_live_projection_gradient(self) -> None: - torch.manual_seed(0) - attention = TemporalAttention( - n_embd=12, - n_head=3, - use_time_rope=False, - use_rbf_bias=True, - ) - features = torch.randn(2, 4, 4, 16) - target = torch.randn(2, 4, 4, 3) - - initial_bias = ( - attention.time_bias_scale.tanh() - * attention.rbf_proj(features) - ) - torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias)) - - (initial_bias * target).sum().backward() - projection_grad = attention.rbf_proj.weight.grad - self.assertIsNotNone(projection_grad) - self.assertGreater(projection_grad.abs().sum().item(), 0.0) - - with torch.no_grad(): - attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3) - attention.zero_grad(set_to_none=True) - updated_bias = ( - attention.time_bias_scale.tanh() - * attention.rbf_proj(features) - ) - (updated_bias * target).sum().backward() - - scale_grad = attention.time_bias_scale.grad - self.assertIsNotNone(scale_grad) - self.assertGreater(scale_grad.abs().item(), 0.0) - - -if __name__ == "__main__": - unittest.main() diff --git a/test_traj_mixer.py b/test_traj_mixer.py deleted file mode 100644 index 4f886e3..0000000 --- a/test_traj_mixer.py +++ /dev/null @@ -1,166 +0,0 @@ -import unittest - -import torch - -from backbones import TrajMixer - - -class TrajMixerTest(unittest.TestCase): - def test_default_shape_parameters_and_initialization(self) -> None: - mixer = TrajMixer( - n_embd=120, - n_head=10, - dropout=0.0, - ) - - x = torch.randn(2, 7, 120) - self.assertEqual(mixer(x).shape, x.shape) - self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040) - self.assertEqual(tuple(mixer.norm.normalized_shape), (120,)) - self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12)) - torch.testing.assert_close( - torch.sigmoid(mixer.intra_gate_logits.detach()), - torch.full((10, 12), 0.1), - ) - self.assertEqual(mixer.intra_hidden, 48) - self.assertEqual( - tuple(mixer.intra_gate_proj.shape), - (10, 12, 48), - ) - self.assertEqual( - tuple(mixer.intra_value_proj.shape), - (10, 12, 48), - ) - self.assertEqual( - tuple(mixer.intra_output_proj.shape), - (10, 48, 12), - ) - self.assertEqual(mixer.hidden_group, 40) - self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40)) - self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40)) - self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10)) - - def test_zero_final_output_projection_makes_mixer_identity(self) -> None: - torch.manual_seed(0) - mixer = TrajMixer(12, n_head=3, dropout=0.0) - with torch.no_grad(): - mixer.output_proj.zero_() - x = torch.randn(2, 5, 12) - torch.testing.assert_close(mixer(x), x) - - def test_forward_matches_single_outer_residual_formula(self) -> None: - torch.manual_seed(0) - mixer = TrajMixer(12, n_head=3, dropout=0.0) - mixer.eval() - x = torch.randn(2, 5, 12) - - grouped = mixer.norm(x).reshape(2, 5, 3, 4) - intra_output = mixer._intra_mix(grouped) - static_gate = torch.sigmoid(mixer.intra_gate_logits).view( - 1, 1, 3, 4 - ) - mixed_input = grouped + static_gate * intra_output - update = mixer._cross_mix(mixed_input).reshape(2, 5, 12) - - torch.testing.assert_close(mixer(x), x + update) - - def test_intra_stage_is_independent_across_groups(self) -> None: - torch.manual_seed(0) - mixer = TrajMixer(12, n_head=3, dropout=0.0) - mixer.eval() - - grouped = torch.randn(2, 4, 3, 4) - changed = grouped.clone() - changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :]) - - original_out = mixer._intra_mix(grouped) - changed_out = mixer._intra_mix(changed) - unchanged_groups = torch.tensor([0, 2]) - torch.testing.assert_close( - original_out.index_select(2, unchanged_groups), - changed_out.index_select(2, unchanged_groups), - ) - - def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None: - mixer = TrajMixer(6, n_head=3, dropout=0.0) - mixer.eval() - with torch.no_grad(): - mixer.gate_proj.zero_() - mixer.value_proj.zero_() - mixer.output_proj.zero_() - - # Coordinate 0 reads group 0 through hidden unit 0 and writes it - # into group 1. Coordinate 1 must remain independent. - mixer.gate_proj[0, 0, 0] = 1.0 - mixer.value_proj[0, 0, 0] = 1.0 - mixer.output_proj[0, 0, 1] = 1.0 - - grouped = torch.tensor( - [[[ - [-1.0, 4.0], - [0.0, 5.0], - [1.0, 6.0], - ]]] - ) - changed = grouped.clone() - changed[0, 0, 0, 0] = 2.0 - - original_out = mixer._cross_mix(grouped) - changed_out = mixer._cross_mix(changed) - - self.assertNotEqual( - original_out[0, 0, 1, 0].item(), - changed_out[0, 0, 1, 0].item(), - ) - torch.testing.assert_close( - original_out[..., 1], - changed_out[..., 1], - ) - - def test_mixer_does_not_mix_sequence_positions(self) -> None: - torch.manual_seed(0) - mixer = TrajMixer(12, n_head=3, dropout=0.0) - mixer.eval() - x = torch.randn(2, 5, 12) - changed = x.clone() - changed[:, 3, :] += torch.randn_like(changed[:, 3, :]) - - original_out = mixer(x) - changed_out = mixer(changed) - unchanged_positions = torch.tensor([0, 1, 2, 4]) - torch.testing.assert_close( - original_out.index_select(1, unchanged_positions), - changed_out.index_select(1, unchanged_positions), - ) - - def test_gradients_reach_every_projection_family(self) -> None: - torch.manual_seed(1) - mixer = TrajMixer(12, n_head=3, dropout=0.0) - x = torch.randn(2, 4, 12, requires_grad=True) - - mixer(x).square().mean().backward() - - self.assertIsNotNone(x.grad) - self.assertTrue(torch.isfinite(x.grad).all()) - self.assertGreater(x.grad.abs().sum().item(), 0.0) - projection_names = ( - "intra_gate_proj", - "intra_value_proj", - "intra_output_proj", - "gate_proj", - "value_proj", - "output_proj", - ) - for name in projection_names: - parameter = getattr(mixer, name) - self.assertIsNotNone(parameter.grad, name) - self.assertTrue(torch.isfinite(parameter.grad).all(), name) - self.assertGreater(parameter.grad.abs().sum().item(), 0.0, name) - - def test_invalid_group_partition_is_rejected(self) -> None: - with self.assertRaisesRegex(ValueError, "divisible"): - TrajMixer(n_embd=121, n_head=10) - - -if __name__ == "__main__": - unittest.main() diff --git a/train_all_future.py b/train_all_future.py index 9a861f0..6bc5200 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -37,8 +37,10 @@ from train_util import ( configure_torch_for_training, create_unique_run_dir, format_extra_info_types, + get_lr, get_model_parameter_counts, load_extra_info_types_file, + move_batch_to_device, resolve_device, save_checkpoint, save_config, @@ -135,24 +137,6 @@ def parse_args() -> argparse.Namespace: return args -def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: - if epoch < args.warmup_epochs: - return adaptive_lr * (epoch + 1) / args.warmup_epochs - progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) - cosine = 0.5 * (1 + math.cos(math.pi * progress)) - return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) - - -def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: - non_blocking = device.type == "cuda" - return { - key: value.to(device, non_blocking=non_blocking) - if isinstance(value, torch.Tensor) - else value - for key, value in batch.items() - } - - def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth: return DeepHealth( vocab_size=dataset.vocab_size, @@ -214,7 +198,6 @@ def compute_all_future_loss( other_value=batch["other_value"], other_value_kind=batch["other_value_kind"], other_time=batch["other_time"], - target_mode="all_future", ) logits = model.calc_risk(hidden) diff --git a/train_next_step.py b/train_next_step.py index bf242a6..e6c8b76 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -1,10 +1,4 @@ -""" -Train DeepHealth with next-token / next-time-point supervision. - -The next-step dataset uses observed event histories, including CHECKUP state -tokens, plus optional gap imputation. UTS training reads out only -same-time group ends. -""" +"""Reproduce Delphi2M with absolute-time next-token supervision.""" from __future__ import annotations import argparse @@ -29,14 +23,15 @@ from model_architectures import ( SUPPORTED_MODEL_ARCHITECTURES, ) from models import DeepHealth, DeepHealthOutput -from readouts import build_readout -from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX +from targets import CHECKUP_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, format_extra_info_types, + get_lr, get_model_parameter_counts, load_extra_info_types_file, + move_batch_to_device, resolve_device, save_checkpoint, save_config, @@ -70,7 +65,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=42) parser.add_argument("--extra_info_types_file", type=str, default=None) parser.add_argument("--no_event_interval_years", type=float, default=5.0) - parser.add_argument("--include_no_event_in_uts_target", action="store_true") parser.add_argument("--train_ratio", type=float, default=0.7) parser.add_argument("--val_ratio", type=float, default=0.15) @@ -85,8 +79,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--n_bins", type=int, default=16) parser.add_argument("--extra_pool_reduce", type=str, default="mean", choices=["mean", "sum"]) - parser.add_argument("--time_mode", type=str, default="relative", - choices=["relative", "absolute"]) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument( "--model_architecture", @@ -95,17 +87,10 @@ def parse_args() -> argparse.Namespace: choices=SUPPORTED_MODEL_ARCHITECTURES, ) - parser.add_argument("--target_mode", type=str, default="uts", - choices=["delphi2m", "uts"]) - parser.add_argument("--readout_name", type=str, default=None, - choices=["token", "same_time_group_end", "last_valid"]) - parser.add_argument("--readout_reduce", type=str, default="mean", - choices=["mean", "sum"]) parser.add_argument("--t_min", type=float, default=0.0027378507871321013) parser.add_argument("--max_exp_input", type=float, default=60.0) parser.add_argument("--ce_weight", type=float, default=1.0) parser.add_argument("--time_weight", type=float, default=1.0) - parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true") parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) @@ -127,11 +112,6 @@ def parse_args() -> argparse.Namespace: ) if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0): raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") - if args.target_mode == "uts": - args.readout_name = args.readout_name or "same_time_group_end" - args.include_no_event_in_uts_target = True - else: - args.readout_name = args.readout_name or "token" args.extra_info_types = ( load_extra_info_types_file(args.extra_info_types_file) if args.extra_info_types_file is not None @@ -140,24 +120,6 @@ def parse_args() -> argparse.Namespace: return args -def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: - if epoch < args.warmup_epochs: - return adaptive_lr * (epoch + 1) / args.warmup_epochs - progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) - cosine = 0.5 * (1 + math.cos(math.pi * progress)) - return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) - - -def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: - non_blocking = device.type == "cuda" - return { - key: value.to(device, non_blocking=non_blocking) - if isinstance(value, torch.Tensor) - else value - for key, value in batch.items() - } - - def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: return DeepHealth( vocab_size=dataset.vocab_size, @@ -171,44 +133,27 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: n_bins=args.n_bins, extra_pool_reduce=args.extra_pool_reduce, target_mode="next_token", - time_mode=args.time_mode, + time_mode="absolute", dist_mode="exponential", dropout=args.dropout, model_architecture=args.model_architecture, ) -def build_next_step_readout(args: argparse.Namespace): - if args.readout_name == "same_time_group_end": - return build_readout("same_time_group_end", reduce=args.readout_reduce) - return build_readout(args.readout_name) - - def build_next_step_loss(args: argparse.Namespace): - if args.target_mode == "delphi2m": - ignored_tokens = {PAD_IDX, CHECKUP_IDX} - if args.ignore_no_event_in_delphi2m: - ignored_tokens.add(NO_EVENT_IDX) - return build_loss( - "delphi2m", - ignored_tokens=ignored_tokens, - t_min=args.t_min, - max_exp_input=args.max_exp_input, - ce_weight=args.ce_weight, - time_weight=args.time_weight, - ) return build_loss( - "uts", - ignored_idx={PAD_IDX, CHECKUP_IDX}, + "delphi2m", + ignored_tokens={PAD_IDX, CHECKUP_IDX}, t_min=args.t_min, max_exp_input=args.max_exp_input, + ce_weight=args.ce_weight, + time_weight=args.time_weight, ) def build_augmented_next_step_targets( batch_cpu: Dict[str, torch.Tensor], model_out: DeepHealthOutput, - include_uts_targets: bool, ) -> Dict[str, torch.Tensor]: hidden_len = model_out.hidden.size(1) event_len = int(model_out.event_len) @@ -216,26 +161,12 @@ def build_augmented_next_step_targets( device = model_out.hidden.device non_blocking = device.type == "cuda" if extra_len <= 0: - targets = { + return { "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking), "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking), - "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking), } - if include_uts_targets: - targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to( - device, non_blocking=non_blocking - ) - targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to( - device, non_blocking=non_blocking - ) - return targets bsz = batch_cpu["target_event_seq"].size(0) - vocab_size = ( - batch_cpu["target_multi_hot"].size(2) - if include_uts_targets - else None - ) other_valid = batch_cpu["other_type"] > 0 extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len) extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool) @@ -268,34 +199,6 @@ def build_augmented_next_step_targets( ], dim=1, ) - readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1) - target_dt_unique = None - target_multi_hot = None - if include_uts_targets: - target_dt_unique = torch.cat( - [ - batch_cpu["target_dt_unique"], - torch.zeros( - bsz, - extra_len, - dtype=batch_cpu["target_dt_unique"].dtype, - ), - ], - dim=1, - ) - target_multi_hot = torch.cat( - [ - batch_cpu["target_multi_hot"], - torch.zeros( - bsz, - extra_len, - vocab_size, - dtype=batch_cpu["target_multi_hot"].dtype, - ), - ], - dim=1, - ) - for b in range(bsz): valid_event = batch_cpu["padding_mask"][b].bool() if not valid_event.any(): @@ -326,7 +229,6 @@ def build_augmented_next_step_targets( t = extra_time[b, j] future = times > t if not future.any(): - readout_mask[b, pos] = False continue first_idx = int(torch.nonzero(future, as_tuple=False)[0].item()) @@ -335,35 +237,15 @@ def build_augmented_next_step_targets( target_event_seq[b, pos] = next_event target_time_seq[b, pos] = next_time - if not include_uts_targets: - continue - - same_next_time = times == next_time - next_events = events[same_next_time] - valid_next_events = next_events[ - (next_events > PAD_IDX) & (next_events < vocab_size) - ].long() - if valid_next_events.numel() == 0: - readout_mask[b, pos] = False - continue - target_multi_hot[b, pos, valid_next_events] = True - target_dt_unique[b, pos] = next_time - t - - targets = { + return { "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking), "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking), - "readout_mask": readout_mask.to(device, non_blocking=non_blocking), } - if include_uts_targets: - targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking) - targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking) - return targets def compute_next_step_loss( args: argparse.Namespace, model: DeepHealth, - readout, criterion, batch: Dict[str, torch.Tensor], device: torch.device, @@ -382,7 +264,6 @@ def compute_next_step_loss( other_value=batch["other_value"], other_value_kind=batch["other_value_kind"], other_time=batch["other_time"], - target_mode="next_token", return_output=True, ) if not isinstance(model_out, DeepHealthOutput): @@ -390,35 +271,17 @@ def compute_next_step_loss( targets = build_augmented_next_step_targets( batch_cpu=batch_cpu, model_out=model_out, - include_uts_targets=args.target_mode == "uts", ) - readout_out = readout( - hidden=model_out.hidden, - time_seq=model_out.time_seq, - padding_mask=model_out.padding_mask, - readout_mask=targets["readout_mask"] - if args.readout_name == "same_time_group_end" - else None, - ) - logits = model.calc_risk(readout_out.hidden) + logits = model.calc_risk(model_out.hidden) - if args.target_mode == "delphi2m": - loss, parts = criterion( - logits=logits, - target_events=targets["target_event_seq"], - target_times=targets["target_time_seq"], - current_times=model_out.time_seq, - padding_mask=readout_out.readout_mask, - return_components=True, - ) - else: - loss, parts = criterion( - logits=logits, - target_multi_hot=targets["target_multi_hot"], - target_dt_unique=targets["target_dt_unique"], - readout_mask=readout_out.readout_mask, - return_components=True, - ) + loss, parts = criterion( + logits=logits, + target_events=targets["target_event_seq"], + target_times=targets["target_time_seq"], + current_times=model_out.time_seq, + padding_mask=model_out.padding_mask, + return_components=True, + ) if not torch.isfinite(loss): raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") return loss, parts @@ -428,7 +291,6 @@ def run_epoch( logger: logging.Logger, args: argparse.Namespace, model: DeepHealth, - readout, criterion, loader: DataLoader, optimizer: AdamW | None, @@ -436,7 +298,6 @@ def run_epoch( is_train: bool, ) -> float: model.train(is_train) - readout.train(is_train) total = torch.zeros((), device=device) n_batches = 0 skipped = 0 @@ -447,7 +308,9 @@ def run_epoch( progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) for batch_idx, batch in enumerate(progress): try: - loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device) + loss, parts = compute_next_step_loss( + args, model, criterion, batch, device + ) if is_train: if optimizer is None: raise ValueError("optimizer is required for training") @@ -497,7 +360,8 @@ def build_metadata( "model_class": "DeepHealth", "model_architecture": args.model_architecture, "model_target_mode": "next_token", - "target_mode": args.target_mode, + "target_mode": "delphi2m", + "time_mode": "absolute", "dist_mode": "exponential", "extra_info_types_file": ( Path(args.extra_info_types_file).name @@ -518,8 +382,8 @@ def build_metadata( "val": int(len(val_subset)), "test": int(len(test_subset)), }, - "resolved_readout_name": args.readout_name, - "resolved_loss_name": args.target_mode, + "resolved_readout_name": "token", + "resolved_loss_name": "delphi2m", } @@ -531,7 +395,7 @@ def main() -> None: run_dir, run_name = create_unique_run_dir( lambda timestamp: ( - f"{args.time_mode}_exponential_next_token_{args.target_mode}_" + "absolute_exponential_next_token_delphi2m_" f"gap_{args.no_event_interval_years:g}y_{timestamp}" ), runs_root=Path(args.runs_root) / args.model_architecture, @@ -542,13 +406,12 @@ def main() -> None: logger.info(f"Device: {device}") logger.info(f"Model architecture: {args.model_architecture}") logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") - logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}") + logger.info("time_mode=absolute, readout=token, target_mode=delphi2m") dataset = HealthDataset( data_prefix=args.data_prefix, labels_file=args.labels_file, no_event_interval_years=args.no_event_interval_years, - include_no_event_in_uts_target=args.include_no_event_in_uts_target, extra_info_types=args.extra_info_types, ) if args.train_eid_file and args.val_eid_file and args.test_eid_file: @@ -616,7 +479,6 @@ def main() -> None: f"total={parameter_counts['model_parameter_count']:,}, " f"trainable={parameter_counts['trainable_parameter_count']:,}" ) - readout = build_next_step_readout(args).to(device) criterion = build_next_step_loss(args) optimizer = AdamW( model.parameters(), @@ -646,9 +508,13 @@ def main() -> None: lr = get_lr(epoch, args, adaptive_lr) set_optimizer_lr(optimizer, lr) - train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True) + train_loss = run_epoch( + logger, args, model, criterion, train_loader, optimizer, device, True + ) with torch.no_grad(): - val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False) + val_loss = run_epoch( + logger, args, model, criterion, val_loader, None, device, False + ) is_best = val_loss < best_val if is_best: @@ -682,7 +548,9 @@ def main() -> None: logger.info("Evaluating best model on next-step test split...") model.load_state_dict(torch.load(best_model_path, map_location=device)) with torch.no_grad(): - test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False) + test_loss = run_epoch( + logger, args, model, criterion, test_loader, None, device, False + ) logger.info(f"Test loss: {test_loss:.6f}") logger.info(f"Best checkpoint: {best_model_path}") diff --git a/train_util.py b/train_util.py index 5c7a806..c1cf1ac 100644 --- a/train_util.py +++ b/train_util.py @@ -6,6 +6,7 @@ import sys import time import csv from datetime import datetime +import math from pathlib import Path from typing import Any, Dict, Iterable, Tuple @@ -141,6 +142,31 @@ def resolve_device(device_arg: str) -> torch.device: raise ValueError(f"Unsupported device: {device_arg}") +def get_lr(epoch: int, args: Any, adaptive_lr: float) -> float: + if epoch < args.warmup_epochs: + return adaptive_lr * (epoch + 1) / args.warmup_epochs + progress = (epoch - args.warmup_epochs) / max( + 1, args.max_epochs - args.warmup_epochs + ) + cosine = 0.5 * (1 + math.cos(math.pi * progress)) + return adaptive_lr * ( + args.min_lr_ratio + cosine * (1 - args.min_lr_ratio) + ) + + +def move_batch_to_device( + batch: Dict[str, torch.Tensor], + device: torch.device, +) -> Dict[str, torch.Tensor]: + non_blocking = device.type == "cuda" + return { + key: value.to(device, non_blocking=non_blocking) + if isinstance(value, torch.Tensor) + else value + for key, value in batch.items() + } + + def split_dataset( dataset: HealthDataset, train_ratio: float, @@ -291,15 +317,6 @@ def split_all_future_datasets_by_eid_files( ) -def build_optimizer(args: Any, model: DeepHealth) -> AdamW: - return AdamW( - model.parameters(), - lr=args.base_lr, - betas=tuple(args.betas), - weight_decay=args.weight_decay, - ) - - def get_model_parameter_counts(model: torch.nn.Module) -> Dict[str, int]: """Return stable total and trainable parameter counts.""" return {