startrail-org--pixelrag
wehub-resource-sync
1275 行
42 KiB
Python
1275 行
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
|