eval: export the LoRA cells as a frozen-retrieval benchmark

eval/README.md

@@ -172,6 +172,24 @@ text-only text retrieval.
 PYTHONPATH=. .venv/bin/python -m lib.grader <task> <responses.jsonl>
 ```
 
+## 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 <nq|nqt|sqa|mms|evqa> <out_root>
+# 2. pack into one Parquet config per bench, images embedded
+.venv/bin/python pack_bench.py <out_root> <pack_dir> --commit <sha>
+```
+
+`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 `<bench>/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.

eval/dump_bench.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# Export the frozen-retrieval PixelRAG reader benchmark for one task.
+#
+#   bash dump_bench.sh <nq|nqt|sqa|mms|evqa> <out_root>
+#
+# 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 <nq|nqt|sqa|mms|evqa> <out_root>}"
+OUT_ROOT="${2:?Usage: dump_bench.sh <bench> <out_root>}"
+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"

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/<article_id>/<chunk file>   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/<id>.<ext>  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()

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 <dump_root> <out_dir> [--commit SHA] [--allow-missing FILE]
+
+Reads <dump_root>/<bench>/records.jsonl plus the tiles and query images beside it and
+writes <out_dir>/<config>/test-*.parquet, one config per bench, images embedded.
+A missing or empty image stops the pack, unless its tile is listed (one "<article_id>/<chunk
+file>" 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()

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(