Add dist_mode parameter to evaluate_landmark_auc for improved flexibility in evaluation
This commit is contained in:
@@ -1138,6 +1138,7 @@ def evaluate_landmark_auc(
|
||||
disease_ids: Sequence[int],
|
||||
disease_chunk_size: int,
|
||||
score_mode: str,
|
||||
dist_mode: str,
|
||||
horizons: np.ndarray,
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
@@ -1208,7 +1209,8 @@ def evaluate_landmark_auc(
|
||||
death_token_ids=np.asarray(
|
||||
landmark_dataset.death_token_ids, dtype=np.int64),
|
||||
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(
|
||||
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,
|
||||
dtype=np.int64),
|
||||
dist_mode,
|
||||
int(getattr(model, "death_idx", dataset.vocab_size - 1)),
|
||||
int(getattr(
|
||||
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
|
||||
),
|
||||
) as ex:
|
||||
nested = list(
|
||||
@@ -1607,6 +1610,7 @@ def main() -> None:
|
||||
disease_ids=disease_ids,
|
||||
disease_chunk_size=disease_chunk_size,
|
||||
score_mode=score_mode,
|
||||
dist_mode=dist_mode,
|
||||
horizons=horizons,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
|
||||
Reference in New Issue
Block a user