move collect

Author: mertalev

machine-learning/export/immich_model_exporter/parse_eval_data.py

@@ -80,8 +80,7 @@ def collapsed_table(language: str, df: pl.DataFrame) -> str:
         (pl.col("image_retrieval_recall@1") + pl.col("image_retrieval_recall@5") + pl.col("image_retrieval_recall@10"))
         * (100 / 3)
     ).round(2)
-).collect()
-eval_df.write_parquet("model_info.parquet")
+)
 
 pareto_front = eval_df.join_where(
     eval_df.select("language", "peak_rss", "exec_time_ms", "recall").rename(
@@ -104,9 +103,12 @@ def collapsed_table(language: str, df: pl.DataFrame) -> str:
 )
 eval_df = eval_df.join(pareto_front, on=["pretrained_model", "language"], how="left")
 eval_df = eval_df.with_columns(is_pareto=pl.col("recall_other").is_null())
-eval_df = eval_df.drop("peak_rss_other", "exec_time_ms_other", "recall_other", "language_other").unique(
-    subset=["pretrained_model", "language"]
+eval_df = (
+    eval_df.drop("peak_rss_other", "exec_time_ms_other", "recall_other", "language_other")
+    .unique(subset=["pretrained_model", "language"])
+    .collect()
 )
+eval_df.write_parquet("model_info.parquet")
 
 eval_df = eval_df.filter(pl.col("recall") >= 20)
 eval_df = eval_df.select(