Add disease history ablation modes

evaluate_auc.py

@@ -39,7 +39,12 @@ import torch
 from torch.utils.data import DataLoader, Subset
 from tqdm.auto import tqdm
 
-from dataset import HealthDataset
+from dataset import (
+    DISEASE_HISTORY_MODE_TIMED,
+    HealthDataset,
+    normalize_disease_history_mode,
+    transform_disease_history_batch_at_position,
+)
 from delphi2m_auc_report import (
     DEFAULT_DELPHI2M_PERIODS_YEARS,
     build_delphi2m_auc_report,
@@ -333,6 +338,7 @@ def infer_readout_hidden(
     loader: DataLoader,
     device: torch.device,
     model_target_mode: str,
+    disease_history_mode: str,
     use_amp: bool,
     hidden_cache_dtype: str = "float16",
 ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
@@ -342,6 +348,16 @@ def infer_readout_hidden(
         raise ValueError(
             f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
         )
+    disease_history_mode = normalize_disease_history_mode(
+        disease_history_mode
+    )
+    if (
+        model_target_mode != "all_future"
+        and disease_history_mode != DISEASE_HISTORY_MODE_TIMED
+    ):
+        raise ValueError(
+            "ordered/set disease history is only supported for all_future models"
+        )
 
     hidden_parts: List[np.ndarray] = []
     arrays: Dict[str, List[np.ndarray]] = {
@@ -385,12 +401,25 @@ def infer_readout_hidden(
                     active = padding_mask[:, pos].bool()
                     if not active.any():
                         continue
-                    hidden_pos = model(
+                    (
+                        model_event_seq,
+                        model_time_seq,
+                        model_padding_mask,
+                        model_t_query,
+                    ) = transform_disease_history_batch_at_position(
                         event_seq=event_seq[active],
-                        time_seq=time_seq[active],
-                        sex=batch_dev["sex"][active],
+                        actual_time_seq=time_seq[active],
                         padding_mask=padding_mask[active],
-                        t_query=time_seq[active, pos],
+                        query_position=pos,
+                        disease_history_mode=disease_history_mode,
+                        vocab_size=model.vocab_size,
+                    )
+                    hidden_pos = model(
+                        event_seq=model_event_seq,
+                        time_seq=model_time_seq,
+                        sex=batch_dev["sex"][active],
+                        padding_mask=model_padding_mask,
+                        t_query=model_t_query,
                         other_type=batch_dev["other_type"][active],
                         other_value=batch_dev["other_value"][active],
                         other_value_kind=batch_dev["other_value_kind"][active],
@@ -899,6 +928,7 @@ def evaluate_auc_pipeline(
     offsets: Sequence[float],
     device: torch.device,
     model_target_mode: str,
+    disease_history_mode: str,
     num_workers_auc: int,
     use_amp: bool,
     auc_task_chunk_size: int = 0,
@@ -957,6 +987,7 @@ def evaluate_auc_pipeline(
         loader=loader,
         device=device,
         model_target_mode=model_target_mode,
+        disease_history_mode=disease_history_mode,
         use_amp=use_amp,
         hidden_cache_dtype=hidden_cache_dtype,
     )
@@ -1165,6 +1196,9 @@ def main() -> None:
             f"got {model_target_mode!r}"
         )
     dist_mode_cfg = cfg.get("dist_mode", "exponential")
+    disease_history_mode = normalize_disease_history_mode(
+        cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
+    )
 
     device = resolve_eval_device(args.device)
     if device.type == "cuda":
@@ -1179,6 +1213,7 @@ def main() -> None:
         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)),
+        disease_history_mode=disease_history_mode,
     )
     validate_dataset_metadata(dataset, cfg)
 
@@ -1209,6 +1244,7 @@ def main() -> None:
     print(f"Resolved dist_mode for evaluation: {dist_mode}")
     print(f"Resolved model architecture: {model_architecture}")
     print(f"Model target mode for AUC: {model_target_mode}")
+    print(f"Disease history mode for AUC: {disease_history_mode}")
     print(
         "AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
         "dist_mode affects model loading but is not converted to horizon-specific risk probability."
@@ -1253,6 +1289,7 @@ def main() -> None:
         offsets=auc_offsets,
         device=device,
         model_target_mode=model_target_mode,
+        disease_history_mode=disease_history_mode,
         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)),