Enhance DeepHealth model to incorporate CHECKUP state tokens in next-step training and evaluation, update dataset cache versioning, and improve handling of observed event histories.

dataset.py

@@ -110,7 +110,7 @@ def _cache_file_path(
                 selected.append(type_id)
         selected_types = ",".join(str(t) for t in selected)
     signature_parts = [
-        "deephealthnew_dataset_cache_v2",
+        "deephealthnew_dataset_cache_v3_checkup_state",
         dataset_kind,
         split or "",
         event_path,
@@ -320,9 +320,6 @@ class _ExpoBaseDataset(Dataset):
             times_days_raw = rows[:, 1].astype(np.float32)
             labels_raw = rows[:, 2].astype(np.int64)
 
-            disease_mask = labels_raw != CHECKUP_IDX
-            times_days_raw = times_days_raw[disease_mask]
-            labels_raw = labels_raw[disease_mask]
             if len(labels_raw) == 0:
                 yield eid, times_days_raw, labels_raw
                 continue
@@ -392,7 +389,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot
     """
 
-    CACHE_VERSION = 2
+    CACHE_VERSION = 3
 
     def __init__(
         self,
@@ -488,7 +485,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
     time range, with at least one future event after every query.
     """
 
-    CACHE_VERSION = 4
+    CACHE_VERSION = 5
 
     def __init__(
         self,
@@ -582,8 +579,13 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
 
     def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
         times = patient["times"]
+        labels = patient["labels"]
+        real_event_mask = ~np.isin(
+            labels,
+            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
+        )
         n_hist = int((times <= t_query).sum())
-        n_future = int((times > t_query).sum())
+        n_future = int(((times > t_query) & real_event_mask).sum())
         return (
             n_hist >= self.min_history_events
             and n_future >= self.min_future_events