diff --git a/eval/README.md b/eval/README.md index bfede20..88d4238 100644 --- a/eval/README.md +++ b/eval/README.md @@ -172,6 +172,24 @@ text-only text retrieval. PYTHONPATH=. .venv/bin/python -m lib.grader ``` +## 6. Exporting the frozen-retrieval benchmark + +[`StarTrail-org/pixelrag-bench`](https://huggingface.co/datasets/StarTrail-org/pixelrag-bench) +holds the LoRA cells' questions with their top-5 retrieved tiles, so a reader can be scored +without a search serve. It is produced in two steps: + +```bash +# 1. retrieve only (no reader); resumable, one directory per bench +LORA_URL=http://localhost:30096/search bash dump_bench.sh +# 2. pack into one Parquet config per bench, images embedded +.venv/bin/python pack_bench.py --commit +``` + +`dump_bench.sh` runs `run_bench.py --dump-retrieval` with the same example sets, query +images, instruction and nprobe as the LoRA cells of `reproduce.sh`. A hit whose tile is +absent on the serve's disk is still recorded and listed in `/missing_tiles.jsonl`; +`pack_bench.py` refuses to pack until each is filled, or listed in `--allow-missing` as lost. + ## Atlas Cloud readers Select Atlas Cloud explicitly with `--atlascloud`; other readers are unchanged. diff --git a/eval/dump_bench.sh b/eval/dump_bench.sh new file mode 100644 index 0000000..151a640 --- /dev/null +++ b/eval/dump_bench.sh @@ -0,0 +1,43 @@ +#!/bin/bash +# Export the frozen-retrieval PixelRAG reader benchmark for one task. +# +# bash dump_bench.sh +# +# Uses the same example sets, query instruction, query images and nprobe as the LoRA +# cells of reproduce.sh, against a serve holding search_index_lora_vit_ckpt200_v2, and +# keeps the top 5 tiles per question (the paper's reader reads the top 3). +# LORA_URL lora pixel search serve [http://localhost:30096/search] +# TILES_DIR wiki tiles dir; EVQA landmark query images are found beside it +# [/mnt/data/yichuan/kiwix_tiles] +# NUM overrides the example count (smoke tests) +set -euo pipefail +cd "$(dirname "$0")" + +BENCH="${1:?Usage: dump_bench.sh }" +OUT_ROOT="${2:?Usage: dump_bench.sh }" +LORA_URL="${LORA_URL:-http://localhost:30096/search}" +TILES_DIR="${TILES_DIR:-/mnt/data/yichuan/kiwix_tiles}" + +case "$BENCH" in + nq) TASK=nq; N=1000; EXTRA=() ;; + nqt) TASK=nq_tables; N=""; EXTRA=() ;; + sqa) TASK=simpleqa; N=1000; EXTRA=(--nprobe 2000) ;; + mms) TASK=mmsearch; N=""; EXTRA=() ;; + evqa) TASK=encyclopedic_vqa; N=""; + EXTRA=(--evqa-dataset-filter landmarks --evqa-question-type-filter automatic) ;; + *) echo "unknown bench: $BENCH" >&2; exit 1 ;; +esac +N="${NUM:-$N}" +NUMFLAG=(); [ -n "$N" ] && NUMFLAG=(--num-examples "$N") + +OUT="$OUT_ROOT/$BENCH" +mkdir -p "$OUT" +# The reader is never called in dump mode; --model/--api-base only satisfy the parser. +.venv/bin/python run_bench.py --task "$TASK" --model Qwen/Qwen3.5-4B \ + --api-base http://localhost:1/v1 --api-key dummy --no-think \ + --retrieval-top-k 5 --reader-top-k 3 "${NUMFLAG[@]}" \ + --tiles-dir "$TILES_DIR" --output "$OUT/unused_reader_output.jsonl" --force \ + --local-api --local-api-url "$LORA_URL" \ + --query-instruction "Retrieve images or text relevant to the user's query." \ + --dump-retrieval "$OUT" "${EXTRA[@]}" +rm -f "$OUT/unused_reader_output.jsonl" diff --git a/eval/lib/dump_retrieval.py b/eval/lib/dump_retrieval.py new file mode 100644 index 0000000..b975853 --- /dev/null +++ b/eval/lib/dump_retrieval.py @@ -0,0 +1,172 @@ +"""Export a frozen-retrieval benchmark: questions + the top-k tiles a pixel serve returned. + +`run_bench.py --dump-retrieval DIR` loads examples and builds the retriever exactly as a +normal run does, then writes what retrieval produced instead of calling a reader. The +result lets a reader model be evaluated without hosting the index or the tile corpus. + +Layout under DIR: + records.jsonl one line per example (see _record); resumable by example_id + tiles// each retrieved tile, decoded from the serve's image_base64 + missing_tiles.jsonl hits the serve returned without an image (its tile is absent on + the serve's disk); these are filled from the tile corpus later + query_images/. the query image, for multimodal tasks + dump.log per-example start/end lines +""" + +from __future__ import annotations + +import base64 +import json +import logging +import os +import shutil +import time +from datetime import datetime, timezone + +logger = logging.getLogger(__name__) + +# Examples sent to the serve per prefetch call; also the checkpoint granularity. +CHUNK = 32 + + +def _now() -> str: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _tile_rel(hit: dict) -> str: + # Keyed by article, not the serve's path: the serve prints an unresolved sub-shard as "?". + return f"tiles/{hit['article_id']}/{os.path.basename(hit['path'])}" + + +def _record( + example: dict, hits: list[dict], query_image: str | None, run_metadata: dict +) -> dict: + return { + "example_id": example["id"], + "load_index": example.get("_load_index"), + "problem": example["problem"], + # Same filter run_bench applies to original_data, so lib.grader reads these + # records unchanged once a "final_response" field is added. + "original_data": { + k: v + for k, v in example.items() + if not hasattr(v, "save") and not k.startswith("_") + }, + "query_image": query_image, + "retrieved_images": [ + { + "rank": rank, + "path": _tile_rel(h), + "serve_path": h["path"], + "score": h.get("score"), + "url": h.get("url"), + "article_id": h.get("article_id"), + "chunk_index": h.get("chunk_index"), + "y_offset": h.get("y_offset"), + } + for rank, h in enumerate(hits) + ], + "run_metadata": run_metadata, + } + + +async def dump_retrieval( + args, examples: list[dict], retriever, run_metadata: dict +) -> None: + out_dir = args.dump_retrieval + os.makedirs(os.path.join(out_dir, "tiles"), exist_ok=True) + os.makedirs(os.path.join(out_dir, "query_images"), exist_ok=True) + records_path = os.path.join(out_dir, "records.jsonl") + + done: set[str] = set() + if os.path.exists(records_path): + with open(records_path) as f: + done = {json.loads(line)["example_id"] for line in f if line.strip()} + + log = open(os.path.join(out_dir, "dump.log"), "a", buffering=1) + log.write( + f"[{_now()}] run_start task={args.task} n={len(examples)} already_done={len(done)} " + f"api={args.local_api_url} top_k={args.retrieval_top_k} nprobe={args.nprobe} " + f"commit={run_metadata.get('git_commit')}\n" + ) + + todo = [ex for ex in examples if ex["id"] not in done] + for ex in examples: + if ex["id"] in done: + log.write( + f"[{_now()}] item={ex['id']} status=skip reason=already_in_records\n" + ) + + with ( + open(records_path, "a", buffering=1) as rec_f, + open(os.path.join(out_dir, "missing_tiles.jsonl"), "a", buffering=1) as miss_f, + ): + for start in range(0, len(todo), CHUNK): + chunk = todo[start : start + CHUNK] + for ex in chunk: + log.write(f"[{_now()}] item={ex['id']} status=start\n") + t0 = time.time() + await retriever.prefetch(chunk) + elapsed = time.time() - t0 + + for ex in chunk: + # A task can repeat an example id (NQ-Tables does); the retriever cache then + # holds both copies' hits back to back, so keep only the first top_k. + hits = retriever._cache.get(ex["id"], [])[: args.retrieval_top_k] + if len(hits) < args.retrieval_top_k: + # Not written, so a rerun retries it. + log.write( + f"[{_now()}] item={ex['id']} status=fail hits={len(hits)} " + f"chunk_elapsed={elapsed:.1f}s\n" + ) + continue + + missing = [] + for h in hits: + dst = os.path.join(out_dir, _tile_rel(h)) + if not h.get("image_base64"): + missing.append(h) + continue + if not os.path.exists(dst): + os.makedirs(os.path.dirname(dst), exist_ok=True) + with open(dst + ".part", "wb") as f: + f.write(base64.b64decode(h["image_base64"])) + os.replace(dst + ".part", dst) + + query_image = None + src = retriever.query_image_fn(ex) if retriever.query_image_fn else None + if src and os.path.exists(src): + ext = os.path.splitext(src)[1] or ".png" + query_image = f"query_images/{ex['id']}{ext}" + shutil.copyfile(src, os.path.join(out_dir, query_image)) + + for h in missing: + miss_f.write( + json.dumps( + { + "example_id": ex["id"], + "article_id": h["article_id"], + "serve_path": h["path"], + "tile": _tile_rel(h), + } + ) + + "\n" + ) + rec_f.write( + json.dumps( + _record(ex, hits, query_image, run_metadata), ensure_ascii=False + ) + + "\n" + ) + log.write( + f"[{_now()}] item={ex['id']} status=ok hits={len(hits)} " + f"missing_tiles={len(missing)}" + f"{' ' + ','.join(h['path'] for h in missing) if missing else ''} " + f"query_image={'yes' if query_image else 'no'} chunk_elapsed={elapsed:.1f}s\n" + ) + print( + f"[dump] {min(start + CHUNK, len(todo))}/{len(todo)} examples ({elapsed:.1f}s last chunk)" + ) + + log.write(f"[{_now()}] run_end task={args.task}\n") + log.close() diff --git a/eval/pack_bench.py b/eval/pack_bench.py new file mode 100644 index 0000000..f9e8fe8 --- /dev/null +++ b/eval/pack_bench.py @@ -0,0 +1,124 @@ +"""Pack dump_bench.sh output into Parquet configs for a Hugging Face dataset repo. + + python pack_bench.py [--commit SHA] [--allow-missing FILE] + +Reads //records.jsonl plus the tiles and query images beside it and +writes //test-*.parquet, one config per bench, images embedded. +A missing or empty image stops the pack, unless its tile is listed (one "/" per line) in --allow-missing: those tiles exist nowhere any more and are packed as null. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +from datasets import Dataset, Features, Image, Sequence, Value + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from lib.grader import build_ground_truth + +# dump_bench.sh bench name -> (config name, grader task) +BENCHES = { + "nq": ("nq", "nq"), + "nqt": ("nq_tables", "nq_tables"), + "sqa": ("simpleqa", "simpleqa"), + "mms": ("mmsearch", "mmsearch"), + "evqa": ("encyclopedic_vqa_landmarks", "encyclopedic_vqa"), +} + +FEATURES = Features( + { + "id": Value("string"), + "question": Value("string"), + "instructions": Value("string"), + "answer": Value("string"), + "query_image": Image(), + "retrieved_images": Sequence(Image()), + "retrieved_scores": Sequence(Value("float32")), + "retrieved_urls": Sequence(Value("string")), + "original_data": Value("string"), + } +) + + +def _img( + bench_dir: str, rel: str, allow_missing: frozenset = frozenset() +) -> dict | None: + path = os.path.join(bench_dir, rel) + if not os.path.exists(path) and rel.removeprefix("tiles/") in allow_missing: + return None + # Bytes, not paths: to_parquet stores a path-only Image as a dangling reference. + with open(path, "rb") as f: + data = f.read() + if not data: + raise SystemExit(f"empty image {rel} in {bench_dir}") + return {"bytes": data, "path": rel} + + +def rows( + bench_dir: str, grader_task: str, commit: str | None, allow_missing: frozenset +): + with open(os.path.join(bench_dir, "records.jsonl")) as f: + recs = sorted((json.loads(line) for line in f), key=lambda r: r["load_index"]) + for r in recs: + if commit: + r["run_metadata"]["git_commit"] = commit + od = r["original_data"] + hits = r["retrieved_images"] + yield { + "id": r["example_id"], + "question": r["problem"], + "instructions": od.get("additional_instructions") or "", + "answer": build_ground_truth(grader_task, od), + "query_image": _img(bench_dir, r["query_image"]) + if r["query_image"] + else None, + "retrieved_images": [ + _img(bench_dir, h["path"], allow_missing) for h in hits + ], + "retrieved_scores": [h["score"] for h in hits], + "retrieved_urls": [h["url"] for h in hits], + "original_data": json.dumps( + {"original_data": od, "run_metadata": r["run_metadata"]}, + ensure_ascii=False, + ), + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("dump_root") + ap.add_argument("out_dir") + ap.add_argument("--commit", default=None) + ap.add_argument("--allow-missing", default=None) + args = ap.parse_args() + allow_missing = frozenset() + if args.allow_missing: + allow_missing = frozenset( + line.strip() for line in open(args.allow_missing) if line.strip() + ) + + for bench, (config, grader_task) in BENCHES.items(): + bench_dir = os.path.join(args.dump_root, bench) + if not os.path.exists(os.path.join(bench_dir, "records.jsonl")): + print(f"[pack] {bench}: no records, skipped") + continue + ds = Dataset.from_list( + list(rows(bench_dir, grader_task, args.commit, allow_missing)), + features=FEATURES, + ) + out = os.path.join(args.out_dir, config) + os.makedirs(out, exist_ok=True) + n_shards = max(1, round(ds.data.nbytes / 500e6) or 1) + for i in range(n_shards): + ds.shard(n_shards, i, contiguous=True).to_parquet( + os.path.join(out, f"test-{i:05d}-of-{n_shards:05d}.parquet") + ) + print(f"[pack] {bench} -> {config}: {len(ds)} rows, {n_shards} shard(s)") + + +if __name__ == "__main__": + main() diff --git a/eval/run_bench.py b/eval/run_bench.py index 05361db..3195624 100644 --- a/eval/run_bench.py +++ b/eval/run_bench.py @@ -1134,6 +1134,12 @@ async def run_async(args): retriever, mode = build_retriever(args, examples, model, api_base, api_key) # (retriever selection logic moved to simpleqa/retriever_factory.py) + + if args.dump_retrieval: + from lib.dump_retrieval import dump_retrieval + + await dump_retrieval(args, examples, retriever, run_metadata) + return # 3. Initialize LLM client llm_client = LLMClient( model=model, @@ -1339,6 +1345,14 @@ def main(): action="store_true", help="Overwrite output file if it exists", ) + parser.add_argument( + "--dump-retrieval", + type=str, + default=None, + metavar="DIR", + help="Retrieve only (no reader) and write records + retrieved tiles to DIR " + "(local-api retrieval only; resumable). See lib/dump_retrieval.py.", + ) # API args parser.add_argument(