Use precomputed exposure embeddings

eval_data.py

@@ -8,8 +8,6 @@ import torch
 from torch.nn.utils.rnn import pad_sequence
 
 from dataset import (
-    DAILY_EXPOSURE_SHAPE,
-    MONTHLY_EXPOSURE_SHAPE,
     AllFutureHealthDataset,
     HealthDataset,
 )
@@ -32,7 +30,7 @@ class AllFutureSequenceEvalDataset:
         min_history_events: int = 1,
         min_future_events: int = 1,
         exposure_cache_dir: str | Path | None = None,
-        mask_onset_exposure: bool = False,
+        exposure_embeddings_file: str | Path | None = None,
     ) -> None:
         base = AllFutureHealthDataset(
             data_prefix=data_prefix,
@@ -41,7 +39,7 @@ class AllFutureSequenceEvalDataset:
             min_history_events=min_history_events,
             min_future_events=min_future_events,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
 
         self.base = base
@@ -91,9 +89,9 @@ class AllFutureSequenceEvalDataset:
             "sex": torch.tensor(s["sex"], dtype=torch.long),
         }
         if "exposure_index" in s:
-            daily, monthly = self.base._load_exposure_windows(s["exposure_index"])
-            out["exposure_daily"] = daily
-            out["exposure_monthly"] = monthly
+            out["exposure_embedding"] = self.base._load_exposure_embeddings(
+                s["exposure_index"]
+            )
         return out
 
 
@@ -107,7 +105,7 @@ def load_sequence_eval_dataset(
     min_history_events: int,
     min_future_events: int,
     exposure_cache_dir: str | Path | None = None,
-    mask_onset_exposure: bool = False,
+    exposure_embeddings_file: str | Path | None = None,
 ):
     mode = str(model_target_mode).lower()
     if mode == "next_token":
@@ -117,7 +115,7 @@ def load_sequence_eval_dataset(
             no_event_interval_years=no_event_interval_years,
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
     if mode == "all_future":
         return AllFutureSequenceEvalDataset(
@@ -126,7 +124,7 @@ def load_sequence_eval_dataset(
             min_history_events=min_history_events,
             min_future_events=min_future_events,
             exposure_cache_dir=exposure_cache_dir,
-            mask_onset_exposure=mask_onset_exposure,
+            exposure_embeddings_file=exposure_embeddings_file,
         )
     raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
 
@@ -156,34 +154,17 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
         "readout_mask": readout_mask,
         "sex": torch.stack([s["sex"] for s in batch]),
     }
-    if any("exposure_daily" in s for s in batch):
-        out["exposure_daily"] = _pad_eval_exposure(
-            batch,
-            "exposure_daily",
-            DAILY_EXPOSURE_SHAPE,
+    if any("exposure_embedding" in s for s in batch):
+        embedding_dim = next(
+            int(s["exposure_embedding"].size(1))
+            for s in batch if "exposure_embedding" in s
         )
-        out["exposure_monthly"] = _pad_eval_exposure(
-            batch,
-            "exposure_monthly",
-            MONTHLY_EXPOSURE_SHAPE,
+        encoded = torch.zeros(
+            len(batch), event_seq.size(1), embedding_dim, dtype=torch.float32
         )
-    return out
-
-
-def _pad_eval_exposure(
-    batch: List[Dict[str, torch.Tensor]],
-    key: str,
-    shape: tuple[int, int],
-) -> torch.Tensor:
-    max_len = max(int(s["event_seq"].numel()) for s in batch)
-    out = torch.full(
-        (len(batch), max_len, shape[0], shape[1]),
-        float("nan"),
-        dtype=torch.float32,
-    )
-    for idx, sample in enumerate(batch):
-        value = sample.get(key)
-        if value is None:
-            continue
-        out[idx, : int(value.size(0))] = value
+        for idx, sample in enumerate(batch):
+            value = sample.get("exposure_embedding")
+            if value is not None:
+                encoded[idx, :value.size(0)] = value
+        out["exposure_embedding"] = encoded
     return out