From 1a7b580f79f803dfef21bfa72d0019289dbc0484 Mon Sep 17 00:00:00 2001 From: Martin Bernstorff Date: Wed, 15 May 2024 12:28:19 +0200 Subject: [PATCH] update mlflow_data_extraction.py --- psycop/common/global_utils/mlflow/mlflow_data_extraction.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/psycop/common/global_utils/mlflow/mlflow_data_extraction.py b/psycop/common/global_utils/mlflow/mlflow_data_extraction.py index a8aa43e2e..21d4326cf 100644 --- a/psycop/common/global_utils/mlflow/mlflow_data_extraction.py +++ b/psycop/common/global_utils/mlflow/mlflow_data_extraction.py @@ -46,6 +46,10 @@ def get_config(self) -> Config: cfg_path = self.download_artifact(artifact_name="config.cfg", save_location=None) return Config().from_disk(cfg_path) + def eval_df(self) -> pl.DataFrame: + eval_df_path = self.download_artifact(artifact_name="eval_df.parquet", save_location=None) + return pl.read_parquet(eval_df_path) + def download_artifact(self, artifact_name: str, save_location: str | None = None) -> Path: """Download an artifact from a run. Returns the path to the downloaded artifact. If save_location is None, will save to temporary directory"""