Add time attention mask handling and baseline class time computation to DeepHealth model

backbones.py

@@ -395,11 +395,28 @@ class BaselineEncoder(nn.Module):
             dtype=dtype,
         ).masked_fill(~mask[:, None, None, :], -1e4)
 
+    def _make_time_attn_mask(
+        self,
+        mask: torch.Tensor,
+        time: torch.Tensor,
+        dtype: torch.dtype,
+    ):
+        valid_key = mask[:, None, :]
+        visible_by_time = time[:, None, :] <= time[:, :, None]
+        valid = valid_key & visible_by_time
+        return torch.zeros(
+            valid.shape,
+            device=valid.device,
+            dtype=dtype,
+        ).masked_fill(~valid, -1e4)[:, None, :, :]
+
     def forward(
         self,
         other_type: torch.LongTensor,        # (B, K), 0 = padding
         other_value: torch.Tensor,           # (B, K), cate stores global id
         other_value_kind: torch.LongTensor,  # (B, K), 0=PAD, 1=CONT, 2=CATE
+        other_time: torch.Tensor | None = None,  # (B, K)
+        cls_time: torch.Tensor | None = None,    # (B,)
     ):
         if other_type.shape != other_value.shape:
             raise ValueError(
@@ -451,7 +468,24 @@ class BaselineEncoder(nn.Module):
         )
         full_valid = torch.cat([cls_valid, other_valid], dim=1)
 
+        if other_time is None:
             attn_mask = self._make_attn_mask(full_valid, f.dtype)
+        else:
+            if other_time.shape != other_type.shape:
+                raise ValueError(
+                    "other_time must have the same shape as other_type, got "
+                    f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
+                )
+            if cls_time is None:
+                raise ValueError("cls_time is required when other_time is provided")
+            full_time = torch.cat(
+                [
+                    cls_time.to(device=other_time.device, dtype=other_time.dtype)[:, None],
+                    other_time,
+                ],
+                dim=1,
+            )
+            attn_mask = self._make_time_attn_mask(full_valid, full_time, f.dtype)
         for block in self.blocks:
             f = block(f, attn_mask=attn_mask)
             f = f * full_valid.unsqueeze(-1).to(f.dtype)

clean_dataset_cache.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import argparse
+import os
+import pickle
+from pathlib import Path
+from typing import Any
+
+
+CURRENT_CACHE_VERSIONS = {
+    "next_step": 3,
+    "all_future": 5,
+}
+
+
+def infer_cache_kind(path: Path) -> str | None:
+    name = path.name
+    for kind in CURRENT_CACHE_VERSIONS:
+        marker = f"_{kind}_cache_"
+        if marker in name:
+            return kind
+    return None
+
+
+def read_cache_version(path: Path) -> int | None:
+    try:
+        with path.open("rb") as f:
+            payload: Any = pickle.load(f)
+    except Exception:
+        return None
+    if not isinstance(payload, dict):
+        return None
+    version = payload.get("_cache_version")
+    try:
+        return int(version)
+    except (TypeError, ValueError):
+        return None
+
+
+def should_remove(path: Path, remove_all: bool) -> tuple[bool, str]:
+    kind = infer_cache_kind(path)
+    if kind is None:
+        return False, "not a DeepHealth dataset cache"
+
+    if remove_all:
+        return True, "remove all dataset caches"
+
+    version = read_cache_version(path)
+    expected = CURRENT_CACHE_VERSIONS[kind]
+    if version is None:
+        return True, f"{kind} cache is unreadable or missing _cache_version"
+    if version != expected:
+        return True, f"{kind} cache version {version} != current {expected}"
+    return False, f"{kind} cache version {version} is current"
+
+
+def iter_cache_files(root: Path, recursive: bool):
+    pattern = "**/*_cache_*.pkl" if recursive else "*_cache_*.pkl"
+    yield from root.glob(pattern)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Remove obsolete DeepHealth dataset cache files. Defaults to dry-run; "
+            "pass --apply to delete."
+        )
+    )
+    parser.add_argument(
+        "--data_dir",
+        type=Path,
+        default=Path("."),
+        help="Directory containing dataset cache files, usually the data_prefix directory.",
+    )
+    parser.add_argument(
+        "--recursive",
+        action="store_true",
+        help="Search recursively under data_dir.",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Remove all recognized DeepHealth dataset caches, including current ones.",
+    )
+    parser.add_argument(
+        "--apply",
+        action="store_true",
+        help="Actually delete files. Without this flag, only prints what would be removed.",
+    )
+    args = parser.parse_args()
+
+    root = args.data_dir.expanduser().resolve()
+    if not root.exists() or not root.is_dir():
+        raise SystemExit(f"data_dir is not a directory: {root}")
+
+    files = sorted(p for p in iter_cache_files(root, args.recursive) if p.is_file())
+    if not files:
+        print(f"No *_cache_*.pkl files found under {root}")
+        return
+
+    kept = 0
+    removed = 0
+    candidates = 0
+    for path in files:
+        remove, reason = should_remove(path, remove_all=bool(args.all))
+        rel = os.path.relpath(path, root)
+        if remove:
+            candidates += 1
+            action = "DELETE" if args.apply else "WOULD_DELETE"
+            print(f"{action}\t{rel}\t{reason}")
+            if args.apply:
+                path.unlink()
+                removed += 1
+        else:
+            kept += 1
+            print(f"KEEP\t{rel}\t{reason}")
+
+    if args.apply:
+        print(f"Removed {removed} cache file(s); kept {kept}.")
+    else:
+        print(
+            f"Dry run: {candidates} cache file(s) would be removed; kept {kept}. "
+            "Re-run with --apply to delete."
+        )
+
+
+if __name__ == "__main__":
+    main()

