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

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"]