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

67
evaluate_auc.py
View File

@@ -403,6 +403,32 @@ def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str
return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)
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)
)
# ---------------------------------------------------------------------------
# Batched inference + cached hidden states
# ---------------------------------------------------------------------------
@@ -426,7 +452,6 @@ def infer_readout_hidden(
model: DeepHealth,
loader: DataLoader,
device: torch.device,
attn_mask_mode: str,
readout_name: str,
readout_reduce: str,
use_amp: bool,
@@ -736,7 +761,6 @@ def _init_auc_worker_flat(
p_sex: np.ndarray,
age_groups: np.ndarray,
n_patients: int,
use_delong: bool,
):
# Prevent BLAS/OpenMP oversubscription when many worker processes are active.
os.environ.setdefault("OMP_NUM_THREADS", "1")
@@ -759,7 +783,6 @@ def _init_auc_worker_flat(
"p_sex": p_sex,
"age_groups": age_groups,
"n_patients": int(n_patients),
"use_delong": bool(use_delong),
})
@@ -793,7 +816,6 @@ def _calibration_auc_one_disease_flat(task: Tuple[int, int]) -> List[Dict[str, A
p_sex = _WORKER["p_sex"]
age_groups = _WORKER["age_groups"]
n_patients = _WORKER["n_patients"]
use_delong = _WORKER["use_delong"]
case_idx = _case_indices_for_token(int(token))
if case_idx.size < 2:
@@ -838,12 +860,7 @@ def _calibration_auc_one_disease_flat(task: Tuple[int, int]) -> List[Dict[str, A
if case_scores.size == 0 or control_scores.size == 0:
continue
if use_delong:
auc_value, auc_var = get_auc_delong_var(
control_scores, case_scores)
else:
auc_value, auc_var = get_auc_delong_var(
control_scores, case_scores)
auc_value, auc_var = get_auc_delong_var(control_scores, case_scores)
out.append({
"token": int(token),
@@ -890,7 +907,6 @@ def compute_auc_chunk_parallel(
offset: float,
valid_target_min_id: int,
num_workers: int,
use_delong: bool,
auc_task_chunk_size: int = 0,
) -> List[Dict[str, Any]]:
sex_mask = arrays["sex"] == sex_value
@@ -928,8 +944,7 @@ def compute_auc_chunk_parallel(
flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"],
flat["target_time"], flat["sort_order"], flat["sorted_target_event"],
flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"],
flat["p_sex"], flat["age_groups"], int(
flat["n_patients"]), use_delong,
flat["p_sex"], flat["age_groups"], int(flat["n_patients"]),
)
nested = [_calibration_auc_one_disease_flat(t) for t in tqdm(
tasks, desc=f"AUC {sex_name}", leave=False, dynamic_ncols=True)]
@@ -946,8 +961,7 @@ def compute_auc_chunk_parallel(
flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"],
flat["target_time"], flat["sort_order"], flat["sorted_target_event"],
flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"],
flat["p_sex"], flat["age_groups"], int(
flat["n_patients"]), use_delong,
flat["p_sex"], flat["age_groups"], int(flat["n_patients"]),
),
) as ex:
nested = list(tqdm(
@@ -985,7 +999,6 @@ def evaluate_auc_pipeline(
age_groups: np.ndarray,
offsets: Sequence[float],
device: torch.device,
attn_mask_mode: str,
readout_name: str,
readout_reduce: str,
num_workers_auc: int,
@@ -1045,7 +1058,6 @@ def evaluate_auc_pipeline(
model=model,
loader=loader,
device=device,
attn_mask_mode=attn_mask_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
use_amp=use_amp,
@@ -1075,7 +1087,6 @@ def evaluate_auc_pipeline(
offset=float(offset),
valid_target_min_id=valid_target_min_id,
num_workers=num_workers_auc,
use_delong=True,
auc_task_chunk_size=auc_task_chunk_size,
)
for r in rows:
@@ -1131,6 +1142,14 @@ def parse_int_list(s: Any) -> Optional[List[int]]:
text = str(s).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()]
@@ -1142,6 +1161,14 @@ def parse_float_list(s: Any) -> Optional[List[float]]:
text = str(s).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()]
@@ -1231,8 +1258,6 @@ def main() -> None:
target_mode = cfg.get("target_mode", "uts")
dist_mode_cfg = cfg.get("dist_mode", "exponential")
attn_mask_mode = cfg.get(
"attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
readout_name = cfg.get(
"readout_name", "same_time_group_end" if target_mode == "uts" else "token")
readout_reduce = cfg.get("readout_reduce", "mean")
@@ -1250,6 +1275,7 @@ def main() -> None:
include_no_event_in_uts_target=include_no_event,
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
)
validate_dataset_metadata(dataset, cfg)
subset, subset_indices = make_eval_subset(dataset, args, cfg)
print(f"Dataset: {len(dataset)} samples, vocab_size={dataset.vocab_size}")
@@ -1309,7 +1335,6 @@ def main() -> None:
age_groups=age_groups,
offsets=auc_offsets,
device=device,
attn_mask_mode=attn_mask_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(