feat: add benchmark support to PrunaDataModule and implement PartiPrompts

src/pruna/data/__init__.py

@@ -26,8 +26,15 @@
     setup_mnist_dataset,
 )
 from pruna.data.datasets.prompt import (
+    setup_dpg_dataset,
     setup_drawbench_dataset,
+    setup_gedit_dataset,
     setup_genai_bench_dataset,
+    setup_geneval_dataset,
+    setup_hps_dataset,
+    setup_imgedit_dataset,
+    setup_long_text_bench_dataset,
+    setup_oneig_dataset,
     setup_parti_prompts_dataset,
 )
 from pruna.data.datasets.question_answering import setup_polyglot_dataset
@@ -97,8 +104,19 @@
         {"img_size": 224},
     ),
     "DrawBench": (setup_drawbench_dataset, "prompt_collate", {}),
-    "PartiPrompts": (setup_parti_prompts_dataset, "prompt_collate", {}),
+    "PartiPrompts": (
+        setup_parti_prompts_dataset,
+        "prompt_with_auxiliaries_collate",
+        {},
+    ),
     "GenAIBench": (setup_genai_bench_dataset, "prompt_collate", {}),
+    "GenEval": (setup_geneval_dataset, "prompt_with_auxiliaries_collate", {}),
+    "HPS": (setup_hps_dataset, "prompt_with_auxiliaries_collate", {}),
+    "ImgEdit": (setup_imgedit_dataset, "prompt_with_auxiliaries_collate", {}),
+    "LongTextBench": (setup_long_text_bench_dataset, "prompt_with_auxiliaries_collate", {}),
+    "GEditBench": (setup_gedit_dataset, "prompt_with_auxiliaries_collate", {}),
+    "OneIG": (setup_oneig_dataset, "prompt_with_auxiliaries_collate", {}),
+    "DPG": (setup_dpg_dataset, "prompt_with_auxiliaries_collate", {}),
     "TinyIMDB": (setup_tiny_imdb_dataset, "text_generation_collate", {}),
     "VBench": (setup_vbench_dataset, "prompt_with_auxiliaries_collate", {}),
 }

src/pruna/data/datasets/image.py

@@ -19,7 +19,6 @@
 from datasets import Dataset, load_dataset
 
 from pruna.data.utils import (
-    define_sample_size_for_dataset,
     split_train_into_train_val,
     split_val_into_val_test,
     stratify_dataset,
@@ -54,11 +53,8 @@ def setup_mnist_dataset(
     train_ds = cast(Dataset, train_ds)
     test_ds = cast(Dataset, test_ds)
 
-    train_sample_size = define_sample_size_for_dataset(train_ds, fraction, train_sample_size)
-    test_sample_size = define_sample_size_for_dataset(test_ds, fraction, test_sample_size)
-
-    train_ds = stratify_dataset(train_ds, train_sample_size, seed)
-    test_ds = stratify_dataset(test_ds, test_sample_size, seed)
+    train_ds = stratify_dataset(train_ds, sample_size=train_sample_size, fraction=fraction, seed=seed)
+    test_ds = stratify_dataset(test_ds, sample_size=test_sample_size, fraction=fraction, seed=seed)
 
     train_ds, val_ds = split_train_into_train_val(train_ds, seed)
     val_ds, test_ds = split_val_into_val_test(val_ds, seed)
@@ -93,11 +89,9 @@ def setup_imagenet_dataset(
     train_ds, val = load_dataset("zh-plus/tiny-imagenet", split=["train", "valid"])  # type: ignore[misc]
     train_ds = cast(Dataset, train_ds)
     val = cast(Dataset, val)
-    train_sample_size = define_sample_size_for_dataset(train_ds, fraction, train_sample_size)
-    train_ds = stratify_dataset(train_ds, train_sample_size, seed)
+    train_ds = stratify_dataset(train_ds, sample_size=train_sample_size, fraction=fraction, seed=seed)
     val_ds, test_ds = split_val_into_val_test(val, seed)
-    test_sample_size = define_sample_size_for_dataset(test_ds, fraction, test_sample_size)
-    test_ds = stratify_dataset(test_ds, test_sample_size, seed)
+    test_ds = stratify_dataset(test_ds, sample_size=test_sample_size, fraction=fraction, seed=seed)
     return train_ds, val_ds, test_ds  # type: ignore[return-value]
 
 
@@ -136,11 +130,8 @@ def setup_cifar10_dataset(
     train_ds = train_ds.rename_column("img", "image")
     test_ds = test_ds.rename_column("img", "image")
 
-    train_sample_size = define_sample_size_for_dataset(train_ds, fraction, train_sample_size)
-    test_sample_size = define_sample_size_for_dataset(test_ds, fraction, test_sample_size)
-
-    train_ds = stratify_dataset(train_ds, train_sample_size, seed)
-    test_ds = stratify_dataset(test_ds, test_sample_size, seed)
+    train_ds = stratify_dataset(train_ds, sample_size=train_sample_size, fraction=fraction, seed=seed)
+    test_ds = stratify_dataset(test_ds, sample_size=test_sample_size, fraction=fraction, seed=seed)
 
     train_ds, val_ds = split_train_into_train_val(train_ds, seed)
     return train_ds, val_ds, test_ds  # type: ignore[return-value]