feat: evaluate on xtd-10 and flickr30k datasets

immich_model_exporter/parse_eval_data.py

@@ -58,19 +58,28 @@ def collapsed_table(language: str, df: pl.DataFrame) -> str:
 eval_df = eval_df.with_columns(
     model=pl.col("model")
     .str.replace("xlm-roberta-base", "XLM-Roberta-Base")
-    .str.replace("xlm-roberta-large", "XLM-Roberta-Large")
+    .str.replace("xlm-roberta-large", "XLM-Roberta-Large"),
+    language=pl.when(pl.col("language") == "jp").then(pl.lit("ja")).otherwise("language"),
 )
-eval_df = eval_df.with_columns(pretrained_model=pl.concat_str(pl.col("model"), pl.col("pretrained"), separator="__"))
+eval_df = eval_df.with_columns(pretrained_model=pl.concat_str("model", "pretrained", separator="__"))
 eval_df = eval_df.drop("model", "pretrained")
+eval_df = eval_df.unnest("metrics")
 eval_df = eval_df.join(profile_df, on="pretrained_model")
 
+eval_df = eval_df.group_by("pretrained_model", "language").agg(
+    [
+        pl.mean("peak_rss"),
+        pl.mean("exec_time_ms"),
+        pl.mean("image_retrieval_recall@1"),
+        pl.mean("image_retrieval_recall@5"),
+        pl.mean("image_retrieval_recall@10"),
+    ]
+)
 eval_df = eval_df.with_columns(
     recall=(
-        pl.col("metrics").struct.field("image_retrieval_recall@1")
-        + pl.col("metrics").struct.field("image_retrieval_recall@5")
-        + pl.col("metrics").struct.field("image_retrieval_recall@10")
-    )
-    * (100 / 3)
+        (pl.col("image_retrieval_recall@1") + pl.col("image_retrieval_recall@5") + pl.col("image_retrieval_recall@10"))
+        * (100 / 3)
+    ).round(2)
 )
 
 pareto_front = eval_df.join_where(
@@ -101,17 +110,19 @@ def collapsed_table(language: str, df: pl.DataFrame) -> str:
 )
 eval_df.write_parquet("model_info.parquet")
 
-eval_df = eval_df.drop("metrics")
 eval_df = eval_df.filter(pl.col("recall") >= 20)
-eval_df = eval_df.sort("recall", descending=True)
 eval_df = eval_df.select(
     pl.col("pretrained_model").alias("Model"),
     (pl.col("peak_rss") / 1024).round().cast(pl.UInt32).alias("Memory (MiB)"),
     pl.col("exec_time_ms").round(2).alias("Execution Time (ms)"),
     pl.col("language").alias("Language"),
-    pl.col("recall").round(2).alias("Recall (%)"),
-    pl.when(pl.col("is_pareto")).then(pl.lit("✅")).otherwise(pl.lit("❌")).alias("Pareto Optimal"),
+    # pl.col("image_retrieval_recall@1").mul(100).round(2).alias("Recall@1 (%)"),
+    # pl.col("image_retrieval_recall@5").mul(100).round(2).alias("Recall@5 (%)"),
+    # pl.col("image_retrieval_recall@10").mul(100).round(2).alias("Recall@10 (%)"),
+    pl.col("recall").alias("Recall (%)"),
+    pl.when("is_pareto").then(pl.lit("✅")).otherwise(pl.lit("❌")).alias("Pareto Optimal"),
 )
+eval_df = eval_df.sort("Recall (%)", "Memory (MiB)", descending=[True, False])
 
 
 for language in languages:

immich_model_exporter/run.py

@@ -110,9 +110,12 @@ def export_models(models: list[str], source: ModelSource) -> None:
     export_models(insightface, ModelSource.INSIGHTFACE)
 
     Path("results").mkdir(exist_ok=True)
+    dataset_root = Path("datasets")
+    dataset_root.mkdir(exist_ok=True)
+
+    crossmodal3600_root = dataset_root / "crossmodal3600"
     subprocess.check_call(
         [
-            "python",
             "clip_benchmark",
             "eval",
             "--pretrained_model",
@@ -121,6 +124,8 @@ def export_models(models: list[str], source: ModelSource) -> None:
             "zeroshot_retrieval",
             "--dataset",
             "crossmodal3600",
+            "--dataset_root",
+            crossmodal3600_root.as_posix(),
             "--batch_size",
             "64",
             "--language",
@@ -169,3 +174,69 @@ def export_models(models: list[str], source: ModelSource) -> None:
             "results/{dataset}_{language}_{model}_{pretrained}.json",
         ]
     )
+
+    xtd10_root = dataset_root / "xtd10"
+    subprocess.check_call(
+        [
+            "clip_benchmark",
+            "eval",
+            "--pretrained_model",
+            *[name.replace("__", ",") for name in openclip],
+            "--task",
+            "zeroshot_retrieval",
+            "--dataset",
+            "xtd10",
+            "--dataset_root",
+            xtd10_root.as_posix(),
+            "--batch_size",
+            "64",
+            "--language",
+            "de",
+            "en",
+            "es",
+            "fr",
+            "it",
+            "jp",
+            "ko",
+            "pl",
+            "ru",
+            "tr",
+            "zh",
+            "--recall_k",
+            "1",
+            "5",
+            "10",
+            "--no_amp",
+            "--output",
+            "results/{dataset}_{language}_{model}_{pretrained}.json",
+        ]
+    )
+
+    flickr30k_root = dataset_root / "flickr30k"
+    # note: need ~/.kaggle/kaggle.json to download the dataset automatically
+    subprocess.check_call(
+        [
+            "clip_benchmark",
+            "eval",
+            "--pretrained_model",
+            *[name.replace("__", ",") for name in openclip],
+            "--task",
+            "zeroshot_retrieval",
+            "--dataset",
+            "flickr30k",
+            "--dataset_root",
+            flickr30k_root.as_posix(),
+            "--batch_size",
+            "64",
+            "--language",
+            "en",
+            "zh",
+            "--recall_k",
+            "1",
+            "5",
+            "10",
+            "--no_amp",
+            "--output",
+            "results/{dataset}_{language}_{model}_{pretrained}.json",
+        ]
+    )

pyproject.toml

@@ -14,8 +14,9 @@ dependencies = [
   "rknn-toolkit2>=2.3.0",
   "transformers>=4.49.0",
   "tenacity>=9.0.0",
-  "clip-benchmark>=1.6.1",
   "polars>=1.25.2",
+  "kaggle>=1.7.4.2",
+  "clip-benchmark",
 ]
 
 [dependency-groups]
@@ -29,6 +30,9 @@ override-dependencies = [
   "torchvision>=0.21",
 ]
 
+[tool.uv.sources]
+clip-benchmark = { git = "https://github.com/mertalev/CLIP_benchmark.git", rev = "1770a603e0d37dfc44ddda5094a93c113b3528f6" }
+
 [tool.hatch.build.targets.sdist]
 include = ["immich_model_exporter"]