refactor: isolate Delphi2M next-token pipeline

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,