wehub-resource-sync
1275 lines
42 KiB
Python
1275 lines
42 KiB
Python
"""Data loading and preprocessing for SimpleQA evaluation.
|
||
|
||
This module handles all data preparation:
|
||
- Loading the SimpleQA dataset
|
||
- Extracting URLs from metadata
|
||
- Capturing screenshots from URLs
|
||
- Fetching text content from URLs
|
||
"""
|
||
|
||
import ast
|
||
import asyncio
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import urllib.parse
|
||
|
||
import pandas as pd
|
||
import trafilatura
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ============================================================================
|
||
# Data Loading
|
||
# ============================================================================
|
||
|
||
|
||
def load_simpleqa_data(num_examples: int | None = None) -> list[dict]:
|
||
"""Load SimpleQA dataset.
|
||
|
||
Args:
|
||
num_examples: Optional limit on number of examples to load.
|
||
|
||
Returns:
|
||
List of example dictionaries with 'id', 'problem', 'answer', etc.
|
||
"""
|
||
logger.info("Loading SimpleQA dataset...")
|
||
try:
|
||
local_path = "evaluation/simple_qa_eval/data/simple_qa_test_set.csv"
|
||
if os.path.exists(local_path):
|
||
df = pd.read_csv(local_path)
|
||
else:
|
||
url = "https://openaipublic.blob.core.windows.net/simple-evals/simple_qa_test_set.csv"
|
||
df = pd.read_csv(url)
|
||
except Exception as e:
|
||
logger.error(f"Failed to load dataset: {e}")
|
||
url = "https://openaipublic.blob.core.windows.net/simple-evals/simple_qa_test_set.csv"
|
||
df = pd.read_csv(url)
|
||
|
||
# Ensure stable ordering: reset index to maintain original CSV row order
|
||
df = df.reset_index(drop=True)
|
||
|
||
# Generate unique ID from problem text
|
||
df["id"] = df["problem"].apply(
|
||
lambda problem: hashlib.md5(problem.encode()).hexdigest()
|
||
)
|
||
|
||
# Convert to list of dicts, maintaining original CSV order
|
||
data = [row.to_dict() for _, row in df.iterrows()]
|
||
|
||
if num_examples:
|
||
logger.info(f"Limiting to first {num_examples} examples.")
|
||
data = data[:num_examples]
|
||
|
||
logger.info(f"Loaded {len(data)} examples.")
|
||
return data
|
||
|
||
|
||
def load_simpleqa_verified_data(num_examples: int | None = None) -> list[dict]:
|
||
"""Load SimpleQA Verified dataset from Hugging Face.
|
||
|
||
Args:
|
||
num_examples: Optional limit on number of examples to load.
|
||
|
||
Returns:
|
||
List of example dictionaries with 'id', 'problem', 'answer', etc.
|
||
Compatible format with SimpleQA dataset.
|
||
"""
|
||
logger.info("Loading SimpleQA Verified dataset...")
|
||
try:
|
||
# Try using datasets library first (recommended)
|
||
try:
|
||
from datasets import load_dataset
|
||
|
||
logger.info("Using Hugging Face datasets library...")
|
||
dataset = load_dataset("google/simpleqa-verified", split="eval")
|
||
df = dataset.to_pandas()
|
||
except ImportError:
|
||
logger.warning("datasets library not available, trying alternative methods")
|
||
# Fallback: try Hugging Face datasets-server API
|
||
try:
|
||
import requests
|
||
|
||
logger.info("Trying Hugging Face datasets-server API...")
|
||
api_url = "https://datasets-server.huggingface.co/parquet?dataset=google%2Fsimpleqa-verified&config=simpleqa_verified&split=eval"
|
||
response = requests.get(api_url, timeout=60)
|
||
if response.status_code == 200:
|
||
import io
|
||
|
||
df = pd.read_parquet(io.BytesIO(response.content))
|
||
logger.info("Successfully loaded via datasets-server API")
|
||
else:
|
||
raise Exception(
|
||
f"Failed to download dataset: HTTP {response.status_code}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Failed to load via API: {e}")
|
||
# Last resort: try direct file download
|
||
try:
|
||
logger.info("Trying direct file download...")
|
||
# Try parquet file
|
||
parquet_url = "https://huggingface.co/datasets/google/simpleqa-verified/resolve/main/data/eval-00000-of-00001.parquet"
|
||
df = pd.read_parquet(parquet_url)
|
||
logger.info("Successfully loaded via direct file download")
|
||
except Exception as e2:
|
||
logger.error(f"Failed to load via direct download: {e2}")
|
||
raise Exception(
|
||
"All methods failed. Please install 'datasets' library: pip install datasets"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Failed to load SimpleQA Verified dataset: {e}")
|
||
raise
|
||
|
||
# Ensure stable ordering: reset index to maintain original order
|
||
df = df.reset_index(drop=True)
|
||
|
||
# Convert to compatible format with SimpleQA
|
||
# SimpleQA Verified has: original_index, problem, answer, topic, answer_type, multi_step, requires_reasoning, urls
|
||
# SimpleQA has: metadata (with urls), problem, answer, id
|
||
|
||
# Generate unique ID from problem text (same as SimpleQA)
|
||
df["id"] = df["problem"].apply(
|
||
lambda problem: hashlib.md5(problem.encode()).hexdigest()
|
||
)
|
||
|
||
# Convert urls to list format if it's a string
|
||
def normalize_urls(urls):
|
||
"""Normalize URLs to list format."""
|
||
if isinstance(urls, str):
|
||
# Try to parse as list string
|
||
try:
|
||
import ast
|
||
|
||
return ast.literal_eval(urls)
|
||
except Exception:
|
||
# Split by comma if it's a comma-separated string
|
||
return [u.strip() for u in urls.split(",") if u.strip()]
|
||
elif isinstance(urls, list):
|
||
return urls
|
||
else:
|
||
return []
|
||
|
||
# Normalize URLs column
|
||
if "urls" in df.columns:
|
||
df["urls"] = df["urls"].apply(normalize_urls)
|
||
else:
|
||
df["urls"] = [[]] * len(df)
|
||
|
||
# Convert to metadata format compatible with SimpleQA
|
||
def create_metadata(row):
|
||
"""Create metadata dict compatible with SimpleQA format."""
|
||
metadata = {
|
||
"topic": str(row.get("topic", "")),
|
||
"answer_type": str(row.get("answer_type", "")),
|
||
"urls": row.get("urls", []),
|
||
}
|
||
if "multi_step" in row and pd.notna(row["multi_step"]):
|
||
metadata["multi_step"] = bool(row["multi_step"])
|
||
if "requires_reasoning" in row and pd.notna(row["requires_reasoning"]):
|
||
metadata["requires_reasoning"] = bool(row["requires_reasoning"])
|
||
if "original_index" in row and pd.notna(row["original_index"]):
|
||
metadata["original_index"] = int(row["original_index"])
|
||
# Convert to string format similar to SimpleQA (using single quotes for Python dict string)
|
||
return str(metadata)
|
||
|
||
df["metadata"] = df.apply(create_metadata, axis=1)
|
||
|
||
# Convert to list of dicts, maintaining original order
|
||
data = [row.to_dict() for _, row in df.iterrows()]
|
||
|
||
if num_examples:
|
||
logger.info(f"Limiting to first {num_examples} examples.")
|
||
data = data[:num_examples]
|
||
|
||
logger.info(f"Loaded {len(data)} SimpleQA Verified examples.")
|
||
return data
|
||
|
||
|
||
def load_text_cache(cache_path: str) -> dict:
|
||
"""Load pre-fetched text from JSONL file.
|
||
|
||
Args:
|
||
cache_path: Path to JSONL file with cached text.
|
||
|
||
Returns:
|
||
Dict mapping example ID to cached item.
|
||
"""
|
||
logger.info(f"Loading text cache from {cache_path}...")
|
||
cache = {}
|
||
with open(cache_path, "r") as f:
|
||
for line in f:
|
||
item = json.loads(line)
|
||
cache[item["id"]] = item
|
||
logger.info(f"Loaded {len(cache)} cached items.")
|
||
return cache
|
||
|
||
|
||
# ============================================================================
|
||
# URL Extraction
|
||
# ============================================================================
|
||
|
||
|
||
def extract_url_from_metadata(example: dict) -> str | None:
|
||
"""Extract URL from example metadata.
|
||
|
||
Args:
|
||
example: Example dict with 'metadata' field.
|
||
|
||
Returns:
|
||
Extracted URL or None.
|
||
"""
|
||
meta = example.get("metadata")
|
||
if isinstance(meta, str):
|
||
try:
|
||
meta = json.loads(meta)
|
||
except json.JSONDecodeError:
|
||
try:
|
||
meta = ast.literal_eval(meta)
|
||
except (ValueError, SyntaxError):
|
||
pass
|
||
|
||
target_url = None
|
||
if isinstance(meta, dict):
|
||
if "url" in meta:
|
||
target_url = meta["url"]
|
||
elif (
|
||
"urls" in meta and isinstance(meta["urls"], list) and len(meta["urls"]) > 0
|
||
):
|
||
# Flatten URLs: some entries have multiple URLs concatenated in a single string
|
||
# (separated by newlines OR directly joined like "https://a.comhttps://b.com")
|
||
all_urls = []
|
||
for url_entry in meta["urls"]:
|
||
if isinstance(url_entry, str):
|
||
# Split on "https://" boundaries to handle concatenated URLs
|
||
parts = re.split(r"(?=https?://)", url_entry)
|
||
for part in parts:
|
||
part = part.strip().rstrip(",'\"").strip("- ").strip()
|
||
if part and re.match(r"https?://", part):
|
||
all_urls.append(part)
|
||
|
||
# Prefer en.wikipedia.org article URLs (exclude non-English and Category pages)
|
||
wikipedia_urls = [
|
||
u
|
||
for u in all_urls
|
||
if "en.wikipedia.org/wiki/" in u
|
||
and "/Category:" not in u
|
||
and "wikipedia-on-ipfs" not in u.lower()
|
||
]
|
||
if wikipedia_urls:
|
||
target_url = wikipedia_urls[0]
|
||
else:
|
||
# Secondary: wikimedia.org URLs (e.g., commons.wikimedia.org)
|
||
wikimedia_urls = [u for u in all_urls if "wikimedia.org" in u.lower()]
|
||
target_url = (
|
||
wikimedia_urls[0]
|
||
if wikimedia_urls
|
||
else (all_urls[0] if all_urls else None)
|
||
)
|
||
|
||
# Extract first valid URL from the string
|
||
if target_url:
|
||
url_match = re.search(r"https?://[^\s<>\"{}|\\^`\[\]]+", target_url)
|
||
target_url = url_match.group(0) if url_match else None
|
||
|
||
# Note by Yichuan: strip URL fragment (#section) so that URLs differing
|
||
# only by anchor are treated as the same page for deduplication and
|
||
# retrieval-accuracy matching.
|
||
if target_url and "#" in target_url:
|
||
target_url = target_url.split("#")[0]
|
||
|
||
return target_url
|
||
|
||
|
||
# ============================================================================
|
||
# Screenshot Capture
|
||
# ============================================================================
|
||
|
||
# Lazy import screenshot utilities
|
||
_capture_screenshot = None
|
||
_encode_image = None
|
||
_encode_image_for_vlm = None
|
||
|
||
|
||
def _init_screenshot_utils():
|
||
"""Initialize screenshot utilities (lazy import)."""
|
||
global _capture_screenshot, _encode_image, _encode_image_for_vlm
|
||
if _capture_screenshot is not None:
|
||
return True
|
||
|
||
try:
|
||
from .screenshot import capture_screenshot, encode_image, encode_image_for_vlm
|
||
|
||
_capture_screenshot = capture_screenshot
|
||
_encode_image = encode_image
|
||
_encode_image_for_vlm = encode_image_for_vlm
|
||
return True
|
||
except ImportError:
|
||
logger.warning("Screenshot utilities not available")
|
||
return False
|
||
|
||
|
||
def capture_screenshot_for_example(
|
||
example: dict, screenshot_dir: str = "screenshots"
|
||
) -> str | None:
|
||
"""Capture screenshot for a single example.
|
||
|
||
Args:
|
||
example: Example dict with metadata containing URL.
|
||
screenshot_dir: Directory to save screenshots.
|
||
|
||
Returns:
|
||
Path to screenshot file, or None if failed.
|
||
"""
|
||
if not _init_screenshot_utils():
|
||
return None
|
||
|
||
target_url = extract_url_from_metadata(example)
|
||
if not target_url:
|
||
return None
|
||
|
||
os.makedirs(screenshot_dir, exist_ok=True)
|
||
screenshot_filename = f"{example['id']}_fullhd.png"
|
||
screenshot_path = os.path.join(screenshot_dir, screenshot_filename)
|
||
|
||
# Check if valid screenshot already exists
|
||
if os.path.exists(screenshot_path) and os.path.getsize(screenshot_path) > 0:
|
||
logger.debug(f"Screenshot exists: {screenshot_path}")
|
||
return screenshot_path
|
||
|
||
# Capture screenshot
|
||
try:
|
||
if _capture_screenshot is None:
|
||
return None
|
||
success = _capture_screenshot(target_url, screenshot_path, True)
|
||
if (
|
||
success
|
||
and os.path.exists(screenshot_path)
|
||
and os.path.getsize(screenshot_path) > 0
|
||
):
|
||
file_size = os.path.getsize(screenshot_path) // 1024
|
||
logger.info(f"Screenshot saved: {screenshot_path} ({file_size}KB)")
|
||
return screenshot_path
|
||
else:
|
||
logger.warning(
|
||
f"Screenshot failed (no output): {target_url} -> {screenshot_path}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Screenshot error for {target_url}: {e}")
|
||
|
||
return None
|
||
|
||
|
||
async def capture_screenshot_async(
|
||
example: dict, screenshot_dir: str = "screenshots"
|
||
) -> str | None:
|
||
"""Async wrapper for screenshot capture."""
|
||
loop = asyncio.get_event_loop()
|
||
return await loop.run_in_executor(
|
||
None, capture_screenshot_for_example, example, screenshot_dir
|
||
)
|
||
|
||
|
||
def encode_screenshot(screenshot_path: str) -> str | None:
|
||
"""Encode screenshot to base64.
|
||
|
||
Args:
|
||
screenshot_path: Path to screenshot file, or already-encoded base64 string.
|
||
|
||
Returns:
|
||
Base64 encoded string, or None if failed.
|
||
"""
|
||
if not screenshot_path:
|
||
return None
|
||
|
||
if not os.path.exists(screenshot_path):
|
||
if len(screenshot_path) > 500 and "/" not in screenshot_path[:20]:
|
||
return screenshot_path
|
||
return None
|
||
|
||
if not _init_screenshot_utils():
|
||
return None
|
||
|
||
try:
|
||
if _encode_image is None:
|
||
return None
|
||
return _encode_image(screenshot_path)
|
||
except Exception as e:
|
||
logger.error(f"Image encoding failed for {screenshot_path}: {e}")
|
||
return None
|
||
|
||
|
||
async def encode_screenshot_async(screenshot_path: str) -> str | None:
|
||
"""Async wrapper for screenshot encoding."""
|
||
loop = asyncio.get_event_loop()
|
||
return await loop.run_in_executor(None, encode_screenshot, screenshot_path)
|
||
|
||
|
||
def encode_screenshot_for_vlm(
|
||
screenshot_path: str, max_pixels: int | None = None
|
||
) -> str | None:
|
||
"""Encode screenshot for VLM ground truth with configurable max_pixels.
|
||
|
||
Unlike encode_screenshot(), this function does NOT apply max_height limit.
|
||
You can control max_pixels to study the effect of resize on VLM performance.
|
||
|
||
Args:
|
||
screenshot_path: Path to screenshot file.
|
||
max_pixels: Maximum pixels before resize. If None, uses default (89M).
|
||
Common values:
|
||
- 16_777_216 (16M): Qwen3-VL default
|
||
- 12_845_056 (12.8M): Qwen2-VL default
|
||
- 4_000_000 (4M): ~4000 tokens
|
||
- 1_000_000 (1M): ~1000 tokens
|
||
|
||
Returns:
|
||
Base64 encoded string, or None if failed.
|
||
"""
|
||
if not _init_screenshot_utils():
|
||
return None
|
||
|
||
if not screenshot_path or not os.path.exists(screenshot_path):
|
||
return None
|
||
|
||
try:
|
||
if _encode_image_for_vlm is None:
|
||
return None
|
||
if max_pixels is not None:
|
||
return _encode_image_for_vlm(screenshot_path, max_pixels=max_pixels)
|
||
return _encode_image_for_vlm(screenshot_path)
|
||
except Exception as e:
|
||
logger.error(f"Image encoding (VLM) failed for {screenshot_path}: {e}")
|
||
return None
|
||
|
||
|
||
async def encode_screenshot_for_vlm_async(
|
||
screenshot_path: str, max_pixels: int | None = None
|
||
) -> str | None:
|
||
"""Async wrapper for VLM screenshot encoding."""
|
||
loop = asyncio.get_event_loop()
|
||
return await loop.run_in_executor(
|
||
None, encode_screenshot_for_vlm, screenshot_path, max_pixels
|
||
)
|
||
|
||
|
||
# ============================================================================
|
||
# Pixel-Compressed Encoding for Generation
|
||
# ============================================================================
|
||
|
||
|
||
def make_compressed_encoder(compress_ratio: int, save_dir: str | None = None):
|
||
"""Create an image encoder that downscales images before encoding to base64.
|
||
|
||
The compression ratio N divides the total pixel count by N, i.e. each
|
||
dimension is scaled by 1/sqrt(N). For a 1024x1024 tile:
|
||
- ratio 1 -> 1024x1024 (no compression, baseline)
|
||
- ratio 4 -> 512x512
|
||
- ratio 9 -> ~341x341
|
||
- ratio 16 -> 256x256
|
||
- ratio 25 -> ~205x205
|
||
|
||
Uses LANCZOS resampling (best quality for downscaling).
|
||
|
||
Compressed images are saved to ``save_dir`` (if provided) so they can
|
||
be visually inspected later. The mapping from original path to saved
|
||
compressed path is recorded in ``encoder.compressed_paths`` (a dict
|
||
attached to the returned function object).
|
||
|
||
Args:
|
||
compress_ratio: Pixel compression ratio (1 = no compression).
|
||
save_dir: Directory to save compressed images. If None, a default
|
||
directory ``compressed_tiles_{ratio}x`` is used.
|
||
|
||
Returns:
|
||
A function with the same signature as ``encode_screenshot`` that
|
||
first downscales the image, then encodes it to base64.
|
||
The function has an attribute ``compressed_paths: dict[str, str]``
|
||
mapping original_path -> compressed_path.
|
||
"""
|
||
if compress_ratio <= 1:
|
||
# No compression – use the normal encoder
|
||
return encode_screenshot
|
||
|
||
import math
|
||
|
||
scale_factor = 1.0 / math.sqrt(compress_ratio)
|
||
|
||
# Set up save directory
|
||
if save_dir is None:
|
||
save_dir = f"compressed_tiles_{compress_ratio}x"
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
logger.info(
|
||
f"Pixel compression enabled: ratio={compress_ratio}, "
|
||
f"scale_factor={scale_factor:.4f} per dimension, "
|
||
f"saving compressed images to {save_dir}"
|
||
)
|
||
|
||
# Shared dict to track original -> compressed path mapping
|
||
_compressed_paths: dict[str, str] = {}
|
||
|
||
def _compressed_encode(screenshot_path: str) -> str | None:
|
||
"""Encode image with pixel compression and save to disk."""
|
||
import base64 as _b64
|
||
from io import BytesIO
|
||
from PIL import Image as _Image
|
||
|
||
if not screenshot_path or not os.path.exists(screenshot_path):
|
||
return None
|
||
|
||
try:
|
||
_Image.MAX_IMAGE_PIXELS = 300_000_000
|
||
with _Image.open(screenshot_path) as img:
|
||
new_w = max(1, int(img.width * scale_factor))
|
||
new_h = max(1, int(img.height * scale_factor))
|
||
|
||
if img.mode != "RGB":
|
||
img = img.convert("RGB")
|
||
|
||
img_resized = img.resize((new_w, new_h), _Image.Resampling.LANCZOS)
|
||
|
||
# Save compressed image to disk
|
||
basename = os.path.splitext(os.path.basename(screenshot_path))[0]
|
||
compressed_filename = f"{basename}_compress{compress_ratio}x.png"
|
||
compressed_path = os.path.join(save_dir, compressed_filename)
|
||
img_resized.save(compressed_path, format="PNG")
|
||
_compressed_paths[screenshot_path] = compressed_path
|
||
|
||
# Encode to base64 from the saved file
|
||
buf = BytesIO()
|
||
img_resized.save(buf, format="PNG")
|
||
return _b64.b64encode(buf.getvalue()).decode("utf-8")
|
||
except Exception as e:
|
||
logger.error(f"Compressed encode failed for {screenshot_path}: {e}")
|
||
return None
|
||
|
||
# Attach the path mapping dict to the function so callers can access it
|
||
_compressed_encode.compressed_paths = _compressed_paths
|
||
_compressed_encode.compress_ratio = compress_ratio
|
||
_compressed_encode.save_dir = save_dir
|
||
|
||
return _compressed_encode
|
||
|
||
|
||
# ============================================================================
|
||
# Text Fetching
|
||
# ============================================================================
|
||
|
||
|
||
def fetch_webpage_text(url: str, max_chars: int = 50000) -> str | None:
|
||
"""Fetch webpage and extract clean text content using trafilatura.
|
||
|
||
Args:
|
||
url: URL to fetch.
|
||
max_chars: Maximum characters to return.
|
||
|
||
Returns:
|
||
Extracted text content, or None if failed.
|
||
"""
|
||
try:
|
||
downloaded = trafilatura.fetch_url(url)
|
||
if downloaded is None:
|
||
logger.warning(f"Failed to download {url}")
|
||
return None
|
||
|
||
text = trafilatura.extract(
|
||
downloaded,
|
||
include_comments=False,
|
||
include_tables=True,
|
||
no_fallback=False,
|
||
)
|
||
|
||
if text is None:
|
||
logger.warning(f"Failed to extract text from {url}")
|
||
return None
|
||
|
||
# Clean up excessive newlines
|
||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||
|
||
# Truncate if needed
|
||
if max_chars and len(text) > max_chars:
|
||
text = text[:max_chars] + "\n...[truncated]"
|
||
|
||
return text
|
||
except Exception as e:
|
||
logger.warning(f"Fetch failed for {url}: {e}")
|
||
return None
|
||
|
||
|
||
def fetch_text_for_example(
|
||
example: dict, max_chars: int = 50000, text_cache: dict | None = None
|
||
) -> tuple[str | None, str | None]:
|
||
"""Fetch text content for a single example.
|
||
|
||
Args:
|
||
example: Example dict with metadata containing URL.
|
||
max_chars: Maximum characters to return.
|
||
text_cache: Optional pre-fetched text cache.
|
||
|
||
Returns:
|
||
Tuple of (text_content, source_url).
|
||
"""
|
||
example_id = example.get("id")
|
||
|
||
# Check cache first
|
||
if text_cache and example_id in text_cache:
|
||
cached = text_cache[example_id]
|
||
text = cached.get("text")
|
||
url = cached.get("extracted_url")
|
||
if text:
|
||
return text, url
|
||
|
||
# Extract URL and fetch
|
||
target_url = extract_url_from_metadata(example)
|
||
if not target_url:
|
||
return None, None
|
||
|
||
text = fetch_webpage_text(target_url, max_chars)
|
||
return text, target_url
|
||
|
||
|
||
async def fetch_text_async(
|
||
example: dict, max_chars: int = 50000, text_cache: dict | None = None
|
||
) -> tuple[str | None, str | None]:
|
||
"""Async wrapper for text fetching."""
|
||
loop = asyncio.get_event_loop()
|
||
return await loop.run_in_executor(
|
||
None, fetch_text_for_example, example, max_chars, text_cache
|
||
)
|
||
|
||
|
||
# ============================================================================
|
||
# Image Tiling
|
||
# ============================================================================
|
||
|
||
|
||
def split_image_to_tiles(
|
||
image_path: str,
|
||
output_dir: str,
|
||
tile_size: int | tuple[int, int] = 512,
|
||
overlap: int = 0,
|
||
) -> list[str]:
|
||
"""Split an image into fixed-size tiles.
|
||
|
||
Args:
|
||
image_path: Path to the source image.
|
||
output_dir: Directory to save tiles.
|
||
tile_size: Size of each tile. Can be int (square) or tuple (width, height).
|
||
overlap: Overlap between tiles in pixels.
|
||
|
||
Returns:
|
||
List of tile file paths.
|
||
"""
|
||
from PIL import Image
|
||
import glob
|
||
|
||
if not os.path.exists(image_path):
|
||
return []
|
||
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# Get base name without extension
|
||
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
||
|
||
# Check if tiles already exist for this image
|
||
existing_tiles = sorted(
|
||
glob.glob(os.path.join(output_dir, f"{base_name}_tile_*.png"))
|
||
)
|
||
if existing_tiles:
|
||
# Tiles already exist, return them
|
||
return existing_tiles
|
||
|
||
# Support both square and rectangular tiles
|
||
if isinstance(tile_size, tuple):
|
||
tile_w, tile_h = tile_size
|
||
else:
|
||
tile_w = tile_h = tile_size
|
||
|
||
try:
|
||
Image.MAX_IMAGE_PIXELS = 300_000_000
|
||
img = Image.open(image_path)
|
||
width, height = img.size
|
||
|
||
tile_paths = []
|
||
step_x = tile_w - overlap
|
||
step_y = tile_h - overlap
|
||
|
||
row = 0
|
||
y = 0
|
||
while y < height:
|
||
col = 0
|
||
x = 0
|
||
while x < width:
|
||
# Calculate tile boundaries
|
||
x2 = min(x + tile_w, width)
|
||
y2 = min(y + tile_h, height)
|
||
|
||
# Calculate tile dimensions
|
||
tile_width = x2 - x
|
||
tile_height = y2 - y
|
||
|
||
# Skip tiles with extreme aspect ratios (> 10:1)
|
||
# This prevents issues with ColQwen which requires aspect ratio < 200
|
||
if tile_width > 0 and tile_height > 0:
|
||
aspect_ratio = max(
|
||
tile_width / tile_height, tile_height / tile_width
|
||
)
|
||
if aspect_ratio > 10:
|
||
col += 1
|
||
x += step_x
|
||
if x >= width:
|
||
break
|
||
continue
|
||
|
||
# Crop tile
|
||
tile = img.crop((x, y, x2, y2))
|
||
|
||
# Save tile
|
||
tile_filename = f"{base_name}_tile_{row}_{col}.png"
|
||
tile_path = os.path.join(output_dir, tile_filename)
|
||
tile.save(tile_path)
|
||
tile_paths.append(tile_path)
|
||
|
||
col += 1
|
||
x += step_x
|
||
if x >= width:
|
||
break
|
||
|
||
row += 1
|
||
y += step_y
|
||
if y >= height:
|
||
break
|
||
|
||
img.close()
|
||
return tile_paths
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Failed to split image {image_path}: {e}")
|
||
return []
|
||
|
||
|
||
def prepare_tiles_for_screenshots(
|
||
screenshot_dir: str, tiles_dir: str, tile_size: int = 512, overlap: int = 0
|
||
) -> dict[str, list[str]]:
|
||
"""Split all screenshots in a directory into tiles.
|
||
|
||
Args:
|
||
screenshot_dir: Directory containing full screenshots.
|
||
tiles_dir: Directory to save tiles.
|
||
tile_size: Size of each tile.
|
||
overlap: Overlap between tiles.
|
||
|
||
Returns:
|
||
Dict mapping original image path to list of tile paths.
|
||
"""
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
|
||
result = {}
|
||
for filename in os.listdir(screenshot_dir):
|
||
if not filename.endswith(".png"):
|
||
continue
|
||
|
||
image_path = os.path.join(screenshot_dir, filename)
|
||
tile_paths = split_image_to_tiles(image_path, tiles_dir, tile_size, overlap)
|
||
|
||
if tile_paths:
|
||
result[image_path] = tile_paths
|
||
logger.info(f"Split {filename} into {len(tile_paths)} tiles")
|
||
|
||
logger.info(
|
||
f"Total: {sum(len(v) for v in result.values())} tiles from {len(result)} images"
|
||
)
|
||
return result
|
||
|
||
|
||
# ============================================================================
|
||
# NQ (Natural Questions) Data Loading
|
||
# ============================================================================
|
||
|
||
|
||
def load_nq_data(
|
||
num_examples: int | None = 1000, split: str = "validation"
|
||
) -> list[dict]:
|
||
"""Load Natural Questions (full) split.
|
||
|
||
For validation, follows the short-answer protocol used by our NQ eval:
|
||
keep only examples where >=2 of 5 annotators marked a non-null short
|
||
answer. The train split has a single annotation per example, so train keeps
|
||
examples with a non-null short answer.
|
||
|
||
Source: HuggingFace google-research-datasets/natural_questions.
|
||
Reference: Kwiatkowski et al. (2019).
|
||
|
||
Args:
|
||
num_examples: Number of examples to return. Default 1000.
|
||
split: HuggingFace split to stream ("train" or "validation").
|
||
|
||
Returns:
|
||
List of dicts with id, problem, gold_answers, metadata.
|
||
"""
|
||
from datasets import load_dataset
|
||
import html as _html
|
||
|
||
if split not in {"train", "validation"}:
|
||
raise ValueError(
|
||
f"Unsupported NQ split: {split!r}. Expected 'train' or 'validation'."
|
||
)
|
||
|
||
logger.info(f"Loading NQ {split} split (streaming)...")
|
||
ds = load_dataset(
|
||
"google-research-datasets/natural_questions",
|
||
split=split,
|
||
streaming=True,
|
||
)
|
||
|
||
data = []
|
||
for ex in ds:
|
||
# Extract short answers from all 5 annotators
|
||
annotations = ex["annotations"]
|
||
short_answer_texts = set()
|
||
non_null_annotators = 0
|
||
|
||
# annotations is a dict with list values (one per annotator)
|
||
num_annotators = len(annotations["id"])
|
||
for i in range(num_annotators):
|
||
texts = annotations["short_answers"][i].get("text", [])
|
||
if texts:
|
||
non_null_annotators += 1
|
||
for t in texts:
|
||
if t.strip():
|
||
short_answer_texts.add(t.strip())
|
||
|
||
min_non_null = 2 if split == "validation" else 1
|
||
if non_null_annotators < min_non_null:
|
||
continue
|
||
|
||
if not short_answer_texts:
|
||
continue
|
||
|
||
question_text = ex["question"]["text"]
|
||
doc_url = ex["document"]["url"]
|
||
# Clean up HTML entities in URL (e.g., & -> &)
|
||
doc_url = _html.unescape(doc_url)
|
||
# Normalize NQ URL format: /w/index.php?title=Foo&oldid=123 -> /wiki/Foo
|
||
_title_match = re.search(r"[?&]title=([^&]+)", doc_url)
|
||
if _title_match:
|
||
doc_url = f"https://en.wikipedia.org/wiki/{urllib.parse.quote(_title_match.group(1), safe='/:(),-')}"
|
||
|
||
example = {
|
||
"id": hashlib.md5(question_text.encode()).hexdigest(),
|
||
"problem": question_text,
|
||
"gold_answers": sorted(short_answer_texts),
|
||
"metadata": {
|
||
"urls": [doc_url],
|
||
"dataset": "nq",
|
||
"document_title": ex["document"]["title"],
|
||
},
|
||
}
|
||
data.append(example)
|
||
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
filter_desc = (
|
||
">=2 annotator agreement" if split == "validation" else "non-null short answer"
|
||
)
|
||
logger.info(f"Loaded {len(data)} NQ {split} examples (filtered by {filter_desc}).")
|
||
return data
|
||
|
||
|
||
# ============================================================================
|
||
# TriviaQA Data Loading
|
||
# ============================================================================
|
||
|
||
|
||
def load_triviaqa_data(num_examples: int | None = 1000) -> list[dict]:
|
||
"""Load TriviaQA rc.wikipedia validation split.
|
||
|
||
Uses entity_pages.title to construct ground truth Wikipedia URLs.
|
||
gold_answers includes answer.value + answer.aliases (following TriviaQA official eval).
|
||
|
||
Source: HuggingFace mandarjoshi/trivia_qa, config rc.wikipedia, validation split.
|
||
Reference: Joshi et al. (2017).
|
||
|
||
Args:
|
||
num_examples: Number of examples to return. Default 1000.
|
||
|
||
Returns:
|
||
List of dicts with id, problem, gold_answers, metadata.
|
||
"""
|
||
from datasets import load_dataset
|
||
import ast as _ast
|
||
from urllib.parse import quote as _url_quote
|
||
|
||
logger.info("Loading TriviaQA rc.wikipedia validation split (streaming)...")
|
||
ds = load_dataset(
|
||
"mandarjoshi/trivia_qa",
|
||
"rc.wikipedia",
|
||
split="validation",
|
||
streaming=True,
|
||
)
|
||
|
||
data = []
|
||
for ex in ds:
|
||
question = ex["question"]
|
||
answer_obj = ex["answer"]
|
||
|
||
# Extract gold answers: value + aliases
|
||
gold_answers = set()
|
||
value = answer_obj.get("value", "")
|
||
if value:
|
||
gold_answers.add(value)
|
||
|
||
# aliases is stored as a string repr of a list
|
||
aliases_raw = answer_obj.get("aliases", "")
|
||
if isinstance(aliases_raw, str) and aliases_raw:
|
||
try:
|
||
aliases = _ast.literal_eval(aliases_raw)
|
||
if isinstance(aliases, list):
|
||
for a in aliases:
|
||
if a and a.strip():
|
||
gold_answers.add(a.strip())
|
||
except (ValueError, SyntaxError):
|
||
pass
|
||
elif isinstance(aliases_raw, list):
|
||
for a in aliases_raw:
|
||
if a and a.strip():
|
||
gold_answers.add(a.strip())
|
||
|
||
if not gold_answers:
|
||
continue
|
||
|
||
# Construct Wikipedia URL from entity_pages.title
|
||
urls = []
|
||
entity_titles = ex.get("entity_pages", {}).get("title", [])
|
||
if entity_titles:
|
||
for title in entity_titles:
|
||
if title:
|
||
wiki_url = f"https://en.wikipedia.org/wiki/{_url_quote(title.replace(' ', '_'))}"
|
||
urls.append(wiki_url)
|
||
|
||
example = {
|
||
"id": hashlib.md5(question.encode()).hexdigest(),
|
||
"problem": question,
|
||
"gold_answers": sorted(gold_answers),
|
||
"question_type": ex.get("question_source", ""),
|
||
"metadata": {
|
||
"urls": urls,
|
||
"dataset": "triviaqa",
|
||
"question_id": ex.get("question_id", ""),
|
||
},
|
||
}
|
||
data.append(example)
|
||
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} TriviaQA examples.")
|
||
return data
|
||
|
||
|
||
# ============================================================================
|
||
# NQ-Tables Data Loading
|
||
# ============================================================================
|
||
|
||
|
||
def load_nq_tables_data(num_examples: int | None = 1000) -> list[dict]:
|
||
"""Load NQ-Tables dev split (table subset of Natural Questions).
|
||
|
||
NQ-Tables filters Natural Questions to only keep examples where the gold
|
||
answer resides inside a Wikipedia HTML table. Each example includes the
|
||
full table content (columns + rows) and the Wikipedia URL.
|
||
|
||
Source: GCS gs://tapas_models/2021_07_22/nq_tables/interactions/dev.jsonl
|
||
Reference: Herzig et al. (2021), "Open Domain Question Answering over
|
||
Tables via Dense Retrieval" (NAACL 2021).
|
||
|
||
Args:
|
||
num_examples: Number of examples to return. Default 1000.
|
||
|
||
Returns:
|
||
List of dicts with id, problem, gold_answers, metadata.
|
||
"""
|
||
import html as _html
|
||
|
||
data_path = os.path.join(
|
||
os.path.dirname(__file__), "..", "..", "data", "nq_tables", "dev.jsonl"
|
||
)
|
||
data_path = os.path.abspath(data_path)
|
||
|
||
if not os.path.exists(data_path):
|
||
raise FileNotFoundError(
|
||
f"NQ-Tables data not found at {data_path}. "
|
||
"Download with: gsutil cp gs://tapas_models/2021_07_22/nq_tables/interactions/dev.jsonl data/nq_tables/"
|
||
)
|
||
|
||
logger.info(f"Loading NQ-Tables dev split from {data_path}...")
|
||
|
||
import json as _json
|
||
|
||
data = []
|
||
with open(data_path) as f:
|
||
for line in f:
|
||
ex = _json.loads(line)
|
||
questions = ex.get("questions", [])
|
||
if not questions:
|
||
continue
|
||
|
||
q = questions[0]
|
||
question_text = q.get("originalText", "")
|
||
answer_texts = q.get("answer", {}).get("answerTexts", [])
|
||
|
||
if not question_text or not answer_texts:
|
||
continue
|
||
|
||
gold_answers = [a.strip() for a in answer_texts if a.strip()]
|
||
if not gold_answers:
|
||
continue
|
||
|
||
# Extract Wikipedia URL from table metadata
|
||
table = ex.get("table", {})
|
||
doc_url = table.get("documentUrl", "")
|
||
doc_url = _html.unescape(doc_url)
|
||
# Normalize NQ URL format: /w/index.php?title=Foo&oldid=123 -> /wiki/Foo
|
||
_title_match = re.search(r"[?&]title=([^&]+)", doc_url)
|
||
if _title_match:
|
||
doc_url = f"https://en.wikipedia.org/wiki/{urllib.parse.quote(_title_match.group(1), safe='/:(),-')}"
|
||
|
||
example = {
|
||
"id": ex.get("id", hashlib.md5(question_text.encode()).hexdigest()),
|
||
"problem": question_text,
|
||
"gold_answers": gold_answers,
|
||
"metadata": {
|
||
"urls": [doc_url] if doc_url else [],
|
||
"dataset": "nq_tables",
|
||
"document_title": table.get("documentTitle", ""),
|
||
"table_id": table.get("tableId", ""),
|
||
},
|
||
}
|
||
data.append(example)
|
||
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} NQ-Tables examples.")
|
||
return data
|
||
|
||
|
||
# ============================================================================
|
||
# Multiple-Choice Reasoning Benchmarks
|
||
# ============================================================================
|
||
|
||
LETTERS = ["A", "B", "C", "D", "E"]
|
||
|
||
|
||
def _format_mc_options(labels: list[str], texts: list[str]) -> str:
|
||
"""Format MC options as 'A. text1\nB. text2\n...'"""
|
||
return "\n".join(f"{label}. {text}" for label, text in zip(labels, texts))
|
||
|
||
|
||
MC_INSTRUCTION = "Choose the best answer from the options above. Reply with ONLY the letter (e.g. A, B, C, or D)."
|
||
|
||
|
||
def load_piqa_data(num_examples: int | None = None) -> list[dict]:
|
||
"""Load PIQA (Physical Intuition QA) validation split.
|
||
|
||
2-choice physical commonsense benchmark. Label is 0 or 1.
|
||
Source: HuggingFace `ybisk/piqa`, validation split.
|
||
|
||
Returns list of dicts with problem (question only), gold_answers (letter),
|
||
additional_instructions (options + MC instruction), metadata.
|
||
"""
|
||
from datasets import load_dataset
|
||
|
||
logger.info("Loading PIQA validation split...")
|
||
ds = load_dataset("ybisk/piqa", split="validation", revision="refs/convert/parquet")
|
||
|
||
data = []
|
||
for ex in ds:
|
||
question = ex["goal"]
|
||
options = [ex["sol1"], ex["sol2"]]
|
||
label = int(ex["label"])
|
||
gold_letter = LETTERS[label]
|
||
|
||
options_text = _format_mc_options(LETTERS[:2], options)
|
||
example = {
|
||
"id": hashlib.md5(question.encode()).hexdigest(),
|
||
"problem": question,
|
||
"gold_answers": [gold_letter],
|
||
"additional_instructions": f"{options_text}\n\n{MC_INSTRUCTION}",
|
||
"metadata": {"dataset": "piqa", "urls": [], "gold_letter": gold_letter},
|
||
}
|
||
data.append(example)
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} PIQA examples.")
|
||
return data
|
||
|
||
|
||
def load_hellaswag_data(num_examples: int | None = None) -> list[dict]:
|
||
"""Load HellaSwag validation split.
|
||
|
||
4-choice sentence completion benchmark. Label is "0"-"3".
|
||
Source: HuggingFace `Rowan/hellaswag`, validation split.
|
||
"""
|
||
from datasets import load_dataset
|
||
|
||
logger.info("Loading HellaSwag validation split...")
|
||
ds = load_dataset(
|
||
"Rowan/hellaswag", split="validation", revision="refs/convert/parquet"
|
||
)
|
||
|
||
data = []
|
||
for ex in ds:
|
||
question = ex["ctx"]
|
||
options = ex["endings"]
|
||
label = int(ex["label"])
|
||
gold_letter = LETTERS[label]
|
||
|
||
options_text = _format_mc_options(LETTERS[: len(options)], options)
|
||
example = {
|
||
"id": hashlib.md5(question.encode()).hexdigest(),
|
||
"problem": question,
|
||
"gold_answers": [gold_letter],
|
||
"additional_instructions": f"{options_text}\n\n{MC_INSTRUCTION}",
|
||
"metadata": {
|
||
"dataset": "hellaswag",
|
||
"urls": [],
|
||
"gold_letter": gold_letter,
|
||
},
|
||
}
|
||
data.append(example)
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} HellaSwag examples.")
|
||
return data
|
||
|
||
|
||
def load_commonsenseqa_data(num_examples: int | None = None) -> list[dict]:
|
||
"""Load CommonsenseQA validation split.
|
||
|
||
5-choice commonsense reasoning benchmark. answerKey is A-E.
|
||
Source: HuggingFace `tau/commonsense_qa`, validation split.
|
||
"""
|
||
from datasets import load_dataset
|
||
|
||
logger.info("Loading CommonsenseQA validation split...")
|
||
ds = load_dataset("tau/commonsense_qa", split="validation")
|
||
|
||
data = []
|
||
for ex in ds:
|
||
question = ex["question"]
|
||
labels = ex["choices"]["label"]
|
||
texts = ex["choices"]["text"]
|
||
gold_letter = ex["answerKey"]
|
||
|
||
options_text = _format_mc_options(labels, texts)
|
||
example = {
|
||
"id": hashlib.md5(question.encode()).hexdigest(),
|
||
"problem": question,
|
||
"gold_answers": [gold_letter],
|
||
"additional_instructions": f"{options_text}\n\n{MC_INSTRUCTION}",
|
||
"metadata": {
|
||
"dataset": "commonsense_qa",
|
||
"urls": [],
|
||
"gold_letter": gold_letter,
|
||
},
|
||
}
|
||
data.append(example)
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} CommonsenseQA examples.")
|
||
return data
|
||
|
||
|
||
def load_openbookqa_data(num_examples: int | None = None) -> list[dict]:
|
||
"""Load OpenBookQA test split.
|
||
|
||
4-choice science QA benchmark. answerKey is A-D.
|
||
Source: HuggingFace `allenai/openbookqa`, main config, test split.
|
||
"""
|
||
from datasets import load_dataset
|
||
|
||
logger.info("Loading OpenBookQA test split...")
|
||
ds = load_dataset("allenai/openbookqa", "main", split="test")
|
||
|
||
data = []
|
||
for ex in ds:
|
||
question = ex["question_stem"]
|
||
labels = ex["choices"]["label"]
|
||
texts = ex["choices"]["text"]
|
||
gold_letter = ex["answerKey"]
|
||
|
||
options_text = _format_mc_options(labels, texts)
|
||
example = {
|
||
"id": hashlib.md5(question.encode()).hexdigest(),
|
||
"problem": question,
|
||
"gold_answers": [gold_letter],
|
||
"additional_instructions": f"{options_text}\n\n{MC_INSTRUCTION}",
|
||
"metadata": {
|
||
"dataset": "openbookqa",
|
||
"urls": [],
|
||
"gold_letter": gold_letter,
|
||
},
|
||
}
|
||
data.append(example)
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} OpenBookQA examples.")
|
||
return data
|
||
|
||
|
||
def load_arc_data(
|
||
config: str = "ARC-Challenge", num_examples: int | None = None
|
||
) -> list[dict]:
|
||
"""Load ARC (AI2 Reasoning Challenge) test split.
|
||
|
||
3-5 choice science exam benchmark. answerKey is A-E or 1-5 (normalized to letters).
|
||
Source: HuggingFace `allenai/ai2_arc`, ARC-Challenge or ARC-Easy config, test split.
|
||
|
||
Args:
|
||
config: "ARC-Challenge" or "ARC-Easy"
|
||
num_examples: Max examples to return. None = all.
|
||
"""
|
||
from datasets import load_dataset
|
||
|
||
dataset_name = config.lower().replace("-", "_")
|
||
logger.info(f"Loading ARC {config} test split...")
|
||
ds = load_dataset("allenai/ai2_arc", config, split="test")
|
||
|
||
# ARC answerKey can be "1","2","3","4","5" instead of letters
|
||
DIGIT_TO_LETTER = {"1": "A", "2": "B", "3": "C", "4": "D", "5": "E"}
|
||
|
||
data = []
|
||
for ex in ds:
|
||
question = ex["question"]
|
||
labels = ex["choices"]["label"]
|
||
texts = ex["choices"]["text"]
|
||
gold_letter = ex["answerKey"]
|
||
gold_letter = DIGIT_TO_LETTER.get(gold_letter, gold_letter)
|
||
|
||
options_text = _format_mc_options(labels, texts)
|
||
example = {
|
||
"id": hashlib.md5(question.encode()).hexdigest(),
|
||
"problem": question,
|
||
"gold_answers": [gold_letter],
|
||
"additional_instructions": f"{options_text}\n\n{MC_INSTRUCTION}",
|
||
"metadata": {
|
||
"dataset": dataset_name,
|
||
"urls": [],
|
||
"gold_letter": gold_letter,
|
||
},
|
||
}
|
||
data.append(example)
|
||
if num_examples and len(data) >= num_examples:
|
||
break
|
||
|
||
logger.info(f"Loaded {len(data)} ARC {config} examples.")
|
||
return data
|