Refactor dataset.py to remove caching functionality and related methods for improved simplicity

This commit is contained in:
2026-06-18 16:25:55 +08:00
parent aa8ec5c3ac
commit c3cac6dcea

View File

@@ -1,9 +1,6 @@
# dataset.py
from __future__ import annotations
import hashlib
import os
import pickle
from typing import Dict, Iterable, List, Literal, Optional, Tuple
import numpy as np
@@ -83,61 +80,6 @@ def _insert_gap_no_event_tokens(
return all_times[order], all_labels[order]
def _cache_file_path(
data_prefix: str,
labels_file: str,
no_event_interval_years: float,
include_no_event_in_uts_target: bool,
dataset_kind: str,
extra_info_types: Iterable[int] | None = None,
split: str | None = None,
min_history_events: int | None = None,
min_future_events: int | None = None,
validation_query_seed: int | None = None,
) -> str:
event_path = f"{data_prefix}_event_data.npy"
basic_path = f"{data_prefix}_basic_info.csv"
other_path = f"{data_prefix}_other_info.npy"
cate_types_path = "cate_types.csv"
selected_types = ""
if extra_info_types is not None:
seen_types: set[int] = set()
selected = []
for raw_type in extra_info_types:
type_id = int(raw_type)
if type_id not in seen_types:
seen_types.add(type_id)
selected.append(type_id)
selected_types = ",".join(str(t) for t in selected)
signature_parts = [
"deephealthnew_dataset_cache_v3_checkup_state",
dataset_kind,
split or "",
event_path,
basic_path,
other_path,
cate_types_path,
selected_types,
labels_file,
f"{no_event_interval_years:.8f}",
str(int(include_no_event_in_uts_target)),
"" if min_history_events is None else str(int(min_history_events)),
"" if min_future_events is None else str(int(min_future_events)),
"" if validation_query_seed is None else str(int(validation_query_seed)),
]
for path in (event_path, basic_path, other_path, cate_types_path, labels_file):
try:
stat = os.stat(path)
signature_parts.append(f"{path}:{stat.st_mtime_ns}:{stat.st_size}")
except OSError:
signature_parts.append(f"{path}:missing")
digest = hashlib.sha1("|".join(signature_parts).encode("utf-8")).hexdigest()
cache_dir = os.path.dirname(event_path) or "."
return os.path.join(cache_dir, f"{data_prefix}_{dataset_kind}_cache_{digest}.pkl")
class _ExpoBaseDataset(Dataset):
def __init__(
self,
@@ -345,40 +287,6 @@ class _ExpoBaseDataset(Dataset):
**other_info,
}
@staticmethod
def _load_cache(cache_path: str, cache_version: int) -> Optional[Dict]:
try:
with open(cache_path, "rb") as f:
payload = pickle.load(f)
except OSError:
return None
except Exception:
return None
if not isinstance(payload, dict):
return None
if payload.get("_cache_version") != cache_version:
return None
state = payload.get("state")
if not isinstance(state, dict):
return None
return state
def _save_cache(self, cache_path: str, cache_version: int) -> None:
payload = {
"_cache_version": cache_version,
"state": {key: value for key, value in self.__dict__.items()},
}
try:
cache_dir = os.path.dirname(cache_path)
if cache_dir:
os.makedirs(cache_dir, exist_ok=True)
with open(cache_path, "wb") as f:
pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
except OSError:
return
class NextStepHealthDataset(_ExpoBaseDataset):
"""
Dataset for next-token and next-time-point losses with unified other-info
@@ -399,19 +307,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
include_no_event_in_uts_target: bool = False,
extra_info_types: Iterable[int] | None = None,
) -> None:
cache_path = _cache_file_path(
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
dataset_kind="next_step",
extra_info_types=extra_info_types,
)
cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
if cached_state is not None:
self.__dict__.update(cached_state)
return
super().__init__(
data_prefix=data_prefix,
labels_file=labels_file,
@@ -451,8 +346,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
**features,
})
self._save_cache(cache_path, self.CACHE_VERSION)
def __len__(self) -> int:
return len(self.samples)
@@ -502,23 +395,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
if split not in {"train", "valid", "test"}:
raise ValueError(f"split must be train/valid/test, got {split!r}")
cache_path = _cache_file_path(
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
dataset_kind="all_future",
extra_info_types=extra_info_types,
split=split,
min_history_events=min_history_events,
min_future_events=min_future_events,
validation_query_seed=validation_query_seed if split in {"valid", "test"} else None,
)
cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
if cached_state is not None:
self.__dict__.update(cached_state)
return
super().__init__(
data_prefix=data_prefix,
labels_file=labels_file,
@@ -575,8 +451,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
if split in {"valid", "test"} and not self.valid_queries:
raise ValueError("No random-but-fixed validation query points were built.")
self._save_cache(cache_path, self.CACHE_VERSION)
def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
times = patient["times"]
labels = patient["labels"]