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

Remove legacy event and mixed distribution paths

Browse Source
This commit is contained in:
Jiarui Li
2026-08-01 14:23:18 +08:00
parent dfb22adf2d
commit de6f9b75b9
22 changed files with 370 additions and 463 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
evaluate_auc.py
View File

@@ -63,7 +63,7 @@ from eval_data import (
)
from model_architectures import resolve_model_architecture
from models import DeepHealth
from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
# ---------------------------------------------------------------------------
@@ -168,7 +168,7 @@ def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> T
# Disease selection
# ---------------------------------------------------------------------------
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
@@ -249,27 +249,20 @@ def load_checkpoint_state_dict(checkpoint_path: str, map_location: str | torch.d
def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
mode = str(cfg_dist_mode).lower()
if mode not in {"exponential", "weibull"}:
raise ValueError(
f"Unsupported dist_mode={mode!r}; expected exponential or weibull."
)
has_rho_head = any(str(k).startswith("rho_head.")
for k in state_dict.keys())
has_rho_death_head = any(str(k).startswith("rho_death_head.")
for k in state_dict.keys())
if has_rho_head and mode != "weibull":
print(
"[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
return "weibull"
if has_rho_death_head and mode != "mixed":
print(
"[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
return "mixed"
if (not has_rho_head) and mode == "weibull":
print(
"[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
return "exponential"
if (not has_rho_death_head) and mode == "mixed":
print(
"[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
return "exponential"
if mode == "weibull" and not has_rho_head:
raise RuntimeError(
"Weibull checkpoint is missing rho_head parameters."
)
if mode == "exponential" and has_rho_head:
raise RuntimeError(
"Exponential checkpoint unexpectedly contains rho_head parameters."
)
return mode
@@ -977,7 +970,7 @@ def evaluate_auc_pipeline(
sex_items = [("female", 0), ("male", 1)]
all_rows: List[Dict[str, Any]] = []
valid_target_min_id = CHECKUP_IDX if NO_EVENT_IDX >= dataset.vocab_size else CHECKUP_IDX
valid_target_min_id = RESERVED_IDX
# If NO_EVENT exists and should not be a disease/control target, require target > NO_EVENT_IDX.
if NO_EVENT_IDX in dataset.label_id_to_code and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>":
valid_target_min_id = NO_EVENT_IDX