models.py

@@ -141,16 +141,42 @@ class DeepHealth(nn.Module):
         summary = baseline_summary.to(device=h_disease.device, dtype=h_disease.dtype)
         return torch.where(checkup_mask.unsqueeze(-1), summary[:, None, :], h_disease)
 
+    def _baseline_cls_time(
+        self,
+        event_seq: torch.Tensor,
+        time_seq: torch.Tensor,
+        padding_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        checkup_mask = event_seq == CHECKUP_IDX
+        inf = torch.full_like(time_seq, float("inf"))
+        first_checkup = torch.where(checkup_mask, time_seq, inf).min(dim=1).values
+        has_checkup = torch.isfinite(first_checkup)
+        fallback_time = torch.where(
+            padding_mask,
+            time_seq,
+            torch.full_like(time_seq, float("-inf")),
+        ).max(dim=1).values
+        fallback_time = torch.where(
+            torch.isfinite(fallback_time),
+            fallback_time,
+            torch.zeros_like(fallback_time),
+        )
+        return torch.where(has_checkup, first_checkup, fallback_time)
+
     def _encode_other_tokens(
         self,
         other_type: torch.LongTensor,
         other_value: torch.Tensor,
         other_value_kind: torch.LongTensor,
+        other_time: torch.Tensor,
+        cls_time: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return self.token_encoder(
             other_type=other_type,
             other_value=other_value,
             other_value_kind=other_value_kind,
+            other_time=other_time,
+            cls_time=cls_time,
         )
 
     def _forward_shared(
@@ -193,16 +219,24 @@ class DeepHealth(nn.Module):
         h_disease = self.token_embedding(event_seq)
         t_disease = time_seq
 
-        h_token, token_mask, baseline_summary = self._encode_other_tokens(
-            other_type=other_type,
-            other_value=other_value,
-            other_value_kind=other_value_kind,
-        )
         if other_time.shape != other_type.shape:
             raise ValueError(
                 "other_time must have the same shape as other_type, got "
                 f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
             )
+        other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
+        cls_time = self._baseline_cls_time(
+            event_seq=event_seq,
+            time_seq=time_seq,
+            padding_mask=padding_mask,
+        )
+        h_token, token_mask, baseline_summary = self._encode_other_tokens(
+            other_type=other_type,
+            other_value=other_value,
+            other_value_kind=other_value_kind,
+            other_time=other_time,
+            cls_time=cls_time,
+        )
         token_time = other_time.to(device=h_token.device, dtype=time_seq.dtype)
 
         h_disease = self.cross_attention(