Use precomputed exposure embeddings

train_next_step.py

@@ -58,8 +58,7 @@ MODEL_INPUT_KEYS = (
 )
 
 EXPOSURE_INPUT_KEYS = (
-    "exposure_daily",
-    "exposure_monthly",
+    "exposure_embedding",
 )
 
 
@@ -174,8 +173,7 @@ class NextStepTrainingModel(nn.Module):
         sex: torch.Tensor,
         padding_mask: torch.Tensor,
         readout_mask: torch.Tensor | None = None,
-        exposure_daily: torch.Tensor | None = None,
-        exposure_monthly: torch.Tensor | None = None,
+        exposure_embedding: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         hidden = self.model(
             event_seq=event_seq,
@@ -183,8 +181,7 @@ class NextStepTrainingModel(nn.Module):
             sex=sex,
             padding_mask=padding_mask,
             target_mode="next_token",
-            exposure_daily=exposure_daily,
-            exposure_monthly=exposure_monthly,
+            exposure_embedding=exposure_embedding,
         )
         if not isinstance(hidden, torch.Tensor):
             raise TypeError("DeepHealth forward must return a hidden-state tensor")
@@ -224,19 +221,15 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_hist_layer", type=int, default=12)
     parser.add_argument("--dropout", type=float, default=0.0)
     parser.add_argument("--exposure_cache_dir", type=str, default=None)
-    parser.add_argument("--mask_onset_exposure", action="store_true")
     parser.add_argument(
-        "--d_model",
-        type=int,
-        default=64,
-        help="Internal TimesNet channel dimension for exposure encoding.",
+        "--exposure_embeddings_file",
+        type=str,
+        default=None,
+        help=(
+            "Precomputed exposure embeddings. Defaults to "
+            "<exposure_cache_dir>/exposure_embeddings.npy."
+        ),
     )
-    parser.add_argument("--exposure_n_layers", type=int, default=2)
-    parser.add_argument("--exposure_top_k", type=int, default=2)
-    parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1)
-    parser.add_argument("--exposure_backbone_kernel_size", type=int, default=5)
-    parser.add_argument("--exposure_backbone_expansion", type=float, default=2.0)
-    parser.add_argument("--no_exposure_gate", action="store_true")
     parser.add_argument("--target_mode", type=str, default="uts",
                         choices=["delphi2m", "uts"])
     parser.add_argument("--readout_name", type=str, default=None,
@@ -311,8 +304,10 @@ def parse_args() -> argparse.Namespace:
         raise ValueError("prefetch_factor must be positive when num_workers > 0")
     if args.exposure_locality_buffer_size < 0:
         raise ValueError("exposure_locality_buffer_size must be non-negative")
-    if args.d_model <= 0:
-        raise ValueError("--d_model must be positive")
+    if args.exposure_embeddings_file and not args.exposure_cache_dir:
+        raise ValueError(
+            "--exposure_cache_dir is required with --exposure_embeddings_file"
+        )
     if args.target_mode == "uts":
         args.readout_name = args.readout_name or "same_time_group_end"
         args.include_no_event_in_uts_target = True
@@ -440,14 +435,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         target_mode="next_token",
         dist_mode="exponential",
         dropout=args.dropout,
-        use_exposure_encoder=args.exposure_cache_dir is not None,
-        exposure_d_model=args.d_model,
-        exposure_n_layers=args.exposure_n_layers,
-        exposure_top_k=args.exposure_top_k,
-        exposure_n_backbone_blocks=args.exposure_n_backbone_blocks,
-        exposure_backbone_kernel_size=args.exposure_backbone_kernel_size,
-        exposure_backbone_expansion=args.exposure_backbone_expansion,
-        exposure_use_gate=not args.no_exposure_gate,
+        use_exposure_embeddings=args.exposure_embeddings_file is not None,
     )
 
 
@@ -499,9 +487,8 @@ def compute_next_step_loss(
         "padding_mask": batch["padding_mask"],
         "readout_mask": batch["readout_mask"],
     }
-    if "exposure_daily" in batch:
-        model_kwargs["exposure_daily"] = batch["exposure_daily"]
-        model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
+    if "exposure_embedding" in batch:
+        model_kwargs["exposure_embedding"] = batch["exposure_embedding"]
     logits, current_times, output_readout_mask = model(**model_kwargs)
     non_blocking = device.type == "cuda"
     targets = {
@@ -631,16 +618,9 @@ def build_metadata(
         "dataset_metadata": {
             "vocab_size": int(dataset.vocab_size),
         },
-        "use_exposure_encoder": args.exposure_cache_dir is not None,
+        "use_exposure_embeddings": args.exposure_embeddings_file is not None,
         "exposure_cache_dir": args.exposure_cache_dir,
-        "mask_onset_exposure": bool(args.mask_onset_exposure),
-        "d_model": int(args.d_model),
-        "exposure_n_layers": int(args.exposure_n_layers),
-        "exposure_top_k": int(args.exposure_top_k),
-        "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks),
-        "exposure_backbone_kernel_size": int(args.exposure_backbone_kernel_size),
-        "exposure_backbone_expansion": float(args.exposure_backbone_expansion),
-        "exposure_use_gate": not bool(args.no_exposure_gate),
+        "exposure_embeddings_file": args.exposure_embeddings_file,
         "num_workers": int(args.num_workers),
         "prefetch_factor": int(args.prefetch_factor),
         "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
@@ -659,6 +639,10 @@ def build_metadata(
 
 def main() -> None:
     args = parse_args()
+    if args.exposure_cache_dir and args.exposure_embeddings_file is None:
+        args.exposure_embeddings_file = str(
+            Path(args.exposure_cache_dir) / "exposure_embeddings.npy"
+        )
     device, rank, local_rank, world_size = init_distributed(args)
     set_seed(args.seed + rank)
     configure_torch_for_training(device)
@@ -670,6 +654,7 @@ def main() -> None:
     logger.info(f"Device: {device}")
     logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
     logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
+    logger.info(f"exposure_embeddings_file={args.exposure_embeddings_file}")
     logger.info(
         "DataLoader IO: "
         f"num_workers={args.num_workers}, "
@@ -683,8 +668,15 @@ def main() -> None:
         no_event_interval_years=args.no_event_interval_years,
         include_no_event_in_uts_target=args.include_no_event_in_uts_target,
         exposure_cache_dir=args.exposure_cache_dir,
-        mask_onset_exposure=args.mask_onset_exposure,
+        exposure_embeddings_file=args.exposure_embeddings_file,
     )
+    if dataset.exposure_cache is not None:
+        embedding_dim = int(dataset.exposure_cache.embeddings.shape[1])
+        if embedding_dim != args.n_embd:
+            raise ValueError(
+                f"Exposure embedding dim {embedding_dim} must equal "
+                f"--n_embd={args.n_embd}"
+            )
     if args.train_eid_file and args.val_eid_file and args.test_eid_file:
         train_subset, val_subset, test_subset = split_dataset_by_eid_files(
             dataset=dataset,