JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhance data parsing and validation, add extra info types files

Browse Source 
- Improved `parse_int_list` and `parse_float_list` functions to support JSON list input.
- Introduced `validate_dataset_metadata` function to ensure dataset metadata consistency with training configuration.
- Added multiple new files for extra information types, categorizing them into assessment-only, exposure-only, and combined types.
- Removed deprecated `merge_extra_info_types` function and adjusted related logic in `train.py`.
- Updated `save_config` function to accept additional metadata for training runs.
- Refactored model and training scripts for better clarity and maintainability.
This commit is contained in:
Jiarui Li
2026-06-12 11:16:19 +08:00
parent fc8c7b7177
commit 0fa8bbbb9a
9 changed files with 818 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
evaluate_auc_v2.py
View File

@@ -56,6 +56,14 @@ def parse_int_list(value: Any) -> Optional[List[int]]:
text = str(value).strip()
if text == "":
return None
if text.startswith("["):
try:
values = json.loads(text)
except json.JSONDecodeError as exc:
raise ValueError(f"Invalid integer list: {text!r}") from exc
if not isinstance(values, list):
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
return [int(x) for x in values]
return [int(x.strip()) for x in text.split(",") if x.strip()]
@@ -67,6 +75,14 @@ def parse_float_list(value: Any) -> Optional[List[float]]:
text = str(value).strip()
if text == "":
return None
if text.startswith("["):
try:
values = json.loads(text)
except json.JSONDecodeError as exc:
raise ValueError(f"Invalid float list: {text!r}") from exc
if not isinstance(values, list):
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
return [float(x) for x in values]
return [float(x.strip()) for x in text.split(",") if x.strip()]
@@ -155,6 +171,32 @@ def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None
f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}")
def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
meta = cfg.get("dataset_metadata")
if not isinstance(meta, dict):
return
actual: Dict[str, Any] = {
"vocab_size": int(dataset.vocab_size),
"n_types": int(dataset.n_types),
"n_cont_types": int(dataset.n_cont_types),
"n_categories": int(dataset.n_categories),
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
"extra_info_types": [int(x) for x in dataset.extra_info_types],
}
mismatches = [
f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
for key, value in actual.items()
if key in meta and meta.get(key) != value
]
if mismatches:
raise RuntimeError(
"Current dataset metadata does not match train_config.json. "
"Use the same prepared data and extra_info_types as training. "
+ "; ".join(mismatches)
)
# ---------------------------------------------------------------------------
# DeLong AUC utilities
# ---------------------------------------------------------------------------
@@ -525,10 +567,8 @@ class LandmarkDataset(Dataset):
np.array([np.float32(landmark_age)], dtype=np.float32),
]
)
target_time_seq = time_seq_landmark.copy()
if self.attn_mask_mode in _TARGET_AWARE_MODES:
target_time_seq[-1] = np.nextafter(
time_seq_landmark[-1] = np.nextafter(
np.float32(landmark_age), np.float32(np.inf), dtype=np.float32
)
@@ -546,7 +586,6 @@ class LandmarkDataset(Dataset):
"landmark_pos": int(len(event_seq_landmark) - 1),
"event_seq": event_seq_landmark,
"time_seq": time_seq_landmark,
"target_time_seq": target_time_seq,
"readout_mask": readout_mask,
"other_type": np.asarray(s["other_type"], dtype=np.int64),
"other_value": np.asarray(s["other_value"], dtype=np.float32),
@@ -570,7 +609,6 @@ class LandmarkDataset(Dataset):
return {
"event_seq": torch.from_numpy(s["event_seq"]).long(),
"time_seq": torch.from_numpy(s["time_seq"]).float(),
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
"readout_mask": torch.from_numpy(s["readout_mask"]),
"sex": torch.tensor(s["sex"], dtype=torch.long),
"other_type": torch.from_numpy(s["other_type"]).long(),
@@ -590,8 +628,6 @@ def collate_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX)
time_seq = pad_sequence([x["time_seq"]
for x in batch], batch_first=True, padding_value=0.0)
target_time_seq = pad_sequence(
[x["target_time_seq"] for x in batch], batch_first=True, padding_value=0.0)
readout_mask = pad_sequence(
[x["readout_mask"] for x in batch], batch_first=True, padding_value=False)
other_type = pad_sequence(
@@ -606,7 +642,6 @@ def collate_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch
return {
"event_seq": event_seq,
"time_seq": time_seq,
"target_time_seq": target_time_seq,
"padding_mask": event_seq > PAD_IDX,
"readout_mask": readout_mask,
"sex": torch.stack([x["sex"] for x in batch]),
@@ -637,7 +672,6 @@ def infer_landmark_hidden(
model: DeepHealth,
loader: DataLoader,
device: torch.device,
attn_mask_mode: str,
readout_name: str,
readout_reduce: str,
use_amp: bool,
@@ -939,7 +973,6 @@ def evaluate_landmark_auc(
score_mode: str,
horizons: np.ndarray,
device: torch.device,
attn_mask_mode: str,
readout_name: str,
readout_reduce: str,
num_workers_auc: int,
@@ -957,7 +990,6 @@ def evaluate_landmark_auc(
model=model,
loader=loader,
device=device,
attn_mask_mode=attn_mask_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
use_amp=use_amp,
@@ -1193,6 +1225,7 @@ def main() -> None:
include_no_event_in_uts_target=bool(include_no_event_in_uts_target),
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
)
validate_dataset_metadata(dataset, cfg)
has_no_event = (
NO_EVENT_IDX in dataset.label_id_to_code
@@ -1381,7 +1414,6 @@ def main() -> None:
score_mode=score_mode,
horizons=horizons,
device=device,
attn_mask_mode=attn_mask_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
num_workers_auc=num_workers_auc,