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

2026-06-15 14:35:10 +08:00
parent c3e49db859
commit 36ec36c8a8
3 changed files with 202 additions and 6 deletions
--- a/clean_dataset_cache.py
+++ 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()