diff --git a/backbones.py b/backbones.py index 2db4c1c..899e0bf 100644 --- a/backbones.py +++ b/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) - attn_mask = self._make_attn_mask(full_valid, f.dtype) + 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) diff --git a/clean_dataset_cache.py b/clean_dataset_cache.py new file mode 100644 index 0000000..98d5437 --- /dev/null +++ b/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() diff --git a/models.py b/models.py index 2abcd96..bbac0c3 100644 --- a/models.py +++ b/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(