fixed a bug in reading the predicted_boolean_probability for computin…

…g auc
cumc-dbmi · Dec 9, 2024 · 562698c · 562698c
1 parent 07292d0
commit 562698c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/cehrbert/runners/hf_cehrbert_finetune_runner.py b/src/cehrbert/runners/hf_cehrbert_finetune_runner.py
@@ -377,7 +377,7 @@ def do_predict(test_dataloader: DataLoader, model_args: ModelArguments, training
     test_prediction_pd = pd.read_parquet(test_prediction_folder)
     # Compute metrics and save results
     metrics = compute_metrics(
-        references=test_prediction_pd.boolean_value, probs=test_prediction_pd.boolean_prediction_probability
+        references=test_prediction_pd.boolean_value, probs=test_prediction_pd.predicted_boolean_probability
     )
     metrics["test_loss"] = np.mean(test_losses)