From f5ac42b50a8233154bbad7ff935b6194fb6efa1c Mon Sep 17 00:00:00 2001 From: davidberenstein1957 Date: Tue, 28 Apr 2026 15:01:40 +0200 Subject: [PATCH 1/4] feat(vision-metrics): split vie_score into dedicated branch Adds VieScoreMetric with GEditBench benchmark wiring and focused unit coverage while keeping image-edit scoring changes for the next stacked PR. Made-with: Cursor --- src/pruna/evaluation/benchmarks.py | 4 +- .../evaluation/metrics/metric_vie_score.py | 363 ++++++++++++++++++ tests/evaluation/test_vision_metrics.py | 20 +- 3 files changed, 384 insertions(+), 3 deletions(-) create mode 100644 src/pruna/evaluation/metrics/metric_vie_score.py diff --git a/src/pruna/evaluation/benchmarks.py b/src/pruna/evaluation/benchmarks.py index 4319226c..99db12db 100644 --- a/src/pruna/evaluation/benchmarks.py +++ b/src/pruna/evaluation/benchmarks.py @@ -267,8 +267,8 @@ def list(cls, task_type: str | None = None) -> list[str]: "material alter, motion change, style change, subject add/remove/replace, text change, " "tone transfer, and human retouching." ), - metrics=[], # Paper uses VIEScore; not in Pruna - task_type="text_to_image", + metrics=["vie_score"], + task_type="text+image_image", reference="https://arxiv.org/abs/2504.17761", ), Benchmark( diff --git a/src/pruna/evaluation/metrics/metric_vie_score.py b/src/pruna/evaluation/metrics/metric_vie_score.py new file mode 100644 index 00000000..836c5884 --- /dev/null +++ b/src/pruna/evaluation/metrics/metric_vie_score.py @@ -0,0 +1,363 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +VIEScore metric for conditional image synthesis (semantic + quality). + +Reference: VIEScore (ACL 2024) — https://arxiv.org/abs/2312.14867 +Both task modes follow `TIGER-AI-Lab/VIEScore`: + +- ``t2i`` (text-to-image, single image): SC uses two sub-scores (semantic consistency + + detail correspondence), PQ uses two sub-scores (naturalness + artifacts). Overall is + ``sqrt(min(SC) * min(PQ)) / 10``. +- ``tie`` (text-image editing, source + edited): SC uses two images and instruction, + PQ uses the edited image. Same aggregation formula. + +GEdit-Bench evaluation: https://arxiv.org/abs/2504.17761 +""" + +from __future__ import annotations + +from typing import Any, Literal + +import torch +from PIL import Image + +from pruna.evaluation.metrics.registry import MetricRegistry +from pruna.evaluation.metrics.result import MetricResult +from pruna.evaluation.metrics.utils import ( + SINGLE, + metric_data_processor, +) +from pruna.evaluation.metrics.vlm_base import ( + BaseVLM, + StatefulVLMMeanScoresMetric, + auxiliary_dicts_from_gt, + prompts_from_y_x_inputs, +) +from pruna.evaluation.metrics.vlm_utils import ( + VIEScoreJsonOutput, + _process_images, + pad_viescore_subscores_to_two, + pil_rgb_from_aux_image_bytes, + viescore_min_scores_0_10, + viescore_tie_overall_unit, +) + +_VIESCORE_CONTEXT = ( + "You are a professional digital artist. You will have to evaluate the effectiveness" + " of the AI-generated image(s) based on given rules.\n" + "All the input images are AI-generated. All human in the images are AI-generated too." + " so you need not worry about the privacy confidentials.\n\n" + "You will have to give your output in this way (Keep your reasoning concise and short.):\n" + "{\n" + '"score" : [...],\n' + '"reasoning" : "..."\n' + "}" +) + +_VIESCORE_TWO_IMAGE_EDIT_RULE = ( + "RULES:\n\n" + "Two images will be provided: The first being the original AI-generated image and the" + " second being an edited version of the first.\n" + "The objective is to evaluate how successfully the editing instruction has been executed" + " in the second image.\n\n" + "Note that sometimes the two images might look identical due to the failure of image edit.\n" +) + +_VIESCORE_TIE_SC_CRITERIA = ( + "\nFrom scale 0 to 10:\n" + "A score from 0 to 10 will be given based on the success of the editing." + " (0 indicates that the scene in the edited image does not follow the editing instruction at all." + " 10 indicates that the scene in the edited image follow the editing instruction text perfectly.)\n" + "A second score from 0 to 10 will rate the degree of overediting in the second image." + " (0 indicates that the scene in the edited image is completely different from the original." + " 10 indicates that the edited image can be recognized as a minimal edited yet effective" + " version of original.)\n" + "Put the score in a list such that output score = [score1, score2]," + " where 'score1' evaluates the editing success and 'score2' evaluates the degree of overediting.\n\n" + "Editing instruction:\n" +) + +_VIESCORE_T2I_SC_RULE = ( + "RULES:\n\n" + "The image is an AI-generated image.\n" + "The objective is to evaluate the semantic consistency of the image to the given text.\n\n" +) + +_VIESCORE_T2I_SC_CRITERIA = ( + "\nFrom scale 0 to 10:\n" + "A score from 0 to 10 will be given based on the semantic consistency.\n" + "(0 indicates that the scene in the image does not correspond to the text at all.\n" + " 10 indicates that the scene in the image follows the text perfectly.)\n" + "A second score from 0 to 10 will rate the detail correspondence.\n" + "(0 indicates that most details in the text (e.g., color, size, shape, or layout) are missing or" + " incorrect in the image.\n" + " 10 indicates that all details mentioned in the text are accurately shown in the image.)\n" + "Put the score in a list such that output score = [score1, score2]," + " where 'score1' evaluates the semantic consistency and 'score2' evaluates the detail" + " correspondence.\n\n" + "Text prompt:\n" +) + +_VIESCORE_PQ_SINGLE_IMAGE = ( + "RULES:\n\n" + "The image is an AI-generated image.\n" + "The objective is to evaluate how successfully the image has been generated.\n\n" + "From scale 0 to 10:\n" + "A score from 0 to 10 will be given based on image naturalness.\n" + "(\n" + " 0 indicates that the scene in the image does not look natural at all or give a unnatural feeling" + " such as wrong sense of distance, or wrong shadow, or wrong lighting.\n" + " 10 indicates that the image looks natural.\n" + ")\n" + "A second score from 0 to 10 will rate the image artifacts.\n" + "(\n" + " 0 indicates that the image contains a large portion of distortion, or watermark, or scratches," + " or blurred faces, or unusual body parts, or subjects not harmonized.\n" + " 10 indicates the image has no artifacts.\n" + ")\n" + "Put the score in a list such that output score = [naturalness, artifacts]\n" +) + + +def _build_viescore_tie_sc_prompt(instruction: str) -> str: + """Build the VIEScore ``tie`` semantic-criteria prompt (source + edited images). + + Args: + instruction: Editing instruction embedded in the prompt. + + Returns: + ------- + Full prompt aligned with TIGER-AI-Lab/VIEScore ``tie`` SC. + """ + return "\n".join( + [ + _VIESCORE_CONTEXT, + _VIESCORE_TWO_IMAGE_EDIT_RULE, + _VIESCORE_TIE_SC_CRITERIA.strip(), + instruction.strip(), + ] + ) + + +def _build_viescore_t2i_sc_prompt(prompt: str) -> str: + """Build the VIEScore ``t2i`` semantic-consistency prompt for one generated image. + + Args: + prompt: Text prompt used to generate the image. + + Returns: + ------- + Full prompt aligned with TIGER-AI-Lab/VIEScore ``t2i`` SC. + """ + return "\n".join( + [ + _VIESCORE_CONTEXT, + _VIESCORE_T2I_SC_RULE.strip(), + _VIESCORE_T2I_SC_CRITERIA.strip(), + prompt.strip(), + ] + ) + + +def _build_viescore_pq_prompt() -> str: + """Build the VIEScore perceptual-quality prompt for one image (SC or edited).""" + return "\n".join([_VIESCORE_CONTEXT, _VIESCORE_PQ_SINGLE_IMAGE]) + + +@MetricRegistry.register("vie_score") +class VieScoreMetric(StatefulVLMMeanScoresMetric): + """ + VIEScore: semantic + perceptual quality with geometric-mean overall. + + **Text-to-image (one generated image):** uses the VIEScore ``t2i`` SC prompt (semantic + consistency + detail correspondence, 0--10 each) and the shared PQ prompt (naturalness + + artifacts, 0--10 each). Overall is ``sqrt(min(SC) * min(PQ)) / 10`` in ``[0, 1]``. + + **Text--image editing (source + edited available):** matches the VIEScore ``tie`` setup + used in GEdit-Bench: semantic criteria use **two** images (source then edited) and the + editing instruction; perceptual criteria use the **edited** image only. Overall is + ``sqrt(min(SC) * min(PQ)) / 10`` in ``[0, 1]``, with ``min`` taken over the sub-scores in + each JSON ``score`` list, consistent with `VIEScore`_. + + .. _VIEScore: https://github.com/TIGER-AI-Lab/VIEScore + + Parameters + ---------- + *args : Any + Additional positional arguments. + vlm : BaseVLM | None, optional + Custom VLM instance. If provided, vlm_type and model_name are ignored. + vlm_type : {"litellm", "transformers"}, optional + VLM backend. Default is "litellm". + model_name : str | None, optional + Litellm model id or HuggingFace checkpoint id. **Required** when ``vlm`` is not + provided (e.g. ``openai/gpt-4o``). + vlm_kwargs : dict, optional + Forwarded by ``get_vlm`` to ``LitellmVLM`` or ``TransformersVLM``. For local models, + set ``model_load_kwargs`` for ``from_pretrained``; for litellm, pass extra API options. + structured_output : bool, optional + Use structured generation (litellm pydantic; transformers may use plain generation for + multi-image). Default is True. + device : str | torch.device | None, optional + Device for transformers VLM. + api_key : str | None, optional + API key for litellm. + call_type : str, optional + Call type for the metric. + **kwargs : Any + Additional arguments. + + References + ---------- + VIEScore: Towards Explainable Metrics for Conditional Image Synthesis Evaluation (ACL 2024) + https://arxiv.org/abs/2312.14867 + https://github.com/TIGER-AI-Lab/VIEScore + + GEdit-Bench (image editing evaluation) + https://arxiv.org/abs/2504.17761 + + Examples + -------- + Same ``hosted`` / ``local`` pattern as :func:`~pruna.evaluation.metrics.vlm_base.get_vlm``. + Multi-image ``tie`` paths call ``generate_with_image_lists`` on ``self.vlm`` internally. + + .. code-block:: python + + import torch + + from pruna.evaluation.metrics import VieScoreMetric + + hosted = VieScoreMetric(vlm_type="litellm", model_name="openai/gpt-4o") + local = VieScoreMetric( + vlm_type="transformers", + model_name="HuggingFaceTB/SmolVLM-256M-Instruct", + device="cpu", + vlm_kwargs={"model_load_kwargs": {"torch_dtype": torch.float32}}, + ) + """ + + scores: list[float] + default_call_type: str = "y_x" + higher_is_better: bool = True + metric_name: str = "vie_score" + + def __init__( + self, + *args, + vlm: BaseVLM | None = None, + vlm_type: Literal["litellm", "transformers"] = "litellm", + model_name: str | None = None, + vlm_kwargs: dict | None = None, + structured_output: bool = True, + device: str | torch.device | None = None, + api_key: str | None = None, + call_type: str = SINGLE, + **kwargs: Any, + ) -> None: + super().__init__(device=device) + self.structured_output = structured_output + self.response_format = VIEScoreJsonOutput if structured_output else None + + self._init_vlm_scores( + vlm=vlm, + vlm_type=vlm_type, + model_name=model_name, + vlm_kwargs=vlm_kwargs, + structured_output=structured_output, + device=device, + api_key=api_key, + call_type=call_type, + ) + + def _score_single_image_t2i(self, image: Image.Image, prompt: str) -> float: + """VIEScore ``t2i``: single-image SC (semantic + detail) and PQ (naturalness + artifacts). + + Matches the VIEScore paper's t2i evaluation: two SC sub-scores on 0--10 and two PQ + sub-scores on 0--10, aggregated as ``sqrt(min(SC) * min(PQ)) / 10``. + """ + sc_prompt = _build_viescore_t2i_sc_prompt(prompt) + pq_prompt = _build_viescore_pq_prompt() + + rf = self.response_format if self.structured_output else None + + sc_raw = self.vlm.generate([image], [sc_prompt], response_format=rf)[0] + pq_raw = self.vlm.generate([image], [pq_prompt], response_format=rf)[0] + + sc_list = pad_viescore_subscores_to_two(viescore_min_scores_0_10(sc_raw)) + pq_list = pad_viescore_subscores_to_two(viescore_min_scores_0_10(pq_raw)) + return viescore_tie_overall_unit(sc_list, pq_list) + + def _score_tie_gedit(self, source: Image.Image, edited: Image.Image, instruction: str) -> float: + """VIEScore ``tie``: two-image SC, single-image PQ, overall geometric mean on 0--10 mins.""" + sc_prompt = _build_viescore_tie_sc_prompt(instruction) + pq_prompt = _build_viescore_pq_prompt() + + rf = self.response_format if self.structured_output else None + + if hasattr(self.vlm, "generate_with_image_lists"): + sc_raw = self.vlm.generate_with_image_lists( + [[source, edited]], + [sc_prompt], + response_format=rf, + )[0] + else: + raise RuntimeError("VLM backend must implement generate_with_image_lists for editing parity.") + + pq_raw = self.vlm.generate([edited], [pq_prompt], response_format=rf)[0] + + sc_list = pad_viescore_subscores_to_two(viescore_min_scores_0_10(sc_raw)) + pq_list = pad_viescore_subscores_to_two(viescore_min_scores_0_10(pq_raw)) + return viescore_tie_overall_unit(sc_list, pq_list) + + def update(self, x: list[Any] | torch.Tensor, gt: Any, outputs: torch.Tensor) -> None: + """ + Update the metric with new batch data. + + Parameters + ---------- + x : List[Any] | torch.Tensor + The input data (prompts). + gt : Any + Per-sample auxiliary dicts (``prompt_with_auxiliaries_collate``), or tensor placeholders + when aux is unused. + outputs : torch.Tensor + The output images. + """ + inputs = metric_data_processor(x, gt, outputs, self.call_type) + images = _process_images(inputs[0]) + prompts = prompts_from_y_x_inputs(inputs, len(images)) + aux_list = auxiliary_dicts_from_gt(gt, len(images)) + + for i, image in enumerate(images): + prompt = prompts[i] if i < len(prompts) else "" + aux = aux_list[i] + source = pil_rgb_from_aux_image_bytes(aux, min_bytes_in_value_scan=100) + + if source is not None: + self.scores.append(self._score_tie_gedit(source, image, prompt)) + else: + self.scores.append(self._score_single_image_t2i(image, prompt)) + + def compute(self) -> MetricResult: + """ + Compute the VIEScore metric. + + Returns + ------- + MetricResult + The mean VIEScore across all updates. + """ + return self.compute_mean_of_scores() diff --git a/tests/evaluation/test_vision_metrics.py b/tests/evaluation/test_vision_metrics.py index bba4b227..e7284eee 100644 --- a/tests/evaluation/test_vision_metrics.py +++ b/tests/evaluation/test_vision_metrics.py @@ -5,10 +5,15 @@ import pytest import torch +from pruna.evaluation.metrics.metric_vie_score import VieScoreMetric from pruna.evaluation.metrics.metric_vqa import VQAMetric from pruna.evaluation.metrics.vlm_base import BaseVLM +def _dummy_image(batch: int = 1, size: int = 64) -> torch.Tensor: + return torch.rand(batch, 3, size, size) + + @pytest.mark.cpu def test_vqa_uses_prompt_question_and_scores_yes_probability() -> None: """VQA asks prompt-grounded yes/no question and stores returned score.""" @@ -16,7 +21,7 @@ def test_vqa_uses_prompt_question_and_scores_yes_probability() -> None: mock_vlm.score.return_value = [0.7] metric = VQAMetric(vlm=mock_vlm, vlm_type="litellm", device="cpu", use_probability=True) - images = torch.rand(1, 3, 64, 64) + images = _dummy_image() metric.update(["a cat"], images, images) result = metric.compute() @@ -24,3 +29,16 @@ def test_vqa_uses_prompt_question_and_scores_yes_probability() -> None: assert result.result == 0.7 call = mock_vlm.score.call_args assert call[0][1] == ['Does this image show "a cat"?'] + + +@pytest.mark.cpu +def test_vie_score_uses_json_score_lists() -> None: + """VieScoreMetric parses JSON score lists and returns normalized value.""" + mock_vlm = MagicMock(spec=BaseVLM) + mock_vlm.generate.return_value = ['{"score": [8.0, 8.0], "reasoning": ""}'] + + metric = VieScoreMetric(vlm=mock_vlm, device="cpu", structured_output=True) + metric.update(["a cat on a sofa"], _dummy_image(), _dummy_image()) + result = metric.compute() + + assert abs(result.result - 0.8) < 0.01 From 63e96aa963218e7242d7488669878eff97e9fbe1 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Tue, 2 Jun 2026 19:28:56 +0200 Subject: [PATCH 2/4] chore(metrics): export VieScoreMetric from __init__ Co-authored-by: Cursor --- src/pruna/evaluation/metrics/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pruna/evaluation/metrics/__init__.py b/src/pruna/evaluation/metrics/__init__.py index c52f354d..4f18b9f8 100644 --- a/src/pruna/evaluation/metrics/__init__.py +++ b/src/pruna/evaluation/metrics/__init__.py @@ -30,6 +30,7 @@ from pruna.evaluation.metrics.metric_sharpness import SharpnessMetric from pruna.evaluation.metrics.metric_text_score import OneIGTextScoreMetric, TextScoreMetric from pruna.evaluation.metrics.metric_torch import TorchMetricWrapper +from pruna.evaluation.metrics.metric_vie_score import VieScoreMetric from pruna.evaluation.metrics.metric_vqa import VQAMetric from pruna.evaluation.metrics.vlm_base import ( BaseVLM, @@ -65,6 +66,7 @@ "RapidataMetric", "TextScoreMetric", "VQAMetric", + "VieScoreMetric", "BaseVLM", "LitellmVLM", "StatefulVLMMeanScoresMetric", From 0beeaba90f7f23dd7d24fec100a36892513735f7 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Thu, 2 Jul 2026 15:50:32 +0200 Subject: [PATCH 3/4] feat(vision-metrics): split vie_score into dedicated branch Co-authored-by: Cursor --- src/pruna/evaluation/benchmarks.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pruna/evaluation/benchmarks.py b/src/pruna/evaluation/benchmarks.py index 99db12db..36e50c5e 100644 --- a/src/pruna/evaluation/benchmarks.py +++ b/src/pruna/evaluation/benchmarks.py @@ -174,7 +174,7 @@ def list(cls, task_type: str | None = None) -> list[str]: "Covers basic skills (scene, attributes, spatial relationships) to advanced reasoning " "(counting, comparison, logic/negation) with over 24k human ratings." ), - metrics=["vqa", "clip_score"], + metrics=[], # Paper uses VQAScore only; not in Pruna task_type="text_to_image", reference="https://arxiv.org/abs/2406.13743", ), @@ -267,8 +267,8 @@ def list(cls, task_type: str | None = None) -> list[str]: "material alter, motion change, style change, subject add/remove/replace, text change, " "tone transfer, and human retouching." ), - metrics=["vie_score"], - task_type="text+image_image", + metrics=[], # Paper uses VIEScore; not in Pruna + task_type="text_to_image", reference="https://arxiv.org/abs/2504.17761", ), Benchmark( From 2f708d886153d91e825bb70f1c615621537dd5cd Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Thu, 2 Jul 2026 16:19:32 +0200 Subject: [PATCH 4/4] fix(ci): lint, format, and benchmark registration - fix benchmarks UTF-8 corruption (1-5 ratings) - sync OneIG subset dataset loaders for benchmark registration - ruff check/format on changed VLM src files --- src/pruna/data/__init__.py | 24 +- src/pruna/data/datasets/prompt.py | 610 ++++++++++++++++++++++++++--- src/pruna/evaluation/benchmarks.py | 6 +- 3 files changed, 578 insertions(+), 62 deletions(-) diff --git a/src/pruna/data/__init__.py b/src/pruna/data/__init__.py index fd14a496..1f0ed5f6 100644 --- 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", {}), diff --git a/src/pruna/data/datasets/prompt.py b/src/pruna/data/datasets/prompt.py index 7764d23b..b18d03d4 100644 --- a/src/pruna/data/datasets/prompt.py +++ b/src/pruna/data/datasets/prompt.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from pathlib import Path from typing import Literal, Tuple, get_args from datasets import Dataset, load_dataset @@ -123,21 +124,95 @@ DPGCategory = Literal["entity", "attribute", "relation", "global", "other"] -def _to_oneig_record(row: dict, questions_by_key: dict[str, dict]) -> dict: - """Convert OneIG row to unified record format.""" +def _warn_ignored_benchmark_seed(seed: int | None, *, dataset: str) -> None: + if seed is not None: + pruna_logger.warning( + "%s: `seed` is ignored for this test-only benchmark; sampling does not shuffle the test split.", + dataset, + ) + + +def _oneig_alignment_language_zh(row: dict) -> bool: + """Return True when the official Q_D file for this row should use the ``*_zh`` graphs.""" + if row.get("category", "") == "Multilingualism": + return True + lang = row.get("language") or row.get("lang") + return isinstance(lang, str) and lang.lower() in {"zh", "zh-cn", "zh_cn", "chinese", "cn"} + + +def _oneig_qd_prefix(row: dict) -> str: + """Map dataset ``category`` (+ language) to Q_D JSON stem (e.g. ``object``, ``anime_zh``).""" + row_category = row.get("category", "") + use_zh = _oneig_alignment_language_zh(row) + if row_category == "Multilingualism": + return "multilingualism_zh" + base = _CATEGORY_TO_QD.get(row_category, "") + if not base: + return "" + return f"{base}_zh" if use_zh else base + + +def _to_oneig_record( + row: dict, + questions_by_key: dict[str, dict], + reasoning_gt_en: dict[str, str], + reasoning_gt_zh: dict[str, str], + reasoning_language: str = "EN", +) -> dict: + """Convert OneIG row to unified record format. + + Parameters + ---------- + row : dict + Raw Hugging Face row (``category``, ``id``, ``class``). EN configs use ``prompt_en``; the + ``OneIG-Bench-ZH`` **Multilingualism** split uses ``prompt_cn`` instead of ``prompt_en``. + questions_by_key : dict[str, dict] + Merged Q_D index keyed as ``{qd_stem}_{prompt_id}`` (see ``_fetch_oneig_alignment``). + reasoning_gt_en : dict[str, str] + Official ``gt_answer.json`` keyed by prompt id (e.g. ``"000"``). + reasoning_gt_zh : dict[str, str] + Official ``gt_answer_zh.json`` keyed by prompt id. + reasoning_language : str, optional + Which reasoning GT to use: ``"EN"`` or ``"ZH"``. Default is ``"EN"``. + + Returns + ------- + dict + Unified record including ``questions``, ``dependencies``, and ``reasoning_gt_answer`` when + applicable (Knowledge_Reasoning only). + """ row_category = row.get("category", "") row_class = row.get("class", "None") or "None" - qd_name = _CATEGORY_TO_QD.get(row_category, "") - lookup_key = f"{qd_name}_{row.get('id', '')}" if qd_name else "" + prompt_id = str(row.get("id", "")) + qd_prefix = _oneig_qd_prefix(row) + lookup_key = f"{qd_prefix}_{prompt_id}" if qd_prefix else "" q_info = questions_by_key.get(lookup_key, {}) + text = row.get("prompt") or row.get("prompt_en") or row.get("prompt_cn") or "" + reasoning_gt_answer: str | None = None + if row_category == "Knowledge_Reasoning": + if reasoning_language.upper() == "ZH": + reasoning_gt_answer = reasoning_gt_zh.get(prompt_id) + else: + reasoning_gt_answer = reasoning_gt_en.get(prompt_id) + is_text_rendering = row_category in ("Text_Rendering", "Text Rendering") + if is_text_rendering and text: + import re as _re + + quoted = _re.findall(r'"([^"]+)"', text) + text_content: str | None = " ".join(quoted) if quoted else (row_class if row_class != "None" else None) + else: + text_content = row_class if row_class != "None" else None + questions = {k: v for k, v in q_info.get("questions", {}).items() if v is not None} + dependencies = {k: v for k, v in q_info.get("dependencies", {}).items() if v is not None} return { - "text": row.get("prompt_en", row.get("prompt", "")), - "subset": "Text_Rendering" if row_category in ("Text_Rendering", "Text Rendering") else row_category, - "text_content": row_class if row_class != "None" else None, + "text": text, + "subset": "Text_Rendering" if is_text_rendering else row_category, + "text_content": text_content, "category": row_category, "class": row_class, - "questions": q_info.get("questions", {}), - "dependencies": q_info.get("dependencies", {}), + "questions": questions, + "dependencies": dependencies, + "reasoning_gt_answer": reasoning_gt_answer, } @@ -159,7 +234,7 @@ def setup_drawbench_dataset() -> Tuple[Dataset, Dataset, Dataset]: def setup_parti_prompts_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -172,8 +247,8 @@ def setup_parti_prompts_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -188,6 +263,7 @@ def setup_parti_prompts_dataset( Tuple[Dataset, Dataset, Dataset] The Parti Prompts dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="PartiPrompts") ds = load_dataset("nateraw/parti-prompts")["train"] # type: ignore[index] if category is not None: @@ -226,7 +302,7 @@ def _generate_geneval_question(entry: dict) -> list[str]: def setup_geneval_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -239,8 +315,8 @@ def setup_geneval_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -255,6 +331,7 @@ def setup_geneval_dataset( Tuple[Dataset, Dataset, Dataset] The GenEval dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="GenEval") import json import requests @@ -286,7 +363,7 @@ def setup_geneval_dataset( def setup_hps_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -299,8 +376,8 @@ def setup_hps_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -315,6 +392,7 @@ def setup_hps_dataset( Tuple[Dataset, Dataset, Dataset] The HPD dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="HPS") import json from huggingface_hub import hf_hub_download @@ -338,7 +416,7 @@ def setup_hps_dataset( def setup_long_text_bench_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -350,8 +428,8 @@ def setup_long_text_bench_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -364,6 +442,7 @@ def setup_long_text_bench_dataset( Tuple[Dataset, Dataset, Dataset] The Long Text Bench dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="LongTextBench") ds = load_dataset("X-Omni/LongText-Bench")["train"] # type: ignore[index] ds = ds.rename_column("text", "text_content") ds = ds.rename_column("prompt", "text") @@ -389,8 +468,74 @@ def setup_genai_bench_dataset() -> Tuple[Dataset, Dataset, Dataset]: return ds.select([0]), ds.select([0]), ds +def ensure_imgedit_benchmark_images_extracted() -> Path: + """ + Download ``Benchmark.tar`` (if needed), extract it, and return the ``singleturn`` folder. + + Returns + ------- + Path + ``Benchmark/singleturn`` directory whose files are addressed by ``image_id``. + + Raises + ------ + RuntimeError + If the archive cannot be downloaded or extracted. + """ + import tarfile + + from huggingface_hub import hf_hub_download + + tar_path = Path(hf_hub_download(repo_id="sysuyy/ImgEdit", filename="Benchmark.tar", repo_type="dataset")) + extract_dir = tar_path.parent / "imgedit_singleturn" + candidate = extract_dir / "Benchmark" / "singleturn" + if not candidate.is_dir() or not any(candidate.iterdir()): + extract_dir.mkdir(parents=True, exist_ok=True) + with tarfile.open(tar_path, "r") as tar: + tar.extractall(path=extract_dir) + if not candidate.is_dir() or not any(candidate.iterdir()): + raise RuntimeError(f"ImgEdit: failed to extract Benchmark.tar to {candidate}") + return candidate + + +def load_imgedit_source_image_bytes(image_id: str, *, image_folder: Path | None = None) -> bytes: + """ + Read one ImgEdit source image as JPEG bytes (RGB). + + Parameters + ---------- + image_id : str + Path relative to the singleturn folder (from the official ``basic_edit.json`` ``id``). + image_folder : Path | None, optional + ``Benchmark/singleturn`` directory; when ``None``, calls + :func:`ensure_imgedit_benchmark_images_extracted`. + + Returns + ------- + bytes + JPEG-encoded bytes for the source image. + + Raises + ------ + FileNotFoundError + If ``image_id`` does not exist under ``image_folder``. + Exception + If PIL cannot open or convert the image. + """ + from io import BytesIO + + from PIL import Image + + folder = image_folder if image_folder is not None else ensure_imgedit_benchmark_images_extracted() + img_path = folder / image_id + pil = Image.open(img_path).convert("RGB") + buf = BytesIO() + pil.save(buf, format="JPEG") + return buf.getvalue() + + def setup_imgedit_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -403,8 +548,8 @@ def setup_imgedit_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -420,6 +565,7 @@ def setup_imgedit_dataset( Tuple[Dataset, Dataset, Dataset] The ImgEdit dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="ImgEdit") import json import requests @@ -433,6 +579,8 @@ def setup_imgedit_dataset( instructions: dict = json.loads(response_instructions.text) judge_prompts: dict = json.loads(response_judge_prompts.text) + image_folder = ensure_imgedit_benchmark_images_extracted() + categories = [category] if category is not None and not isinstance(category, list) else category records = [] for _, instruction in instructions.items(): @@ -441,14 +589,17 @@ def setup_imgedit_dataset( if categories is not None and edit_type not in categories: continue - records.append( - { - "text": instruction.get("prompt", ""), - "category": edit_type, - "image_id": instruction.get("id", ""), - "judge_prompt": judge_prompts.get(edit_type, ""), - } - ) + image_id = instruction.get("id", "") + record: dict = { + "text": instruction.get("prompt", ""), + "category": edit_type, + "image_id": image_id, + "judge_prompt": judge_prompts.get(edit_type, ""), + } + src = load_imgedit_source_image_bytes(image_id, image_folder=image_folder) + if src is not None: + record["source_image_bytes"] = src + records.append(record) ds = Dataset.from_list(records) ds = stratify_dataset(ds, sample_size=test_sample_size, fraction=fraction) @@ -466,18 +617,47 @@ def setup_imgedit_dataset( "General_Object": "object", } -_ONEIG_ALIGNMENT_BASE = "https://raw.githubusercontent.com/OneIG-Bench/OneIG-Benchmark/41b49831e79e6dde5323618c164da1c4cf0f699d/scripts/alignment/Q_D" +_ONEIG_BENCHMARK_REF = "41b49831e79e6dde5323618c164da1c4cf0f699d" +_ONEIG_RAW_BASE = f"https://raw.githubusercontent.com/OneIG-Bench/OneIG-Benchmark/{_ONEIG_BENCHMARK_REF}" +_ONEIG_ALIGNMENT_QD_URL = f"{_ONEIG_RAW_BASE}/scripts/alignment/Q_D" +_ONEIG_REASONING_GT_URL_EN = f"{_ONEIG_RAW_BASE}/scripts/reasoning/gt_answer.json" +_ONEIG_REASONING_GT_URL_ZH = f"{_ONEIG_RAW_BASE}/scripts/reasoning/gt_answer_zh.json" + +_ONEIG_QD_JSON_STEMS: tuple[str, ...] = ( + "anime", + "human", + "object", + "anime_zh", + "human_zh", + "object_zh", + "multilingualism_zh", +) def _fetch_oneig_alignment() -> dict[str, dict]: - """Fetch alignment questions from per-category Q_D files (InferBench-style).""" + """Load OneIG question/dependency graphs from the official repo (HTTP, no on-disk cache). + + Fetches every ``scripts/alignment/Q_D/*.json`` file used by upstream ``alignment_score.py`` (EN + ZH), + including ``multilingualism_zh.json``. Keys in the returned map are ``{stem}_{prompt_id}`` matching + upstream file stems (e.g. ``object_012``, ``multilingualism_zh_000``). + + Returns + ------- + dict[str, dict] + ``prompt_id``-level ``questions`` and ``dependencies`` dicts (parsed from JSON strings when needed). + + Raises + ------ + requests.HTTPError + If any asset URL is missing or the response is not successful. + """ import json import requests questions_by_key: dict[str, dict] = {} - for qd_name in ("anime", "human", "object"): - url = f"{_ONEIG_ALIGNMENT_BASE}/{qd_name}.json" + for stem in _ONEIG_QD_JSON_STEMS: + url = f"{_ONEIG_ALIGNMENT_QD_URL}/{stem}.json" resp = requests.get(url, timeout=30) resp.raise_for_status() data = json.loads(resp.text) @@ -488,16 +668,55 @@ def _fetch_oneig_alignment() -> dict[str, dict]: q = json.loads(q) if isinstance(d, str): d = json.loads(d) - questions_by_key[f"{qd_name}_{prompt_id}"] = {"questions": q, "dependencies": d} + questions_by_key[f"{stem}_{prompt_id}"] = {"questions": q, "dependencies": d} return questions_by_key +def _fetch_oneig_reasoning_gt() -> tuple[dict[str, str], dict[str, str]]: + """Load official knowledge-reasoning reference answers (HTTP, no on-disk cache). + + Mirrors ``scripts/reasoning/gt_answer.json`` and ``gt_answer_zh.json`` from the same pinned commit as Q_D. + Keys are prompt ids (``str``), values are answer strings; downstream metrics may slice filenames to the + first three characters like ``reasoning_score.py``. + + Returns + ------- + tuple[dict[str, str], dict[str, str]] + ``(en_by_id, zh_by_id)``. + + Raises + ------ + requests.HTTPError + If any asset URL is missing or the response is not successful. + """ + import json + + import requests + + def _load(url: str) -> dict[str, str]: + resp = requests.get(url, timeout=60) + resp.raise_for_status() + raw = json.loads(resp.text) + return {str(k): str(v) for k, v in raw.items()} + + return _load(_ONEIG_REASONING_GT_URL_EN), _load(_ONEIG_REASONING_GT_URL_ZH) + + +def _oneig_needs_zh_multilingualism_hub(category: OneIGCategory | list[OneIGCategory] | None) -> bool: + """Whether ``OneIG-Bench-ZH`` must be loaded for ``Multilingualism`` rows.""" + if category is None: + return True + categories = [category] if not isinstance(category, list) else category + return "Multilingualism" in categories + + def setup_oneig_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, category: OneIGCategory | list[OneIGCategory] | None = None, + reasoning_language: str = "EN", ) -> Tuple[Dataset, Dataset, Dataset]: """ Setup the OneIG benchmark dataset. @@ -506,8 +725,8 @@ def setup_oneig_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -517,16 +736,43 @@ def setup_oneig_dataset( category : OneIGCategory | list[OneIGCategory] | None Filter by dataset category (Anime_Stylization, Portrait, etc.) or class (fauvism, watercolor, etc.). If None, returns all subsets. + reasoning_language : str, optional + Which reasoning GT to use for Knowledge_Reasoning rows: ``"EN"`` or ``"ZH"``. Default is ``"EN"``. Returns ------- Tuple[Dataset, Dataset, Dataset] - The OneIG dataset (dummy train, dummy val, test). + The OneIG dataset (dummy train, dummy val, test). Rows include ``questions`` and + ``dependencies`` from official Q_D JSON (EN + ZH stems, including ``multilingualism_zh``), + plus ``reasoning_gt_answer`` for ``Knowledge_Reasoning`` (language chosen by ``reasoning_language``). + Rows cover EN categories from ``OneIG-Bench`` plus ``Multilingualism`` from ``OneIG-Bench-ZH``. + Assets are downloaded over HTTP on each call (pinned commit ``_ONEIG_BENCHMARK_REF``); there is + no local disk cache. + + Notes + ----- + Non-multilingual prompts are loaded from the Hub config ``OneIG-Bench``; **Multilingualism** rows + are taken only from ``OneIG-Bench-ZH`` (they use ``prompt_cn``). The ZH config is fetched only when + the requested ``category`` is ``None`` (full suite) or explicitly includes ``Multilingualism``. + Q_D / reasoning JSON URLs are defined next to ``_fetch_oneig_alignment`` and + ``_fetch_oneig_reasoning_gt``. """ + _warn_ignored_benchmark_seed(seed, dataset="OneIG") questions_by_key = _fetch_oneig_alignment() - - ds_raw = load_dataset("OneIG-Bench/OneIG-Bench", "OneIG-Bench")["train"] # type: ignore[index] - records = [_to_oneig_record(dict(row), questions_by_key) for row in ds_raw] + reasoning_gt_en, reasoning_gt_zh = _fetch_oneig_reasoning_gt() + + ds_en = load_dataset("OneIG-Bench/OneIG-Bench", "OneIG-Bench")["train"] # type: ignore[index] + records = [ + _to_oneig_record(dict(row), questions_by_key, reasoning_gt_en, reasoning_gt_zh, reasoning_language) + for row in ds_en + ] + if _oneig_needs_zh_multilingualism_hub(category): + ds_zh = load_dataset("OneIG-Bench/OneIG-Bench", "OneIG-Bench-ZH")["train"] # type: ignore[index] + ds_zh_ml = ds_zh.filter(lambda r: r["category"] == "Multilingualism") + records.extend( + _to_oneig_record(dict(row), questions_by_key, reasoning_gt_en, reasoning_gt_zh, reasoning_language) + for row in ds_zh_ml + ) ds = Dataset.from_list(records) if category is not None: @@ -544,8 +790,252 @@ def setup_oneig_dataset( return ds.select([0]), ds.select([0]), ds +# functools.partial is not used for these wrappers: get_literal_values_from_param would unwrap +# partial objects back to setup_oneig_dataset and expose every OneIGCategory instead of one. + + +def setup_oneig_anime_stylization_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Anime_Stylization``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Anime_Stylization subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Anime_Stylization", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_general_object_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``General_Object``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the General_Object subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="General_Object", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_knowledge_reasoning_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Knowledge_Reasoning``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Knowledge_Reasoning subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Knowledge_Reasoning", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_multilingualism_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Multilingualism``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Multilingualism subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Multilingualism", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_portrait_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Portrait``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Portrait subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Portrait", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_text_rendering_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Text_Rendering``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Text_Rendering subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Text_Rendering", + reasoning_language=reasoning_language, + ) + + def setup_gedit_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -558,8 +1048,8 @@ def setup_gedit_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -576,6 +1066,7 @@ def setup_gedit_dataset( Tuple[Dataset, Dataset, Dataset] The GEditBench dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="GEditBench") task_type_map = { "subject_add": "subject-add", "subject_remove": "subject-remove", @@ -595,12 +1086,18 @@ def setup_gedit_dataset( for row in ds: task_type = row.get("task_type", "") category_name = task_type_to_category.get(task_type, task_type) - records.append( - { - "text": row.get("instruction", ""), - "category": category_name, - } - ) + record: dict = { + "text": row.get("instruction", ""), + "category": category_name, + } + src = row.get("input_image_raw") + if src is not None: + from io import BytesIO + + buf = BytesIO() + src.save(buf, format="JPEG") + record["source_image_bytes"] = buf.getvalue() + records.append(record) ds = Dataset.from_list(records) ds = stratify_dataset(ds, sample_size=test_sample_size, fraction=fraction) @@ -613,7 +1110,7 @@ def setup_gedit_dataset( def setup_dpg_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -626,8 +1123,8 @@ def setup_dpg_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -642,6 +1139,7 @@ def setup_dpg_dataset( Tuple[Dataset, Dataset, Dataset] The DPG dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="DPG") import csv import io from collections import defaultdict diff --git a/src/pruna/evaluation/benchmarks.py b/src/pruna/evaluation/benchmarks.py index 36e50c5e..009fe410 100644 --- a/src/pruna/evaluation/benchmarks.py +++ b/src/pruna/evaluation/benchmarks.py @@ -75,7 +75,7 @@ class BenchmarkRegistry: - WikiText (1609.07843 ?5): perplexity on validation/test. - GenEval (2310.11513 ?3.2): Mask2Former + CLIP color pipeline, binary score. - HPS (2306.09341): HPS v2 scoring model (CLIP fine-tuned on HPD v2). - - ImgEdit (2505.20275 ?4.2): GPT-4o 1ÿÿÿ5 ratings and ImgEdit-Judge. + - ImgEdit (2505.20275 ?4.2): GPT-4o 1-5 ratings and ImgEdit-Judge. - Long Text Bench (2507.22058 ?4): Text Accuracy (OCR, Qwen2.5-VL-7B). - GEditBench (2504.17761 ?4.2): VIEScore (SQ, PQ, O via GPT-4.1/Qwen2.5-VL). - OneIG (2506.07977 ?4.1): per-dimension metrics (semantic alignment, ED, etc.). @@ -88,9 +88,7 @@ class BenchmarkRegistry: def _register(cls, benchmark: Benchmark) -> None: missing = [m for m in benchmark.metrics if not MetricRegistry.has_metric(m)] if missing: - raise ValueError( - f"Benchmark '{benchmark.name}' references metrics not in MetricRegistry: {missing}." - ) + raise ValueError(f"Benchmark '{benchmark.name}' references metrics not in MetricRegistry: {missing}.") if benchmark.lookup_key not in base_datasets: available = ", ".join(base_datasets.keys()) raise ValueError(