From 7ec21389aab1b37b9c913a0b2fcaafba1de6f28d Mon Sep 17 00:00:00 2001 From: Zhifei Li Date: Wed, 30 Sep 2026 16:47:21 -0400 Subject: [PATCH 1/6] eval: --dump-retrieval exports questions + retrieved tiles as a frozen-retrieval benchmark --- eval/dump_bench.sh | 43 +++++++++++++ eval/lib/dump_retrieval.py | 127 +++++++++++++++++++++++++++++++++++++ eval/run_bench.py | 14 ++++ 3 files changed, 184 insertions(+) create mode 100644 eval/dump_bench.sh create mode 100644 eval/lib/dump_retrieval.py 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..d241bfe --- /dev/null +++ b/eval/lib/dump_retrieval.py @@ -0,0 +1,127 @@ +"""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 + 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 _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": f"tiles/{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: + 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: + hits = retriever._cache.get(ex["id"], []) + missing = [h.get("path") for h in hits if not h.get("image_base64")] + if len(hits) < args.retrieval_top_k or missing: + # Not written, so a rerun retries it. + log.write( + f"[{_now()}] item={ex['id']} status=fail hits={len(hits)} " + f"no_image={len(missing)} chunk_elapsed={elapsed:.1f}s\n" + ) + continue + + for h in hits: + dst = os.path.join(out_dir, "tiles", h["path"]) + 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)) + + 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"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/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( From f15e4c73e1b4cc89fa7f3eb80b0cb79d7d55963c Mon Sep 17 00:00:00 2001 From: Zhifei Li Date: Wed, 30 Sep 2026 16:52:40 -0400 Subject: [PATCH 2/6] eval: keep records whose tiles the serve lacks; key tiles by article; pack_bench.py --- eval/lib/dump_retrieval.py | 34 ++++++++++--- eval/pack_bench.py | 102 +++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 7 deletions(-) create mode 100644 eval/pack_bench.py diff --git a/eval/lib/dump_retrieval.py b/eval/lib/dump_retrieval.py index d241bfe..9300502 100644 --- a/eval/lib/dump_retrieval.py +++ b/eval/lib/dump_retrieval.py @@ -6,7 +6,9 @@ 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 + 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 """ @@ -31,6 +33,11 @@ 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"], @@ -45,7 +52,8 @@ def _record(example: dict, hits: list[dict], query_image: str | None, run_metada "retrieved_images": [ { "rank": rank, - "path": f"tiles/{h['path']}", + "path": _tile_rel(h), + "serve_path": h["path"], "score": h.get("score"), "url": h.get("url"), "article_id": h.get("article_id"), @@ -81,7 +89,9 @@ async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: di 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: + 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: @@ -92,17 +102,20 @@ async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: di for ex in chunk: hits = retriever._cache.get(ex["id"], []) - missing = [h.get("path") for h in hits if not h.get("image_base64")] - if len(hits) < args.retrieval_top_k or missing: + 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"no_image={len(missing)} chunk_elapsed={elapsed:.1f}s\n" + f"chunk_elapsed={elapsed:.1f}s\n" ) continue + missing = [] for h in hits: - dst = os.path.join(out_dir, "tiles", h["path"]) + 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: @@ -116,9 +129,16 @@ async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: di 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)") diff --git a/eval/pack_bench.py b/eval/pack_bench.py new file mode 100644 index 0000000..f50301a --- /dev/null +++ b/eval/pack_bench.py @@ -0,0 +1,102 @@ +"""Pack dump_bench.sh output into Parquet configs for a Hugging Face dataset repo. + + python pack_bench.py [--commit SHA] + +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: that bench must be re-dumped. +""" + +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 # noqa: E402 + +# 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) -> dict: + # Bytes, not paths: to_parquet stores a path-only Image as a dangling reference. + with open(os.path.join(bench_dir, rel), "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): + 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"]) 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) + args = ap.parse_args() + + 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)), 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() From fb274ca464e3ef8e301a6a1da606691e868bfbfd Mon Sep 17 00:00:00 2001 From: Zhifei Li Date: Wed, 30 Sep 2026 18:57:17 -0400 Subject: [PATCH 3/6] eval: backfill_tiles.py fills tiles the serve lacks from pixelrag-tiles --- eval/backfill_tiles.py | 84 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 eval/backfill_tiles.py diff --git a/eval/backfill_tiles.py b/eval/backfill_tiles.py new file mode 100644 index 0000000..5c63c50 --- /dev/null +++ b/eval/backfill_tiles.py @@ -0,0 +1,84 @@ +"""Fill tiles the search serve returned without an image, from StarTrail-org/pixelrag-tiles. + + python backfill_tiles.py + +Reads /*/missing_tiles.jsonl (written by --dump-retrieval), downloads each +needed shard tar once, and writes the tile as PNG at //. The HF +tiles are lossless WebP, so the PNG is pixel-identical to the original render. Tiles not +found in the tar are listed at the end; pack_bench.py refuses a bench that still lacks one. +""" + +from __future__ import annotations + +import glob +import io +import json +import os +import sys +import tarfile +from collections import defaultdict +from datetime import datetime, timezone + +from huggingface_hub import hf_hub_download +from PIL import Image + +ARTICLES_PER_SHARD = 8284 # shard = article_id // 8284, as in the corpus layout + + +def _log(msg: str) -> None: + print(f"[{datetime.now(timezone.utc):%Y-%m-%dT%H:%M:%SZ}] {msg}", flush=True) + + +def main(): + dump_root, cache = sys.argv[1], sys.argv[2] + # (article_id, chunk stem) -> destination paths across benches + wanted: dict[tuple[str, str], list[str]] = defaultdict(list) + for f in glob.glob(os.path.join(dump_root, "*", "missing_tiles.jsonl")): + bench_dir = os.path.dirname(f) + for line in open(f): + m = json.loads(line) + dst = os.path.join(bench_dir, m["tile"]) + if not os.path.exists(dst): + stem = os.path.splitext(os.path.basename(m["tile"]))[0] + wanted[(str(m["article_id"]), stem)].append(dst) + + by_shard: dict[int, set] = defaultdict(set) + for aid, stem in wanted: + by_shard[int(aid) // ARTICLES_PER_SHARD].add((aid, stem)) + _log(f"need {len(wanted)} tiles from {len(by_shard)} shards: {sorted(by_shard)}") + + found = set() + for shard, keys in sorted(by_shard.items()): + _log(f"shard={shard:03d} download_start tiles={len(keys)}") + tar_path = hf_hub_download( + "StarTrail-org/pixelrag-tiles", f"shard_{shard:03d}.tar", + repo_type="dataset", local_dir=cache, + ) + with tarfile.open(tar_path) as tar: + for member in tar: + parts = member.name.split("/") + if len(parts) < 2 or not parts[-2].endswith(".png.tiles"): + continue + key = (parts[-2].removesuffix(".png.tiles"), os.path.splitext(parts[-1])[0]) + if key not in keys: + continue + img = Image.open(io.BytesIO(tar.extractfile(member).read())) + for dst in wanted[key]: + os.makedirs(os.path.dirname(dst), exist_ok=True) + img.save(dst + ".part", format="PNG") + os.replace(dst + ".part", dst) + found.add(key) + _log(f"shard={shard:03d} tile={key[0]}/{key[1]} status=ok member={member.name}") + for key in sorted(keys - found): + _log(f"shard={shard:03d} tile={key[0]}/{key[1]} status=not_in_tar") + os.remove(tar_path) # re-downloadable; 1-4 GB each + _log(f"shard={shard:03d} done found={len(keys & found)}/{len(keys)}") + + missing = sorted(set(wanted) - found) + _log(f"backfill_end filled={len(found)} still_missing={len(missing)}") + for aid, stem in missing: + print(f" still missing: {aid}/{stem} -> {wanted[(aid, stem)]}") + + +if __name__ == "__main__": + main() From eb9325a92f51c03593faa87b189cdb32a8e4b1c3 Mon Sep 17 00:00:00 2001 From: Zhifei Li Date: Wed, 30 Sep 2026 19:26:38 -0400 Subject: [PATCH 4/6] eval: truncate hits of repeated example ids; pack --allow-missing --- eval/lib/dump_retrieval.py | 4 +++- eval/pack_bench.py | 22 +++++++++++++++------- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/eval/lib/dump_retrieval.py b/eval/lib/dump_retrieval.py index 9300502..a435a47 100644 --- a/eval/lib/dump_retrieval.py +++ b/eval/lib/dump_retrieval.py @@ -101,7 +101,9 @@ async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: di elapsed = time.time() - t0 for ex in chunk: - hits = retriever._cache.get(ex["id"], []) + # 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( diff --git a/eval/pack_bench.py b/eval/pack_bench.py index f50301a..744c472 100644 --- a/eval/pack_bench.py +++ b/eval/pack_bench.py @@ -1,10 +1,11 @@ """Pack dump_bench.sh output into Parquet configs for a Hugging Face dataset repo. - python pack_bench.py [--commit SHA] + 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: that bench must be re-dumped. +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 @@ -43,16 +44,19 @@ ) -def _img(bench_dir: str, rel: str) -> dict: +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(os.path.join(bench_dir, rel), "rb") as f: + 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): +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: @@ -66,7 +70,7 @@ def rows(bench_dir: str, grader_task: str, commit: str | None): "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"]) for h in hits], + "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( @@ -80,14 +84,18 @@ def main(): 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)), features=FEATURES) + 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) From 8b57b4e13c33fc5137e8bdea9ff3b5118f20237b Mon Sep 17 00:00:00 2001 From: Zhifei Li Date: Wed, 30 Sep 2026 19:29:05 -0400 Subject: [PATCH 5/6] eval: document the frozen-retrieval benchmark export; drop the HF backfill helper --- eval/README.md | 18 +++++++++ eval/backfill_tiles.py | 84 ------------------------------------------ 2 files changed, 18 insertions(+), 84 deletions(-) delete mode 100644 eval/backfill_tiles.py 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/backfill_tiles.py b/eval/backfill_tiles.py deleted file mode 100644 index 5c63c50..0000000 --- a/eval/backfill_tiles.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Fill tiles the search serve returned without an image, from StarTrail-org/pixelrag-tiles. - - python backfill_tiles.py - -Reads /*/missing_tiles.jsonl (written by --dump-retrieval), downloads each -needed shard tar once, and writes the tile as PNG at //. The HF -tiles are lossless WebP, so the PNG is pixel-identical to the original render. Tiles not -found in the tar are listed at the end; pack_bench.py refuses a bench that still lacks one. -""" - -from __future__ import annotations - -import glob -import io -import json -import os -import sys -import tarfile -from collections import defaultdict -from datetime import datetime, timezone - -from huggingface_hub import hf_hub_download -from PIL import Image - -ARTICLES_PER_SHARD = 8284 # shard = article_id // 8284, as in the corpus layout - - -def _log(msg: str) -> None: - print(f"[{datetime.now(timezone.utc):%Y-%m-%dT%H:%M:%SZ}] {msg}", flush=True) - - -def main(): - dump_root, cache = sys.argv[1], sys.argv[2] - # (article_id, chunk stem) -> destination paths across benches - wanted: dict[tuple[str, str], list[str]] = defaultdict(list) - for f in glob.glob(os.path.join(dump_root, "*", "missing_tiles.jsonl")): - bench_dir = os.path.dirname(f) - for line in open(f): - m = json.loads(line) - dst = os.path.join(bench_dir, m["tile"]) - if not os.path.exists(dst): - stem = os.path.splitext(os.path.basename(m["tile"]))[0] - wanted[(str(m["article_id"]), stem)].append(dst) - - by_shard: dict[int, set] = defaultdict(set) - for aid, stem in wanted: - by_shard[int(aid) // ARTICLES_PER_SHARD].add((aid, stem)) - _log(f"need {len(wanted)} tiles from {len(by_shard)} shards: {sorted(by_shard)}") - - found = set() - for shard, keys in sorted(by_shard.items()): - _log(f"shard={shard:03d} download_start tiles={len(keys)}") - tar_path = hf_hub_download( - "StarTrail-org/pixelrag-tiles", f"shard_{shard:03d}.tar", - repo_type="dataset", local_dir=cache, - ) - with tarfile.open(tar_path) as tar: - for member in tar: - parts = member.name.split("/") - if len(parts) < 2 or not parts[-2].endswith(".png.tiles"): - continue - key = (parts[-2].removesuffix(".png.tiles"), os.path.splitext(parts[-1])[0]) - if key not in keys: - continue - img = Image.open(io.BytesIO(tar.extractfile(member).read())) - for dst in wanted[key]: - os.makedirs(os.path.dirname(dst), exist_ok=True) - img.save(dst + ".part", format="PNG") - os.replace(dst + ".part", dst) - found.add(key) - _log(f"shard={shard:03d} tile={key[0]}/{key[1]} status=ok member={member.name}") - for key in sorted(keys - found): - _log(f"shard={shard:03d} tile={key[0]}/{key[1]} status=not_in_tar") - os.remove(tar_path) # re-downloadable; 1-4 GB each - _log(f"shard={shard:03d} done found={len(keys & found)}/{len(keys)}") - - missing = sorted(set(wanted) - found) - _log(f"backfill_end filled={len(found)} still_missing={len(missing)}") - for aid, stem in missing: - print(f" still missing: {aid}/{stem} -> {wanted[(aid, stem)]}") - - -if __name__ == "__main__": - main() From ae501cbfce24820e68c6aa11f551d635469b4194 Mon Sep 17 00:00:00 2001 From: Aakash Suresh Date: Wed, 30 Sep 2026 22:35:59 -0700 Subject: [PATCH 6/6] style(eval): satisfy the lint gate on the retrieval-dump tooling `Python lint` fails on this branch. Two things, both mechanical: - `pack_bench.py` carries `# noqa: E402` on the post-`sys.path` import, but E402 is already in the repo-wide ruff ignore list (pyproject.toml), so the suppression suppresses nothing and RUF100 reports it unused. - `ruff format` reflows two of the new files. No behaviour change: both files still compile and the 28 existing eval tests pass. Also merges current main into the branch. Co-Authored-By: Claude Opus 5 --- eval/lib/dump_retrieval.py | 47 ++++++++++++++++++++++++++++---------- eval/pack_bench.py | 30 +++++++++++++++++------- 2 files changed, 57 insertions(+), 20 deletions(-) diff --git a/eval/lib/dump_retrieval.py b/eval/lib/dump_retrieval.py index a435a47..b975853 100644 --- a/eval/lib/dump_retrieval.py +++ b/eval/lib/dump_retrieval.py @@ -38,7 +38,9 @@ def _tile_rel(hit: dict) -> str: 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: +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"), @@ -46,7 +48,9 @@ def _record(example: dict, hits: list[dict], query_image: str | None, run_metada # 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("_") + k: v + for k, v in example.items() + if not hasattr(v, "save") and not k.startswith("_") }, "query_image": query_image, "retrieved_images": [ @@ -66,7 +70,9 @@ def _record(example: dict, hits: list[dict], query_image: str | None, run_metada } -async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: dict) -> None: +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) @@ -87,11 +93,14 @@ async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: di 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: + 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: @@ -133,17 +142,31 @@ async def dump_retrieval(args, examples: list[dict], retriever, run_metadata: di 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" + 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 ) - rec_f.write(json.dumps(_record(ex, hits, query_image, run_metadata), ensure_ascii=False) + "\n") + + "\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)") + 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 index 744c472..f9e8fe8 100644 --- a/eval/pack_bench.py +++ b/eval/pack_bench.py @@ -18,7 +18,7 @@ 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 # noqa: E402 +from lib.grader import build_ground_truth # dump_bench.sh bench name -> (config name, grader task) BENCHES = { @@ -44,7 +44,9 @@ ) -def _img(bench_dir: str, rel: str, allow_missing: frozenset = frozenset()) -> dict | None: +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 @@ -56,7 +58,9 @@ def _img(bench_dir: str, rel: str, allow_missing: frozenset = frozenset()) -> di return {"bytes": data, "path": rel} -def rows(bench_dir: str, grader_task: str, commit: str | None, allow_missing: frozenset): +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: @@ -69,12 +73,17 @@ def rows(bench_dir: str, grader_task: str, commit: str | None, allow_missing: fr "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], + "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 + {"original_data": od, "run_metadata": r["run_metadata"]}, + ensure_ascii=False, ), } @@ -88,14 +97,19 @@ def main(): 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()) + 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) + 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)