Add dist_mode parameter to evaluate_landmark_auc for improved flexibility in evaluation

This commit is contained in:
2026-06-22 09:40:08 +08:00
parent 615f395058
commit 56421eb3f5

View File

@@ -1138,6 +1138,7 @@ def evaluate_landmark_auc(
disease_ids: Sequence[int], disease_ids: Sequence[int],
disease_chunk_size: int, disease_chunk_size: int,
score_mode: str, score_mode: str,
dist_mode: str,
horizons: np.ndarray, horizons: np.ndarray,
device: torch.device, device: torch.device,
model_target_mode: str, model_target_mode: str,
@@ -1208,7 +1209,8 @@ def evaluate_landmark_auc(
death_token_ids=np.asarray( death_token_ids=np.asarray(
landmark_dataset.death_token_ids, dtype=np.int64), landmark_dataset.death_token_ids, dtype=np.int64),
dist_mode=dist_mode, dist_mode=dist_mode,
model_death_idx=int(getattr(model, "death_idx", dataset.vocab_size - 1)), model_death_idx=int(getattr(
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
) )
nested = [_eval_token(t) for t in tqdm( nested = [_eval_token(t) for t in tqdm(
tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)] tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
@@ -1236,7 +1238,8 @@ def evaluate_landmark_auc(
np.asarray(landmark_dataset.death_token_ids, np.asarray(landmark_dataset.death_token_ids,
dtype=np.int64), dtype=np.int64),
dist_mode, dist_mode,
int(getattr(model, "death_idx", dataset.vocab_size - 1)), int(getattr(
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
), ),
) as ex: ) as ex:
nested = list( nested = list(
@@ -1607,6 +1610,7 @@ def main() -> None:
disease_ids=disease_ids, disease_ids=disease_ids,
disease_chunk_size=disease_chunk_size, disease_chunk_size=disease_chunk_size,
score_mode=score_mode, score_mode=score_mode,
dist_mode=dist_mode,
horizons=horizons, horizons=horizons,
device=device, device=device,
model_target_mode=model_target_mode, model_target_mode=model_target_mode,