From 111a2abf5e10ffaeb5a5cfaead6f8cbd2bc093d4 Mon Sep 17 00:00:00 2001 From: lee-groq Date: Wed, 17 Dec 2025 20:52:21 -0800 Subject: [PATCH] feat(mathvista): add image resizing flags --- src/openbench/datasets/mathvista.py | 199 +++++++++++++++++++--------- src/openbench/evals/mathvista.py | 10 ++ 2 files changed, 143 insertions(+), 66 deletions(-) diff --git a/src/openbench/datasets/mathvista.py b/src/openbench/datasets/mathvista.py index 0954640c..c07aaa2d 100644 --- a/src/openbench/datasets/mathvista.py +++ b/src/openbench/datasets/mathvista.py @@ -8,10 +8,12 @@ GitHub: https://github.com/lupantech/MathVista """ -from typing import Any, Dict, List, Optional, Union, cast +import io +from typing import Any, Callable, Dict, List, Optional, Union, cast from inspect_ai.dataset import Dataset, MemoryDataset, Sample from inspect_ai.model import ChatMessageUser, ContentImage, ContentText +from PIL import Image from openbench.utils.image import ( compress_image, @@ -20,76 +22,130 @@ ) -def record_to_sample(record: Dict[str, Any]) -> Sample: - """Convert a MathVista record to an Inspect Sample. +def record_to_sample( + max_dimension: Optional[int] = 1536, + quality: int = 75, + max_size_mb: float = 5.0, +) -> Callable[[Dict[str, Any]], Sample]: + """Creates a record-to-sample converter with specified image parameters. Args: - record: A record from the MathVista dataset + max_dimension: Maximum width/height in pixels for image resizing. + If None, disables dimension-based resizing. (default: 1536) + quality: JPEG quality (1-100) for image compression (default: 75) + max_size_mb: Maximum allowed size in MB before compression (default: 5.0) Returns: - An Inspect AI Sample with properly formatted input and metadata + A function that converts a MathVista record to an Inspect Sample """ - # Extract core fields - pid = str(record["pid"]) - question = record["question"] - answer = record["answer"] - question_type = record["question_type"] # "multi_choice" or "free_form" - answer_type = record["answer_type"] # "text", "integer", or "float" - query = record.get("query", "") # Pre-formatted query with hints - - # Use pre-formatted query if available (faithful to original) - prompt_text = query if query else question - - # Build input content with text first - input_content: List[Union[ContentText, ContentImage]] = [ - ContentText(text=prompt_text) - ] - - # Add the image if present - if "decoded_image" in record and record["decoded_image"] is not None: - image_data = record["decoded_image"] - - # Extract bytes from various image formats (HF dict, raw bytes, or PIL) - image_bytes = extract_image_bytes(image_data) - - # Compress and encode image to data URI - compressed_bytes = compress_image( - image_bytes, max_size_mb=5.0, quality=75, max_dimension=1536 + + def _convert(record: Dict[str, Any]) -> Sample: + """Convert a MathVista record to an Inspect Sample. + + Args: + record: A record from the MathVista dataset + + Returns: + An Inspect AI Sample with properly formatted input and metadata + """ + # Extract core fields + pid = str(record["pid"]) + question = record["question"] + answer = record["answer"] + question_type = record["question_type"] # "multi_choice" or "free_form" + answer_type = record["answer_type"] # "text", "integer", or "float" + query = record.get("query", "") # Pre-formatted query with hints + + # Use pre-formatted query if available (faithful to original) + prompt_text = query if query else question + + # Build input content with text first + input_content: List[Union[ContentText, ContentImage]] = [ + ContentText(text=prompt_text) + ] + + # Add the image if present + if "decoded_image" in record and record["decoded_image"] is not None: + image_data = record["decoded_image"] + + # Extract bytes from various image formats (HF dict, raw bytes, or PIL) + image_bytes = extract_image_bytes(image_data) + + # Always process images through quality/dimension pipeline + try: + with Image.open(io.BytesIO(image_bytes)) as img: + # Convert to RGB if necessary (for JPEG compatibility) + if img.mode in ("RGBA", "LA", "P"): + background = Image.new("RGB", img.size, (255, 255, 255)) + if img.mode == "P": + img = img.convert("RGBA") + background.paste( + img, + mask=img.split()[-1] + if img.mode in ("RGBA", "LA") + else None, + ) + img = background + elif img.mode != "RGB": + img = img.convert("RGB") + + # Resize if max_dimension is set and image exceeds it + if max_dimension is not None and max(img.size) > max_dimension: + img.thumbnail( + (max_dimension, max_dimension), Image.Resampling.LANCZOS + ) + + # Always re-encode at specified quality level + output = io.BytesIO() + img.save(output, format="JPEG", quality=quality, optimize=True) + image_bytes = output.getvalue() + except Exception: + # If processing fails, use original bytes + pass + + # Then apply additional size-based compression if needed + compressed_bytes = compress_image( + image_bytes, + max_size_mb=max_size_mb, + quality=quality, + max_dimension=100000, # Use very large value to skip dimension check + ) + data_uri = image_bytes_to_data_uri(compressed_bytes) + + # Add the image to input content + input_content.append(ContentImage(image=data_uri)) + + # Extract metadata + record_metadata = record.get("metadata", {}) + + # Build comprehensive metadata (faithful to original structure) + metadata = { + "pid": pid, + "question": question, + "question_type": question_type, + "answer_type": answer_type, + "category": record.get("category", ""), + "task": record.get("task", ""), + "context": record.get("context", ""), + "grade": record.get("grade", ""), + "skills": record.get("skills", []), + "unit": record.get("unit"), + "precision": record.get("precision"), + "choices": record.get("choices"), + } + + # Add any additional metadata from the record + if record_metadata: + metadata["original_metadata"] = record_metadata + + return Sample( + id=pid, + input=[ChatMessageUser(content=cast(Any, input_content))], + target=str(answer), + metadata=metadata, ) - data_uri = image_bytes_to_data_uri(compressed_bytes) - - # Add the image to input content - input_content.append(ContentImage(image=data_uri)) - - # Extract metadata - record_metadata = record.get("metadata", {}) - - # Build comprehensive metadata (faithful to original structure) - metadata = { - "pid": pid, - "question": question, - "question_type": question_type, - "answer_type": answer_type, - "category": record.get("category", ""), - "task": record.get("task", ""), - "context": record.get("context", ""), - "grade": record.get("grade", ""), - "skills": record.get("skills", []), - "unit": record.get("unit"), - "precision": record.get("precision"), - "choices": record.get("choices"), - } - - # Add any additional metadata from the record - if record_metadata: - metadata["original_metadata"] = record_metadata - - return Sample( - id=pid, - input=[ChatMessageUser(content=cast(Any, input_content))], - target=str(answer), - metadata=metadata, - ) + + return _convert def get_dataset( @@ -97,6 +153,9 @@ def get_dataset( question_type: Optional[str] = None, shuffle: bool = True, seed: int = 42, + max_dimension: Optional[int] = 1536, + quality: int = 75, + max_size_mb: float = 5.0, ) -> Dataset: """Load the MathVista dataset from HuggingFace. @@ -105,6 +164,10 @@ def get_dataset( question_type: Optional filter by question type ("multi_choice" or "free_form") shuffle: Whether to shuffle the dataset seed: Random seed for shuffling + max_dimension: Maximum width/height in pixels for image resizing. + If None, disables dimension-based resizing. (default: 1536) + quality: JPEG quality (1-100) for image compression (default: 75) + max_size_mb: Maximum allowed size in MB before compression (default: 5.0) Returns: An Inspect AI Dataset @@ -115,7 +178,11 @@ def get_dataset( dataset = hf_dataset( path="AI4Math/MathVista", split=split, - sample_fields=record_to_sample, + sample_fields=record_to_sample( + max_dimension=max_dimension, + quality=quality, + max_size_mb=max_size_mb, + ), shuffle=shuffle, seed=seed, ) diff --git a/src/openbench/evals/mathvista.py b/src/openbench/evals/mathvista.py index 4fca2720..cf8855cc 100644 --- a/src/openbench/evals/mathvista.py +++ b/src/openbench/evals/mathvista.py @@ -26,6 +26,9 @@ def mathvista( shuffle: bool = True, seed: int = 42, grader_model: str = "openai/gpt-4-turbo", + max_dimension: Optional[int] = 1536, + quality: int = 75, + max_size_mb: float = 5.0, ) -> Task: """MathVista: Mathematical Reasoning in Visual Contexts. @@ -38,6 +41,10 @@ def mathvista( seed: Random seed for shuffling grader_model: Model to use for LLM-based answer extraction fallback (default: gpt-4-turbo, matching original paper) + max_dimension: Maximum width/height in pixels for image resizing. + If None, disables dimension-based resizing. (default: 1536) + quality: JPEG quality (1-100) for image compression (default: 75) + max_size_mb: Maximum allowed size in MB before compression (default: 5.0) Returns: Task for evaluation @@ -47,6 +54,9 @@ def mathvista( question_type=question_type, shuffle=shuffle, seed=seed, + max_dimension=max_dimension, + quality=quality, + max_size_mb=max_size_mb, ) return Task(