wehub-resource-sync
4357 lines
167 KiB
Python
4357 lines
167 KiB
Python
"""Retrieval strategies for SimpleQA evaluation.
|
||
|
||
This module defines retrieval strategies that work with prepared data:
|
||
- NaiveRetriever: No retrieval, just pass query to LLM
|
||
- ScreenshotRetriever: Use pre-captured screenshot for the example
|
||
- TextRetriever: Use pre-fetched or cached text for the example
|
||
- VectorRetriever: Search across all screenshots using vector similarity
|
||
"""
|
||
|
||
import base64
|
||
import io
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
from abc import ABC, abstractmethod
|
||
from dataclasses import dataclass, field
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class RetrievalResult:
|
||
"""Result from a retrieval operation."""
|
||
|
||
# Text content (for text-based retrieval)
|
||
text: str | None = None
|
||
|
||
# Image paths with scores (for vector retrieval)
|
||
images: list[tuple[str, float]] = field(default_factory=list)
|
||
|
||
# Per-image source URLs, aligned with ``images`` when provided.
|
||
image_urls: list[str | None] = field(default_factory=list)
|
||
|
||
# Base64 encoded image (for screenshot)
|
||
base64_image: str | None = None
|
||
|
||
# Source URL
|
||
source_url: str | None = None
|
||
|
||
# Which retrieval type was used
|
||
retrieval_type: str = "naive"
|
||
|
||
# Path to pixel query image used for retrieval embedding (rendered card or raw photo)
|
||
pixel_query_path: str | None = None
|
||
|
||
# Path to raw species/landmark photo for generation (always the original photo,
|
||
# never the rendered card). If None, falls back to pixel_query_path in build_messages.
|
||
query_image_path: str | None = None
|
||
|
||
@property
|
||
def has_content(self) -> bool:
|
||
"""Check if retrieval found any content."""
|
||
return bool(self.text or self.images or self.base64_image)
|
||
|
||
|
||
class BaseRetriever(ABC):
|
||
"""Base class for retrieval strategies."""
|
||
|
||
@abstractmethod
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
"""Retrieve relevant content for the query.
|
||
|
||
Args:
|
||
query: The question/query text.
|
||
example: The full example dict (may contain metadata, prepared data, etc.).
|
||
|
||
Returns:
|
||
RetrievalResult with retrieved content.
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
|
||
# EVQA query image data dirs (iNaturalist 2021, Google Landmarks v2)
|
||
_INAT2021_DATA_DIR = os.path.join(
|
||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
||
"data",
|
||
"inat2021",
|
||
)
|
||
_LANDMARK_V2_DATA_DIR = os.path.join(
|
||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
||
"data",
|
||
"landmark_v2",
|
||
)
|
||
|
||
# Local kiwix tile store (pre-rendered Wikipedia pages)
|
||
_WIKI_SCREENSHOT_DIR = "/path/to/project"
|
||
_KIWIX_OUTPUT_DIR = "/path/to/data"
|
||
_KIWIX_ARTICLES_JSON = "/path/to/data"
|
||
_KIWIX_REDIRECTS_JSON = "/path/to/data"
|
||
|
||
|
||
def _lookup_and_copy_local_wiki_tiles(
|
||
ex_id: str,
|
||
url: str,
|
||
tiles_dir: str,
|
||
wiki_cache_dir: str,
|
||
cut_height: int,
|
||
) -> list[str]:
|
||
"""Look up a Wikipedia URL in the local kiwix tile store, copy raw tiles, cut into strips.
|
||
|
||
Args:
|
||
ex_id: Example ID (used for output tile naming).
|
||
url: Wikipedia URL.
|
||
tiles_dir: Directory where cut tile strips are written ({ex_id}_tile_*.png).
|
||
wiki_cache_dir: Directory where raw kiwix tile pages are cached ({ex_id}/).
|
||
cut_height: Height of each output strip in pixels.
|
||
|
||
Returns:
|
||
Sorted list of cut tile paths.
|
||
|
||
Raises:
|
||
RuntimeError: If kiwix index unavailable, URL not found, or no tiles produced.
|
||
"""
|
||
import glob as _glob
|
||
import shutil
|
||
import sys as _sys
|
||
from PIL import Image
|
||
|
||
# Return cached tiles if already cut
|
||
existing = sorted(_glob.glob(os.path.join(tiles_dir, f"{ex_id}_tile_*.png")))
|
||
if existing:
|
||
return existing
|
||
|
||
if not url or "wikipedia.org" not in url:
|
||
raise RuntimeError(f"Not a Wikipedia URL: {url!r}")
|
||
|
||
if not os.path.isdir(_KIWIX_OUTPUT_DIR) or not os.path.isfile(_KIWIX_ARTICLES_JSON):
|
||
raise RuntimeError(f"kiwix tiles unavailable at {_KIWIX_OUTPUT_DIR}")
|
||
|
||
if _WIKI_SCREENSHOT_DIR not in _sys.path:
|
||
_sys.path.insert(0, _WIKI_SCREENSHOT_DIR)
|
||
from scripts.build_index import batch_query_by_url as _batch_query
|
||
|
||
redirects = _KIWIX_REDIRECTS_JSON if os.path.isfile(_KIWIX_REDIRECTS_JSON) else None
|
||
results = _batch_query(
|
||
_KIWIX_OUTPUT_DIR, [url], _KIWIX_ARTICLES_JSON, redirects_json=redirects
|
||
)
|
||
result = results.get(url)
|
||
if result is None:
|
||
raise RuntimeError(f"URL not found in local kiwix: {url}")
|
||
|
||
# Copy raw kiwix tiles to wiki_cache_dir/{ex_id}/
|
||
src_dir = os.path.join(_KIWIX_OUTPUT_DIR, result["tiles_dir"])
|
||
article_cache = os.path.join(wiki_cache_dir, str(ex_id))
|
||
if not os.path.exists(article_cache):
|
||
if not os.path.isdir(src_dir):
|
||
raise RuntimeError(f"kiwix tiles dir not on disk: {src_dir}")
|
||
shutil.copytree(src_dir, article_cache)
|
||
|
||
# Cut raw tiles into height=cut_height strips
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
raw_tiles = sorted(
|
||
f
|
||
for f in os.listdir(article_cache)
|
||
if f.endswith(".png") and f.startswith("tile_")
|
||
)
|
||
if not raw_tiles:
|
||
raise RuntimeError(f"No tile PNGs found in {article_cache}")
|
||
|
||
global_row = 0
|
||
for raw_name in raw_tiles:
|
||
raw_path = os.path.join(article_cache, raw_name)
|
||
if os.path.getsize(raw_path) == 0:
|
||
continue
|
||
img = Image.open(raw_path)
|
||
img.load()
|
||
w, h = img.size
|
||
y = 0
|
||
while y < h:
|
||
y2 = min(y + cut_height, h)
|
||
strip = img.crop((0, y, w, y2))
|
||
strip.save(os.path.join(tiles_dir, f"{ex_id}_tile_{global_row}_0.png"))
|
||
strip.close()
|
||
global_row += 1
|
||
y += cut_height
|
||
img.close()
|
||
|
||
tile_paths = sorted(_glob.glob(os.path.join(tiles_dir, f"{ex_id}_tile_*.png")))
|
||
if not tile_paths:
|
||
raise RuntimeError(f"No strips cut for {ex_id} (source: {article_cache})")
|
||
return tile_paths
|
||
|
||
|
||
def _get_inat_image_path_for_example(example: dict, tiles_dir: str) -> str | None:
|
||
"""Get iNaturalist 2021 query image path. dataset_name must be 'inaturalist'."""
|
||
inat_ids = example.get("inat_image_ids", [])
|
||
if not inat_ids:
|
||
return None
|
||
cache_dir = os.path.join(os.path.dirname(tiles_dir), "inat_images")
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
example_id = example.get("id", "unknown")
|
||
local_path = os.path.join(cache_dir, f"{example_id}.jpg")
|
||
if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
|
||
return local_path
|
||
import shutil
|
||
|
||
id_map = TiledQwen3VLEmbeddingRetriever._load_inat2021_mapping()
|
||
for str_id in inat_ids:
|
||
try:
|
||
img_id = int(str_id)
|
||
except ValueError:
|
||
continue
|
||
file_name = id_map.get(img_id)
|
||
if not file_name:
|
||
continue
|
||
src_path = os.path.join(_INAT2021_DATA_DIR, file_name)
|
||
if os.path.exists(src_path) and os.path.getsize(src_path) > 0:
|
||
shutil.copy2(src_path, local_path)
|
||
return local_path
|
||
logger.warning(f"Failed to find iNaturalist image for {example_id}")
|
||
return None
|
||
|
||
|
||
def _get_landmark_image_path_for_example(
|
||
example: dict, tiles_dir: str, quiet: bool = False
|
||
) -> str | None:
|
||
"""Get Google Landmarks v2 query image path. dataset_name must be 'landmarks'.
|
||
|
||
GLDv2 stores images as {split}/{a}/{b}/{c}/{id}.jpg (a,b,c = first 3 chars of id).
|
||
Searches train, index, test in order.
|
||
"""
|
||
ids = example.get("dataset_image_ids_parsed", [])
|
||
if not ids:
|
||
return None
|
||
cache_dir = os.path.join(os.path.dirname(tiles_dir), "landmark_images")
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
example_id = example.get("id", "unknown")
|
||
local_path = os.path.join(cache_dir, f"{example_id}.jpg")
|
||
if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
|
||
return local_path
|
||
import shutil
|
||
|
||
data_dir = _LANDMARK_V2_DATA_DIR
|
||
for img_id in ids:
|
||
if len(img_id) < 3:
|
||
continue
|
||
# GLDv2 path: {split}/{a}/{b}/{c}/{id}.jpg
|
||
subpath = f"{img_id[0]}/{img_id[1]}/{img_id[2]}/{img_id}.jpg"
|
||
for split in ("train", "index", "test"):
|
||
src_path = os.path.join(data_dir, split, subpath)
|
||
if os.path.exists(src_path) and os.path.getsize(src_path) > 0:
|
||
shutil.copy2(src_path, local_path)
|
||
return local_path
|
||
# Fallback: download from train.csv URL (requires data/landmark_v2/train.csv)
|
||
# Try each img_id in order; first URL may be 404, others might work
|
||
for img_id in ids:
|
||
if _try_download_landmark_from_url(example_id, img_id, local_path):
|
||
return local_path
|
||
if not quiet:
|
||
logger.warning(
|
||
f"Failed to find Landmark image for {example_id} (data in {data_dir}?)"
|
||
)
|
||
return None
|
||
|
||
|
||
def _try_download_landmark_from_url(
|
||
example_id: str, img_id: str, local_path: str
|
||
) -> bool:
|
||
"""Try to download landmark image from train.csv URL. Used when GLDv2 TARs unavailable.
|
||
|
||
Returns True if download succeeded and file is valid, False otherwise.
|
||
"""
|
||
import urllib.request
|
||
|
||
train_csv = os.path.join(_LANDMARK_V2_DATA_DIR, "train.csv")
|
||
if not os.path.exists(train_csv):
|
||
return False
|
||
import csv
|
||
|
||
with open(train_csv) as f:
|
||
for row in csv.DictReader(f):
|
||
if row.get("id") == img_id:
|
||
url = row.get("url", "")
|
||
if url:
|
||
try:
|
||
req = urllib.request.Request(
|
||
url, headers={"User-Agent": "PixelRAG-Bot/1.0"}
|
||
)
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
data = resp.read()
|
||
if len(data) >= 1000:
|
||
with open(local_path, "wb") as out:
|
||
out.write(data)
|
||
return True
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"URL download failed for {example_id} (img_id={img_id}): {e}"
|
||
)
|
||
return False
|
||
return False
|
||
|
||
|
||
def _get_query_image_path_for_example(
|
||
example: dict, tiles_dir: str, quiet: bool = False
|
||
) -> str | None:
|
||
"""Get EVQA query image path. Dispatches by dataset_name: inaturalist | landmarks."""
|
||
ds = (example.get("dataset_name") or "").lower()
|
||
if ds == "inaturalist":
|
||
return _get_inat_image_path_for_example(example, tiles_dir)
|
||
if ds == "landmarks":
|
||
return _get_landmark_image_path_for_example(example, tiles_dir, quiet=quiet)
|
||
# Fallback: try inaturalist (backward compat when dataset_name missing)
|
||
return _get_inat_image_path_for_example(example, tiles_dir)
|
||
|
||
|
||
def _get_all_inat_image_paths(example: dict, tiles_dir: str) -> list[str]:
|
||
"""Get ALL iNaturalist query image paths for an example (not just the first)."""
|
||
inat_ids = example.get("inat_image_ids", [])
|
||
if not inat_ids:
|
||
return []
|
||
cache_dir = os.path.join(os.path.dirname(tiles_dir), "inat_images_multi")
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
example_id = example.get("id", "unknown")
|
||
import shutil
|
||
|
||
id_map = TiledQwen3VLEmbeddingRetriever._load_inat2021_mapping()
|
||
paths = []
|
||
for i, str_id in enumerate(inat_ids):
|
||
local_path = os.path.join(cache_dir, f"{example_id}_{i}.jpg")
|
||
if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
|
||
paths.append(local_path)
|
||
continue
|
||
try:
|
||
img_id = int(str_id)
|
||
except ValueError:
|
||
continue
|
||
file_name = id_map.get(img_id)
|
||
if not file_name:
|
||
continue
|
||
src_path = os.path.join(_INAT2021_DATA_DIR, file_name)
|
||
if os.path.exists(src_path) and os.path.getsize(src_path) > 0:
|
||
shutil.copy2(src_path, local_path)
|
||
paths.append(local_path)
|
||
return paths
|
||
|
||
|
||
_landmark_url_map_cache: dict[str, str] | None = None
|
||
|
||
|
||
def _load_landmark_url_map() -> dict[str, str]:
|
||
"""Load GLDv2 train.csv: img_id -> url. Cached after first call."""
|
||
global _landmark_url_map_cache
|
||
if _landmark_url_map_cache is not None:
|
||
return _landmark_url_map_cache
|
||
import csv
|
||
|
||
train_csv = os.path.join(_LANDMARK_V2_DATA_DIR, "train.csv")
|
||
if not os.path.exists(train_csv):
|
||
_landmark_url_map_cache = {}
|
||
return _landmark_url_map_cache
|
||
url_map = {}
|
||
with open(train_csv, encoding="utf-8") as f:
|
||
for row in csv.DictReader(f):
|
||
img_id = row.get("id", "").strip()
|
||
url = row.get("url", "").strip()
|
||
if img_id and url:
|
||
url_map[img_id] = url
|
||
_landmark_url_map_cache = url_map
|
||
logger.info(f"Loaded landmark URL map: {len(url_map)} entries")
|
||
return url_map
|
||
|
||
|
||
def _download_landmark_image_by_id(img_id: str, local_path: str) -> bool:
|
||
"""Download a landmark image by its GLDv2 ID. Returns True on success."""
|
||
import urllib.request
|
||
|
||
url_map = _load_landmark_url_map()
|
||
url = url_map.get(img_id)
|
||
if not url:
|
||
return False
|
||
try:
|
||
req = urllib.request.Request(url, headers={"User-Agent": "PixelRAG-Bot/1.0"})
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
data = resp.read()
|
||
if len(data) >= 1000:
|
||
with open(local_path, "wb") as out:
|
||
out.write(data)
|
||
return True
|
||
except Exception as e:
|
||
logger.debug(f"Download failed for landmark {img_id}: {e}")
|
||
return False
|
||
|
||
|
||
def _get_all_landmark_image_paths(example: dict, tiles_dir: str) -> list[str]:
|
||
"""Get ALL Google Landmarks query image paths for an example (not just the first)."""
|
||
ids = example.get("dataset_image_ids_parsed", [])
|
||
if not ids:
|
||
return []
|
||
cache_dir = os.path.join(os.path.dirname(tiles_dir), "landmark_images_multi")
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
example_id = example.get("id", "unknown")
|
||
import shutil
|
||
|
||
data_dir = _LANDMARK_V2_DATA_DIR
|
||
paths = []
|
||
for i, img_id in enumerate(ids):
|
||
local_path = os.path.join(cache_dir, f"{example_id}_{i}.jpg")
|
||
if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
|
||
paths.append(local_path)
|
||
continue
|
||
if len(img_id) < 3:
|
||
continue
|
||
subpath = f"{img_id[0]}/{img_id[1]}/{img_id[2]}/{img_id}.jpg"
|
||
found = False
|
||
for split in ("train", "index", "test"):
|
||
src_path = os.path.join(data_dir, split, subpath)
|
||
if os.path.exists(src_path) and os.path.getsize(src_path) > 0:
|
||
shutil.copy2(src_path, local_path)
|
||
paths.append(local_path)
|
||
found = True
|
||
break
|
||
if not found:
|
||
if _download_landmark_image_by_id(img_id, local_path):
|
||
paths.append(local_path)
|
||
return paths
|
||
|
||
|
||
def _get_all_query_image_paths(example: dict, tiles_dir: str) -> list[str]:
|
||
"""Get ALL query image paths for an EVQA example (all available images, not just the first).
|
||
|
||
Falls back to the single ``query_image_path`` / ``_get_query_image_path_for_example``
|
||
when the multi-image helpers return nothing (e.g. ``dataset_image_ids_parsed`` lives
|
||
inside ``original_data`` rather than at top level).
|
||
"""
|
||
ds = (example.get("dataset_name") or "").lower()
|
||
if ds not in ("inaturalist", "landmarks"):
|
||
od = example.get("original_data", {})
|
||
if isinstance(od, str):
|
||
import ast
|
||
|
||
try:
|
||
od = ast.literal_eval(od)
|
||
except Exception:
|
||
od = {}
|
||
ds = (od.get("dataset_name") or "").lower()
|
||
if ds == "inaturalist":
|
||
paths = _get_all_inat_image_paths(example, tiles_dir)
|
||
elif ds == "landmarks":
|
||
paths = _get_all_landmark_image_paths(example, tiles_dir)
|
||
else:
|
||
paths = _get_all_inat_image_paths(example, tiles_dir)
|
||
if not paths:
|
||
single = example.get("query_image_path") or _get_query_image_path_for_example(
|
||
example, tiles_dir, quiet=True
|
||
)
|
||
if single and os.path.exists(single):
|
||
paths = [single]
|
||
return paths
|
||
|
||
|
||
class NaiveRetriever(BaseRetriever):
|
||
"""No retrieval - returns empty result, LLM answers from its own knowledge."""
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
return RetrievalResult(retrieval_type="naive")
|
||
|
||
|
||
class EVQANoRetrievalRetriever(BaseRetriever):
|
||
"""EVQA without retrieval: query + iNaturalist image only, no Wikipedia tiles.
|
||
|
||
Used to test VLM's ability to answer from the species image alone.
|
||
"""
|
||
|
||
def __init__(self, tiles_dir: str = "tiles/evqa"):
|
||
self.tiles_dir = tiles_dir
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
inat_image_path = _get_query_image_path_for_example(example, self.tiles_dir)
|
||
return RetrievalResult(
|
||
images=[],
|
||
retrieval_type="evqa_no_retrieval_multimodal",
|
||
pixel_query_path=inat_image_path,
|
||
query_image_path=inat_image_path,
|
||
)
|
||
|
||
|
||
def _save_task_query_image(
|
||
example: dict, task_name: str, base_dir: str = "tiles"
|
||
) -> str | None:
|
||
"""Save query image from any task to disk. Returns path or None.
|
||
Images saved to {base_dir}/{task_name}_images/{example_id}.png
|
||
Works with PIL images, base64 strings, or dict with 'bytes' key.
|
||
"""
|
||
img = example.get("image")
|
||
if img is None:
|
||
return None
|
||
example_id = example.get("id", "unknown")
|
||
save_dir = os.path.join(base_dir, f"{task_name}_images")
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
out_path = os.path.join(save_dir, f"{example_id}.png")
|
||
if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
|
||
return out_path
|
||
try:
|
||
if hasattr(img, "save"):
|
||
img.save(out_path, format="PNG")
|
||
return out_path
|
||
if isinstance(img, str):
|
||
raw = (
|
||
img.split(",", 1)[1] if img.startswith("data:") and "," in img else img
|
||
)
|
||
data = base64.b64decode(raw)
|
||
with open(out_path, "wb") as f:
|
||
f.write(data)
|
||
return out_path
|
||
if isinstance(img, dict) and "bytes" in img:
|
||
b = img["bytes"]
|
||
if b:
|
||
with open(out_path, "wb") as f:
|
||
f.write(b)
|
||
return out_path
|
||
except Exception as e:
|
||
logger.warning(f"Failed to save {task_name} image for {example_id}: {e}")
|
||
return None
|
||
|
||
|
||
def _save_worldvqa_query_image(example: dict, base_dir: str = "tiles") -> str | None:
|
||
"""Save WorldVQA query image to disk. Returns path or None.
|
||
Images saved to {base_dir}/worldvqa_images/{example_id}.png
|
||
"""
|
||
img = example.get("image")
|
||
if img is None:
|
||
return None
|
||
example_id = example.get("id", "unknown")
|
||
save_dir = os.path.join(base_dir, "worldvqa_images")
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
out_path = os.path.join(save_dir, f"{example_id}.png")
|
||
|
||
try:
|
||
if hasattr(img, "save"):
|
||
img.save(out_path, format="PNG")
|
||
return out_path
|
||
if isinstance(img, str):
|
||
raw = (
|
||
img.split(",", 1)[1] if img.startswith("data:") and "," in img else img
|
||
)
|
||
data = base64.b64decode(raw)
|
||
ext = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
|
||
out_path = os.path.join(save_dir, f"{example_id}{ext}")
|
||
with open(out_path, "wb") as f:
|
||
f.write(data)
|
||
return out_path
|
||
if isinstance(img, dict) and "bytes" in img:
|
||
b = img["bytes"]
|
||
if b:
|
||
ext = ".png" if b[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
|
||
out_path = os.path.join(save_dir, f"{example_id}{ext}")
|
||
with open(out_path, "wb") as f:
|
||
f.write(b)
|
||
return out_path
|
||
except Exception as e:
|
||
logger.warning(f"Failed to save WorldVQA image for {example_id}: {e}")
|
||
return None
|
||
|
||
|
||
def _worldvqa_image_to_base64(img) -> str | None:
|
||
"""Convert WorldVQA image (PIL, base64 str, or dict) to base64 string."""
|
||
if img is None:
|
||
return None
|
||
if isinstance(img, str):
|
||
if img.startswith("data:"):
|
||
if "," in img:
|
||
return img.split(",", 1)[1]
|
||
return img
|
||
if hasattr(img, "save"):
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
return base64.b64encode(buf.getvalue()).decode()
|
||
if isinstance(img, dict) and "bytes" in img:
|
||
b = img["bytes"]
|
||
return base64.b64encode(b).decode() if b else None
|
||
return None
|
||
|
||
|
||
class WorldVQANoRetrievalRetriever(BaseRetriever):
|
||
"""WorldVQA without retrieval: query + image from dataset only.
|
||
|
||
WorldVQA images are embedded in the HuggingFace dataset (PIL or base64).
|
||
"""
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
img = example.get("image")
|
||
base64_img = _worldvqa_image_to_base64(img)
|
||
return RetrievalResult(
|
||
base64_image=base64_img,
|
||
retrieval_type="worldvqa_no_retrieval",
|
||
)
|
||
|
||
|
||
class ScreenshotRetriever(BaseRetriever):
|
||
"""Use screenshot that was prepared in data layer.
|
||
|
||
Expects screenshot to be captured beforehand. This retriever just
|
||
loads and encodes the existing screenshot.
|
||
|
||
For ground truth evaluation, uses encode_screenshot_for_vlm_async which
|
||
does NOT apply max_height limit. You can control max_pixels to study
|
||
the effect of resize on VLM performance.
|
||
|
||
Args:
|
||
screenshot_dir: Directory containing screenshots.
|
||
max_pixels: Maximum pixels before resize. If None, no resize (89M limit).
|
||
Common values:
|
||
- None: No resize (let VLM handle it)
|
||
- 16_777_216 (16M): Qwen3-VL default, ~16K tokens
|
||
- 4_000_000 (4M): ~4K tokens
|
||
- 1_000_000 (1M): ~1K tokens
|
||
"""
|
||
|
||
def __init__(
|
||
self, screenshot_dir: str = "screenshots", max_pixels: int | None = None
|
||
):
|
||
self.screenshot_dir = screenshot_dir
|
||
self.max_pixels = max_pixels
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
from .simpleqa_data import (
|
||
capture_screenshot_async,
|
||
encode_screenshot_for_vlm_async,
|
||
extract_url_from_metadata,
|
||
)
|
||
|
||
# Get or capture screenshot
|
||
screenshot_path = await capture_screenshot_async(example, self.screenshot_dir)
|
||
|
||
if not screenshot_path:
|
||
return RetrievalResult(
|
||
retrieval_type="screenshot",
|
||
source_url=extract_url_from_metadata(example),
|
||
)
|
||
|
||
# Encode to base64 with configurable max_pixels
|
||
base64_image = await encode_screenshot_for_vlm_async(
|
||
screenshot_path, max_pixels=self.max_pixels
|
||
)
|
||
|
||
return RetrievalResult(
|
||
base64_image=base64_image,
|
||
source_url=extract_url_from_metadata(example),
|
||
retrieval_type="screenshot",
|
||
)
|
||
|
||
|
||
class TiledScreenshotRetriever(BaseRetriever):
|
||
"""Use tiled screenshot from ground truth URL.
|
||
|
||
Captures screenshot for the example's URL, splits it into tiles,
|
||
and returns tiles. This is ground truth (not vector search).
|
||
|
||
Args:
|
||
max_tiles: Maximum number of tiles to return. If None, returns all tiles.
|
||
For context-aware limiting, calculate based on model context length.
|
||
Rough estimate: max_tiles = (context_length - 2000) / tokens_per_tile
|
||
where tokens_per_tile ≈ 1500-2000 for most VLMs.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
screenshot_dir: str = "screenshots",
|
||
tiles_dir: str = "tiles",
|
||
tile_size: int = 512,
|
||
overlap: int = 0,
|
||
max_tiles: int | None = None,
|
||
):
|
||
self.screenshot_dir = screenshot_dir
|
||
self.tiles_dir = tiles_dir
|
||
self.tile_size = tile_size
|
||
self.overlap = overlap
|
||
self.max_tiles = max_tiles
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
from .simpleqa_data import (
|
||
capture_screenshot_async,
|
||
encode_screenshot_async,
|
||
extract_url_from_metadata,
|
||
split_image_to_tiles,
|
||
)
|
||
|
||
# Get or capture screenshot
|
||
screenshot_path = await capture_screenshot_async(example, self.screenshot_dir)
|
||
|
||
if not screenshot_path:
|
||
return RetrievalResult(
|
||
retrieval_type="tiled_screenshot",
|
||
source_url=extract_url_from_metadata(example),
|
||
)
|
||
|
||
# Split into tiles
|
||
example_id = example.get("id", "unknown")
|
||
example_tiles_dir = os.path.join(self.tiles_dir, example_id)
|
||
tile_paths = split_image_to_tiles(
|
||
screenshot_path,
|
||
example_tiles_dir,
|
||
tile_size=self.tile_size,
|
||
overlap=self.overlap,
|
||
)
|
||
|
||
if not tile_paths:
|
||
# Fall back to full screenshot
|
||
base64_image = await encode_screenshot_async(screenshot_path)
|
||
return RetrievalResult(
|
||
base64_image=base64_image,
|
||
source_url=extract_url_from_metadata(example),
|
||
retrieval_type="tiled_screenshot",
|
||
)
|
||
|
||
# Limit tiles if max_tiles is set
|
||
if self.max_tiles is not None and len(tile_paths) > self.max_tiles:
|
||
logger.info(f"Limiting tiles from {len(tile_paths)} to {self.max_tiles}")
|
||
tile_paths = tile_paths[: self.max_tiles]
|
||
|
||
# Return tiles as images list (path, score=1.0 for ground truth)
|
||
images = [(path, 1.0) for path in tile_paths]
|
||
|
||
return RetrievalResult(
|
||
images=images,
|
||
source_url=extract_url_from_metadata(example),
|
||
retrieval_type="tiled_screenshot",
|
||
)
|
||
|
||
|
||
class LocalWikiTiledScreenshotRetriever(BaseRetriever):
|
||
"""Ground-truth tiled retriever using pre-rendered Wikipedia tiles from local kiwix.
|
||
|
||
For each example, looks up the Wikipedia URL in the local kiwix tile store,
|
||
copies raw tiles to a local cache, cuts into tile_height strips, and passes
|
||
all tiles to the VLM as context. No Selenium, no SSH.
|
||
|
||
Args:
|
||
tiles_dir: Directory for cut tile strips (output).
|
||
wiki_cache_dir: Directory for raw kiwix tile copies.
|
||
tile_height: Height of each strip in pixels (default 1024).
|
||
max_tiles: Maximum tiles to pass to VLM (None = all).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
tiles_dir: str = "tiles-local-wiki",
|
||
wiki_cache_dir: str = "screenshots-localwiki",
|
||
tile_height: int = 1024,
|
||
max_tiles: int | None = None,
|
||
):
|
||
self.tiles_dir = tiles_dir
|
||
self.wiki_cache_dir = wiki_cache_dir
|
||
self.tile_height = tile_height
|
||
self.max_tiles = max_tiles
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
os.makedirs(wiki_cache_dir, exist_ok=True)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
ex_id = example.get("id", "unknown")
|
||
url = extract_url_from_metadata(example) or ""
|
||
|
||
loop = asyncio.get_event_loop()
|
||
try:
|
||
tile_paths = await loop.run_in_executor(
|
||
None,
|
||
lambda: _lookup_and_copy_local_wiki_tiles(
|
||
ex_id, url, self.tiles_dir, self.wiki_cache_dir, self.tile_height
|
||
),
|
||
)
|
||
except RuntimeError as e:
|
||
logger.error(f"local-wiki [{ex_id}]: {e}")
|
||
return RetrievalResult(retrieval_type="local_wiki_tiled", source_url=url)
|
||
|
||
if self.max_tiles is not None and len(tile_paths) > self.max_tiles:
|
||
tile_paths = tile_paths[: self.max_tiles]
|
||
|
||
images = [(path, 1.0) for path in tile_paths]
|
||
return RetrievalResult(
|
||
images=images,
|
||
source_url=url,
|
||
retrieval_type="local_wiki_tiled",
|
||
)
|
||
|
||
|
||
class TextRetriever(BaseRetriever):
|
||
"""Use text content fetched from URL.
|
||
|
||
Can use pre-cached text or fetch on demand.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
max_chars: int = 50000,
|
||
text_cache: dict | None = None,
|
||
cache_path: str | None = None,
|
||
):
|
||
self.max_chars = max_chars
|
||
self.text_cache = text_cache
|
||
self.cache_path = cache_path
|
||
self._cache_lock = asyncio.Lock()
|
||
|
||
async def _save_to_cache(self, example_id: str, text: str, url: str):
|
||
"""Append result to cache file."""
|
||
if not self.cache_path:
|
||
return
|
||
try:
|
||
import json
|
||
|
||
async with self._cache_lock:
|
||
with open(self.cache_path, "a") as f:
|
||
cache_entry = {"id": example_id, "text": text, "url": url}
|
||
f.write(json.dumps(cache_entry) + "\n")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to save to cache: {e}")
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
from .simpleqa_data import fetch_text_async
|
||
|
||
example_id = example.get("id", "")
|
||
was_cached = self.text_cache and example_id in self.text_cache
|
||
|
||
text, source_url = await fetch_text_async(
|
||
example, self.max_chars, self.text_cache
|
||
)
|
||
|
||
# Save to cache if not already cached
|
||
if not was_cached and text and source_url:
|
||
await self._save_to_cache(example_id, text, source_url)
|
||
|
||
return RetrievalResult(
|
||
text=text, source_url=source_url, retrieval_type="text_rag"
|
||
)
|
||
|
||
|
||
class JinaReaderRetriever(BaseRetriever):
|
||
"""Use Jina Reader API to fetch clean markdown text from URL.
|
||
|
||
Jina Reader (r.jina.ai) converts any URL to LLM-friendly markdown text.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
max_chars: int = 50000,
|
||
api_key: str | None = None,
|
||
text_cache: dict | None = None,
|
||
cache_path: str | None = None,
|
||
):
|
||
self.max_chars = max_chars
|
||
self.api_key = api_key
|
||
self.text_cache = text_cache
|
||
self.cache_path = cache_path
|
||
self._cache_lock = asyncio.Lock()
|
||
|
||
async def _save_to_cache(self, example_id: str, text: str, url: str):
|
||
"""Append result to cache file."""
|
||
if not self.cache_path:
|
||
return
|
||
try:
|
||
import json
|
||
|
||
async with self._cache_lock:
|
||
with open(self.cache_path, "a") as f:
|
||
cache_entry = {"id": example_id, "text": text, "url": url}
|
||
f.write(json.dumps(cache_entry) + "\n")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to save to cache: {e}")
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
import aiohttp
|
||
import asyncio
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
# Check cache first
|
||
example_id = example.get("id", "")
|
||
if self.text_cache and example_id in self.text_cache:
|
||
cached = self.text_cache[example_id]
|
||
text = cached.get("text", "")
|
||
source_url = cached.get("url", "")
|
||
if text:
|
||
if len(text) > self.max_chars:
|
||
text = text[: self.max_chars] + "\n\n[Content truncated...]"
|
||
return RetrievalResult(
|
||
text=text, source_url=source_url, retrieval_type="jina_reader"
|
||
)
|
||
|
||
target_url = extract_url_from_metadata(example)
|
||
if not target_url:
|
||
return RetrievalResult(
|
||
text="No URL found in metadata.", retrieval_type="jina_reader"
|
||
)
|
||
|
||
# Use Jina Reader API with retry logic
|
||
reader_url = f"https://r.jina.ai/{target_url}"
|
||
headers = {}
|
||
if self.api_key:
|
||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||
|
||
max_retries = 5
|
||
for attempt in range(max_retries):
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(
|
||
reader_url,
|
||
headers=headers,
|
||
timeout=aiohttp.ClientTimeout(total=60),
|
||
) as response:
|
||
# Handle rate limiting (429) with exponential backoff
|
||
if response.status == 429:
|
||
if attempt < max_retries - 1:
|
||
wait_time = min(2**attempt * 2, 30) # Max 30 seconds
|
||
logger.warning(
|
||
f"Rate limited (429) for {target_url}, waiting {wait_time}s before retry ({attempt + 1}/{max_retries})"
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
else:
|
||
error_msg = f"Jina Reader API rate limited (429) after {max_retries} retries"
|
||
logger.error(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg,
|
||
source_url=target_url,
|
||
retrieval_type="jina_reader",
|
||
)
|
||
|
||
# Handle server errors (5xx) with retry
|
||
if 500 <= response.status < 600:
|
||
if attempt < max_retries - 1:
|
||
wait_time = min(2**attempt, 10) # Max 10 seconds
|
||
logger.warning(
|
||
f"Server error ({response.status}) for {target_url}, waiting {wait_time}s before retry ({attempt + 1}/{max_retries})"
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
else:
|
||
error_msg = (
|
||
f"Jina Reader API server error: {response.status}"
|
||
)
|
||
logger.error(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg,
|
||
source_url=target_url,
|
||
retrieval_type="jina_reader",
|
||
)
|
||
|
||
# Handle client errors (4xx) - don't retry for most
|
||
if response.status == 200:
|
||
text = await response.text()
|
||
# Save to cache before truncation
|
||
await self._save_to_cache(example_id, text, target_url)
|
||
# Truncate if too long
|
||
if len(text) > self.max_chars:
|
||
text = (
|
||
text[: self.max_chars]
|
||
+ "\n\n[Content truncated...]"
|
||
)
|
||
return RetrievalResult(
|
||
text=text,
|
||
source_url=target_url,
|
||
retrieval_type="jina_reader",
|
||
)
|
||
else:
|
||
# Other 4xx errors (403, 404, etc.) - don't retry
|
||
error_msg = f"Jina Reader API error: {response.status}"
|
||
logger.warning(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg,
|
||
source_url=target_url,
|
||
retrieval_type="jina_reader",
|
||
)
|
||
except asyncio.TimeoutError:
|
||
if attempt < max_retries - 1:
|
||
wait_time = min(2**attempt, 10) # Max 10 seconds
|
||
logger.warning(
|
||
f"Timeout for {target_url}, waiting {wait_time}s before retry ({attempt + 1}/{max_retries})"
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
else:
|
||
error_msg = f"Jina Reader fetch timeout after {max_retries} retries"
|
||
logger.error(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg,
|
||
source_url=target_url,
|
||
retrieval_type="jina_reader",
|
||
)
|
||
except aiohttp.ClientError as e:
|
||
if attempt < max_retries - 1:
|
||
wait_time = min(2**attempt, 10) # Max 10 seconds
|
||
logger.warning(
|
||
f"Client error for {target_url}: {e}, waiting {wait_time}s before retry ({attempt + 1}/{max_retries})"
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
else:
|
||
error_msg = f"Jina Reader fetch failed: {e}"
|
||
logger.error(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg,
|
||
source_url=target_url,
|
||
retrieval_type="jina_reader",
|
||
)
|
||
except Exception as e:
|
||
error_msg = f"Jina Reader fetch failed: {e}"
|
||
logger.error(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg, source_url=target_url, retrieval_type="jina_reader"
|
||
)
|
||
|
||
# Should not reach here, but just in case
|
||
error_msg = f"Jina Reader fetch failed after {max_retries} retries"
|
||
return RetrievalResult(
|
||
text=error_msg, source_url=target_url, retrieval_type="jina_reader"
|
||
)
|
||
|
||
|
||
class WikipediaAPIRetriever(BaseRetriever):
|
||
"""Use Wikipedia API to fetch clean article text.
|
||
|
||
Extracts Wikipedia page title from URL and fetches content via API.
|
||
Much cleaner and faster than web scraping.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
max_chars: int = 50000,
|
||
text_cache: dict | None = None,
|
||
cache_path: str | None = None,
|
||
):
|
||
self.max_chars = max_chars
|
||
self.text_cache = text_cache
|
||
self.cache_path = cache_path
|
||
self._cache_lock = asyncio.Lock()
|
||
|
||
async def _save_to_cache(self, example_id: str, text: str, url: str):
|
||
"""Append result to cache file."""
|
||
if not self.cache_path:
|
||
return
|
||
try:
|
||
import json
|
||
|
||
async with self._cache_lock:
|
||
with open(self.cache_path, "a") as f:
|
||
cache_entry = {"id": example_id, "text": text, "url": url}
|
||
f.write(json.dumps(cache_entry) + "\n")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to save to cache: {e}")
|
||
|
||
def _extract_wiki_title(self, url: str) -> str | None:
|
||
"""Extract Wikipedia page title from URL."""
|
||
import re
|
||
from urllib.parse import unquote
|
||
|
||
# Match patterns like:
|
||
# https://en.wikipedia.org/wiki/Python_(programming_language)
|
||
# https://zh.wikipedia.org/wiki/Artificial_intelligence
|
||
pattern = r"https?://[a-z]{2,3}\.wikipedia\.org/wiki/(.+?)(?:#.*)?$"
|
||
match = re.match(pattern, url)
|
||
if match:
|
||
title = unquote(match.group(1))
|
||
# Replace underscores with spaces
|
||
title = title.replace("_", " ")
|
||
return title
|
||
return None
|
||
|
||
def _get_wiki_lang(self, url: str) -> str:
|
||
"""Extract Wikipedia language code from URL."""
|
||
import re
|
||
|
||
match = re.match(r"https?://([a-z]{2,3})\.wikipedia\.org", url)
|
||
return match.group(1) if match else "en"
|
||
|
||
def _html_to_text(self, html: str) -> str:
|
||
"""Convert Wikipedia HTML to plain text, preserving table content."""
|
||
import re
|
||
from bs4 import BeautifulSoup
|
||
|
||
soup = BeautifulSoup(html, "html.parser")
|
||
|
||
# Remove unwanted elements
|
||
for tag in soup.find_all(["script", "style", "link", "meta"]):
|
||
tag.decompose()
|
||
|
||
# Remove edit section links
|
||
for tag in soup.find_all("span", class_="mw-editsection"):
|
||
tag.decompose()
|
||
|
||
# Remove reference numbers [1], [2], etc.
|
||
for tag in soup.find_all("sup", class_="reference"):
|
||
tag.decompose()
|
||
|
||
# Get text
|
||
text = soup.get_text(separator="\n", strip=True)
|
||
|
||
# Clean up excessive newlines
|
||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||
|
||
return text
|
||
|
||
def _parse_infobox(self, wikitext: str) -> str:
|
||
"""Parse infobox from wikitext and convert to plain text."""
|
||
import re
|
||
|
||
# Find infobox start
|
||
start = wikitext.find("{{Infobox")
|
||
if start == -1:
|
||
start = wikitext.find("{{infobox")
|
||
if start == -1:
|
||
return ""
|
||
|
||
# Count braces to find matching end
|
||
depth = 0
|
||
end = start
|
||
for i in range(start, len(wikitext)):
|
||
if wikitext[i : i + 2] == "{{":
|
||
depth += 1
|
||
elif wikitext[i : i + 2] == "}}":
|
||
depth -= 1
|
||
if depth == 0:
|
||
end = i + 2
|
||
break
|
||
|
||
infobox_raw = wikitext[start:end]
|
||
|
||
# Parse fields
|
||
lines = []
|
||
for match in re.finditer(
|
||
r"\|\s*([^=|]+?)\s*=\s*([^|]*?)(?=\n\s*\||\}\})", infobox_raw, re.DOTALL
|
||
):
|
||
key = match.group(1).strip()
|
||
value = match.group(2).strip()
|
||
|
||
# Skip image-related fields
|
||
if key.lower() in (
|
||
"image",
|
||
"caption",
|
||
"alt",
|
||
"width",
|
||
"height",
|
||
"image_size",
|
||
"imagesize",
|
||
):
|
||
continue
|
||
|
||
# Clean up wikitext markup
|
||
value = re.sub(
|
||
r"\{\{[^}|]*\|([^}]*)\}\}", r"\1", value
|
||
) # {{template|value}} -> value
|
||
value = re.sub(
|
||
r"\[\[([^|\]]*\|)?([^\]]*)\]\]", r"\2", value
|
||
) # [[link|text]] -> text
|
||
value = re.sub(r"'''?", "", value) # bold/italic
|
||
value = re.sub(r"<[^>]+>", "", value) # HTML tags
|
||
value = re.sub(r"\{\{[^}]*\}\}", "", value) # remaining templates
|
||
value = " ".join(value.split()) # normalize whitespace
|
||
|
||
if value:
|
||
lines.append(f"{key}: {value}")
|
||
|
||
return "\n".join(lines)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
import aiohttp
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
# Check cache first
|
||
example_id = example.get("id", "")
|
||
if self.text_cache and example_id in self.text_cache:
|
||
cached = self.text_cache[example_id]
|
||
text = cached.get("text", "")
|
||
source_url = cached.get("url", "")
|
||
if text:
|
||
if len(text) > self.max_chars:
|
||
text = text[: self.max_chars] + "\n\n[Content truncated...]"
|
||
return RetrievalResult(
|
||
text=text, source_url=source_url, retrieval_type="wikipedia_api"
|
||
)
|
||
|
||
target_url = extract_url_from_metadata(example)
|
||
if not target_url:
|
||
return RetrievalResult(
|
||
text="No URL found in metadata.", retrieval_type="wikipedia_api"
|
||
)
|
||
|
||
# Check if it's a Wikipedia URL
|
||
if "wikipedia.org" not in target_url.lower():
|
||
return RetrievalResult(
|
||
text=f"URL is not a Wikipedia page: {target_url}",
|
||
source_url=target_url,
|
||
retrieval_type="wikipedia_api",
|
||
)
|
||
|
||
title = self._extract_wiki_title(target_url)
|
||
if not title:
|
||
return RetrievalResult(
|
||
text=f"Could not extract Wikipedia title from: {target_url}",
|
||
source_url=target_url,
|
||
retrieval_type="wikipedia_api",
|
||
)
|
||
|
||
lang = self._get_wiki_lang(target_url)
|
||
api_url = f"https://{lang}.wikipedia.org/w/api.php"
|
||
|
||
headers = {
|
||
"User-Agent": "SimpleQA-Evaluation/1.0 (https://github.com/example; contact@example.com)"
|
||
}
|
||
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
# Use action=parse to get full HTML (includes tables)
|
||
parse_params = {
|
||
"action": "parse",
|
||
"page": title,
|
||
"prop": "text",
|
||
"format": "json",
|
||
"redirects": "1",
|
||
}
|
||
|
||
timeout = aiohttp.ClientTimeout(total=30)
|
||
async with session.get(
|
||
api_url, params=parse_params, headers=headers, timeout=timeout
|
||
) as resp:
|
||
if resp.status != 200:
|
||
error_msg = f"Wikipedia API error: {resp.status}"
|
||
logger.warning(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg,
|
||
source_url=target_url,
|
||
retrieval_type="wikipedia_api",
|
||
)
|
||
|
||
data = await resp.json()
|
||
|
||
# Check for error
|
||
if "error" in data:
|
||
error_msg = data["error"].get("info", "Unknown error")
|
||
return RetrievalResult(
|
||
text=f"Wikipedia page not found: {title}",
|
||
source_url=target_url,
|
||
retrieval_type="wikipedia_api",
|
||
)
|
||
|
||
html = data.get("parse", {}).get("text", {}).get("*", "")
|
||
if not html:
|
||
return RetrievalResult(
|
||
text=f"No content found for Wikipedia page: {title}",
|
||
source_url=target_url,
|
||
retrieval_type="wikipedia_api",
|
||
)
|
||
|
||
# Parse HTML to text (includes tables)
|
||
text = self._html_to_text(html)
|
||
|
||
# Save to cache before truncation
|
||
await self._save_to_cache(example_id, text, target_url)
|
||
|
||
# Truncate if too long
|
||
if len(text) > self.max_chars:
|
||
text = text[: self.max_chars] + "\n\n[Content truncated...]"
|
||
|
||
return RetrievalResult(
|
||
text=text, source_url=target_url, retrieval_type="wikipedia_api"
|
||
)
|
||
except Exception as e:
|
||
error_msg = f"Wikipedia API fetch failed: {e}"
|
||
logger.warning(f"{error_msg} for {target_url}")
|
||
return RetrievalResult(
|
||
text=error_msg, source_url=target_url, retrieval_type="wikipedia_api"
|
||
)
|
||
|
||
|
||
class VectorRetriever(BaseRetriever):
|
||
"""Retrieve similar screenshots using vector similarity search.
|
||
|
||
Uses Jina API for embedding and retrieval across dataset screenshots only.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str,
|
||
screenshot_dir: str = "screenshots",
|
||
cache_path: str | None = None,
|
||
use_multivector: bool = True,
|
||
top_k: int = 3,
|
||
examples: list[dict] | None = None,
|
||
):
|
||
self.top_k = top_k
|
||
self.screenshot_dir = screenshot_dir
|
||
self.examples = examples or []
|
||
os.makedirs(screenshot_dir, exist_ok=True)
|
||
|
||
# Prepare missing screenshots and get file paths
|
||
screenshot_paths = self._prepare_screenshots()
|
||
|
||
# Import retrieval system
|
||
try:
|
||
from scripts.jina_retrieval import JinaAPIRetrievalSystem
|
||
except ImportError:
|
||
try:
|
||
from jina_retrieval import JinaAPIRetrievalSystem
|
||
except ImportError:
|
||
raise ImportError("JinaAPIRetrievalSystem not available")
|
||
|
||
vector_type = "single vector" if not use_multivector else "multivector"
|
||
logger.info(f"Initializing VectorRetriever with {vector_type} mode")
|
||
|
||
self.retrieval_system = JinaAPIRetrievalSystem(
|
||
api_key=api_key,
|
||
use_multivector=use_multivector,
|
||
device="cpu", # Use CPU to avoid OOM when VLM is on GPU
|
||
)
|
||
# Only embed screenshots for current dataset
|
||
self.retrieval_system.embed_images(
|
||
file_paths=screenshot_paths, cache_path=cache_path
|
||
)
|
||
logger.info(
|
||
f"VectorRetriever ready with {len(self.retrieval_system.image_paths)} images"
|
||
)
|
||
|
||
def _prepare_screenshots(self) -> list[str]:
|
||
"""Prepare screenshots for dataset and return list of paths."""
|
||
from .simpleqa_data import capture_screenshot_for_example
|
||
|
||
screenshot_paths = []
|
||
missing = []
|
||
|
||
for ex in self.examples:
|
||
screenshot_path = os.path.join(
|
||
self.screenshot_dir, f"{ex['id']}_fullhd.png"
|
||
)
|
||
screenshot_paths.append(screenshot_path)
|
||
if (
|
||
not os.path.exists(screenshot_path)
|
||
or os.path.getsize(screenshot_path) == 0
|
||
):
|
||
missing.append(ex)
|
||
|
||
if missing:
|
||
logger.info(
|
||
f"Found {len(missing)} missing screenshots out of {len(self.examples)} total examples"
|
||
)
|
||
logger.info(f"Preparing {len(missing)} missing screenshots...")
|
||
# Use a more robust approach: continue even if some screenshots fail
|
||
success_count = 0
|
||
for ex in missing:
|
||
try:
|
||
capture_screenshot_for_example(ex, self.screenshot_dir)
|
||
success_count += 1
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"Failed to capture screenshot for {ex.get('id', 'unknown')}: {e}"
|
||
)
|
||
# Continue with next screenshot instead of failing completely
|
||
logger.info(
|
||
f"Screenshots prepared: {success_count}/{len(missing)} successful"
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"All {len(self.examples)} screenshots already exist, skipping preparation"
|
||
)
|
||
|
||
# Return only existing screenshots
|
||
return [
|
||
p for p in screenshot_paths if os.path.exists(p) and os.path.getsize(p) > 0
|
||
]
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
loop = asyncio.get_event_loop()
|
||
|
||
try:
|
||
results = await loop.run_in_executor(
|
||
None, self.retrieval_system.retrieve, query, self.top_k
|
||
)
|
||
|
||
if results:
|
||
return RetrievalResult(images=results, retrieval_type="vector")
|
||
except Exception as e:
|
||
logger.warning(f"Vector retrieval failed: {e}")
|
||
|
||
return RetrievalResult(retrieval_type="vector")
|
||
|
||
|
||
class ColQwenVectorRetriever(BaseRetriever):
|
||
"""Retrieve similar screenshots using ColQwen2 LEANN multi-vector retrieval."""
|
||
|
||
def __init__(
|
||
self,
|
||
index_path: str,
|
||
screenshot_dir: str = "screenshots",
|
||
model_name: str = "colqwen2",
|
||
search_method: str = "ann",
|
||
first_stage_k: int = 500,
|
||
rebuild_index: bool = False,
|
||
recursive: bool = False,
|
||
top_k: int = 3,
|
||
examples: list[dict] | None = None,
|
||
prepare_screenshots: bool = False, # ColQwen2 doesn't need to prepare specific screenshots
|
||
):
|
||
self.top_k = top_k
|
||
self.screenshot_dir = screenshot_dir
|
||
self.examples = examples or []
|
||
os.makedirs(screenshot_dir, exist_ok=True)
|
||
|
||
# Build list of image paths for the specific examples (only Wikipedia samples)
|
||
image_paths = self._get_example_image_paths()
|
||
|
||
if image_paths:
|
||
logger.info(
|
||
f"ColQwen2 will retrieve from {len(image_paths)} images for {len(self.examples)} examples"
|
||
)
|
||
else:
|
||
logger.warning(
|
||
f"No images found for examples, falling back to all images in: {screenshot_dir}"
|
||
)
|
||
|
||
# Import ColQwen2 retrieval system
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# Add scripts directory to path for import
|
||
scripts_dir = Path(__file__).parent.parent
|
||
if str(scripts_dir) not in sys.path:
|
||
sys.path.insert(0, str(scripts_dir))
|
||
|
||
try:
|
||
from colqwen_leann_retrieval import ColQwenLEANNRetrievalSystem
|
||
except ImportError:
|
||
try:
|
||
from scripts.colqwen_leann_retrieval import ColQwenLEANNRetrievalSystem
|
||
except ImportError:
|
||
raise ImportError(
|
||
"ColQwenLEANNRetrievalSystem not available. Make sure colqwen_leann_retrieval.py is in the scripts directory."
|
||
)
|
||
|
||
logger.info("Initializing ColQwen2 LEANN retrieval system...")
|
||
logger.info(f"Search method: {search_method}")
|
||
|
||
# Use filtered image paths if available, otherwise fall back to directory scanning
|
||
if image_paths:
|
||
self.retrieval_system = ColQwenLEANNRetrievalSystem(
|
||
index_path=index_path,
|
||
model_name=model_name,
|
||
search_method=search_method,
|
||
first_stage_k=first_stage_k,
|
||
rebuild_index=rebuild_index,
|
||
custom_image_paths=image_paths, # Pass specific image paths
|
||
)
|
||
else:
|
||
self.retrieval_system = ColQwenLEANNRetrievalSystem(
|
||
index_path=index_path,
|
||
model_name=model_name,
|
||
search_method=search_method,
|
||
first_stage_k=first_stage_k,
|
||
rebuild_index=rebuild_index,
|
||
custom_folder_path=screenshot_dir,
|
||
custom_folder_recursive=recursive,
|
||
)
|
||
logger.info("ColQwen2 LEANN retrieval system ready")
|
||
|
||
def _get_example_image_paths(self) -> list[str]:
|
||
"""Get image paths for the specific examples."""
|
||
image_paths = []
|
||
for ex in self.examples:
|
||
example_id = ex.get("id", "")
|
||
if not example_id:
|
||
continue
|
||
path = os.path.join(self.screenshot_dir, f"{example_id}_fullhd.png")
|
||
if os.path.exists(path) and os.path.getsize(path) > 0:
|
||
image_paths.append(path)
|
||
return image_paths
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
loop = asyncio.get_event_loop()
|
||
|
||
try:
|
||
results = await loop.run_in_executor(
|
||
None, self.retrieval_system.retrieve, query, self.top_k
|
||
)
|
||
|
||
if results:
|
||
return RetrievalResult(images=results, retrieval_type="colqwen_vector")
|
||
except Exception as e:
|
||
logger.warning(f"ColQwen2 vector retrieval failed: {e}")
|
||
|
||
return RetrievalResult(retrieval_type="colqwen_vector")
|
||
|
||
|
||
def _filter_tiles_by_aspect_ratio(
|
||
tile_paths: list[str], max_aspect_ratio: float = 100.0
|
||
) -> list[str]:
|
||
"""Filter out tiles with extreme aspect ratios.
|
||
|
||
Args:
|
||
tile_paths: List of tile image paths.
|
||
max_aspect_ratio: Maximum allowed aspect ratio (default 100, ColQwen requires < 200).
|
||
|
||
Returns:
|
||
Filtered list of tile paths.
|
||
"""
|
||
from PIL import Image
|
||
|
||
filtered = []
|
||
for tile_path in tile_paths:
|
||
try:
|
||
with Image.open(tile_path) as img:
|
||
w, h = img.size
|
||
if w > 0 and h > 0:
|
||
aspect_ratio = max(w / h, h / w)
|
||
if aspect_ratio <= max_aspect_ratio:
|
||
filtered.append(tile_path)
|
||
else:
|
||
logger.warning(
|
||
f"Skipping tile with extreme aspect ratio {aspect_ratio:.2f}: {tile_path}"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to check tile {tile_path}: {e}")
|
||
|
||
return filtered
|
||
|
||
|
||
class TiledVectorRetriever(BaseRetriever):
|
||
"""Retrieve similar image tiles using vector similarity search.
|
||
|
||
Splits dataset screenshots into fixed-size tiles, embeds each tile,
|
||
and retrieves the most relevant tiles for a query.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str,
|
||
screenshot_dir: str = "screenshots",
|
||
tiles_dir: str = "tiles",
|
||
tile_size: int = 512,
|
||
overlap: int = 0,
|
||
cache_path: str | None = None,
|
||
use_multivector: bool = True,
|
||
top_k: int = 3,
|
||
examples: list[dict] | None = None,
|
||
):
|
||
self.top_k = top_k
|
||
self.screenshot_dir = screenshot_dir
|
||
self.tiles_dir = tiles_dir
|
||
self.tile_size = tile_size
|
||
self.overlap = overlap
|
||
self.examples = examples or []
|
||
os.makedirs(screenshot_dir, exist_ok=True)
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
|
||
# Build example_id -> URL mapping (prioritize Wikipedia URLs)
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
self.id_to_url = {}
|
||
for ex in self.examples:
|
||
ex_id = ex.get("id", "")
|
||
url = extract_url_from_metadata(ex) # Uses Wikipedia-first priority
|
||
if url:
|
||
self.id_to_url[ex_id] = url
|
||
|
||
# Prepare screenshots and get tile paths
|
||
tile_paths = self._prepare_screenshots_and_tiles()
|
||
|
||
# Import retrieval system
|
||
try:
|
||
from scripts.jina_retrieval import JinaAPIRetrievalSystem
|
||
except ImportError:
|
||
try:
|
||
from jina_retrieval import JinaAPIRetrievalSystem
|
||
except ImportError:
|
||
raise ImportError("JinaAPIRetrievalSystem not available")
|
||
|
||
vector_type = "single vector" if not use_multivector else "multivector"
|
||
logger.info(f"Initializing TiledVectorRetriever with {vector_type} mode")
|
||
|
||
self.retrieval_system = JinaAPIRetrievalSystem(
|
||
api_key=api_key,
|
||
use_multivector=use_multivector,
|
||
device="cpu", # Use CPU to avoid OOM when VLM is on GPU
|
||
)
|
||
# Only embed tiles for current dataset
|
||
self.retrieval_system.embed_images(file_paths=tile_paths, cache_path=cache_path)
|
||
logger.info(
|
||
f"TiledVectorRetriever ready with {len(self.retrieval_system.image_paths)} tiles"
|
||
)
|
||
|
||
def _prepare_screenshots_and_tiles(self) -> list[str]:
|
||
"""Prepare screenshots and tiles for dataset, return tile paths."""
|
||
from .simpleqa_data import capture_screenshot_for_example, split_image_to_tiles
|
||
from tqdm import tqdm
|
||
|
||
screenshot_paths = []
|
||
missing = []
|
||
|
||
# Collect screenshot paths and identify missing
|
||
for ex in self.examples:
|
||
screenshot_path = os.path.join(
|
||
self.screenshot_dir, f"{ex['id']}_fullhd.png"
|
||
)
|
||
screenshot_paths.append(screenshot_path)
|
||
if (
|
||
not os.path.exists(screenshot_path)
|
||
or os.path.getsize(screenshot_path) == 0
|
||
):
|
||
missing.append(ex)
|
||
|
||
# Capture missing screenshots
|
||
if missing:
|
||
logger.info(f"Preparing {len(missing)} missing screenshots...")
|
||
for ex in tqdm(missing, desc="Capturing screenshots"):
|
||
capture_screenshot_for_example(ex, self.screenshot_dir)
|
||
logger.info("Screenshots prepared.")
|
||
|
||
# Split each screenshot into tiles
|
||
all_tile_paths = []
|
||
logger.info(
|
||
f"Splitting {len(screenshot_paths)} screenshots into tiles (output: {self.tiles_dir})..."
|
||
)
|
||
for screenshot_path in tqdm(screenshot_paths, desc="Splitting tiles"):
|
||
if os.path.exists(screenshot_path) and os.path.getsize(screenshot_path) > 0:
|
||
tile_paths = split_image_to_tiles(
|
||
screenshot_path, self.tiles_dir, self.tile_size, self.overlap
|
||
)
|
||
all_tile_paths.extend(tile_paths)
|
||
|
||
# Filter out tiles with extreme aspect ratios
|
||
filtered_tile_paths = _filter_tiles_by_aspect_ratio(all_tile_paths)
|
||
logger.info(
|
||
f"Prepared {len(filtered_tile_paths)} tiles from {len(screenshot_paths)} screenshots (filtered {len(all_tile_paths) - len(filtered_tile_paths)} extreme aspect ratio tiles)"
|
||
)
|
||
return filtered_tile_paths
|
||
|
||
def _extract_urls_from_results(self, results: list) -> str:
|
||
"""Extract source URLs from tile paths in results, preserving retrieval order."""
|
||
urls = []
|
||
seen = set()
|
||
for item in results:
|
||
# item is (path, score) tuple
|
||
path = item[0] if isinstance(item, tuple) else item
|
||
# Extract example_id from tile path: {example_id}_fullhd_tile_{x}_{y}.png
|
||
filename = os.path.basename(path)
|
||
# Split by _fullhd_ or just get the first part before _tile_
|
||
if "_tile_" in filename:
|
||
example_id = filename.split("_tile_")[0]
|
||
# Remove _fullhd suffix if present
|
||
if example_id.endswith("_fullhd"):
|
||
example_id = example_id[:-7]
|
||
if example_id in self.id_to_url:
|
||
url = self.id_to_url[example_id]
|
||
if url not in seen:
|
||
seen.add(url)
|
||
urls.append(url)
|
||
return ", ".join(urls)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
del example # Not used - retrieval is from pre-built index
|
||
loop = asyncio.get_event_loop()
|
||
|
||
try:
|
||
results = await loop.run_in_executor(
|
||
None, self.retrieval_system.retrieve, query, self.top_k
|
||
)
|
||
|
||
if results:
|
||
source_url = self._extract_urls_from_results(results)
|
||
return RetrievalResult(
|
||
images=results, source_url=source_url, retrieval_type="tiled_vector"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Tiled vector retrieval failed: {e}")
|
||
|
||
return RetrievalResult(retrieval_type="tiled_vector")
|
||
|
||
|
||
class TiledColQwenVectorRetriever(BaseRetriever):
|
||
"""Retrieve similar image tiles using ColQwen2 LEANN multi-vector retrieval.
|
||
|
||
Splits dataset screenshots into fixed-size tiles, embeds each tile with ColQwen2,
|
||
and retrieves the most relevant tiles for a query using LEANN.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
index_path: str,
|
||
screenshot_dir: str = "screenshots",
|
||
tiles_dir: str = "tiles",
|
||
tile_size: int = 512,
|
||
overlap: int = 0,
|
||
model_name: str = "colqwen2",
|
||
search_method: str = "ann",
|
||
first_stage_k: int = 500,
|
||
rebuild_index: bool = False,
|
||
top_k: int = 3,
|
||
examples: list[dict] | None = None,
|
||
):
|
||
self.top_k = top_k
|
||
self.screenshot_dir = screenshot_dir
|
||
self.tiles_dir = tiles_dir
|
||
self.tile_size = tile_size
|
||
self.overlap = overlap
|
||
self.examples = examples or []
|
||
os.makedirs(screenshot_dir, exist_ok=True)
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
|
||
# Build example_id -> URL mapping (prioritize Wikipedia URLs)
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
self.id_to_url = {}
|
||
for ex in self.examples:
|
||
ex_id = ex.get("id", "")
|
||
url = extract_url_from_metadata(ex) # Uses Wikipedia-first priority
|
||
if url:
|
||
self.id_to_url[ex_id] = url
|
||
|
||
# Prepare screenshots and get tile paths
|
||
tile_paths = self._prepare_screenshots_and_tiles()
|
||
|
||
# Import ColQwen2 retrieval system
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
scripts_dir = Path(__file__).parent.parent
|
||
if str(scripts_dir) not in sys.path:
|
||
sys.path.insert(0, str(scripts_dir))
|
||
|
||
try:
|
||
from colqwen_leann_retrieval import ColQwenLEANNRetrievalSystem
|
||
except ImportError:
|
||
try:
|
||
from scripts.colqwen_leann_retrieval import ColQwenLEANNRetrievalSystem
|
||
except ImportError:
|
||
raise ImportError("ColQwenLEANNRetrievalSystem not available.")
|
||
|
||
logger.info("Initializing TiledColQwen2 LEANN retrieval system...")
|
||
logger.info(f"Search method: {search_method}, tiles: {len(tile_paths)}")
|
||
|
||
self.retrieval_system = ColQwenLEANNRetrievalSystem(
|
||
index_path=index_path,
|
||
custom_image_paths=tile_paths,
|
||
model_name=model_name,
|
||
search_method=search_method,
|
||
first_stage_k=first_stage_k,
|
||
rebuild_index=rebuild_index,
|
||
)
|
||
|
||
logger.info(
|
||
f"TiledColQwen2 LEANN retrieval system ready with {len(tile_paths)} tiles"
|
||
)
|
||
|
||
def _prepare_screenshots_and_tiles(self) -> list[str]:
|
||
"""Prepare screenshots and tiles for dataset, return tile paths."""
|
||
from .simpleqa_data import capture_screenshot_for_example, split_image_to_tiles
|
||
from tqdm import tqdm
|
||
|
||
screenshot_paths = []
|
||
missing = []
|
||
|
||
# Collect screenshot paths and identify missing
|
||
for ex in self.examples:
|
||
screenshot_path = os.path.join(
|
||
self.screenshot_dir, f"{ex['id']}_fullhd.png"
|
||
)
|
||
screenshot_paths.append(screenshot_path)
|
||
if (
|
||
not os.path.exists(screenshot_path)
|
||
or os.path.getsize(screenshot_path) == 0
|
||
):
|
||
missing.append(ex)
|
||
|
||
# Capture missing screenshots
|
||
if missing:
|
||
logger.info(f"Preparing {len(missing)} missing screenshots...")
|
||
for ex in tqdm(missing, desc="Capturing screenshots"):
|
||
capture_screenshot_for_example(ex, self.screenshot_dir)
|
||
logger.info("Screenshots prepared.")
|
||
|
||
# Split each screenshot into tiles
|
||
all_tile_paths = []
|
||
logger.info(
|
||
f"Splitting {len(screenshot_paths)} screenshots into tiles (output: {self.tiles_dir})..."
|
||
)
|
||
for screenshot_path in tqdm(screenshot_paths, desc="Splitting tiles"):
|
||
if os.path.exists(screenshot_path) and os.path.getsize(screenshot_path) > 0:
|
||
tile_paths = split_image_to_tiles(
|
||
screenshot_path, self.tiles_dir, self.tile_size, self.overlap
|
||
)
|
||
all_tile_paths.extend(tile_paths)
|
||
|
||
# Filter out tiles with extreme aspect ratios
|
||
filtered_tile_paths = _filter_tiles_by_aspect_ratio(all_tile_paths)
|
||
logger.info(
|
||
f"Prepared {len(filtered_tile_paths)} tiles from {len(screenshot_paths)} screenshots (filtered {len(all_tile_paths) - len(filtered_tile_paths)} extreme aspect ratio tiles)"
|
||
)
|
||
return filtered_tile_paths
|
||
|
||
def _extract_urls_from_results(self, results: list) -> str:
|
||
"""Extract source URLs from tile paths in results, preserving retrieval order."""
|
||
urls = []
|
||
seen = set()
|
||
for item in results:
|
||
# item is (path, score) tuple
|
||
path = item[0] if isinstance(item, tuple) else item
|
||
# Extract example_id from tile path: {example_id}_fullhd_tile_{x}_{y}.png
|
||
filename = os.path.basename(path)
|
||
if "_tile_" in filename:
|
||
example_id = filename.split("_tile_")[0]
|
||
if example_id.endswith("_fullhd"):
|
||
example_id = example_id[:-7]
|
||
if example_id in self.id_to_url:
|
||
url = self.id_to_url[example_id]
|
||
if url not in seen:
|
||
seen.add(url)
|
||
urls.append(url)
|
||
return ", ".join(urls)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
del example # Not used - retrieval is from pre-built index
|
||
loop = asyncio.get_event_loop()
|
||
|
||
try:
|
||
results = await loop.run_in_executor(
|
||
None, self.retrieval_system.retrieve, query, self.top_k
|
||
)
|
||
|
||
if results:
|
||
source_url = self._extract_urls_from_results(results)
|
||
return RetrievalResult(
|
||
images=results,
|
||
source_url=source_url,
|
||
retrieval_type="tiled_colqwen_vector",
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"TiledColQwen2 vector retrieval failed: {e}")
|
||
|
||
return RetrievalResult(retrieval_type="tiled_colqwen_vector")
|
||
|
||
|
||
class TextVectorRetriever(BaseRetriever):
|
||
"""Retrieve text using LEANN vector search.
|
||
|
||
Uses LEANN's integrated embedding + indexing system for text retrieval.
|
||
Supports various embedding models (Qwen3, nomic-embed-text, OpenAI, etc.)
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
text_cache: dict,
|
||
index_path: str,
|
||
embedding_model: str = "Qwen/Qwen3-Embedding-0.6B",
|
||
embedding_mode: str = "sentence-transformers",
|
||
embedding_options: dict | None = None,
|
||
top_k: int = 3,
|
||
rebuild_index: bool = False,
|
||
chunk_size: int = 512,
|
||
chunk_overlap: int = 128,
|
||
):
|
||
"""Initialize TextVectorRetriever.
|
||
|
||
Args:
|
||
text_cache: Dict of {id: {"text": ..., "url": ...}}
|
||
index_path: Path to LEANN index
|
||
embedding_model: Embedding model name (default: Qwen/Qwen3-Embedding-0.6B)
|
||
embedding_mode: Embedding mode (sentence-transformers, openai, gemini, ollama)
|
||
embedding_options: Additional options for embedding (e.g., base_url, api_key for OpenAI-compatible APIs)
|
||
top_k: Number of results to retrieve
|
||
rebuild_index: Force rebuild index even if exists
|
||
chunk_size: Max tokens per chunk (default: 512)
|
||
chunk_overlap: Overlap tokens between chunks (default: 128)
|
||
"""
|
||
import sys
|
||
from pathlib import Path as PathLib
|
||
|
||
# Add LEANN to path
|
||
leann_path = (
|
||
PathLib(__file__).parent.parent.parent
|
||
/ "LEANN"
|
||
/ "packages"
|
||
/ "leann-core"
|
||
/ "src"
|
||
)
|
||
if str(leann_path) not in sys.path:
|
||
sys.path.insert(0, str(leann_path))
|
||
|
||
from leann.api import LeannBuilder, LeannSearcher
|
||
|
||
self.text_cache = text_cache
|
||
self.top_k = top_k
|
||
self.index_path = index_path
|
||
self.embedding_model = embedding_model
|
||
self.embedding_mode = embedding_mode
|
||
self.embedding_options = embedding_options or {}
|
||
self.chunk_size = chunk_size
|
||
self.chunk_overlap = chunk_overlap
|
||
|
||
# Check if index exists
|
||
meta_path = f"{index_path}.meta.json"
|
||
index_exists = os.path.exists(meta_path)
|
||
|
||
if rebuild_index or not index_exists:
|
||
logger.info(f"Building LEANN text index at {index_path}...")
|
||
self._build_index(LeannBuilder)
|
||
logger.info(f"LEANN text index built with {len(text_cache)} documents")
|
||
else:
|
||
logger.info(f"Loading existing LEANN text index from {index_path}")
|
||
|
||
# Load searcher
|
||
self.searcher = LeannSearcher(index_path)
|
||
logger.info(
|
||
f"TextVectorRetriever ready with {len(text_cache)} documents, top_k={top_k}"
|
||
)
|
||
|
||
def _build_index(self, LeannBuilder):
|
||
"""Build LEANN index from text_cache with chunking for long texts."""
|
||
builder = LeannBuilder(
|
||
backend_name="hnsw",
|
||
embedding_model=self.embedding_model,
|
||
embedding_mode=self.embedding_mode,
|
||
embedding_options=self.embedding_options,
|
||
is_recompute=False, # Store embeddings to avoid recomputing at search time
|
||
)
|
||
|
||
# Chunking parameters (from CLI or defaults)
|
||
max_tokens = self.chunk_size
|
||
overlap_tokens = self.chunk_overlap
|
||
|
||
# Import tiktoken for accurate chunking
|
||
try:
|
||
import tiktoken
|
||
|
||
enc = tiktoken.get_encoding("cl100k_base")
|
||
except ImportError:
|
||
enc = None
|
||
logger.warning("tiktoken not available, using character-based chunking")
|
||
|
||
chunk_count = 0
|
||
for example_id, data in self.text_cache.items():
|
||
text = data.get("text", "")
|
||
url = data.get("url", "")
|
||
if not text:
|
||
continue
|
||
|
||
if enc:
|
||
# Token-based chunking
|
||
tokens = enc.encode(text)
|
||
if len(tokens) <= max_tokens:
|
||
# Short text, add as single passage
|
||
builder.add_text(text, metadata={"id": example_id, "url": url})
|
||
chunk_count += 1
|
||
else:
|
||
# Long text, chunk it with overlap
|
||
start = 0
|
||
chunk_idx = 0
|
||
while start < len(tokens):
|
||
end = min(start + max_tokens, len(tokens))
|
||
chunk_tokens = tokens[start:end]
|
||
chunk_text = enc.decode(chunk_tokens)
|
||
|
||
chunk_id = f"{example_id}_chunk_{chunk_idx}"
|
||
builder.add_text(
|
||
chunk_text,
|
||
metadata={
|
||
"id": chunk_id,
|
||
"original_id": example_id,
|
||
"url": url,
|
||
"chunk_idx": chunk_idx,
|
||
},
|
||
)
|
||
chunk_count += 1
|
||
chunk_idx += 1
|
||
|
||
if end >= len(tokens):
|
||
break
|
||
start = end - overlap_tokens # Overlap
|
||
else:
|
||
# Fallback: character-based chunking (~4 chars per token)
|
||
max_chars = max_tokens * 4
|
||
overlap_chars = overlap_tokens * 4
|
||
|
||
if len(text) <= max_chars:
|
||
builder.add_text(text, metadata={"id": example_id, "url": url})
|
||
chunk_count += 1
|
||
else:
|
||
start = 0
|
||
chunk_idx = 0
|
||
while start < len(text):
|
||
end = min(start + max_chars, len(text))
|
||
chunk_text = text[start:end]
|
||
|
||
chunk_id = f"{example_id}_chunk_{chunk_idx}"
|
||
builder.add_text(
|
||
chunk_text,
|
||
metadata={
|
||
"id": chunk_id,
|
||
"original_id": example_id,
|
||
"url": url,
|
||
"chunk_idx": chunk_idx,
|
||
},
|
||
)
|
||
chunk_count += 1
|
||
chunk_idx += 1
|
||
|
||
if end >= len(text):
|
||
break
|
||
start = end - overlap_chars
|
||
|
||
logger.info(
|
||
f"Created {chunk_count} chunks from {len(self.text_cache)} documents"
|
||
)
|
||
|
||
# Build index
|
||
builder.build_index(self.index_path)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
"""Retrieve relevant texts using LEANN vector search."""
|
||
del example # Not used - retrieval is from pre-built index
|
||
loop = asyncio.get_event_loop()
|
||
|
||
try:
|
||
# Run search in executor (LEANN search is sync)
|
||
results = await loop.run_in_executor(
|
||
None,
|
||
lambda: self.searcher.search(
|
||
query, top_k=self.top_k, recompute_embeddings=False
|
||
),
|
||
)
|
||
|
||
if results:
|
||
# Combine retrieved texts
|
||
texts = []
|
||
urls = []
|
||
for r in results:
|
||
texts.append(r.text)
|
||
url = r.metadata.get("url", "") if r.metadata else ""
|
||
urls.append(url)
|
||
|
||
combined_text = "\n\n---\n\n".join(texts)
|
||
combined_urls = ", ".join(u for u in urls if u)
|
||
|
||
return RetrievalResult(
|
||
text=combined_text,
|
||
source_url=combined_urls,
|
||
retrieval_type="text_vector",
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Text vector retrieval failed: {e}")
|
||
|
||
return RetrievalResult(retrieval_type="text_vector")
|
||
|
||
|
||
class DsServeRetriever(BaseRetriever):
|
||
"""Use ds-serve API for external text augmentation.
|
||
|
||
Calls ds-serve search API to retrieve relevant text passages for the query.
|
||
"""
|
||
|
||
def __init__(
|
||
self, api_url: str = "http://api.ds-serve.org:30888/search", top_k: int = 3
|
||
):
|
||
self.api_url = api_url
|
||
self.top_k = top_k
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
import aiohttp
|
||
import asyncio
|
||
|
||
max_retries = 3
|
||
for attempt in range(max_retries):
|
||
try:
|
||
headers = {"Content-Type": "application/json"}
|
||
payload = {"query": query}
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
self.api_url,
|
||
headers=headers,
|
||
json=payload,
|
||
timeout=aiohttp.ClientTimeout(total=30),
|
||
) as response:
|
||
if response.status == 200:
|
||
result = await response.json()
|
||
|
||
# Extract passages from response
|
||
passages = []
|
||
if "results" in result and "passages" in result["results"]:
|
||
# passages is a list of lists, get the first list
|
||
passage_list = (
|
||
result["results"]["passages"][0]
|
||
if result["results"]["passages"]
|
||
else []
|
||
)
|
||
|
||
# Take top_k passages
|
||
for i, passage_data in enumerate(
|
||
passage_list[: self.top_k]
|
||
):
|
||
if isinstance(passage_data, dict):
|
||
text = passage_data.get(
|
||
"text", ""
|
||
) or passage_data.get("center_text", "")
|
||
if text:
|
||
passages.append(text)
|
||
|
||
# Combine passages into context text
|
||
if passages:
|
||
combined_text = "\n\n".join(
|
||
[
|
||
f"[Passage {i + 1}]\n{text}"
|
||
for i, text in enumerate(passages)
|
||
]
|
||
)
|
||
|
||
return RetrievalResult(
|
||
text=combined_text,
|
||
source_url=f"ds-serve:{self.api_url}",
|
||
retrieval_type="ds_serve",
|
||
)
|
||
else:
|
||
return RetrievalResult(
|
||
text="No passages found from ds-serve.",
|
||
source_url=f"ds-serve:{self.api_url}",
|
||
retrieval_type="ds_serve",
|
||
)
|
||
elif response.status == 429:
|
||
if attempt < max_retries - 1:
|
||
wait_time = min(2**attempt * 2, 10)
|
||
logger.warning(
|
||
f"Rate limited (429), waiting {wait_time}s before retry ({attempt + 1}/{max_retries})"
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
else:
|
||
error_msg = f"ds-serve API rate limited after {max_retries} retries"
|
||
logger.error(error_msg)
|
||
return RetrievalResult(
|
||
text=error_msg, retrieval_type="ds_serve"
|
||
)
|
||
else:
|
||
error_text = await response.text()
|
||
error_msg = f"ds-serve API error: {response.status} - {error_text[:200]}"
|
||
logger.error(error_msg)
|
||
return RetrievalResult(
|
||
text=error_msg, retrieval_type="ds_serve"
|
||
)
|
||
except asyncio.TimeoutError:
|
||
if attempt < max_retries - 1:
|
||
wait_time = min(2**attempt, 5)
|
||
logger.warning(
|
||
f"Timeout, waiting {wait_time}s before retry ({attempt + 1}/{max_retries})"
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
else:
|
||
error_msg = f"ds-serve API timeout after {max_retries} retries"
|
||
logger.error(error_msg)
|
||
return RetrievalResult(text=error_msg, retrieval_type="ds_serve")
|
||
except Exception as e:
|
||
error_msg = f"ds-serve API call failed: {e}"
|
||
logger.error(error_msg)
|
||
return RetrievalResult(text=error_msg, retrieval_type="ds_serve")
|
||
|
||
return RetrievalResult(
|
||
text="ds-serve API call failed after all retries", retrieval_type="ds_serve"
|
||
)
|
||
|
||
|
||
class LocalAPIRetriever(BaseRetriever):
|
||
"""Retrieve tiles from a local search API (e.g. localhost:30888/search).
|
||
|
||
The API accepts batch queries:
|
||
{"queries": [{"text": "..."}, ...], "n_docs": N}
|
||
and returns:
|
||
{"results": [{"hits": [{"path": ..., "url": ..., "score": ...}, ...]}, ...]}
|
||
|
||
Call prefetch(examples) before the main loop to batch all queries in one API
|
||
call. Individual retrieve() calls then return cached results instantly.
|
||
|
||
When query_rewrite is enabled, uses an LLM to rewrite questions into
|
||
keyword-rich search queries before retrieval.
|
||
"""
|
||
|
||
REWRITE_PROMPT = (
|
||
"You are a search query optimizer. Given a trivia/factual question, "
|
||
"rewrite it as a Wikipedia search query that would find the article "
|
||
"containing the answer. Output ONLY the search query, nothing else.\n\n"
|
||
"Rules:\n"
|
||
"- Focus on the key entity or topic the question is about\n"
|
||
"- Include all specific names, dates, awards, events, or other details mentioned\n"
|
||
"- Remove filler words like 'what is', 'who was', 'in which year'\n"
|
||
"- Preserve all proper nouns and technical terms exactly as written\n\n"
|
||
"Question: {question}\n"
|
||
"Search query:"
|
||
)
|
||
|
||
def __init__(
|
||
self,
|
||
api_url: str = "http://localhost:30888/search",
|
||
top_k: int = 5,
|
||
batch_size: int = 32,
|
||
query_rewrite: bool = False,
|
||
rewrite_model: str | None = None,
|
||
rewrite_api_base: str | None = None,
|
||
rewrite_api_key: str = "dummy",
|
||
nprobe: int | None = None,
|
||
reranker=None,
|
||
rerank_top_k: int = 3,
|
||
query_image_fn=None,
|
||
multi_image_query: bool = False,
|
||
tiles_dir: str = "tiles/evqa",
|
||
lookup_reference_url: bool = False,
|
||
query_instruction: str | None = None,
|
||
):
|
||
self.api_url = api_url
|
||
self.top_k = top_k
|
||
self.batch_size = batch_size
|
||
self.query_rewrite = query_rewrite
|
||
self.rewrite_model = rewrite_model
|
||
self.rewrite_api_base = rewrite_api_base
|
||
self.rewrite_api_key = rewrite_api_key
|
||
self.nprobe = nprobe
|
||
self.reranker = reranker
|
||
self.rerank_top_k = rerank_top_k
|
||
self.query_image_fn = query_image_fn # callable(example) -> image_path or None
|
||
self.multi_image_query = multi_image_query
|
||
self.tiles_dir = tiles_dir
|
||
self.lookup_reference_url = lookup_reference_url
|
||
self.query_instruction = query_instruction
|
||
self._cache: dict[str, list[dict]] = {} # example_id -> hits
|
||
self._rewritten_queries: dict[str, str] = {} # example_id -> rewritten query
|
||
|
||
async def _rewrite_queries(self, examples: list[dict]) -> dict[str, str]:
|
||
"""Batch-rewrite questions into search queries using an LLM."""
|
||
from openai import AsyncOpenAI
|
||
|
||
client = AsyncOpenAI(
|
||
api_key=self.rewrite_api_key,
|
||
base_url=self.rewrite_api_base,
|
||
timeout=60.0,
|
||
)
|
||
|
||
rewritten = {}
|
||
sem = asyncio.Semaphore(20)
|
||
|
||
async def rewrite_one(ex):
|
||
eid = ex.get("id", "unknown")
|
||
prompt = self.REWRITE_PROMPT.format(question=ex["problem"])
|
||
async with sem:
|
||
try:
|
||
resp = await client.chat.completions.create(
|
||
model=self.rewrite_model,
|
||
messages=[{"role": "user", "content": prompt}],
|
||
temperature=0.0,
|
||
max_tokens=200,
|
||
)
|
||
rewritten[eid] = resp.choices[0].message.content.strip()
|
||
except Exception as e:
|
||
logger.warning(f"Query rewrite failed for {eid}: {e}")
|
||
rewritten[eid] = ex["problem"] # fallback to original
|
||
|
||
await asyncio.gather(*[rewrite_one(ex) for ex in examples])
|
||
return rewritten
|
||
|
||
def _lookup_reference_tiles(self, examples: list[dict]) -> dict[str, list[dict]]:
|
||
"""Look up reference URL tiles from kiwix for each example.
|
||
|
||
Returns dict: example_id -> list of hit dicts with path/score/url/is_reference.
|
||
"""
|
||
import sys as _sys
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
if not os.path.isdir(_KIWIX_OUTPUT_DIR) or not os.path.isfile(
|
||
_KIWIX_ARTICLES_JSON
|
||
):
|
||
logger.error(
|
||
f"lookup_reference_url: kiwix tiles unavailable at {_KIWIX_OUTPUT_DIR}"
|
||
)
|
||
return {}
|
||
|
||
if _WIKI_SCREENSHOT_DIR not in _sys.path:
|
||
_sys.path.insert(0, _WIKI_SCREENSHOT_DIR)
|
||
from scripts.build_index import batch_query_by_url as _batch_query
|
||
|
||
# Collect URLs, group by URL to avoid duplicate lookups
|
||
url_to_eids: dict[str, list[str]] = {}
|
||
for ex in examples:
|
||
eid = ex.get("id", "unknown")
|
||
url = extract_url_from_metadata(ex)
|
||
if url and "wikipedia.org" in url:
|
||
url_to_eids.setdefault(url, []).append(eid)
|
||
|
||
if not url_to_eids:
|
||
return {}
|
||
|
||
redirects = (
|
||
_KIWIX_REDIRECTS_JSON if os.path.isfile(_KIWIX_REDIRECTS_JSON) else None
|
||
)
|
||
results = _batch_query(
|
||
_KIWIX_OUTPUT_DIR,
|
||
list(url_to_eids.keys()),
|
||
_KIWIX_ARTICLES_JSON,
|
||
redirects_json=redirects,
|
||
)
|
||
|
||
ref_tiles: dict[str, list[dict]] = {}
|
||
found, missing = 0, 0
|
||
for url, eids in url_to_eids.items():
|
||
result = results.get(url)
|
||
if result is None:
|
||
missing += 1
|
||
logger.warning(f"lookup_reference_url: URL not found in kiwix: {url}")
|
||
continue
|
||
tiles_dir_abs = os.path.join(_KIWIX_OUTPUT_DIR, result["tiles_dir"])
|
||
if not os.path.isdir(tiles_dir_abs):
|
||
missing += 1
|
||
logger.warning(
|
||
f"lookup_reference_url: tiles dir missing: {tiles_dir_abs}"
|
||
)
|
||
continue
|
||
chunks = sorted(
|
||
f
|
||
for f in os.listdir(tiles_dir_abs)
|
||
if f.startswith("chunk_") and f.endswith(".png")
|
||
)
|
||
if not chunks:
|
||
missing += 1
|
||
logger.warning(
|
||
f"lookup_reference_url: no chunk files in {tiles_dir_abs}"
|
||
)
|
||
continue
|
||
found += 1
|
||
hits = [
|
||
{
|
||
"path": os.path.join(tiles_dir_abs, c),
|
||
"score": 0.0,
|
||
"url": url,
|
||
"is_reference": True,
|
||
}
|
||
for c in chunks
|
||
]
|
||
for eid in eids:
|
||
ref_tiles[eid] = hits
|
||
|
||
logger.info(
|
||
f"lookup_reference_url: batch lookup {found} found, {missing} missing "
|
||
f"out of {len(url_to_eids)} unique URLs"
|
||
)
|
||
return ref_tiles
|
||
|
||
async def prefetch(self, examples: list[dict]):
|
||
"""Batch-fetch retrieval results for all examples via the API."""
|
||
import aiohttp
|
||
|
||
# Step 1: Query rewriting (if enabled)
|
||
if self.query_rewrite and self.rewrite_model:
|
||
to_rewrite = [
|
||
ex
|
||
for ex in examples
|
||
if ex.get("id", "unknown") not in self._rewritten_queries
|
||
]
|
||
if to_rewrite:
|
||
logger.info(
|
||
f"LocalAPIRetriever: rewriting {len(to_rewrite)} queries..."
|
||
)
|
||
self._rewritten_queries.update(await self._rewrite_queries(to_rewrite))
|
||
# Log some examples
|
||
for ex in to_rewrite[:3]:
|
||
eid = ex.get("id", "unknown")
|
||
orig = ex["problem"][:60]
|
||
rewr = self._rewritten_queries.get(eid, "")[:60]
|
||
logger.info(f" Rewrite: '{orig}...' -> '{rewr}'")
|
||
|
||
# Step 2: Build query list
|
||
queries = []
|
||
example_ids = []
|
||
|
||
if self.multi_image_query:
|
||
# Multi-image: send one query per image, track which example each belongs to
|
||
# We'll aggregate after receiving results
|
||
multi_image_groups: dict[
|
||
str, list[int]
|
||
] = {} # eid -> list of indices in queries[]
|
||
for ex in examples:
|
||
eid = ex.get("id", "unknown")
|
||
if eid in self._cache:
|
||
continue
|
||
if self.query_rewrite and eid in self._rewritten_queries:
|
||
query_text = self._rewritten_queries[eid]
|
||
else:
|
||
query_text = ex["problem"]
|
||
|
||
all_paths = _get_all_query_image_paths(ex, self.tiles_dir)
|
||
if len(all_paths) <= 1:
|
||
# Single or no image: just use the standard path
|
||
query_dict = {"text": query_text}
|
||
if all_paths:
|
||
import base64
|
||
|
||
with open(all_paths[0], "rb") as f:
|
||
query_dict["image"] = base64.b64encode(f.read()).decode()
|
||
elif self.query_image_fn:
|
||
img_path = self.query_image_fn(ex)
|
||
if img_path and os.path.exists(img_path):
|
||
import base64
|
||
|
||
with open(img_path, "rb") as f:
|
||
query_dict["image"] = base64.b64encode(
|
||
f.read()
|
||
).decode()
|
||
multi_image_groups[eid] = [len(queries)]
|
||
queries.append(query_dict)
|
||
example_ids.append(eid)
|
||
else:
|
||
# Multiple images: one query per image
|
||
group_indices = []
|
||
import base64
|
||
|
||
for img_path in all_paths:
|
||
query_dict = {"text": query_text}
|
||
with open(img_path, "rb") as f:
|
||
query_dict["image"] = base64.b64encode(f.read()).decode()
|
||
group_indices.append(len(queries))
|
||
queries.append(query_dict)
|
||
example_ids.append(eid)
|
||
multi_image_groups[eid] = group_indices
|
||
logger.info(
|
||
f"Multi-image query for {eid[:8]}: {len(all_paths)} images"
|
||
)
|
||
else:
|
||
for ex in examples:
|
||
eid = ex.get("id", "unknown")
|
||
if eid in self._cache:
|
||
continue
|
||
if self.query_rewrite and eid in self._rewritten_queries:
|
||
query_text = self._rewritten_queries[eid]
|
||
else:
|
||
query_text = ex["problem"]
|
||
query_dict = {"text": query_text}
|
||
if self.query_image_fn:
|
||
img_path = self.query_image_fn(ex)
|
||
if img_path and os.path.exists(img_path):
|
||
import base64
|
||
|
||
with open(img_path, "rb") as f:
|
||
query_dict["image"] = base64.b64encode(f.read()).decode()
|
||
queries.append(query_dict)
|
||
example_ids.append(eid)
|
||
|
||
if not queries:
|
||
logger.info("LocalAPIRetriever: all examples already cached")
|
||
return
|
||
|
||
# Use smaller batches when queries contain images (GPU memory)
|
||
has_images = any("image" in q for q in queries)
|
||
batch_size = min(self.batch_size, 16) if has_images else self.batch_size
|
||
logger.info(
|
||
f"LocalAPIRetriever: prefetching {len(queries)} queries in batches of {batch_size}"
|
||
f"{' (multimodal)' if has_images else ''}"
|
||
)
|
||
|
||
for batch_start in range(0, len(queries), batch_size):
|
||
batch_queries = queries[batch_start : batch_start + batch_size]
|
||
batch_ids = example_ids[batch_start : batch_start + batch_size]
|
||
|
||
n_docs = self.top_k * 2 if self.multi_image_query else self.top_k
|
||
payload = {
|
||
"queries": batch_queries,
|
||
"n_docs": n_docs,
|
||
"include_images": True,
|
||
}
|
||
if self.nprobe is not None:
|
||
payload["nprobe"] = self.nprobe
|
||
if self.query_instruction is not None:
|
||
payload["instruction"] = self.query_instruction
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
self.api_url,
|
||
json=payload,
|
||
timeout=aiohttp.ClientTimeout(total=600),
|
||
) as response:
|
||
if response.status != 200:
|
||
error_text = await response.text()
|
||
logger.error(
|
||
f"Local API batch error {response.status}: {error_text[:200]}"
|
||
)
|
||
for eid in batch_ids:
|
||
self._cache[eid] = []
|
||
continue
|
||
result = await response.json()
|
||
except Exception as e:
|
||
logger.error(f"Local API batch call failed: {e}")
|
||
for eid in batch_ids:
|
||
self._cache[eid] = []
|
||
continue
|
||
|
||
results_list = result.get("results", [])
|
||
for i, eid in enumerate(batch_ids):
|
||
if i < len(results_list):
|
||
hits = results_list[i].get("hits", [])
|
||
else:
|
||
hits = []
|
||
if eid not in self._cache:
|
||
self._cache[eid] = hits
|
||
else:
|
||
# Multi-image: accumulate hits from all images for this example
|
||
self._cache[eid].extend(hits)
|
||
|
||
logger.info(
|
||
f" Batch {batch_start // batch_size + 1}/{(len(queries) + batch_size - 1) // batch_size}: "
|
||
f"{len(batch_queries)} queries done"
|
||
)
|
||
|
||
# Multi-image aggregation: deduplicate and keep max score per tile path
|
||
if self.multi_image_query:
|
||
for eid in list(self._cache.keys()):
|
||
hits = self._cache[eid]
|
||
if not hits:
|
||
continue
|
||
# Aggregate by path: keep hit with max score
|
||
best_by_path: dict[str, dict] = {}
|
||
for hit in hits:
|
||
path = hit.get("path", "")
|
||
score = hit.get("score", 0.0)
|
||
if path not in best_by_path or score > best_by_path[path].get(
|
||
"score", 0.0
|
||
):
|
||
best_by_path[path] = hit
|
||
# Sort by score descending, take top_k
|
||
sorted_hits = sorted(
|
||
best_by_path.values(),
|
||
key=lambda h: h.get("score", 0.0),
|
||
reverse=True,
|
||
)
|
||
self._cache[eid] = sorted_hits[: self.top_k]
|
||
|
||
logger.info(f"LocalAPIRetriever: prefetch complete, {len(self._cache)} cached")
|
||
|
||
# Step 2.5: Merge reference URL tiles (if enabled) — chunk-level dedup
|
||
if self.lookup_reference_url:
|
||
ref_tiles = self._lookup_reference_tiles(examples)
|
||
total_added, total_skipped = 0, 0
|
||
for eid, ref_hits in ref_tiles.items():
|
||
existing = self._cache.get(eid, [])
|
||
existing_paths = {hit.get("path", "") for hit in existing}
|
||
new_chunks = [rh for rh in ref_hits if rh["path"] not in existing_paths]
|
||
skipped = len(ref_hits) - len(new_chunks)
|
||
if new_chunks:
|
||
logger.info(
|
||
f" [{eid[:8]}]: adding {len(new_chunks)} reference URL chunks "
|
||
f"({skipped} already in API results)"
|
||
)
|
||
self._cache[eid] = existing + new_chunks
|
||
total_added += len(new_chunks)
|
||
total_skipped += skipped
|
||
logger.info(
|
||
f"lookup_reference_url: added {total_added} chunks, "
|
||
f"skipped {total_skipped} duplicates"
|
||
)
|
||
|
||
# Step 3: Rerank (if reranker provided)
|
||
if self.reranker is not None:
|
||
# Build batch of (query, candidates) for all examples
|
||
batch_inputs = []
|
||
batch_eids = []
|
||
for ex in examples:
|
||
eid = ex.get("id", "unknown")
|
||
hits = self._cache.get(eid, [])
|
||
if not hits:
|
||
continue
|
||
candidates = []
|
||
for hit in hits:
|
||
path = hit.get("path", "")
|
||
score = hit.get("score", 0.0)
|
||
if path and os.path.exists(path):
|
||
candidates.append((path, score))
|
||
if not candidates:
|
||
continue
|
||
batch_inputs.append((ex["problem"], candidates))
|
||
batch_eids.append(eid)
|
||
|
||
if batch_inputs:
|
||
all_reranked = self.reranker.rerank_batch(
|
||
batch_inputs,
|
||
top_k=self.rerank_top_k,
|
||
)
|
||
# Update cache with reranked results
|
||
for eid, reranked_results in zip(batch_eids, all_reranked):
|
||
hits = self._cache[eid]
|
||
path_to_hit = {hit["path"]: hit for hit in hits if "path" in hit}
|
||
new_hits = []
|
||
for path, rerank_score in reranked_results:
|
||
orig_hit = path_to_hit.get(path, {})
|
||
new_hits.append(
|
||
{**orig_hit, "path": path, "score": rerank_score}
|
||
)
|
||
self._cache[eid] = new_hits
|
||
logger.info(
|
||
f"LocalAPIRetriever: reranking complete ({len(batch_inputs)} examples)"
|
||
)
|
||
|
||
@staticmethod
|
||
@staticmethod
|
||
def _resolve_tile_path(hit: dict, tiles_dir: str | None = None) -> str | None:
|
||
"""Resolve tile path from hit, searching local shard dirs if needed."""
|
||
path = hit.get("path", "")
|
||
if path and os.path.exists(path):
|
||
return path
|
||
if not tiles_dir:
|
||
return path if path else None
|
||
article_id = hit.get("article_id")
|
||
tile_index = hit.get("tile_index", 0)
|
||
chunk_index = hit.get("chunk_index", 0)
|
||
if article_id is None:
|
||
return path if path else None
|
||
tiles_dirname = f"{article_id}.png.tiles"
|
||
chunk_name = f"chunk_{tile_index:04d}_{chunk_index:02d}.png"
|
||
shard_size = 8284
|
||
top_shard = article_id // shard_size
|
||
top_shard_dir = os.path.join(tiles_dir, f"shard_{top_shard:03d}")
|
||
if os.path.isdir(top_shard_dir):
|
||
for sub in sorted(os.listdir(top_shard_dir)):
|
||
sub_path = os.path.join(top_shard_dir, sub, tiles_dirname)
|
||
if os.path.isdir(sub_path):
|
||
full = os.path.join(sub_path, chunk_name)
|
||
if os.path.exists(full):
|
||
return full
|
||
flat = os.path.join(tiles_dir, tiles_dirname, chunk_name)
|
||
if os.path.exists(flat):
|
||
return flat
|
||
return path if path else None
|
||
|
||
@staticmethod
|
||
def _hits_to_result(
|
||
hits: list[dict], tiles_dir: str | None = None
|
||
) -> RetrievalResult:
|
||
"""Convert API hits to RetrievalResult."""
|
||
if not hits:
|
||
return RetrievalResult(retrieval_type="local_api")
|
||
|
||
images = []
|
||
image_urls = []
|
||
urls = []
|
||
seen_urls = set()
|
||
for hit in hits:
|
||
score = hit.get("score", 0.0)
|
||
url = hit.get("url", "")
|
||
path = LocalAPIRetriever._resolve_tile_path(hit, tiles_dir)
|
||
if path and os.path.exists(path):
|
||
images.append((path, score))
|
||
image_urls.append(url or None)
|
||
elif hit.get("image_base64"):
|
||
images.append((hit["image_base64"], score))
|
||
image_urls.append(url or None)
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
urls.append(url)
|
||
|
||
return RetrievalResult(
|
||
images=images,
|
||
image_urls=image_urls,
|
||
source_url=", ".join(urls) if urls else None,
|
||
retrieval_type="local_api",
|
||
)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
eid = example.get("id", "unknown")
|
||
|
||
# Return cached result if available (from prefetch)
|
||
if eid in self._cache:
|
||
return self._hits_to_result(self._cache[eid], tiles_dir=self.tiles_dir)
|
||
|
||
# Fallback: single query (if prefetch wasn't called)
|
||
import aiohttp
|
||
|
||
query_dict = {"text": query}
|
||
if self.query_image_fn:
|
||
img_path = self.query_image_fn(example)
|
||
if img_path and os.path.exists(img_path):
|
||
import base64
|
||
|
||
with open(img_path, "rb") as f:
|
||
query_dict["image"] = base64.b64encode(f.read()).decode()
|
||
payload = {"queries": [query_dict], "n_docs": self.top_k}
|
||
if self.nprobe is not None:
|
||
payload["nprobe"] = self.nprobe
|
||
if self.query_instruction is not None:
|
||
payload["instruction"] = self.query_instruction
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
self.api_url,
|
||
json=payload,
|
||
timeout=aiohttp.ClientTimeout(total=300),
|
||
) as response:
|
||
if response.status != 200:
|
||
return RetrievalResult(retrieval_type="local_api")
|
||
result = await response.json()
|
||
except Exception as e:
|
||
logger.error(f"Local API call failed: {e}")
|
||
return RetrievalResult(retrieval_type="local_api")
|
||
|
||
hits = result.get("results", [{}])[0].get("hits", [])
|
||
self._cache[eid] = hits
|
||
return self._hits_to_result(hits, tiles_dir=self.tiles_dir)
|
||
|
||
async def get_hits(self, query: str, example: dict) -> list[dict]:
|
||
"""Return raw per-hit dicts (path/url/score/...) for this example.
|
||
|
||
Used by wrappers that need per-hit granularity (e.g. HybridRetriever).
|
||
Uses the same cache as retrieve().
|
||
"""
|
||
await self.retrieve(query, example)
|
||
return self._cache.get(example.get("id", "unknown"), [])
|
||
|
||
|
||
class TiledQwen3VLEmbeddingRetriever(BaseRetriever):
|
||
"""Retrieves context by searching through image tiles using Qwen3-VL-Embedding.
|
||
|
||
Uses single vector embeddings (2048 dim) with cosine similarity for retrieval.
|
||
|
||
When *pixel_query_map* is provided the retriever embeds the rendered query
|
||
image (pixel query) instead of the raw text, so retrieval happens entirely
|
||
in pixel space.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
screenshot_dir: str = "screenshots",
|
||
tiles_dir: str = "tiles",
|
||
tile_size: int | tuple[int, int] = 512,
|
||
overlap: int = 0,
|
||
cache_path: str | None = None,
|
||
model_name: str = "Qwen/Qwen3-VL-Embedding-2B",
|
||
top_k: int = 3,
|
||
examples: list[dict] | None = None,
|
||
gpu_ids: list[int] | None = None,
|
||
tensor_parallel_size: int = 1,
|
||
pixel_query_map: dict[str, str] | None = None,
|
||
multimodal_query_text_only: bool = False,
|
||
multimodal_query_image_only: bool = False,
|
||
local_wiki: bool = False,
|
||
local_wiki_screenshot_dir: str | None = None,
|
||
multi_image_query: bool = False,
|
||
prebuilt_tiles_dir: str | None = None,
|
||
embedding_backend: str = "vllm", # "vllm", "hf", or "biqwen3"
|
||
peft_adapter: str | None = None,
|
||
):
|
||
self.top_k = top_k
|
||
self.screenshot_dir = screenshot_dir
|
||
self.tiles_dir = tiles_dir
|
||
self.tile_size = tile_size
|
||
self.overlap = overlap
|
||
self.examples = examples or []
|
||
self.pixel_query_map = pixel_query_map # example_id -> pixel query image path
|
||
self.multimodal_query_text_only = multimodal_query_text_only
|
||
self.multimodal_query_image_only = multimodal_query_image_only
|
||
self.local_wiki = local_wiki
|
||
self.local_wiki_screenshot_dir = local_wiki_screenshot_dir
|
||
self.multi_image_query = multi_image_query
|
||
self.prebuilt_tiles_dir = prebuilt_tiles_dir
|
||
self.embedding_backend = embedding_backend
|
||
self.peft_adapter = peft_adapter
|
||
os.makedirs(screenshot_dir, exist_ok=True)
|
||
os.makedirs(tiles_dir, exist_ok=True)
|
||
|
||
# Build example_id -> URL mapping and deduplicate by URL
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
|
||
self.id_to_url = {}
|
||
seen_urls: dict[str, str] = {} # url -> first example_id that uses it
|
||
self.url_to_representative_id: dict[
|
||
str, str
|
||
] = {} # url -> representative example_id
|
||
dedup_examples = []
|
||
for ex in self.examples:
|
||
ex_id = ex.get("id", "")
|
||
url = extract_url_from_metadata(ex)
|
||
if url:
|
||
self.id_to_url[ex_id] = url
|
||
if url not in seen_urls:
|
||
seen_urls[url] = ex_id
|
||
self.url_to_representative_id[url] = ex_id
|
||
dedup_examples.append(ex)
|
||
|
||
logger.info(
|
||
f"Deduplicated {len(self.examples)} examples -> {len(dedup_examples)} unique URLs "
|
||
f"(removed {len(self.examples) - len(dedup_examples)} duplicate pages)"
|
||
)
|
||
self._dedup_examples = dedup_examples
|
||
|
||
# Prepare tile paths: prebuilt dir (hard mini-datastore), local-wiki, or Selenium
|
||
if self.prebuilt_tiles_dir:
|
||
tile_paths = self._load_prebuilt_tiles()
|
||
elif self.local_wiki:
|
||
tile_paths = self._prepare_local_wiki_tiles()
|
||
else:
|
||
tile_paths = self._prepare_screenshots_and_tiles()
|
||
|
||
# Import Qwen3-VL-Embedding retrieval system
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
scripts_dir = Path(__file__).parent.parent
|
||
if str(scripts_dir) not in sys.path:
|
||
sys.path.insert(0, str(scripts_dir))
|
||
|
||
try:
|
||
from qwen3vl_embedding_retrieval import Qwen3VLEmbeddingSystem
|
||
except ImportError:
|
||
try:
|
||
from scripts.qwen3vl_embedding_retrieval import Qwen3VLEmbeddingSystem
|
||
except ImportError:
|
||
raise ImportError("Qwen3VLEmbeddingSystem not available.")
|
||
|
||
logger.info("Initializing Qwen3-VL-Embedding retrieval system...")
|
||
logger.info(f"Model: {model_name}, tiles: {len(tile_paths)}, GPUs: {gpu_ids}")
|
||
if self.pixel_query_map:
|
||
logger.info(
|
||
f"Pixel query mode ENABLED ({len(self.pixel_query_map)} queries)"
|
||
)
|
||
|
||
self.retrieval_system = Qwen3VLEmbeddingSystem(
|
||
model_name=model_name,
|
||
cache_path=cache_path,
|
||
gpu_ids=gpu_ids,
|
||
tensor_parallel_size=tensor_parallel_size,
|
||
backend=self.embedding_backend,
|
||
peft_adapter=self.peft_adapter,
|
||
)
|
||
|
||
# Embed all tiles (batch_size=8 for HF backend to avoid OOM on shared GPUs)
|
||
embed_bs = 8 if self.embedding_backend == "hf" else 32
|
||
self.retrieval_system.embed_images(
|
||
file_paths=tile_paths,
|
||
cache_path=cache_path,
|
||
batch_size=embed_bs,
|
||
)
|
||
logger.info(
|
||
f"Qwen3-VL-Embedding retrieval ready with {len(self.retrieval_system.image_paths)} tiles"
|
||
)
|
||
|
||
def _load_prebuilt_tiles(self) -> list[str]:
|
||
"""Load ALL .png tiles from a prebuilt tile directory (e.g. hard mini-datastore).
|
||
|
||
Unlike _prepare_local_wiki_tiles which only loads golden tiles matching
|
||
example IDs, this loads every tile in the directory — including distractors.
|
||
"""
|
||
import glob as _glob
|
||
|
||
all_tiles = sorted(_glob.glob(os.path.join(self.prebuilt_tiles_dir, "*.png")))
|
||
filtered = _filter_tiles_by_aspect_ratio(all_tiles)
|
||
logger.info(
|
||
f"prebuilt-tiles: loaded {len(filtered)} tiles from {self.prebuilt_tiles_dir} "
|
||
f"(filtered {len(all_tiles) - len(filtered)} extreme aspect ratio tiles)"
|
||
)
|
||
return filtered
|
||
|
||
def _prepare_local_wiki_tiles(self) -> list[str]:
|
||
"""Prepare tiles from local kiwix tile store for all examples in the batch.
|
||
|
||
Does a single batch URL lookup (fast), then copies+cuts tiles per example.
|
||
Reports an error (no fallback) if a URL is not found in kiwix.
|
||
|
||
Returns the list of all cut tile paths ready for embedding.
|
||
"""
|
||
import glob as _glob
|
||
import shutil
|
||
import sys as _sys
|
||
from PIL import Image
|
||
from .simpleqa_data import extract_url_from_metadata
|
||
from tqdm import tqdm
|
||
|
||
cut_height = (
|
||
self.tile_size[1] if isinstance(self.tile_size, tuple) else self.tile_size
|
||
)
|
||
wiki_cache = self.local_wiki_screenshot_dir or os.path.join(
|
||
self.screenshot_dir, "local-wiki"
|
||
)
|
||
os.makedirs(wiki_cache, exist_ok=True)
|
||
os.makedirs(self.tiles_dir, exist_ok=True)
|
||
|
||
# Separate already-cached examples from ones that need processing
|
||
need: list[tuple[str, str]] = [] # (ex_id, url)
|
||
for ex in self._dedup_examples:
|
||
ex_id = ex["id"]
|
||
if not _glob.glob(os.path.join(self.tiles_dir, f"{ex_id}_tile_*.png")):
|
||
url = extract_url_from_metadata(ex) or ""
|
||
need.append((ex_id, url))
|
||
|
||
logger.info(
|
||
f"local-wiki: {len(self._dedup_examples) - len(need)} cached, {len(need)} need processing"
|
||
)
|
||
|
||
if need:
|
||
# Single batch lookup for all URLs at once (loads articles.json once)
|
||
if not os.path.isdir(_KIWIX_OUTPUT_DIR) or not os.path.isfile(
|
||
_KIWIX_ARTICLES_JSON
|
||
):
|
||
logger.error(
|
||
f"local-wiki: kiwix tiles unavailable at {_KIWIX_OUTPUT_DIR}"
|
||
)
|
||
else:
|
||
if _WIKI_SCREENSHOT_DIR not in _sys.path:
|
||
_sys.path.insert(0, _WIKI_SCREENSHOT_DIR)
|
||
from scripts.build_index import batch_query_by_url as _batch_query
|
||
|
||
redirects = (
|
||
_KIWIX_REDIRECTS_JSON
|
||
if os.path.isfile(_KIWIX_REDIRECTS_JSON)
|
||
else None
|
||
)
|
||
urls_to_lookup = [u for _, u in need if u and "wikipedia.org" in u]
|
||
results = _batch_query(
|
||
_KIWIX_OUTPUT_DIR,
|
||
urls_to_lookup,
|
||
_KIWIX_ARTICLES_JSON,
|
||
redirects_json=redirects,
|
||
)
|
||
found = sum(1 for r in results.values() if r is not None)
|
||
logger.info(
|
||
f"local-wiki: batch lookup found {found}/{len(urls_to_lookup)} URLs"
|
||
)
|
||
|
||
# Copy + cut per example
|
||
ok, failed = 0, 0
|
||
for ex_id, url in tqdm(need, desc="local-wiki: copying+cutting tiles"):
|
||
# Check cache again (may have been done by a parallel run)
|
||
if _glob.glob(os.path.join(self.tiles_dir, f"{ex_id}_tile_*.png")):
|
||
ok += 1
|
||
continue
|
||
result = results.get(url)
|
||
if result is None:
|
||
logger.error(
|
||
f"local-wiki [{ex_id}]: URL not found in kiwix: {url}"
|
||
)
|
||
failed += 1
|
||
continue
|
||
src_dir = os.path.join(_KIWIX_OUTPUT_DIR, result["tiles_dir"])
|
||
article_cache = os.path.join(wiki_cache, str(ex_id))
|
||
if not os.path.exists(article_cache):
|
||
if not os.path.isdir(src_dir):
|
||
logger.error(
|
||
f"local-wiki [{ex_id}]: tiles dir not on disk: {src_dir}"
|
||
)
|
||
failed += 1
|
||
continue
|
||
shutil.copytree(src_dir, article_cache)
|
||
# Cut into strips
|
||
raw_tiles = sorted(
|
||
f
|
||
for f in os.listdir(article_cache)
|
||
if f.endswith(".png") and f.startswith("tile_")
|
||
)
|
||
if not raw_tiles:
|
||
logger.error(
|
||
f"local-wiki [{ex_id}]: no tile PNGs in {article_cache}"
|
||
)
|
||
failed += 1
|
||
continue
|
||
global_row = 0
|
||
for raw_name in raw_tiles:
|
||
raw_path = os.path.join(article_cache, raw_name)
|
||
if os.path.getsize(raw_path) == 0:
|
||
continue
|
||
try:
|
||
img = Image.open(raw_path)
|
||
img.load()
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"local-wiki [{ex_id}]: corrupt tile {raw_path}: {e}"
|
||
)
|
||
continue
|
||
w, h = img.size
|
||
y = 0
|
||
while y < h:
|
||
y2 = min(y + cut_height, h)
|
||
img.crop((0, y, w, y2)).save(
|
||
os.path.join(
|
||
self.tiles_dir, f"{ex_id}_tile_{global_row}_0.png"
|
||
)
|
||
)
|
||
global_row += 1
|
||
y += cut_height
|
||
img.close()
|
||
ok += 1
|
||
logger.info(
|
||
f"local-wiki: {ok} articles prepared, {failed} not found/failed"
|
||
)
|
||
|
||
all_tile_paths = []
|
||
for ex in self._dedup_examples:
|
||
ex_id = ex["id"]
|
||
tiles = sorted(
|
||
_glob.glob(os.path.join(self.tiles_dir, f"{ex_id}_tile_*.png"))
|
||
)
|
||
all_tile_paths.extend(tiles)
|
||
|
||
filtered = _filter_tiles_by_aspect_ratio(all_tile_paths)
|
||
logger.info(
|
||
f"local-wiki: {len(filtered)} tiles ready for embedding "
|
||
f"(filtered {len(all_tile_paths) - len(filtered)} extreme aspect ratio tiles)"
|
||
)
|
||
return filtered
|
||
|
||
def _prepare_screenshots_and_tiles(self) -> list[str]:
|
||
"""Prepare screenshots and tiles for dataset, return tile paths.
|
||
|
||
Uses deduplicated examples (one per unique URL) to avoid
|
||
duplicate tiles inflating the retrieval index.
|
||
"""
|
||
from .simpleqa_data import capture_screenshot_for_example, split_image_to_tiles
|
||
from tqdm import tqdm
|
||
|
||
examples_to_process = self._dedup_examples
|
||
screenshot_paths = []
|
||
missing = []
|
||
|
||
# Collect screenshot paths and identify missing (deduplicated)
|
||
for ex in examples_to_process:
|
||
screenshot_path = os.path.join(
|
||
self.screenshot_dir, f"{ex['id']}_fullhd.png"
|
||
)
|
||
screenshot_paths.append(screenshot_path)
|
||
if (
|
||
not os.path.exists(screenshot_path)
|
||
or os.path.getsize(screenshot_path) == 0
|
||
):
|
||
missing.append(ex)
|
||
|
||
# Capture missing screenshots
|
||
if missing:
|
||
logger.info(f"Preparing {len(missing)} missing screenshots...")
|
||
for ex in tqdm(missing, desc="Capturing screenshots"):
|
||
capture_screenshot_for_example(ex, self.screenshot_dir)
|
||
logger.info("Screenshots prepared.")
|
||
|
||
# Split each screenshot into tiles
|
||
all_tile_paths = []
|
||
logger.info(
|
||
f"Splitting {len(screenshot_paths)} unique screenshots into tiles (output: {self.tiles_dir})..."
|
||
)
|
||
for screenshot_path in tqdm(screenshot_paths, desc="Splitting tiles"):
|
||
if os.path.exists(screenshot_path) and os.path.getsize(screenshot_path) > 0:
|
||
tile_paths = split_image_to_tiles(
|
||
screenshot_path, self.tiles_dir, self.tile_size, self.overlap
|
||
)
|
||
all_tile_paths.extend(tile_paths)
|
||
|
||
# Filter out tiles with extreme aspect ratios
|
||
filtered_tile_paths = _filter_tiles_by_aspect_ratio(all_tile_paths)
|
||
logger.info(
|
||
f"Prepared {len(filtered_tile_paths)} tiles from {len(screenshot_paths)} unique screenshots "
|
||
f"(filtered {len(all_tile_paths) - len(filtered_tile_paths)} extreme aspect ratio tiles)"
|
||
)
|
||
return filtered_tile_paths
|
||
|
||
def _extract_urls_from_results(self, results: list) -> str:
|
||
"""Extract source URLs from tile paths in results, preserving retrieval order."""
|
||
urls = []
|
||
seen = set()
|
||
for item in results:
|
||
# item is (path, score) tuple
|
||
path = item[0] if isinstance(item, tuple) else item
|
||
# Extract example_id from tile path: {example_id}_fullhd_tile_{x}_{y}.png
|
||
filename = os.path.basename(path)
|
||
if "_tile_" in filename:
|
||
example_id = filename.split("_tile_")[0]
|
||
if example_id.endswith("_fullhd"):
|
||
example_id = example_id[:-7]
|
||
if example_id in self.id_to_url:
|
||
url = self.id_to_url[example_id]
|
||
if url not in seen:
|
||
seen.add(url)
|
||
urls.append(url)
|
||
return ", ".join(urls)
|
||
|
||
# Class-level cache for iNat 2021 image_id -> file_name mapping
|
||
_inat2021_id_map: dict[int, str] | None = None
|
||
INAT2021_DATA_DIR = _INAT2021_DATA_DIR
|
||
|
||
@classmethod
|
||
def _load_inat2021_mapping(cls) -> dict[int, str]:
|
||
"""Load iNaturalist 2021 competition image_id -> file_name mapping.
|
||
|
||
Downloads val.json from the competition S3 bucket if not cached locally.
|
||
"""
|
||
if cls._inat2021_id_map is not None:
|
||
return cls._inat2021_id_map
|
||
|
||
import json
|
||
import tarfile
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
data_dir = Path(cls.INAT2021_DATA_DIR)
|
||
data_dir.mkdir(parents=True, exist_ok=True)
|
||
val_json = data_dir / "val.json"
|
||
|
||
if not val_json.exists():
|
||
tar_path = data_dir / "val.json.tar.gz"
|
||
if not tar_path.exists():
|
||
logger.info("Downloading iNaturalist 2021 val annotations...")
|
||
urllib.request.urlretrieve(
|
||
"https://ml-inat-competition-datasets.s3.amazonaws.com/2021/val.json.tar.gz",
|
||
str(tar_path),
|
||
)
|
||
with tarfile.open(str(tar_path), "r:gz") as tf:
|
||
tf.extractall(path=str(data_dir))
|
||
logger.info(f"Extracted iNat 2021 val.json to {val_json}")
|
||
|
||
with open(val_json) as f:
|
||
data = json.load(f)
|
||
|
||
cls._inat2021_id_map = {img["id"]: img["file_name"] for img in data["images"]}
|
||
logger.info(f"Loaded iNat 2021 mapping: {len(cls._inat2021_id_map)} images")
|
||
return cls._inat2021_id_map
|
||
|
||
def _get_inat_image_path(self, example: dict) -> str | None:
|
||
"""Get EVQA query image (iNaturalist or Landmarks). Delegates to _get_query_image_path_for_example."""
|
||
return _get_query_image_path_for_example(example, self.tiles_dir)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
# Dispatch to multi-image retrieval if enabled
|
||
if self.multi_image_query:
|
||
return await self.retrieve_multi_image(query, example)
|
||
return await self._retrieve_single(query, example)
|
||
|
||
async def _retrieve_single(self, query: str, example: dict) -> RetrievalResult:
|
||
example_id = example.get("id", "")
|
||
loop = asyncio.get_event_loop()
|
||
|
||
# Priority: pixel_query_map > iNaturalist image > text-only
|
||
pixel_query_path = None
|
||
if self.pixel_query_map and example_id in self.pixel_query_map:
|
||
pixel_query_path = self.pixel_query_map[example_id]
|
||
|
||
# Check for iNaturalist query image (multimodal text+image query)
|
||
inat_image_path = self._get_inat_image_path(example)
|
||
|
||
try:
|
||
# Determine query modality
|
||
query_image = None
|
||
if pixel_query_path and os.path.exists(pixel_query_path):
|
||
# Pixel query: image-only (rendered text as image)
|
||
query_image = pixel_query_path
|
||
query_text = None
|
||
retrieval_type = "tiled_qwen3vl_embedding_pixel_query"
|
||
elif self.multimodal_query_text_only:
|
||
# Ablation: text-only (no image)
|
||
query_image = None
|
||
query_text = query
|
||
retrieval_type = "tiled_qwen3vl_embedding_multimodal_textonly"
|
||
elif self.multimodal_query_image_only and inat_image_path:
|
||
# Ablation: image-only (no text)
|
||
query_image = inat_image_path
|
||
query_text = None
|
||
retrieval_type = "tiled_qwen3vl_embedding_multimodal_imageonly"
|
||
elif inat_image_path:
|
||
# Multimodal: text + image
|
||
query_image = inat_image_path
|
||
query_text = query
|
||
retrieval_type = "tiled_qwen3vl_embedding_multimodal"
|
||
else:
|
||
# Text-only (no query image available)
|
||
query_text = query
|
||
retrieval_type = "tiled_qwen3vl_embedding"
|
||
|
||
results = await loop.run_in_executor(
|
||
None,
|
||
lambda: self.retrieval_system.search(
|
||
text=query_text, image=query_image, top_k=self.top_k
|
||
),
|
||
)
|
||
|
||
if results:
|
||
source_url = self._extract_urls_from_results(results)
|
||
return RetrievalResult(
|
||
images=results,
|
||
source_url=source_url,
|
||
retrieval_type=retrieval_type,
|
||
pixel_query_path=pixel_query_path or inat_image_path,
|
||
query_image_path=inat_image_path,
|
||
)
|
||
else:
|
||
return RetrievalResult(
|
||
text="No relevant tiles found via Qwen3-VL-Embedding search",
|
||
retrieval_type=retrieval_type,
|
||
pixel_query_path=pixel_query_path or inat_image_path,
|
||
query_image_path=inat_image_path,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Qwen3-VL-Embedding search failed: {e}")
|
||
return RetrievalResult(
|
||
text=f"Qwen3-VL-Embedding retrieval error: {e}",
|
||
retrieval_type="tiled_qwen3vl_embedding",
|
||
pixel_query_path=pixel_query_path or inat_image_path,
|
||
query_image_path=inat_image_path,
|
||
)
|
||
|
||
async def retrieve_multi_image(self, query: str, example: dict) -> RetrievalResult:
|
||
"""Multi-image retrieval: search with ALL query images, aggregate scores, return top-K.
|
||
|
||
For each query image, does a multimodal search (text + image), then combines
|
||
scores across all images using max-score aggregation per tile.
|
||
Falls back to single-image retrieve() if only 0-1 images available.
|
||
"""
|
||
all_image_paths = _get_all_query_image_paths(example, self.tiles_dir)
|
||
# Get single image for generation (first available, used in RetrievalResult)
|
||
single_image_path = self._get_inat_image_path(example)
|
||
|
||
if len(all_image_paths) <= 1:
|
||
return await self._retrieve_single(query, example)
|
||
|
||
example_id = example.get("id", "")
|
||
loop = asyncio.get_event_loop()
|
||
logger.info(
|
||
f"Multi-image retrieval for {example_id}: {len(all_image_paths)} query images"
|
||
)
|
||
|
||
try:
|
||
# Score aggregation: for each tile, keep the max score across all query images
|
||
tile_best_score: dict[str, float] = {}
|
||
|
||
for img_path in all_image_paths:
|
||
results = await loop.run_in_executor(
|
||
None,
|
||
lambda p=img_path: self.retrieval_system.search(
|
||
text=query, image=p, top_k=self.top_k * 2
|
||
),
|
||
)
|
||
for tile_path, score in results:
|
||
if (
|
||
tile_path not in tile_best_score
|
||
or score > tile_best_score[tile_path]
|
||
):
|
||
tile_best_score[tile_path] = score
|
||
|
||
# Sort by score descending, take top_k
|
||
sorted_tiles = sorted(
|
||
tile_best_score.items(), key=lambda x: x[1], reverse=True
|
||
)
|
||
top_results = sorted_tiles[: self.top_k]
|
||
|
||
retrieval_type = (
|
||
f"tiled_qwen3vl_embedding_multiimage_{len(all_image_paths)}imgs"
|
||
)
|
||
|
||
if top_results:
|
||
source_url = self._extract_urls_from_results(top_results)
|
||
return RetrievalResult(
|
||
images=top_results,
|
||
source_url=source_url,
|
||
retrieval_type=retrieval_type,
|
||
pixel_query_path=single_image_path,
|
||
query_image_path=single_image_path,
|
||
)
|
||
else:
|
||
return RetrievalResult(
|
||
text="No relevant tiles found via multi-image search",
|
||
retrieval_type=retrieval_type,
|
||
pixel_query_path=single_image_path,
|
||
query_image_path=single_image_path,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Multi-image retrieval failed: {e}")
|
||
return await self._retrieve_single(query, example)
|
||
|
||
|
||
class TextAPIRetriever(BaseRetriever):
|
||
"""Retrieve text chunks from a text search API (wiki-screenshot text_search_api.py).
|
||
|
||
The API accepts:
|
||
POST /search
|
||
{"queries": [{"text": "..."}], "n_docs": N}
|
||
and returns:
|
||
{"results": [{"hits": [{"text": ..., "title": ..., "url": ..., "score": ...}, ...]}]}
|
||
|
||
Supports batch prefetch for efficient evaluation.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_url: str = "http://localhost:30889/search",
|
||
top_k: int = 3,
|
||
batch_size: int = 32,
|
||
nprobe: int | None = None,
|
||
query_instruction: str | None = None,
|
||
reader_top_k: int | None = None,
|
||
query_image_fn=None,
|
||
):
|
||
self.api_url = api_url
|
||
self.top_k = top_k
|
||
# If reader_top_k is set and < top_k, only the first reader_top_k hits are
|
||
# passed to the reader. Mirrors the image-side reader_top_k slicing in
|
||
# run_naive_simpleqa.py so text + image cells are comparable at fixed k.
|
||
self.reader_top_k = reader_top_k
|
||
self.batch_size = batch_size
|
||
self.nprobe = nprobe
|
||
self.query_instruction = query_instruction
|
||
self.query_image_fn = query_image_fn
|
||
self._cache: dict[str, list[dict]] = {}
|
||
|
||
async def prefetch(self, examples: list[dict]):
|
||
"""Batch-fetch retrieval results for all examples."""
|
||
import aiohttp
|
||
|
||
queries = []
|
||
example_ids = []
|
||
for ex in examples:
|
||
eid = ex.get("id", "unknown")
|
||
if eid in self._cache:
|
||
continue
|
||
query_dict = {"text": ex["problem"]}
|
||
if self.query_image_fn:
|
||
img_path = self.query_image_fn(ex)
|
||
if img_path and os.path.exists(img_path):
|
||
import base64
|
||
|
||
with open(img_path, "rb") as f:
|
||
query_dict["image"] = base64.b64encode(f.read()).decode()
|
||
queries.append(query_dict)
|
||
example_ids.append(eid)
|
||
|
||
if not queries:
|
||
logger.info("TextAPIRetriever: all examples already cached")
|
||
return
|
||
|
||
has_images = any("image" in q for q in queries)
|
||
batch_size = min(self.batch_size, 16) if has_images else self.batch_size
|
||
logger.info(
|
||
f"TextAPIRetriever: prefetching {len(queries)} queries in batches of {batch_size}"
|
||
f"{' (multimodal)' if has_images else ''}"
|
||
)
|
||
|
||
for batch_start in range(0, len(queries), batch_size):
|
||
batch_queries = queries[batch_start : batch_start + batch_size]
|
||
batch_ids = example_ids[batch_start : batch_start + batch_size]
|
||
|
||
payload = {"queries": batch_queries, "n_docs": self.top_k}
|
||
if self.nprobe is not None:
|
||
payload["nprobe"] = self.nprobe
|
||
if self.query_instruction is not None:
|
||
payload["instruction"] = self.query_instruction
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
self.api_url,
|
||
json=payload,
|
||
timeout=aiohttp.ClientTimeout(total=600),
|
||
) as response:
|
||
if response.status != 200:
|
||
error_text = await response.text()
|
||
logger.error(
|
||
f"TextAPI batch error {response.status}: {error_text[:200]}"
|
||
)
|
||
for eid in batch_ids:
|
||
self._cache[eid] = []
|
||
continue
|
||
result = await response.json()
|
||
except Exception as e:
|
||
logger.error(f"TextAPI batch call failed: {e}")
|
||
for eid in batch_ids:
|
||
self._cache[eid] = []
|
||
continue
|
||
|
||
results_list = result.get("results", [])
|
||
for i, eid in enumerate(batch_ids):
|
||
if i < len(results_list):
|
||
self._cache[eid] = results_list[i].get("hits", [])
|
||
else:
|
||
self._cache[eid] = []
|
||
|
||
logger.info(
|
||
f" Batch {batch_start // self.batch_size + 1}/"
|
||
f"{(len(queries) + self.batch_size - 1) // self.batch_size}: "
|
||
f"{len(batch_queries)} queries done"
|
||
)
|
||
|
||
logger.info(f"TextAPIRetriever: prefetch complete, {len(self._cache)} cached")
|
||
|
||
@staticmethod
|
||
def _hits_to_result(
|
||
hits: list[dict], max_passages: int | None = None
|
||
) -> RetrievalResult:
|
||
"""Convert text API hits to RetrievalResult.
|
||
|
||
If max_passages is set, only the first max_passages hits are joined into
|
||
the reader prompt. The cache itself is not truncated, so the same cached
|
||
hits can serve multiple reader_top_k values.
|
||
"""
|
||
if not hits:
|
||
return RetrievalResult(retrieval_type="text_api")
|
||
|
||
if max_passages is not None and max_passages < len(hits):
|
||
hits = hits[:max_passages]
|
||
|
||
passages = []
|
||
urls = []
|
||
seen_urls = set()
|
||
for hit in hits:
|
||
text = hit.get("text", "")
|
||
url = hit.get("url", "")
|
||
# Option 1 (2026-04-29): no `[title]` prefix on chunks. Title is leaked
|
||
# metadata for entity-answering tasks (often contains the answer outright).
|
||
# Reader sees only the chunk content. URL lives in retrieval_result.source_url
|
||
# for logging/grading but is not injected into the prompt by build_messages.
|
||
if text:
|
||
passages.append(text)
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
urls.append(url)
|
||
|
||
combined_text = "\n\n".join(passages) if passages else None
|
||
return RetrievalResult(
|
||
text=combined_text,
|
||
source_url=", ".join(urls) if urls else None,
|
||
retrieval_type="text_api",
|
||
)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
eid = example.get("id", "unknown")
|
||
|
||
if eid in self._cache:
|
||
return self._hits_to_result(
|
||
self._cache[eid], max_passages=self.reader_top_k
|
||
)
|
||
|
||
# Fallback: single query
|
||
import aiohttp
|
||
|
||
payload = {"queries": [{"text": query}], "n_docs": self.top_k}
|
||
if self.nprobe is not None:
|
||
payload["nprobe"] = self.nprobe
|
||
if self.query_instruction is not None:
|
||
payload["instruction"] = self.query_instruction
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
self.api_url,
|
||
json=payload,
|
||
timeout=aiohttp.ClientTimeout(total=300),
|
||
) as response:
|
||
if response.status != 200:
|
||
return RetrievalResult(retrieval_type="text_api")
|
||
result = await response.json()
|
||
except Exception as e:
|
||
logger.error(f"TextAPI call failed: {e}")
|
||
return RetrievalResult(retrieval_type="text_api")
|
||
|
||
hits = result.get("results", [{}])[0].get("hits", [])
|
||
self._cache[eid] = hits
|
||
return self._hits_to_result(hits, max_passages=self.reader_top_k)
|
||
|
||
async def get_hits(self, query: str, example: dict) -> list[dict]:
|
||
"""Return raw per-hit dicts (title/text/url/score/...) for this example.
|
||
|
||
Used by wrappers that need per-chunk granularity (e.g. RenderedTextWrapper).
|
||
Uses the same cache as retrieve().
|
||
"""
|
||
await self.retrieve(query, example)
|
||
return self._cache.get(example.get("id", "unknown"), [])
|
||
|
||
|
||
class OCRWrappedRetriever(BaseRetriever):
|
||
"""Wraps an image retriever; OCRs retrieved tiles and returns text.
|
||
|
||
Ablation A pipeline: image retrieve -> OCR -> text to reader.
|
||
Talks to an OpenAI-compatible chat endpoint (PaddleOCR-VL served via vLLM).
|
||
Caches OCR output to a JSONL file keyed by absolute image path so reruns
|
||
reuse prior work.
|
||
"""
|
||
|
||
DEFAULT_PROMPT = "OCR this image. Output only the extracted text verbatim, preserving paragraph and line breaks."
|
||
|
||
def __init__(
|
||
self,
|
||
base: BaseRetriever,
|
||
ocr_url: str = "http://localhost:8202/v1",
|
||
model: str = "PaddlePaddle/PaddleOCR-VL",
|
||
api_key: str = "dummy",
|
||
cache_path: str = "ocr_cache/paddleocr_vl.jsonl",
|
||
concurrency: int = 16,
|
||
prompt: str | None = None,
|
||
timeout: float = 180.0,
|
||
max_tokens: int = 4096,
|
||
reader_top_k: int | None = None,
|
||
):
|
||
self.base = base
|
||
self.ocr_url = ocr_url.rstrip("/")
|
||
self.model = model
|
||
self.api_key = api_key
|
||
self.cache_path = cache_path
|
||
self.concurrency = concurrency
|
||
self.prompt = prompt or self.DEFAULT_PROMPT
|
||
self.timeout = timeout
|
||
self.max_tokens = max_tokens
|
||
self.reader_top_k = reader_top_k
|
||
self._cache: dict[str, str] = {}
|
||
self.tiles_dir = getattr(base, "tiles_dir", None)
|
||
self._load_cache()
|
||
|
||
def _load_cache(self):
|
||
if not os.path.isfile(self.cache_path):
|
||
return
|
||
import json
|
||
|
||
loaded = 0
|
||
try:
|
||
with open(self.cache_path) as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
entry = json.loads(line)
|
||
self._cache[entry["path"]] = entry["text"]
|
||
loaded += 1
|
||
logger.info(
|
||
f"OCRWrappedRetriever: loaded {loaded} cached OCR entries from {self.cache_path}"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"OCRWrappedRetriever: cache load failed ({e}); starting fresh"
|
||
)
|
||
|
||
def _append_cache(self, path: str, text: str):
|
||
import json
|
||
|
||
os.makedirs(os.path.dirname(self.cache_path) or ".", exist_ok=True)
|
||
with open(self.cache_path, "a") as f:
|
||
f.write(json.dumps({"path": path, "text": text}, ensure_ascii=False) + "\n")
|
||
self._cache[path] = text
|
||
|
||
async def _ocr_one(self, path: str, session) -> str:
|
||
if path in self._cache:
|
||
return self._cache[path]
|
||
import aiohttp
|
||
import base64
|
||
|
||
try:
|
||
with open(path, "rb") as f:
|
||
img_b64 = base64.b64encode(f.read()).decode("ascii")
|
||
except Exception as e:
|
||
logger.error(f"OCR read failed for {path}: {e}")
|
||
return ""
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:image/png;base64,{img_b64}"},
|
||
},
|
||
{"type": "text", "text": self.prompt},
|
||
],
|
||
}
|
||
],
|
||
"max_tokens": self.max_tokens,
|
||
"temperature": 0.0,
|
||
}
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
try:
|
||
async with session.post(
|
||
f"{self.ocr_url}/chat/completions",
|
||
json=payload,
|
||
headers=headers,
|
||
timeout=aiohttp.ClientTimeout(total=self.timeout),
|
||
) as resp:
|
||
if resp.status != 200:
|
||
err = await resp.text()
|
||
logger.error(f"OCR HTTP {resp.status} for {path}: {err[:200]}")
|
||
return ""
|
||
result = await resp.json()
|
||
text = result["choices"][0]["message"]["content"]
|
||
except Exception as e:
|
||
logger.error(f"OCR request failed for {path}: {e}")
|
||
return ""
|
||
self._append_cache(path, text)
|
||
return text
|
||
|
||
async def _batch_ocr(self, paths: list[str]) -> dict[str, str]:
|
||
import aiohttp
|
||
|
||
to_fetch = [p for p in paths if p not in self._cache]
|
||
if not to_fetch:
|
||
return {p: self._cache[p] for p in paths}
|
||
sem = asyncio.Semaphore(self.concurrency)
|
||
async with aiohttp.ClientSession() as session:
|
||
|
||
async def _one(p):
|
||
async with sem:
|
||
return await self._ocr_one(p, session)
|
||
|
||
await asyncio.gather(*[_one(p) for p in to_fetch])
|
||
return {p: self._cache.get(p, "") for p in paths}
|
||
|
||
async def prefetch(self, examples: list[dict]):
|
||
"""Forward to base's prefetch, then batch-OCR all tiles up front."""
|
||
if hasattr(self.base, "prefetch"):
|
||
await self.base.prefetch(examples)
|
||
all_paths: set[str] = set()
|
||
for ex in examples:
|
||
r = await self.base.retrieve(ex.get("problem", ""), ex)
|
||
images = (
|
||
r.images[: self.reader_top_k]
|
||
if self.reader_top_k is not None
|
||
else r.images
|
||
)
|
||
for p, _ in images:
|
||
all_paths.add(os.path.abspath(p))
|
||
uncached = [p for p in all_paths if p not in self._cache]
|
||
logger.info(
|
||
f"OCRWrappedRetriever: {len(all_paths)} unique tiles across {len(examples)} examples; "
|
||
f"{len(all_paths) - len(uncached)} cached, OCRing {len(uncached)}"
|
||
)
|
||
if uncached:
|
||
await self._batch_ocr(uncached)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
r = await self.base.retrieve(query, example)
|
||
if not r.images:
|
||
return r
|
||
images = (
|
||
r.images[: self.reader_top_k] if self.reader_top_k is not None else r.images
|
||
)
|
||
image_urls = (
|
||
r.image_urls[: self.reader_top_k]
|
||
if self.reader_top_k is not None and r.image_urls
|
||
else list(r.image_urls or [])
|
||
)
|
||
urls: list[str] = []
|
||
seen_urls: set[str] = set()
|
||
for url in image_urls:
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
urls.append(url)
|
||
paths = [os.path.abspath(p) for p, _ in images]
|
||
ocr_map = await self._batch_ocr(paths)
|
||
passages = [ocr_map[p].strip() for p in paths if ocr_map.get(p, "").strip()]
|
||
combined = "\n\n---\n\n".join(passages) if passages else None
|
||
return RetrievalResult(
|
||
text=combined,
|
||
images=[],
|
||
source_url=", ".join(urls) if urls else r.source_url,
|
||
retrieval_type=f"{r.retrieval_type}+ocr",
|
||
pixel_query_path=r.pixel_query_path,
|
||
query_image_path=r.query_image_path,
|
||
)
|
||
|
||
|
||
class RenderedTextWrapper(BaseRetriever):
|
||
"""Wraps a text retriever; renders each chunk as an image.
|
||
|
||
Ablation B pipeline: text retrieve -> render as Wikipedia-style image -> VLM reader.
|
||
Requires the base retriever to expose get_hits(query, example) returning
|
||
per-hit dicts with keys: title, text, url, score, article_id, chunk_index.
|
||
(TextAPIRetriever satisfies this.)
|
||
|
||
Renders are cached on disk at {render_dir}/{article_id}_{chunk_index}.png
|
||
so repeated eval runs don't re-render.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
base: BaseRetriever,
|
||
render_dir: str = "rendered_chunks",
|
||
reader_top_k: int | None = None,
|
||
):
|
||
if not hasattr(base, "get_hits"):
|
||
raise TypeError(
|
||
f"RenderedTextWrapper requires base retriever with get_hits(); "
|
||
f"got {type(base).__name__}"
|
||
)
|
||
self.base = base
|
||
self.render_dir = render_dir
|
||
self.reader_top_k = reader_top_k
|
||
os.makedirs(self.render_dir, exist_ok=True)
|
||
self.tiles_dir = render_dir
|
||
|
||
async def prefetch(self, examples: list[dict]):
|
||
if hasattr(self.base, "prefetch"):
|
||
await self.base.prefetch(examples)
|
||
|
||
def _render(self, hit: dict) -> str:
|
||
from .text_renderer import render_text_chunk
|
||
|
||
article_id = hit.get("article_id", "unknown")
|
||
chunk_index = hit.get("chunk_index", 0)
|
||
out_path = os.path.join(self.render_dir, f"{article_id}_{chunk_index}.png")
|
||
if os.path.isfile(out_path):
|
||
return out_path
|
||
# No-title policy: mirrors `_hits_to_result` (line ~3035) — title/url are
|
||
# leaked metadata for entity-answering tasks and were stripped from the
|
||
# text→text path on 2026-04-29. Apply the same constraint here so
|
||
# rendered and text→text differ only in modality, not in content.
|
||
render_text_chunk(
|
||
text=hit.get("text", ""),
|
||
title=None,
|
||
url=None,
|
||
output_path=out_path,
|
||
)
|
||
return out_path
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
hits = await self.base.get_hits(query, example)
|
||
if not hits:
|
||
return RetrievalResult(retrieval_type="text_api+rendered")
|
||
if self.reader_top_k is not None:
|
||
hits = hits[: self.reader_top_k]
|
||
images: list[tuple[str, float]] = []
|
||
urls: list[str] = []
|
||
seen_urls: set[str] = set()
|
||
for hit in hits:
|
||
if not hit.get("text"):
|
||
continue
|
||
path = self._render(hit)
|
||
images.append((path, float(hit.get("score", 0.0))))
|
||
url = hit.get("url", "")
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
urls.append(url)
|
||
return RetrievalResult(
|
||
images=images,
|
||
source_url=", ".join(urls) if urls else None,
|
||
retrieval_type="text_api+rendered",
|
||
)
|
||
|
||
|
||
class HybridRetriever(BaseRetriever):
|
||
"""Merge image (LocalAPIRetriever) and text (TextAPIRetriever) hits by raw score.
|
||
|
||
Both underlying retrievers embed with Qwen3-VL-Embedding-2B against L2-normalized
|
||
FAISS IVFFlat (IP metric) indices, so their per-hit scores are cosine similarities
|
||
on the same scale and directly comparable without any normalization step.
|
||
|
||
Each base is called with its own configured top_k, then the combined candidate pool
|
||
is sorted by score desc and the top `top_k` are kept. The reader receives the
|
||
surviving image hits as image inputs and the surviving text hits as a concatenated
|
||
text block in the same prompt — VL-4B handles mixed modality natively.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
image_base: "LocalAPIRetriever",
|
||
text_base: "TextAPIRetriever",
|
||
top_k: int = 3,
|
||
reader_top_k: int | None = None,
|
||
):
|
||
if not hasattr(image_base, "get_hits"):
|
||
raise TypeError(
|
||
f"HybridRetriever.image_base requires get_hits(); got {type(image_base).__name__}"
|
||
)
|
||
if not hasattr(text_base, "get_hits"):
|
||
raise TypeError(
|
||
f"HybridRetriever.text_base requires get_hits(); got {type(text_base).__name__}"
|
||
)
|
||
self.image_base = image_base
|
||
self.text_base = text_base
|
||
self.top_k = top_k
|
||
self.reader_top_k = reader_top_k
|
||
self.tiles_dir = getattr(image_base, "tiles_dir", None)
|
||
|
||
async def prefetch(self, examples: list[dict]):
|
||
if hasattr(self.image_base, "prefetch"):
|
||
await self.image_base.prefetch(examples)
|
||
if hasattr(self.text_base, "prefetch"):
|
||
await self.text_base.prefetch(examples)
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
image_hits = await self.image_base.get_hits(query, example)
|
||
text_hits = await self.text_base.get_hits(query, example)
|
||
|
||
# Tag each hit with its modality, then merge and sort by score desc.
|
||
merged: list[tuple[float, str, dict]] = []
|
||
for h in image_hits:
|
||
score = float(h.get("score", 0.0))
|
||
merged.append((score, "image", h))
|
||
for h in text_hits:
|
||
score = float(h.get("score", 0.0))
|
||
merged.append((score, "text", h))
|
||
|
||
merged.sort(key=lambda x: x[0], reverse=True)
|
||
keep_k = self.reader_top_k if self.reader_top_k is not None else self.top_k
|
||
top = merged[:keep_k]
|
||
|
||
images: list[tuple[str, float]] = []
|
||
passages: list[str] = []
|
||
urls: list[str] = []
|
||
seen_urls: set[str] = set()
|
||
|
||
for score, modality, hit in top:
|
||
url = hit.get("url", "")
|
||
if modality == "image":
|
||
path = hit.get("path", "")
|
||
if path and os.path.exists(path):
|
||
images.append((path, score))
|
||
else: # text
|
||
title = hit.get("title", "")
|
||
text = hit.get("text", "")
|
||
if text:
|
||
header = f"[{title}]" if title else ""
|
||
passages.append(f"{header}\n{text}" if header else text)
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
urls.append(url)
|
||
|
||
return RetrievalResult(
|
||
text="\n\n".join(passages) if passages else None,
|
||
images=images,
|
||
source_url=", ".join(urls) if urls else None,
|
||
retrieval_type="hybrid",
|
||
)
|
||
|
||
|
||
class HTMLDOMLookupRetriever(BaseRetriever):
|
||
"""Text-retrieve → DOM lookup: retrieve text chunks, then find their HTML context.
|
||
|
||
Wraps TextAPIRetriever. For each retrieved text chunk:
|
||
1. Fetches original HTML from kiwix-serve using article_id
|
||
2. Locates the chunk text within the HTML DOM
|
||
3. Extracts the enclosing semantic container (section/table/div)
|
||
4. Returns structured HTML context to the reader
|
||
|
||
This gives the reader table/list structure without needing a separate HTML index.
|
||
Falls back to plain text if DOM lookup fails for a chunk.
|
||
"""
|
||
|
||
KIWIX_BASE = "http://localhost:9454/content/wikipedia_en_all_maxi_2025-08"
|
||
|
||
def __init__(
|
||
self,
|
||
text_api_url: str = "http://localhost:30889/search",
|
||
top_k: int = 3,
|
||
nprobe: int | None = None,
|
||
query_instruction: str | None = None,
|
||
reader_top_k: int | None = None,
|
||
query_image_fn=None,
|
||
kiwix_base: str | None = None,
|
||
articles_json: str = "/path/to/data",
|
||
context_mode: str = "section",
|
||
llm_verify: bool = False,
|
||
llm_verify_model: str = "gpt-4.1-mini",
|
||
):
|
||
import json as _json
|
||
|
||
self._text_retriever = TextAPIRetriever(
|
||
api_url=text_api_url,
|
||
top_k=top_k,
|
||
nprobe=nprobe,
|
||
query_instruction=query_instruction,
|
||
reader_top_k=reader_top_k,
|
||
query_image_fn=query_image_fn,
|
||
)
|
||
if kiwix_base:
|
||
self.KIWIX_BASE = kiwix_base
|
||
self.top_k = top_k
|
||
self.reader_top_k = reader_top_k
|
||
self.context_mode = context_mode
|
||
self.llm_verify = llm_verify
|
||
self.llm_verify_model = llm_verify_model
|
||
|
||
with open(articles_json) as f:
|
||
self._articles: list[str] = _json.load(f)
|
||
|
||
self._html_cache: dict[int, str] = {}
|
||
|
||
async def prefetch(self, examples: list[dict]):
|
||
await self._text_retriever.prefetch(examples)
|
||
|
||
def _fetch_html(self, article_id: int) -> str | None:
|
||
"""Fetch article HTML from kiwix-serve (with caching)."""
|
||
if article_id in self._html_cache:
|
||
return self._html_cache[article_id]
|
||
|
||
if article_id >= len(self._articles):
|
||
return None
|
||
|
||
import requests
|
||
from urllib.parse import quote
|
||
|
||
slug = self._articles[article_id]
|
||
url = f"{self.KIWIX_BASE}/{quote(slug, safe='/:@!$&()*+,;=')}"
|
||
try:
|
||
resp = requests.get(url, timeout=15)
|
||
resp.encoding = "utf-8"
|
||
if resp.status_code != 200:
|
||
return None
|
||
self._html_cache[article_id] = resp.text
|
||
return resp.text
|
||
except Exception:
|
||
return None
|
||
|
||
@staticmethod
|
||
def _normalize(text: str) -> str:
|
||
"""Normalize text for fuzzy DOM matching."""
|
||
import re
|
||
import unicodedata
|
||
|
||
text = re.sub(r"[\xa0 ]", " ", text)
|
||
text = re.sub(r"[‐-―−﹘﹣－—–]", "-", text)
|
||
text = re.sub(r" +", " ", text)
|
||
text = unicodedata.normalize("NFKD", text)
|
||
text = "".join(c for c in text if not unicodedata.combining(c))
|
||
return text.lower()
|
||
|
||
def _dom_lookup(self, html: str, chunk_text: str) -> str | None:
|
||
"""Find the contiguous DOM span covering chunk_text, return its HTML.
|
||
|
||
Strategy:
|
||
1. Extract search keys from chunk text (table cells + prose fragments)
|
||
2. For each key, find the tightest DOM element and walk up to a
|
||
direct child of mw-parser-output
|
||
3. Return ALL direct children from the first match to the last match
|
||
(inclusive), plus everything in between — this preserves the full
|
||
contiguous region the chunk spans.
|
||
"""
|
||
from lxml import html as lxml_html, etree
|
||
|
||
tree = lxml_html.fromstring(html)
|
||
|
||
keys = self._extract_search_keys(chunk_text)
|
||
if not keys:
|
||
return None
|
||
|
||
mw_output = tree.xpath('//div[contains(@class, "mw-parser-output")]')
|
||
if not mw_output:
|
||
return None
|
||
content_root = mw_output[0]
|
||
children = list(content_root)
|
||
if not children:
|
||
return None
|
||
|
||
# For each key, find the tightest match and resolve to a
|
||
# direct-child index of mw-parser-output
|
||
matched_child_indices = set()
|
||
SKIP_TAGS = frozenset(
|
||
("script", "style", "title", "meta", "link", "nav", "header", "footer")
|
||
)
|
||
|
||
for key in keys:
|
||
key_norm = self._normalize(key)
|
||
if len(key_norm) < 4:
|
||
continue
|
||
|
||
best_elem = None
|
||
best_len = float("inf")
|
||
|
||
for elem in content_root.iter():
|
||
if not isinstance(elem, lxml_html.HtmlElement):
|
||
continue
|
||
if elem.tag in SKIP_TAGS:
|
||
continue
|
||
try:
|
||
tc = elem.text_content()
|
||
except Exception:
|
||
continue
|
||
tc_norm = self._normalize(tc)
|
||
if key_norm in tc_norm and len(tc) < best_len:
|
||
best_elem = elem
|
||
best_len = len(tc)
|
||
|
||
if best_elem is None:
|
||
continue
|
||
|
||
# Walk up from best_elem to find which direct child of content_root
|
||
# contains it
|
||
current = best_elem
|
||
while current is not None:
|
||
parent = current.getparent()
|
||
if parent is None:
|
||
break
|
||
if parent == content_root:
|
||
# current is a direct child of mw-parser-output
|
||
try:
|
||
idx = children.index(current)
|
||
matched_child_indices.add(idx)
|
||
except ValueError:
|
||
pass
|
||
break
|
||
current = parent
|
||
|
||
if not matched_child_indices:
|
||
return None
|
||
|
||
# Return contiguous range from first to last matched child (inclusive)
|
||
first = min(matched_child_indices)
|
||
last = max(matched_child_indices)
|
||
|
||
span_elems = children[first : last + 1]
|
||
|
||
# Build result: serialize all elements in the span
|
||
parts = []
|
||
for el in span_elems:
|
||
# Strip style/script/navbox noise
|
||
for tag in ("style", "script"):
|
||
for junk in list(el.iter(tag)):
|
||
if junk.getparent() is not None:
|
||
junk.getparent().remove(junk)
|
||
if hasattr(el, "xpath"):
|
||
for nav in el.xpath('.//*[contains(@class, "navbox")]'):
|
||
if nav.getparent() is not None:
|
||
nav.getparent().remove(nav)
|
||
try:
|
||
parts.append(etree.tostring(el, encoding="unicode", method="html"))
|
||
except Exception:
|
||
continue
|
||
|
||
if not parts:
|
||
return None
|
||
|
||
html_str = "\n".join(parts)
|
||
|
||
# Log oversized results but still return them (caller decides)
|
||
if len(html_str) > self.MAX_CONTAINER_CHARS * 2:
|
||
logger.warning(
|
||
"DOM lookup oversized: %d chars (max %d) for chunk starting with %r",
|
||
len(html_str),
|
||
self.MAX_CONTAINER_CHARS * 2,
|
||
chunk_text[:50],
|
||
)
|
||
|
||
# Minimum useful size
|
||
if len(html_str) < 100 and len(chunk_text) > 200:
|
||
return None
|
||
|
||
return html_str
|
||
|
||
MAX_CONTAINER_CHARS = 8000
|
||
|
||
def _find_semantic_container(self, elem) -> "lxml_html.HtmlElement": # noqa: F821
|
||
"""Walk up from matched element to find a meaningful semantic container.
|
||
|
||
Hard cap: never return a container with text_content > MAX_CONTAINER_CHARS.
|
||
Stops at mw-parser-output boundary (never returns the whole article).
|
||
"""
|
||
|
||
SEMANTIC_TAGS = {
|
||
"section",
|
||
"article",
|
||
"table",
|
||
"blockquote",
|
||
"details",
|
||
"figure",
|
||
}
|
||
STOP_CLASSES = {"mw-parser-output", "mw-body-content", "mw-body"}
|
||
MIN_CONTEXT_LEN = 200
|
||
|
||
if elem.tag in SEMANTIC_TAGS:
|
||
return elem
|
||
|
||
best = elem
|
||
current = elem
|
||
|
||
for _ in range(15):
|
||
parent = current.getparent()
|
||
if parent is None:
|
||
break
|
||
# Hard stop: never go past the article content container
|
||
parent_classes = parent.get("class", "")
|
||
if any(sc in parent_classes for sc in STOP_CLASSES):
|
||
# We've reached the article root — use section gathering instead
|
||
if self.context_mode == "section":
|
||
gathered = self._gather_section_context(current)
|
||
if gathered is not None:
|
||
return gathered
|
||
break
|
||
|
||
try:
|
||
parent_len = len(parent.text_content())
|
||
except Exception:
|
||
break
|
||
|
||
# Prefer semantic tags — even if parent exceeds size cap
|
||
# Bug fix 2: tbody→table jump — don't let size cap block us from
|
||
# reaching a semantic container that's just one level up
|
||
if parent.tag in SEMANTIC_TAGS:
|
||
return parent
|
||
|
||
# Stop if parent is too large (but we already checked semantic tags above)
|
||
if parent_len > self.MAX_CONTAINER_CHARS:
|
||
# One more chance: check if grandparent is a semantic tag
|
||
grandparent = parent.getparent()
|
||
if grandparent is not None and grandparent.tag in SEMANTIC_TAGS:
|
||
return grandparent
|
||
break
|
||
|
||
# Accept block containers that are reasonably sized
|
||
if parent_len >= MIN_CONTEXT_LEN:
|
||
best = parent
|
||
|
||
current = parent
|
||
|
||
return best
|
||
|
||
def _gather_section_context(self, elem) -> "lxml_html.HtmlElement": # noqa: F821
|
||
"""Gather all sibling elements within the same h2/h3 section."""
|
||
from lxml import etree
|
||
|
||
# Walk up to find direct child of mw-parser-output
|
||
current = elem
|
||
mw_output = None
|
||
while current is not None:
|
||
parent = current.getparent()
|
||
if parent is not None:
|
||
classes = parent.get("class", "")
|
||
if "mw-parser-output" in classes:
|
||
mw_output = parent
|
||
break
|
||
current = parent
|
||
|
||
if mw_output is None:
|
||
return elem
|
||
|
||
# Find the element's position among mw-parser-output children
|
||
children = list(mw_output)
|
||
try:
|
||
idx = children.index(current)
|
||
except ValueError:
|
||
return elem
|
||
|
||
# Gather backward until we hit a heading, forward until next heading
|
||
section_elems = [current]
|
||
|
||
# Backward
|
||
for i in range(idx - 1, max(idx - 10, -1), -1):
|
||
child = children[i]
|
||
if hasattr(child, "tag") and child.tag in ("h1", "h2", "h3"):
|
||
section_elems.insert(0, child)
|
||
break
|
||
section_elems.insert(0, child)
|
||
|
||
# Forward
|
||
for i in range(idx + 1, min(idx + 10, len(children))):
|
||
child = children[i]
|
||
if hasattr(child, "tag") and child.tag in ("h1", "h2", "h3"):
|
||
break
|
||
section_elems.append(child)
|
||
|
||
# Build a container div with these elements
|
||
container = etree.Element("div")
|
||
for el in section_elems:
|
||
try:
|
||
container.append(el)
|
||
except Exception:
|
||
pass
|
||
|
||
return container
|
||
|
||
@staticmethod
|
||
def _extract_search_keys(chunk_text: str) -> list[str]:
|
||
"""Extract distinctive search keys from chunk text for DOM matching.
|
||
|
||
Detects chunk type (table-heavy vs prose-heavy) and picks the best strategy.
|
||
Returns keys ordered by distinctiveness — first key is tried first in DOM lookup.
|
||
"""
|
||
import re
|
||
|
||
lines = chunk_text.split("\n")
|
||
# Skip first line if it looks like an article title (short, no pipes, no punctuation)
|
||
# These match <h1> in DOM and cause Bug 1
|
||
if lines and len(lines[0]) < 80 and "|" not in lines[0] and "." not in lines[0]:
|
||
content_lines = lines[1:]
|
||
else:
|
||
content_lines = lines
|
||
table_lines = [l for l in content_lines if "|" in l and "---" not in l]
|
||
prose_lines = [
|
||
l
|
||
for l in content_lines
|
||
if len(l) > 30 and "|" not in l and not l.startswith("- ^")
|
||
]
|
||
is_table_heavy = (
|
||
len(table_lines) > len(content_lines) * 0.4 if content_lines else False
|
||
)
|
||
|
||
keys = []
|
||
|
||
if is_table_heavy:
|
||
# Mixed strategy: include keys from BOTH table cells and prose
|
||
# so coverage scorer can find a container spanning both parts.
|
||
cell_candidates = []
|
||
for tl in table_lines:
|
||
cells = [c.strip() for c in tl.split("|") if c.strip()]
|
||
for cell in cells:
|
||
if len(cell) < 5 or len(cell) > 80:
|
||
continue
|
||
if cell.lower() in ("yes", "no", "n/a", "none", ""):
|
||
continue
|
||
has_code = bool(re.search(r"[A-Z]\d|[a-z]\d{4,}", cell))
|
||
has_mixed = bool(re.search(r"\d.*[a-zA-Z]|[a-zA-Z].*\d", cell))
|
||
has_proper = bool(re.search(r"[A-Z][a-z]+\s+[A-Z]", cell))
|
||
if has_code or has_mixed or has_proper:
|
||
cell_candidates.insert(0, cell)
|
||
elif len(cell) > 12:
|
||
cell_candidates.append(cell)
|
||
|
||
# Table cells first (these anchor to the infobox)
|
||
for cc in cell_candidates[:3]:
|
||
if cc not in keys:
|
||
keys.append(cc)
|
||
|
||
# Then prose keys (these anchor to body paragraphs)
|
||
for line in prose_lines[:3]:
|
||
mid = len(line) // 2
|
||
candidate = line[mid - 15 : mid + 15].strip()
|
||
if len(candidate) >= 10 and re.search(r"[a-zA-Z]{4,}", candidate):
|
||
keys.append(candidate)
|
||
|
||
else:
|
||
# Prose-dominant chunk: use prose fragments as primary keys
|
||
for line in prose_lines[:4]:
|
||
mid = len(line) // 2
|
||
candidate = line[mid - 15 : mid + 15].strip()
|
||
if len(candidate) >= 10 and re.search(r"[a-zA-Z]{4,}", candidate):
|
||
keys.append(candidate)
|
||
|
||
# Add table cell values as secondary
|
||
if table_lines:
|
||
for tl in table_lines[:5]:
|
||
cells = [
|
||
c.strip()
|
||
for c in tl.split("|")
|
||
if c.strip() and len(c.strip()) > 8
|
||
]
|
||
for cell in cells[:1]:
|
||
if cell not in keys:
|
||
keys.append(cell)
|
||
|
||
# List item content
|
||
if not keys:
|
||
list_lines = [
|
||
l[2:]
|
||
for l in lines
|
||
if l.startswith("- ") and len(l) > 20 and not l.startswith("- ^")
|
||
]
|
||
for ll in list_lines[:3]:
|
||
mid = len(ll) // 2
|
||
candidate = ll[mid - 15 : mid + 15].strip()
|
||
if len(candidate) >= 10:
|
||
keys.append(candidate)
|
||
|
||
# Fallback
|
||
if not keys and len(chunk_text) > 40:
|
||
candidate = chunk_text[10:50].strip()
|
||
keys.append(candidate)
|
||
|
||
return keys
|
||
|
||
async def _llm_dom_closure(self, raw_html: str, chunk_text: str) -> str | None:
|
||
"""Use an LLM to find the minimal DOM closure containing the chunk text.
|
||
|
||
Sends the article HTML (truncated) and the chunk text to the model,
|
||
asks it to return the minimal enclosing HTML subtree.
|
||
"""
|
||
import openai
|
||
|
||
# Truncate HTML to avoid context limits — keep first 60K chars
|
||
# (most Wikipedia articles are under 100K)
|
||
html_truncated = raw_html[:60000]
|
||
|
||
prompt = f"""Given this HTML document and a text chunk extracted from it, find the minimal DOM subtree that contains ALL the text in the chunk. Return ONLY the raw HTML of that subtree, no explanation.
|
||
|
||
The text chunk (extracted by Trafilatura, so formatting differs from HTML):
|
||
---
|
||
{chunk_text[:2000]}
|
||
---
|
||
|
||
The HTML document:
|
||
---
|
||
{html_truncated}
|
||
---
|
||
|
||
Return the minimal HTML subtree containing all the information from the text chunk. Include complete table/list structures if the chunk spans table cells. Return ONLY HTML, no markdown fences."""
|
||
|
||
try:
|
||
client = openai.AsyncOpenAI()
|
||
response = await client.chat.completions.create(
|
||
model=self.llm_verify_model,
|
||
messages=[{"role": "user", "content": prompt}],
|
||
max_tokens=8000,
|
||
temperature=0,
|
||
)
|
||
result = response.choices[0].message.content.strip()
|
||
# Strip markdown fences if present
|
||
if result.startswith("```"):
|
||
lines = result.split("\n")
|
||
result = "\n".join(
|
||
lines[1:-1] if lines[-1].startswith("```") else lines[1:]
|
||
)
|
||
return result if "<" in result else None
|
||
except Exception as e:
|
||
logger.warning(f"LLM DOM closure failed: {e}")
|
||
return None
|
||
|
||
async def retrieve(self, query: str, example: dict) -> RetrievalResult:
|
||
"""Retrieve text chunks, then do DOM lookup for HTML context."""
|
||
example.get("id", "unknown")
|
||
|
||
# Get raw hits from text retriever
|
||
hits = await self._text_retriever.get_hits(query, example)
|
||
if not hits:
|
||
return RetrievalResult(retrieval_type="html_dom_lookup")
|
||
|
||
keep_k = self.reader_top_k if self.reader_top_k is not None else self.top_k
|
||
hits = hits[:keep_k]
|
||
|
||
passages = []
|
||
urls = []
|
||
seen_urls: set[str] = set()
|
||
|
||
for hit in hits:
|
||
article_id = hit.get("article_id")
|
||
chunk_text = hit.get("text", "")
|
||
url = hit.get("url", "")
|
||
|
||
html_context = None
|
||
if article_id is not None:
|
||
raw_html = self._fetch_html(int(article_id))
|
||
if raw_html:
|
||
# Heuristic DOM lookup first
|
||
html_context = self._dom_lookup(raw_html, chunk_text)
|
||
|
||
# LLM verification/fallback
|
||
if self.llm_verify and (
|
||
html_context is None or len(chunk_text) > 500
|
||
):
|
||
llm_result = await self._llm_dom_closure(raw_html, chunk_text)
|
||
if llm_result:
|
||
html_context = llm_result
|
||
|
||
if html_context:
|
||
passages.append(html_context)
|
||
else:
|
||
passages.append(chunk_text)
|
||
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
urls.append(url)
|
||
|
||
# Hard cap per passage. HTML is ~2 chars/token; reader has 65K tokens
|
||
# with ~2K for output + system prompt. Budget ~50K tokens for context
|
||
# = ~100K chars across all passages. Per-passage cap avoids one huge
|
||
# article starving the others.
|
||
MAX_PER_PASSAGE = 30000
|
||
passages = [p[:MAX_PER_PASSAGE] for p in passages]
|
||
MAX_TOTAL_CHARS = 90000
|
||
total = sum(len(p) for p in passages)
|
||
if total > MAX_TOTAL_CHARS:
|
||
per_passage = MAX_TOTAL_CHARS // max(len(passages), 1)
|
||
passages = [p[:per_passage] for p in passages]
|
||
logger.warning(
|
||
"Truncated %d passages from %d to %d total chars",
|
||
len(passages),
|
||
total,
|
||
sum(len(p) for p in passages),
|
||
)
|
||
|
||
combined = "\n\n---\n\n".join(passages) if passages else None
|
||
return RetrievalResult(
|
||
text=combined,
|
||
source_url=", ".join(urls) if urls else None,
|
||
retrieval_type="html_dom_lookup",
|
||
)
|