PrunaAI · davidberenstein1957 · Feb 21, 2026 · Feb 21, 2026 · Feb 21, 2026 · Feb 21, 2026
diff --git a/docs/user_manual/configure.rst b/docs/user_manual/configure.rst
@@ -253,7 +253,7 @@ Underneath you can find the list of all the available datasets.
      - ``text: str``
    * - Image Generation
      - `LAION256 <https://huggingface.co/datasets/nannullna/laion_subset>`_, `OpenImage <https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1>`_, `COCO <https://huggingface.co/datasets/phiyodr/coco2017>`_, `DrawBench <https://huggingface.co/datasets/sayakpaul/drawbench>`_, `PartiPrompts <https://huggingface.co/datasets/nateraw/parti-prompts>`_, `GenAIBench <https://huggingface.co/datasets/BaiqiL/GenAI-Bench>`_
-     - ``image_generation_collate``, ``prompt_collate``
+     - ``image_generation_collate``, ``prompt_with_auxiliaries_collate``
      - ``text: str``, ``image: Optional[PIL.Image.Image]``
    * - Image Classification
      - `ImageNet <https://huggingface.co/datasets/zh-plus/tiny-imagenet>`_, `MNIST <https://huggingface.co/datasets/ylecun/mnist>`_, `CIFAR10 <https://huggingface.co/datasets/uoft-cs/cifar10>`_

diff --git a/pyproject.toml b/pyproject.toml
@@ -151,7 +151,7 @@ dependencies = [
     "peft>=0.18.0",
     "trl<=0.21.0",
     "termcolor==2.3.0",
-    "realesrgan"
+    "realesrgan",
 ]
 
 [project.optional-dependencies]
@@ -166,6 +166,10 @@ vllm = [
     "vllm>=0.16.0",
     "ray",
 ]
+evaluation = [
+    "outlines>1.2.0,<2.0.0",
+    "litellm>=1.0.0",
+]
 stable-fast = [
     "xformers>=0.0.30",
     "stable-fast-pruna>=1.0.8,<1.0.9",
@@ -222,6 +226,7 @@ dev = [
     "types-PyYAML",
     "logbar",
     "pytest-xdist>=3.8.0",
+    "pruna[evaluation]",
 ]
 cpu = []
 lmharness = [
@@ -234,6 +239,9 @@ intel = [
     "torch>=2.7.0,<2.9.0",
     "torchvision>=0.22.0,<0.24.0",
 ]
+mine-replicate = [
+    "replicate>=0.26.0",
+]
 
 [build-system]
 requires = ["hatchling"]

diff --git a/src/pruna/data/__init__.py b/src/pruna/data/__init__.py
@@ -34,7 +34,13 @@
     setup_hps_dataset,
     setup_imgedit_dataset,
     setup_long_text_bench_dataset,
+    setup_oneig_anime_stylization_dataset,
     setup_oneig_dataset,
+    setup_oneig_general_object_dataset,
+    setup_oneig_knowledge_reasoning_dataset,
+    setup_oneig_multilingualism_dataset,
+    setup_oneig_portrait_dataset,
+    setup_oneig_text_rendering_dataset,
     setup_parti_prompts_dataset,
 )
 from pruna.data.datasets.question_answering import setup_polyglot_dataset
@@ -103,19 +109,33 @@
         "image_classification_collate",
         {"img_size": 224},
     ),
-    "DrawBench": (setup_drawbench_dataset, "prompt_collate", {}),
+    "DrawBench": (setup_drawbench_dataset, "prompt_with_auxiliaries_collate", {}),
     "PartiPrompts": (
         setup_parti_prompts_dataset,
         "prompt_with_auxiliaries_collate",
         {},
     ),
-    "GenAIBench": (setup_genai_bench_dataset, "prompt_collate", {}),
+    "GenAIBench": (setup_genai_bench_dataset, "prompt_with_auxiliaries_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", {}),
+    "OneIGAnimeStylization": (
+        setup_oneig_anime_stylization_dataset,
+        "prompt_with_auxiliaries_collate",
+        {},
+    ),
+    "OneIGGeneralObject": (setup_oneig_general_object_dataset, "prompt_with_auxiliaries_collate", {}),
+    "OneIGKnowledgeReasoning": (
+        setup_oneig_knowledge_reasoning_dataset,
+        "prompt_with_auxiliaries_collate",
+        {},
+    ),
+    "OneIGMultilingualism": (setup_oneig_multilingualism_dataset, "prompt_with_auxiliaries_collate", {}),
+    "OneIGPortrait": (setup_oneig_portrait_dataset, "prompt_with_auxiliaries_collate", {}),
+    "OneIGTextRendering": (setup_oneig_text_rendering_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", {}),