update mlflow_data_extraction.py

Aarhus-Psychiatry-Research · May 15, 2024 · 1a7b580 · 1a7b580
1 parent a18f876
commit 1a7b580
Showing 1 changed file with 4 additions and 0 deletions.
diff --git 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"""