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"""