updated the prediction output schema to be compatible with the meds-e…

…valuation
cumc-dbmi · Dec 9, 2024 · 181c765 · 181c765
1 parent 2772603
commit 181c765
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/cehrbert/runners/hf_cehrbert_finetune_runner.py b/src/cehrbert/runners/hf_cehrbert_finetune_runner.py
@@ -365,8 +365,8 @@ def do_predict(test_dataloader: DataLoader, model_args: ModelArguments, training
                 {
                     "subject_id": person_ids,
                     "prediction_time": index_dates,
-                    "boolean_prediction_probability": probabilities,
-                    "boolean_prediction": logits,
+                    "predicted_boolean_probability": probabilities,
+                    "predicted_boolean_value": None,
                     "boolean_value": labels,
                 }
             )