wehub-resource-sync
551 lines
21 KiB
Python
551 lines
21 KiB
Python
"""Retriever factory — maps CLI flags to a (retriever, mode_str) pair.
|
|
|
|
Extracted from run_naive_simpleqa.py to keep the orchestrator readable.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
from . import (
|
|
NaiveRetriever,
|
|
ScreenshotRetriever,
|
|
TiledScreenshotRetriever,
|
|
LocalWikiTiledScreenshotRetriever,
|
|
TextRetriever,
|
|
JinaReaderRetriever,
|
|
WikipediaAPIRetriever,
|
|
VectorRetriever,
|
|
ColQwenVectorRetriever,
|
|
TiledVectorRetriever,
|
|
TiledColQwenVectorRetriever,
|
|
TiledQwen3VLEmbeddingRetriever,
|
|
EVQANoRetrievalRetriever,
|
|
WorldVQANoRetrievalRetriever,
|
|
TextVectorRetriever,
|
|
DsServeRetriever,
|
|
LocalAPIRetriever,
|
|
TextAPIRetriever,
|
|
OCRWrappedRetriever,
|
|
RenderedTextWrapper,
|
|
HybridRetriever,
|
|
HTMLDOMLookupRetriever,
|
|
load_text_cache,
|
|
)
|
|
from .retrieval import _get_query_image_path_for_example, _save_task_query_image
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TILE_WIDTH = 1024
|
|
|
|
|
|
def build_retriever(args, examples, model, api_base, api_key):
|
|
"""Build a retriever from CLI args.
|
|
|
|
Args:
|
|
args: Parsed argparse namespace.
|
|
examples: Loaded dataset examples (some retrievers need them for setup).
|
|
model: Reader model name (for query rewrite fallback).
|
|
api_base: Reader API base (for query rewrite fallback).
|
|
api_key: Reader API key (for query rewrite fallback).
|
|
|
|
Returns:
|
|
(retriever, mode_str) tuple.
|
|
"""
|
|
tile_size = (TILE_WIDTH, args.tile_height)
|
|
|
|
retrieval_mode_count = sum(
|
|
[
|
|
args.url_screenshot,
|
|
args.url_tiled_screenshot,
|
|
args.url_text,
|
|
args.url_jina_reader,
|
|
args.retrieval_augment,
|
|
args.use_tiled_retrieval,
|
|
args.text_vector,
|
|
args.local_api,
|
|
args.text_api,
|
|
args.html_dom_lookup,
|
|
args.hybrid,
|
|
]
|
|
)
|
|
|
|
if args.url_screenshot:
|
|
retriever = ScreenshotRetriever(
|
|
screenshot_dir=args.screenshot_dir, max_pixels=args.max_pixels
|
|
)
|
|
mode = f"Screenshot (Ground Truth, max_pixels={args.max_pixels or 'None'})"
|
|
|
|
elif args.url_tiled_screenshot and args.local_wiki:
|
|
retriever = LocalWikiTiledScreenshotRetriever(
|
|
tiles_dir=args.tiles_dir,
|
|
wiki_cache_dir=args.local_wiki_screenshot_dir,
|
|
tile_height=args.tile_height,
|
|
max_tiles=args.max_tiles,
|
|
)
|
|
mode = f"Local-Wiki Tiled Screenshot (Ground Truth, tile_height={args.tile_height}, max_tiles={args.max_tiles})"
|
|
|
|
elif args.url_tiled_screenshot:
|
|
retriever = TiledScreenshotRetriever(
|
|
screenshot_dir=args.screenshot_dir,
|
|
tiles_dir=args.tiles_dir,
|
|
tile_size=tile_size,
|
|
overlap=args.tile_overlap,
|
|
max_tiles=args.max_tiles,
|
|
)
|
|
mode = f"Tiled Screenshot (Ground Truth, max_tiles={args.max_tiles})"
|
|
|
|
elif args.url_text:
|
|
text_cache = None
|
|
if args.text_cache and os.path.exists(args.text_cache):
|
|
text_cache = load_text_cache(args.text_cache)
|
|
logger.info(f"Loaded {len(text_cache)} cached items from {args.text_cache}")
|
|
elif args.text_cache:
|
|
logger.info(
|
|
f"Cache file not found: {args.text_cache} (will fetch from source)"
|
|
)
|
|
if args.text_source == "jina":
|
|
retriever = JinaReaderRetriever(
|
|
max_chars=args.max_context_chars,
|
|
api_key=args.jina_api_key,
|
|
text_cache=text_cache,
|
|
cache_path=args.text_cache,
|
|
)
|
|
mode = "Text RAG (Jina)"
|
|
elif args.text_source == "wikipedia":
|
|
retriever = WikipediaAPIRetriever(
|
|
max_chars=args.max_context_chars,
|
|
text_cache=text_cache,
|
|
cache_path=args.text_cache,
|
|
)
|
|
mode = "Text RAG (Wikipedia API)"
|
|
else:
|
|
retriever = TextRetriever(
|
|
max_chars=args.max_context_chars,
|
|
text_cache=text_cache,
|
|
cache_path=args.text_cache,
|
|
)
|
|
mode = "Text RAG (Crawl)"
|
|
|
|
elif args.url_jina_reader:
|
|
logger.warning(
|
|
"--url-jina-reader is deprecated, use --url-text --text-source jina instead"
|
|
)
|
|
retriever = JinaReaderRetriever(
|
|
max_chars=args.max_context_chars, api_key=args.jina_api_key
|
|
)
|
|
mode = "Jina Reader"
|
|
|
|
elif args.retrieval_augment:
|
|
if args.use_colqwen_retrieval:
|
|
retriever = ColQwenVectorRetriever(
|
|
index_path=args.colqwen_index_path,
|
|
screenshot_dir=args.screenshot_dir,
|
|
model_name=args.colqwen_model,
|
|
search_method=args.colqwen_search_method,
|
|
first_stage_k=args.colqwen_first_stage_k,
|
|
rebuild_index=args.rebuild_colqwen_index,
|
|
recursive=args.colqwen_recursive,
|
|
top_k=args.retrieval_top_k,
|
|
examples=examples,
|
|
)
|
|
mode = "ColQwen Vector Retrieval"
|
|
else:
|
|
retriever = VectorRetriever(
|
|
api_key=args.jina_api_key,
|
|
screenshot_dir=args.screenshot_dir,
|
|
cache_path=args.retrieval_cache,
|
|
use_multivector=not args.single_vector,
|
|
top_k=args.retrieval_top_k,
|
|
examples=examples,
|
|
)
|
|
mode = "Vector Retrieval"
|
|
|
|
elif args.use_tiled_retrieval:
|
|
if args.use_colqwen_retrieval:
|
|
tiled_index_path = args.colqwen_index_path.replace(
|
|
".leann", f"_tiled_{args.tile_height}.leann"
|
|
)
|
|
retriever = TiledColQwenVectorRetriever(
|
|
index_path=tiled_index_path,
|
|
screenshot_dir=args.screenshot_dir,
|
|
tiles_dir=args.tiles_dir,
|
|
tile_size=tile_size,
|
|
overlap=args.tile_overlap,
|
|
model_name=args.colqwen_model,
|
|
search_method=args.colqwen_search_method,
|
|
first_stage_k=args.colqwen_first_stage_k,
|
|
rebuild_index=args.rebuild_colqwen_index,
|
|
top_k=args.retrieval_top_k,
|
|
examples=examples,
|
|
)
|
|
mode = "Tiled ColQwen Vector Retrieval"
|
|
elif args.use_qwen3vl_embedding:
|
|
qwen3vl_cache_path = args.retrieval_cache
|
|
if qwen3vl_cache_path is None:
|
|
task_subset = f"{args.task}_{args.subset}" if args.subset else args.task
|
|
localwiki_suffix = "_localwiki" if args.local_wiki else ""
|
|
qwen3vl_cache_path = f"qwen3vl_tiles_{task_subset}_{TILE_WIDTH}x{args.tile_height}_{args.num_examples}ex{localwiki_suffix}_embeddings.pkl"
|
|
qwen3vl_gpu_ids = [int(x.strip()) for x in args.qwen3vl_gpu_ids.split(",")]
|
|
|
|
pixel_query_map = None
|
|
if (
|
|
args.task == "encyclopedic_vqa"
|
|
and not args.evqa_multimodal_query
|
|
and not args.evqa_multi_image_query
|
|
):
|
|
from .pixel_query import QueryImageTextRenderer
|
|
|
|
tiles_dir = args.tiles_dir or "tiles/evqa"
|
|
renderer = QueryImageTextRenderer(
|
|
output_dir="query_cards/evqa",
|
|
tiles_dir=tiles_dir,
|
|
)
|
|
pixel_query_map = {}
|
|
for ex in examples:
|
|
inat_path = _get_query_image_path_for_example(ex, tiles_dir)
|
|
path = renderer.render(
|
|
ex["id"], ex["problem"], inat_path, force=args.force
|
|
)
|
|
pixel_query_map[ex["id"]] = path
|
|
logger.info(f"EVQA query cards: {len(pixel_query_map)} rendered")
|
|
elif args.pixel_query:
|
|
from .pixel_query import PixelQueryRenderer
|
|
|
|
pq_renderer = PixelQueryRenderer(output_dir=args.pixel_query_dir)
|
|
pixel_query_map = pq_renderer.render_all(examples)
|
|
logger.info(
|
|
f"Pixel query mode: rendered {len(pixel_query_map)} query images"
|
|
)
|
|
|
|
retriever = TiledQwen3VLEmbeddingRetriever(
|
|
screenshot_dir=args.screenshot_dir,
|
|
tiles_dir=args.tiles_dir,
|
|
tile_size=tile_size,
|
|
overlap=args.tile_overlap,
|
|
cache_path=qwen3vl_cache_path,
|
|
model_name=args.qwen3vl_model,
|
|
top_k=args.retrieval_top_k,
|
|
examples=examples,
|
|
gpu_ids=qwen3vl_gpu_ids,
|
|
tensor_parallel_size=args.qwen3vl_tp_size,
|
|
pixel_query_map=pixel_query_map,
|
|
multimodal_query_text_only=args.evqa_multimodal_query_text_only,
|
|
multimodal_query_image_only=args.evqa_multimodal_query_image_only,
|
|
local_wiki=args.local_wiki,
|
|
local_wiki_screenshot_dir=args.local_wiki_screenshot_dir,
|
|
multi_image_query=args.evqa_multi_image_query,
|
|
prebuilt_tiles_dir=getattr(args, "prebuilt_tiles_dir", None),
|
|
embedding_backend=getattr(args, "embedding_backend", "vllm"),
|
|
peft_adapter=getattr(args, "peft_adapter", None),
|
|
)
|
|
mode = "Tiled Qwen3-VL-Embedding Retrieval"
|
|
if getattr(args, "prebuilt_tiles_dir", None):
|
|
mode += " (prebuilt hard-mini)"
|
|
elif args.local_wiki:
|
|
mode += " (local-wiki)"
|
|
if args.task == "encyclopedic_vqa":
|
|
if args.evqa_multi_image_query:
|
|
mode += " (EVQA multi-image query)"
|
|
elif args.evqa_multimodal_query:
|
|
if args.evqa_multimodal_query_text_only:
|
|
mode += " (EVQA multimodal: text-only)"
|
|
elif args.evqa_multimodal_query_image_only:
|
|
mode += " (EVQA multimodal: image-only)"
|
|
else:
|
|
mode += " (EVQA multimodal: text+image)"
|
|
else:
|
|
mode += " (EVQA query card)"
|
|
elif args.pixel_query:
|
|
mode += " (Pixel Query)"
|
|
else:
|
|
tile_cache_path = args.retrieval_cache
|
|
if tile_cache_path is None:
|
|
vector_type = "single" if args.single_vector else "multi"
|
|
task_subset = f"{args.task}_{args.subset}" if args.subset else args.task
|
|
tile_cache_path = f"jina_tiles_{task_subset}_{TILE_WIDTH}x{args.tile_height}_{vector_type}_{args.num_examples}ex_embeddings.pkl"
|
|
retriever = TiledVectorRetriever(
|
|
api_key=args.jina_api_key,
|
|
screenshot_dir=args.screenshot_dir,
|
|
tiles_dir=args.tiles_dir,
|
|
tile_size=tile_size,
|
|
overlap=args.tile_overlap,
|
|
cache_path=tile_cache_path,
|
|
use_multivector=not args.single_vector,
|
|
top_k=args.retrieval_top_k,
|
|
examples=examples,
|
|
)
|
|
mode = "Tiled Jina Vector Retrieval"
|
|
|
|
elif args.local_api:
|
|
rw_model = args.rewrite_model or model
|
|
rw_api_base = args.rewrite_api_base or api_base
|
|
rw_api_key = args.rewrite_api_key or api_key
|
|
reranker_obj = None
|
|
if args.reranker:
|
|
logger.info(f"Loading reranker on GPU {args.reranker_gpu_id}")
|
|
from .reranker import Qwen3VLReranker
|
|
|
|
reranker_obj = Qwen3VLReranker(
|
|
model_name=args.reranker_model,
|
|
gpu_id=args.reranker_gpu_id,
|
|
)
|
|
query_image_fn = None
|
|
if args.no_query_image:
|
|
logger.info(
|
|
"--no-query-image set: retrieval queries will be text-only (reader still sees query image)"
|
|
)
|
|
elif args.task == "encyclopedic_vqa":
|
|
_tiles_dir = args.tiles_dir or "tiles/evqa"
|
|
|
|
def query_image_fn(ex, _td=_tiles_dir):
|
|
return _get_query_image_path_for_example(ex, _td, quiet=True)
|
|
elif args.task in (
|
|
"worldvqa",
|
|
"simplevqa",
|
|
"factualvqa",
|
|
"mmsearch",
|
|
"webqa",
|
|
"multimodalqa",
|
|
):
|
|
_task = args.task
|
|
|
|
def query_image_fn(ex, _t=_task):
|
|
return _save_task_query_image(ex, _t, base_dir="tiles")
|
|
|
|
retriever = LocalAPIRetriever(
|
|
api_url=args.local_api_url,
|
|
top_k=args.retrieval_top_k,
|
|
query_rewrite=args.query_rewrite,
|
|
rewrite_model=rw_model if args.query_rewrite else None,
|
|
rewrite_api_base=rw_api_base if args.query_rewrite else None,
|
|
rewrite_api_key=rw_api_key if args.query_rewrite else "dummy",
|
|
nprobe=args.nprobe,
|
|
reranker=reranker_obj,
|
|
rerank_top_k=args.rerank_top_k,
|
|
query_image_fn=query_image_fn,
|
|
multi_image_query=args.evqa_multi_image_query,
|
|
tiles_dir=args.tiles_dir or "tiles/evqa",
|
|
lookup_reference_url=args.lookup_reference_url,
|
|
query_instruction=args.query_instruction,
|
|
)
|
|
mode = f"Local API Retrieval ({args.local_api_url})"
|
|
if args.query_instruction is not None:
|
|
mode += f" [instr={args.query_instruction!r}]"
|
|
if args.evqa_multi_image_query:
|
|
mode += " (multi-image query)"
|
|
elif query_image_fn:
|
|
mode += " (multimodal query)"
|
|
if args.query_rewrite:
|
|
mode += f" + QueryRewrite({rw_model})"
|
|
if args.lookup_reference_url:
|
|
mode += " + RefURL"
|
|
if args.reranker:
|
|
mode += f" + Reranker({args.reranker_model}, top{args.rerank_top_k})"
|
|
if args.react:
|
|
mode += f" + ReAct({args.react_prompt}, max_turns={args.react_max_turns})"
|
|
|
|
elif args.text_api:
|
|
text_query_image_fn = None
|
|
if not args.no_query_image:
|
|
if args.task == "encyclopedic_vqa":
|
|
_tiles_dir = args.tiles_dir or "tiles/evqa"
|
|
|
|
def text_query_image_fn(ex, _td=_tiles_dir):
|
|
return _get_query_image_path_for_example(ex, _td, quiet=True)
|
|
elif args.task in (
|
|
"worldvqa",
|
|
"simplevqa",
|
|
"factualvqa",
|
|
"mmsearch",
|
|
"webqa",
|
|
"multimodalqa",
|
|
):
|
|
_task = args.task
|
|
|
|
def text_query_image_fn(ex, _t=_task):
|
|
return _save_task_query_image(ex, _t, base_dir="tiles")
|
|
|
|
retriever = TextAPIRetriever(
|
|
api_url=args.text_api_url,
|
|
top_k=args.retrieval_top_k,
|
|
nprobe=args.nprobe,
|
|
query_instruction=args.query_instruction,
|
|
reader_top_k=args.reader_top_k,
|
|
query_image_fn=text_query_image_fn,
|
|
)
|
|
mode = f"Text API Retrieval ({args.text_api_url})"
|
|
if args.query_instruction is not None:
|
|
mode += f" [instr={args.query_instruction!r}]"
|
|
|
|
elif args.html_dom_lookup:
|
|
retriever = HTMLDOMLookupRetriever(
|
|
text_api_url=args.text_api_url,
|
|
top_k=args.retrieval_top_k,
|
|
nprobe=args.nprobe,
|
|
query_instruction=args.query_instruction,
|
|
reader_top_k=args.reader_top_k,
|
|
query_image_fn=None,
|
|
context_mode="section",
|
|
llm_verify=getattr(args, "llm_verify", False),
|
|
)
|
|
mode = f"HTML DOM Lookup (text_api={args.text_api_url}, top_k={args.retrieval_top_k})"
|
|
if args.llm_verify:
|
|
mode += " [llm-verify]"
|
|
|
|
elif args.hybrid:
|
|
if args.read_as_text_ocr or args.render_as_image:
|
|
print(
|
|
"Error: --hybrid is not compatible with --read-as-text-ocr or --render-as-image."
|
|
)
|
|
sys.exit(1)
|
|
image_base = LocalAPIRetriever(
|
|
api_url=args.local_api_url,
|
|
top_k=args.retrieval_top_k,
|
|
nprobe=args.nprobe,
|
|
tiles_dir=args.tiles_dir or "tiles/evqa",
|
|
query_instruction=args.query_instruction,
|
|
)
|
|
text_base = TextAPIRetriever(
|
|
api_url=args.text_api_url,
|
|
top_k=args.retrieval_top_k,
|
|
nprobe=args.nprobe,
|
|
query_instruction=args.query_instruction,
|
|
reader_top_k=args.reader_top_k,
|
|
)
|
|
retriever = HybridRetriever(
|
|
image_base=image_base,
|
|
text_base=text_base,
|
|
top_k=args.retrieval_top_k,
|
|
reader_top_k=args.reader_top_k,
|
|
)
|
|
mode = f"Hybrid Retrieval (image={args.local_api_url}, text={args.text_api_url}, top_k={args.retrieval_top_k})"
|
|
|
|
elif args.text_vector:
|
|
if args.text_source == "ds-serve":
|
|
retriever = DsServeRetriever(
|
|
api_url=args.ds_serve_api_url, top_k=args.retrieval_top_k
|
|
)
|
|
mode = "Text Vector (ds-serve)"
|
|
else:
|
|
text_cache_path = f"text_cache/text_cache_{args.text_source}.jsonl"
|
|
text_cache = load_text_cache(text_cache_path)
|
|
if not text_cache:
|
|
print(f"Error: Text cache not found at {text_cache_path}")
|
|
print(
|
|
f"Run with --url-text --text-source {args.text_source} first to build the cache."
|
|
)
|
|
sys.exit(1)
|
|
|
|
if args.text_embed_preset == "qwen":
|
|
embedding_model = "Qwen/Qwen3-Embedding-0.6B"
|
|
embedding_mode = "sentence-transformers"
|
|
embedding_options = {"batch_size": args.embed_batch_size}
|
|
preset_name = "qwen3-0.6b"
|
|
elif args.text_embed_preset == "jina":
|
|
embedding_model = "jina-embeddings-v4"
|
|
embedding_mode = "openai"
|
|
embedding_options = {
|
|
"base_url": "https://api.jina.ai/v1",
|
|
"api_key": args.jina_api_key,
|
|
}
|
|
preset_name = "jina-v4"
|
|
elif args.text_embed_preset == "contriever":
|
|
embedding_model = "facebook/contriever"
|
|
embedding_mode = "sentence-transformers"
|
|
embedding_options = {"batch_size": args.embed_batch_size}
|
|
preset_name = "contriever"
|
|
else:
|
|
embedding_model = "facebook/contriever"
|
|
embedding_mode = "sentence-transformers"
|
|
embedding_options = {"batch_size": args.embed_batch_size}
|
|
preset_name = "contriever"
|
|
|
|
index_path = (
|
|
f"indexes/text_{args.text_source}_{preset_name}_c{args.chunk_size}"
|
|
)
|
|
retriever = TextVectorRetriever(
|
|
text_cache=text_cache,
|
|
index_path=index_path,
|
|
embedding_model=embedding_model,
|
|
embedding_mode=embedding_mode,
|
|
embedding_options=embedding_options,
|
|
top_k=args.retrieval_top_k,
|
|
rebuild_index=args.rebuild_text_index,
|
|
chunk_size=args.chunk_size,
|
|
chunk_overlap=args.chunk_overlap,
|
|
)
|
|
mode = f"Text Vector ({args.text_source}, {preset_name})"
|
|
|
|
elif args.task == "encyclopedic_vqa" and retrieval_mode_count == 0:
|
|
retriever = EVQANoRetrievalRetriever(tiles_dir=args.tiles_dir or "tiles/evqa")
|
|
mode = "EVQA no retrieval (query + image only)"
|
|
|
|
elif args.task == "worldvqa" and retrieval_mode_count == 0:
|
|
retriever = WorldVQANoRetrievalRetriever()
|
|
mode = "WorldVQA no retrieval (query + image only)"
|
|
|
|
elif (
|
|
args.task in ("simplevqa", "factualvqa", "mmsearch", "webqa", "multimodalqa")
|
|
and retrieval_mode_count == 0
|
|
):
|
|
retriever = WorldVQANoRetrievalRetriever()
|
|
mode = f"{args.task} no retrieval (query + image only)"
|
|
|
|
else:
|
|
retriever = NaiveRetriever()
|
|
mode = "Naive"
|
|
|
|
# Ablation A: wrap image retriever with OCR
|
|
if args.read_as_text_ocr:
|
|
image_modes = (
|
|
args.local_api
|
|
or args.use_tiled_retrieval
|
|
or args.retrieval_augment
|
|
or args.url_screenshot
|
|
or args.url_tiled_screenshot
|
|
)
|
|
if not image_modes:
|
|
print(
|
|
"Error: --read-as-text-ocr requires an image retrieval mode "
|
|
"(--local-api, --use-tiled-retrieval, --retrieval-augment, "
|
|
"--url-screenshot, or --url-tiled-screenshot)."
|
|
)
|
|
sys.exit(1)
|
|
if args.react:
|
|
print(
|
|
"Error: --read-as-text-ocr is not compatible with --react "
|
|
"(react bypasses the retriever wrapper on subsequent turns)."
|
|
)
|
|
sys.exit(1)
|
|
retriever = OCRWrappedRetriever(
|
|
base=retriever,
|
|
ocr_url=args.ocr_url,
|
|
model=args.ocr_model,
|
|
cache_path=args.ocr_cache,
|
|
concurrency=args.ocr_concurrency,
|
|
reader_top_k=args.reader_top_k,
|
|
)
|
|
mode += f" + OCR({args.ocr_url})"
|
|
logger.info(
|
|
f"Ablation A: OCR wrapper enabled ({args.ocr_url}, cache={args.ocr_cache})"
|
|
)
|
|
|
|
# Ablation B: wrap text retriever with renderer
|
|
if args.render_as_image:
|
|
if not args.text_api:
|
|
print(
|
|
"Error: --render-as-image requires --text-api (needs a text retriever "
|
|
"exposing get_hits())."
|
|
)
|
|
sys.exit(1)
|
|
retriever = RenderedTextWrapper(
|
|
base=retriever,
|
|
render_dir=args.render_dir,
|
|
reader_top_k=args.reader_top_k,
|
|
)
|
|
mode += f" + Render({args.render_dir})"
|
|
logger.info(f"Ablation B: text->image renderer enabled (dir={args.render_dir})")
|
|
|
|
return retriever, mode
|