From c3cac6dcea564560caab0782c34fb7882498f6f3 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 18 Jun 2026 16:25:55 +0800 Subject: [PATCH] Refactor dataset.py to remove caching functionality and related methods for improved simplicity --- dataset.py | 126 ----------------------------------------------------- 1 file changed, 126 deletions(-) diff --git a/dataset.py b/dataset.py index 7c45028..2743c73 100644 --- a/dataset.py +++ b/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"]