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( diff --git a/src/pruna/evaluation/metrics/__init__.py b/src/pruna/evaluation/metrics/__init__.py index d17ea6dd..4a08d622 100644 --- a/src/pruna/evaluation/metrics/__init__.py +++ b/src/pruna/evaluation/metrics/__init__.py @@ -23,6 +23,7 @@ from pruna.evaluation.metrics.metric_memory import DiskMemoryMetric, InferenceMemoryMetric, TrainingMemoryMetric from pruna.evaluation.metrics.metric_model_architecture import TotalMACsMetric, TotalParamsMetric from pruna.evaluation.metrics.metric_oneig_alignment import OneIGAlignmentMetric +from pruna.evaluation.metrics.metric_oneig_reasoning import OneIGReasoningMetric from pruna.evaluation.metrics.metric_pairwise_clip import PairwiseClipScore from pruna.evaluation.metrics.metric_qa_accuracy import QAAccuracyMetric from pruna.evaluation.metrics.metric_rapiddata import RapidataMetric as RapidataMetric @@ -57,6 +58,7 @@ "AestheticLAION", "LMEvalMetric", "OneIGAlignmentMetric", + "OneIGReasoningMetric", "OneIGTextScoreMetric", "QAAccuracyMetric", "RapidataMetric", diff --git a/src/pruna/evaluation/metrics/metric_oneig_reasoning.py b/src/pruna/evaluation/metrics/metric_oneig_reasoning.py new file mode 100644 index 00000000..23889f0c --- /dev/null +++ b/src/pruna/evaluation/metrics/metric_oneig_reasoning.py @@ -0,0 +1,420 @@ +# 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. + +"""OneIG reasoning score via LLM2CLIP text-image similarity. + +Llama-derived checkpoints may require ``HF_TOKEN`` and ``huggingface-cli login``. + +Hugging Face download tuning (optional): + +- ``PRUNA_ONEIG_HF_VERBOSE=1`` or ``HF_DEBUG=1`` — hub **debug** logging and tqdm + progress bars (helps when stderr is piped; pair with ``python -u`` or + ``PYTHONUNBUFFERED=1`` for line-buffered output). +- ``PRUNA_ONEIG_HF_FAST_DOWNLOAD=1`` — enable **hf_transfer** multi-part downloads + (requires ``pruna[evaluation]``, which lists ``hf_transfer``). Alternatively, set + ``HF_HUB_ENABLE_HF_TRANSFER=1`` **before** starting Python so the hub picks it up at + import time. + +``transformers`` is pinned to ``<5`` in ``pyproject.toml``. The LLM2CLIP loading path +(``CLIPImageProcessor``, ``AutoModel``, ``LlamaEncoderModel``) is exercised on **4.x** +releases in CI and manual smoke runs. ``transformers`` 5.x has had reports of +``from_pretrained`` not fully initializing some non-persistent buffers (for example +``position_ids``) for certain architectures; the pin avoids that class of failures +until those issues are clearly resolved upstream. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from typing import Any, Iterator + +import torch +from PIL import Image + +from pruna.evaluation.metrics.metric_stateful import StatefulMetric +from pruna.evaluation.metrics.registry import MetricRegistry +from pruna.evaluation.metrics.result import MetricResult +from pruna.evaluation.metrics.utils import ( + SINGLE, + get_call_type_for_single_metric, + metric_data_processor, +) +from pruna.evaluation.metrics.vlm_utils import ( + _process_images, + _tensor_to_pil, + resolve_oneig_reasoning_device, + split_mxn_grid, +) +from pruna.logging.logger import pruna_logger + + +def _is_cuda_device(device: str) -> bool: + return device == "cuda" or device.startswith("cuda:") + + +@contextmanager +def _oneig_hf_download_env() -> Iterator[None]: + """Apply optional OneIG HF hub env tweaks, then restore prior values.""" + keys = ("HF_HUB_ENABLE_HF_TRANSFER",) + saved = {k: os.environ.get(k) for k in keys} + try: + _prepare_huggingface_hub_for_oneig_downloads() + yield + finally: + for key, val in saved.items(): + if val is None: + os.environ.pop(key, None) + else: + os.environ[key] = val + + +def _env_truthy(raw: str | None) -> bool: + if raw is None: + return False + return raw.strip().upper() in {"1", "ON", "YES", "TRUE"} + + +def _prepare_huggingface_hub_for_oneig_downloads() -> None: + """ + Apply Hugging Face Hub verbosity and optional fast downloads before checkpoints load. + + ``HF_HUB_ENABLE_HF_TRANSFER`` is read when ``huggingface_hub`` loads; if it was + false, we flip the in-module flag after importing ``hf_transfer`` when + ``PRUNA_ONEIG_HF_FAST_DOWNLOAD=1``. + """ + if _env_truthy(os.environ.get("PRUNA_ONEIG_HF_VERBOSE")) or _env_truthy(os.environ.get("HF_DEBUG")): + from huggingface_hub.utils import enable_progress_bars + from huggingface_hub.utils.logging import set_verbosity_debug + + set_verbosity_debug() + enable_progress_bars() + + if not _env_truthy(os.environ.get("PRUNA_ONEIG_HF_FAST_DOWNLOAD")): + return + + import hf_transfer # noqa: F401 # type: ignore[import-not-found] + import huggingface_hub.constants as hf_constants + + hf_constants.HF_HUB_ENABLE_HF_TRANSFER = True + pruna_logger.info("oneig_reasoning: enabled hf_transfer downloads (PRUNA_ONEIG_HF_FAST_DOWNLOAD=1).") + + +def _to_pil_list(images: list) -> list: + """Convert images to list of PIL.Image (RGB).""" + import numpy as np + from PIL import Image + + out: list = [] + for img in images: + if isinstance(img, Image.Image): + out.append(img.convert("RGB")) + elif isinstance(img, torch.Tensor): + if img.ndim == 4: + img = img[0] + if img.max() > 1: + img = img / 255.0 + np_img = (img.cpu().numpy() * 255).astype("uint8") + if np_img.shape[0] == 3: + np_img = np_img.transpose(1, 2, 0) + out.append(Image.fromarray(np_img)) + elif hasattr(img, "__array__"): + out.append(Image.fromarray(np.asarray(img)).convert("RGB")) + else: + out.append(img) + return out + + +class _LLM2CLIPScorer: + """ + Thin wrapper around LLM2CLIP text-image similarity. + + Accepts PIL images and a single answer string; returns per-image scores. + Best-effort alignment with OneIG-Benchmark scripts (CUDA + bfloat16). + """ + + def __init__( + self, + processor_model: str = "openai/clip-vit-large-patch14-336", + model_name: str = "microsoft/LLM2CLIP-Openai-L-14-336", + llm_model_name: str = "microsoft/LLM2CLIP-Llama-3-8B-Instruct-CC-Finetuned", + device: str = "cuda", + attn_implementation: str | None = None, + ) -> None: + self.processor_model = processor_model + self.model_name = model_name + self.llm_model_name = llm_model_name + self.device = device + self.attn_implementation = attn_implementation + self._processor = None + self._clip_model = None + self._l2v = None + + def _load_models(self) -> None: + if self._clip_model is not None: + return + with _oneig_hf_download_env(): + self._load_models_inner() + + def _load_models_inner(self) -> None: + from transformers import AutoConfig, AutoModel, AutoTokenizer, CLIPImageProcessor + + from pruna.evaluation.metrics.vendor.oneig_llm2vec import LLM2Vec + from pruna.evaluation.metrics.vendor.oneig_llm2vec.modeling_llama_encoder import LlamaEncoderModel + + pruna_logger.info( + "oneig_reasoning: downloading or loading LLM2CLIP checkpoints " + "(%s, %s). First run can take many minutes and several gigabytes; " + "Hugging Face download progress may look idle when logs are piped.", + self.model_name, + self.llm_model_name, + ) + dtype = torch.bfloat16 if _is_cuda_device(str(self.device)) else torch.float32 + self._processor = CLIPImageProcessor.from_pretrained(self.processor_model) + self._clip_model = AutoModel.from_pretrained( + self.model_name, + dtype=dtype, + trust_remote_code=True, + ).to(self.device) + self._clip_model.train(mode=False) + + config = AutoConfig.from_pretrained(self.llm_model_name, trust_remote_code=True) + if self.attn_implementation is not None: + # User override (e.g. sdpa, flash_attention_2, eager). Upstream OneIG leaves HF default. + config.attn_implementation = self.attn_implementation + if hasattr(config, "_attn_implementation"): + config._attn_implementation = self.attn_implementation + elif _is_cuda_device(str(self.device)): + config.attn_implementation = "sdpa" + if hasattr(config, "_attn_implementation"): + config._attn_implementation = "sdpa" + llm_model = LlamaEncoderModel.from_pretrained( + self.llm_model_name, + dtype=dtype, + config=config, + trust_remote_code=True, + ) + llm_model.config._name_or_path = "meta-llama/Meta-Llama-3-8B-Instruct" + tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name) + self._l2v = LLM2Vec(llm_model, tokenizer, pooling_mode="mean", max_length=512, doc_max_length=512) + + def score(self, images: list, text_prompt: str) -> list[float] | None: + """ + Compute similarity scores between images and text. + + Parameters + ---------- + images : list + List of PIL.Image.Image. + text_prompt : str + Reference text (e.g. ground-truth answer). + + Returns + ------- + list[float] | None + Per-image scores, or None on failure. + """ + self._load_models() + pil_images = _to_pil_list(images) + if not pil_images: + return None + input_pixels = self._processor(images=pil_images, return_tensors="pt").pixel_values.to(self.device) + captions = [text_prompt] + with torch.no_grad(): + text_features = self._l2v.encode(captions, convert_to_tensor=True, device=self.device).to(self.device) + text_features = self._clip_model.get_text_features(text_features) + if _is_cuda_device(str(self.device)): + with torch.amp.autocast(device_type="cuda"): + image_features = self._clip_model.get_image_features(input_pixels) + else: + image_features = self._clip_model.get_image_features(input_pixels.float()) + + image_features = image_features.float() + text_features = text_features.float() + eps = 1e-8 + image_features /= image_features.norm(dim=-1, keepdim=True) + eps + text_features /= text_features.norm(dim=-1, keepdim=True) + eps + + text_probs = (image_features @ text_features.T).cpu().tolist() + return [p[0] for p in text_probs] + + +@MetricRegistry.register("oneig_reasoning") +class OneIGReasoningMetric(StatefulMetric): + """ + OneIG reasoning score: LLM2CLIP similarity between GT answer text and generated image. + + Uses ``reasoning_gt_answer`` from aux (populated by OneIG Knowledge_Reasoning loader; + language is chosen at dataset load via ``reasoning_language``). Splits a ``2 x 2`` grid + by default (OneIG ``split_mxn_grid``) and averages cell scores. Llama-derived checkpoints may require + ``HF_TOKEN`` and ``huggingface-cli login``. + + Parameters + ---------- + processor_model : str, optional + CLIP processor model ID. + model_name : str, optional + LLM2CLIP model ID. + llm_model_name : str, optional + LLM2Vec model ID. + grid_size : tuple[int, int], optional + ``(columns, rows)`` for :func:`~pruna.evaluation.metrics.vlm_utils.split_mxn_grid``. + attn_implementation : str | None, optional + Transformers attention backend for the LLM encoder (``None`` uses HF default; + Pruna defaults to ``sdpa`` on CUDA when unset). + device : str | torch.device | None, optional + Device for inference. + scorer : _LLM2CLIPScorer | None, optional + Optional scorer instance for testing (injected mock). + call_type : str, optional + Call type for the metric. + **kwargs : Any + Additional keyword arguments for :class:`StatefulMetric`. + + Notes + ----- + Prompt benchmarks yield ``(prompts, aux_list)``. With default ``call_type`` + ``y_gt``, ``aux_list`` is the list (or tensor coerced to a list) of per-sample + dicts parallel to generated images. Each dict must include a non-empty + ``reasoning_gt_answer`` for Knowledge/Reasoning samples. Missing GT, scorer + failures, or :meth:`compute` with no scored samples raise ``ValueError`` or + ``RuntimeError`` instead of returning a placeholder score. + """ + + metric_name: str = "oneig_reasoning" + default_call_type: str = "y_gt" + higher_is_better: bool = True + runs_on: list[str] = ["cuda", "cpu"] + + def __init__( + self, + processor_model: str = "openai/clip-vit-large-patch14-336", + model_name: str = "microsoft/LLM2CLIP-Openai-L-14-336", + llm_model_name: str = "microsoft/LLM2CLIP-Llama-3-8B-Instruct-CC-Finetuned", + grid_size: tuple[int, int] = (2, 2), + attn_implementation: str | None = None, + device: str | torch.device | None = None, + scorer: _LLM2CLIPScorer | None = None, + call_type: str | None = None, + **kwargs: Any, + ) -> None: + resolved = resolve_oneig_reasoning_device(device) + super().__init__(device=resolved, **kwargs) + self.device = resolved + self.call_type = get_call_type_for_single_metric( + call_type if call_type is not None else SINGLE, self.default_call_type + ) + self.processor_model = processor_model + self.model_name = model_name + self.llm_model_name = llm_model_name + self.grid_size = (int(grid_size[0]), int(grid_size[1])) + self.attn_implementation = attn_implementation + self._injected_scorer = scorer + self._llm2clip_scorer: _LLM2CLIPScorer | None = None + self.add_state("scores", default=[]) + + def _get_scorer(self) -> _LLM2CLIPScorer: + if self._injected_scorer is not None: + return self._injected_scorer + if self._llm2clip_scorer is None: + self._llm2clip_scorer = _LLM2CLIPScorer( + processor_model=self.processor_model, + model_name=self.model_name, + llm_model_name=self.llm_model_name, + device=str(self.device), + attn_implementation=self.attn_implementation, + ) + return self._llm2clip_scorer + + def reset(self) -> None: + """Clear scores and release cached LLM2CLIP weights.""" + super().reset() + self._llm2clip_scorer = None + + def _get_gt_text(self, aux: dict) -> str: + val = aux.get("reasoning_gt_answer") + if val is None or (isinstance(val, str) and not val.strip()): + raise ValueError( + "oneig_reasoning requires 'reasoning_gt_answer' in aux for Knowledge_Reasoning rows. " + f"Got keys: {list(aux.keys())}." + ) + return str(val).strip() + + def update(self, x: list[Any] | torch.Tensor, gt: torch.Tensor, outputs: torch.Tensor) -> None: + """ + Score each image against its GT answer text via LLM2CLIP similarity. + + Parameters + ---------- + x : list[Any] | torch.Tensor + Unused batch metadata. + gt : torch.Tensor + Ground-truth slot with per-sample aux dicts containing ``reasoning_gt_answer``. + outputs : torch.Tensor + Model outputs (generated images). + + Raises + ------ + ValueError + If a per-sample aux entry is not a dict or lacks a non-empty + ``reasoning_gt_answer``. + RuntimeError + If the LLM2CLIP scorer returns no scores for a sample. + """ + inputs = metric_data_processor(x, gt, outputs, self.call_type) + images = _process_images(inputs[0]) + aux_slot = inputs[1] if len(inputs) > 1 else [] + if isinstance(aux_slot, torch.Tensor): + raise ValueError("oneig_reasoning expects gt as list[dict] with 'reasoning_gt_answer'.") + + scorer = self._get_scorer() + + for i, image in enumerate(images): + aux_row = aux_slot[i] if isinstance(aux_slot, (list, tuple)) and i < len(aux_slot) else {} + if not isinstance(aux_row, dict): + raise ValueError(f"oneig_reasoning requires aux[{i}] to be a dict. Got: {type(aux_row)}.") + text = self._get_gt_text(aux_row) + if isinstance(image, Image.Image): + pil = image.convert("RGB") + elif isinstance(image, torch.Tensor): + pil = _tensor_to_pil(image) + else: + pil = Image.fromarray(image).convert("RGB") + cells = split_mxn_grid(pil, self.grid_size) + result = scorer.score(cells, text) + if result is None or len(result) == 0: + raise RuntimeError(f"oneig_reasoning: LLM2CLIP scorer returned no scores for sample {i}.") + self.scores.append(float(sum(result) / len(result))) + + def compute(self) -> MetricResult: + """ + Compute the mean reasoning score across all samples. + + Returns + ------- + MetricResult + Mean LLM2CLIP similarity. + + Raises + ------ + RuntimeError + If :meth:`update` was not called or scored no samples. + """ + if not self.scores: + raise RuntimeError( + "oneig_reasoning: no samples were scored; call update() with valid " + "batches and non-empty reasoning_gt_answer before compute()." + ) + mean_score = sum(self.scores) / len(self.scores) + return MetricResult(self.metric_name, self.__dict__, float(mean_score)) diff --git a/src/pruna/evaluation/metrics/registry.py b/src/pruna/evaluation/metrics/registry.py index 5efd721a..650f8a76 100644 --- a/src/pruna/evaluation/metrics/registry.py +++ b/src/pruna/evaluation/metrics/registry.py @@ -14,6 +14,7 @@ from __future__ import annotations +import importlib from functools import partial from inspect import isclass from typing import Any, Callable, Dict, Iterable, List @@ -29,9 +30,17 @@ class MetricRegistry: Registry for metrics. The registry is a dictionary that maps metric names to metric classes. + + Notes + ----- + ``_lazy_metrics`` lists names that :meth:`has_metric` treats as registered before the + implementing module is loaded. The ``oneig_reasoning`` metric imports the LLM2CLIP-related + stack (vendored helpers, heavy optional dependencies); it is imported only when + :meth:`get_metric` is called with that name so other code paths avoid that cost. """ _registry: Dict[str, Callable[..., Any]] = {} + _lazy_metrics: frozenset[str] = frozenset({"oneig_reasoning"}) @classmethod def register(cls, name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]: @@ -104,7 +113,7 @@ def has_metric(cls, name: str) -> bool: bool True if the metric is registered, False otherwise. """ - return name in cls._registry + return name in cls._registry or name in cls._lazy_metrics @classmethod def get_metric(cls, name: str, **kwargs) -> BaseMetric | StatefulMetric: @@ -122,6 +131,9 @@ def get_metric(cls, name: str, **kwargs) -> BaseMetric | StatefulMetric: ------- The metric instance. """ + if name in cls._lazy_metrics and name not in cls._registry: + importlib.import_module("pruna.evaluation.metrics.metric_oneig_reasoning") + if name not in cls._registry: raise ValueError(f"Metric '{name}' is not registered.")