Refactor dataset.py to remove caching functionality and related methods for improved simplicity
This commit is contained in:
126
dataset.py
126
dataset.py
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user