Initial release of PixelRAG

Co-Authored-By: Yichuan Wang <yichuan_wang@berkeley.edu>
Co-Authored-By: Zhifei Li <andylizf@outlook.com>
This commit is contained in:
Yichuan Wang
2026-05-29 01:06:45 -07:00
committed by yichuan520030910320
commit 82c5794348
314 changed files with 87135 additions and 0 deletions
+16
View File
@@ -0,0 +1,16 @@
{
"name": "pixelrag-plugins",
"owner": {
"name": "StarTrail-org"
},
"description": "PixelRAG Claude Code plugins",
"plugins": [
{
"name": "pixelbrowse",
"source": "./plugin",
"description": "Give Claude eyes — screenshot any URL with pixelshot and read it visually",
"homepage": "https://github.com/StarTrail-org/PixelRAG",
"repository": "https://github.com/StarTrail-org/PixelRAG"
}
]
}
+9
View File
@@ -0,0 +1,9 @@
# Markdown
*.md linguist-detectable=true
*.md linguist-documentation=false
# JSON
*.json linguist-detectable=true
# YAML
*.yml linguist-detectable=true
+103
View File
@@ -0,0 +1,103 @@
name: CI
on:
push:
branches: [main]
pull_request:
branches: [main]
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
python-lint:
name: Python lint
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v6
with:
version: "latest"
- uses: actions/setup-python@v5
with:
python-version: "3.12"
- name: Lint with ruff
run: uvx ruff check .
- name: Check formatting with ruff
run: uvx ruff format --check .
frontend-build:
name: Frontend build
runs-on: ubuntu-latest
defaults:
run:
working-directory: web
steps:
- uses: actions/checkout@v4
- uses: actions/setup-node@v4
with:
node-version: 22
cache: npm
cache-dependency-path: web/package-lock.json
- name: Install dependencies
run: npm ci
- name: Build
run: npm run build
test:
name: Tests
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v6
with:
version: "latest"
- uses: actions/setup-python@v5
with:
python-version: "3.12"
- name: Install (core + dev)
run: uv sync --extra dev
- name: Chrome runtime libraries
run: npx -y playwright install-deps chromium
- name: Cache patched Chrome
uses: actions/cache@v4
with:
path: ~/.cache/pixelrag/chrome
key: pixelrag-chrome-150.0.7844.0
- name: Download patched Chrome (cached after first run)
run: uv run pixelshot install-chrome
- name: Run tests
run: uv run pytest tests/ -v
plugin-validate:
name: Plugin validate
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-node@v4
with:
node-version: 22
- name: Install Claude Code CLI
run: npm install -g @anthropic-ai/claude-code
- name: Validate marketplace + plugin
run: |
claude plugin validate .
claude plugin validate ./plugin
+45
View File
@@ -0,0 +1,45 @@
name: Release to PyPI
# Publishes the `pixelrag` package to PyPI. Auth is an account-scoped API token stored as
# the `PYPI_API_TOKEN` repo secret (see RELEASING.md). Triggered by publishing a GitHub
# Release (tag like `v0.2.0`), or run manually:
# - dry_run = true → build only (no upload)
# - dry_run = false → build + publish
# `--check-url` makes publishing idempotent: files already on PyPI are skipped.
on:
release:
types: [published]
workflow_dispatch:
inputs:
dry_run:
description: "Build only, do not publish"
type: boolean
default: true
concurrency:
group: release-${{ github.ref }}
cancel-in-progress: false
jobs:
release:
name: pixelrag
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v6
with:
version: "latest"
- name: Build
run: uv build --out-dir dist
- name: List artifacts
run: ls -l dist
- name: Publish to PyPI
if: ${{ github.event_name == 'release' || !inputs.dry_run }}
env:
UV_PUBLISH_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
run: uv publish --check-url https://pypi.org/simple/
+24
View File
@@ -0,0 +1,24 @@
name: Status Site CI
on:
push:
branches: [main]
paths:
- "status-site/**"
- ".github/workflows/status-site.yml"
workflow_dispatch:
permissions:
contents: write
pages: write
jobs:
deploy:
name: Deploy static status page
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v5
- name: Deploy to gh-pages
uses: JamesIves/github-pages-deploy-action@v4
with:
branch: gh-pages
folder: ./status-site
clean: true
@@ -0,0 +1,24 @@
name: Response Time CI
on:
schedule:
- cron: "0 23 * * *"
workflow_dispatch:
permissions:
contents: write
jobs:
response-time:
name: Check response time
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v5
with:
ref: ${{ github.head_ref }}
token: ${{ secrets.GH_PAT || github.token }}
- name: Update response time
uses: upptime/uptime-monitor@v1.41.8
with:
command: "response-time"
env:
GH_PAT: ${{ secrets.GH_PAT || github.token }}
SECRETS_CONTEXT: ${{ toJson(secrets) }}
+25
View File
@@ -0,0 +1,25 @@
name: Summary CI
on:
schedule:
# Every 10 min so history/summary.json (the data source for the static
# status page) stays near-real-time, not just daily.
- cron: "*/10 * * * *"
workflow_dispatch:
permissions:
contents: write
jobs:
summary:
name: Generate summary
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v5
with:
ref: ${{ github.head_ref }}
token: ${{ secrets.GH_PAT || github.token }}
- name: Update summary
uses: upptime/uptime-monitor@v1.41.8
with:
command: "readme"
env:
GH_PAT: ${{ secrets.GH_PAT || github.token }}
+24
View File
@@ -0,0 +1,24 @@
name: Uptime CI
on:
schedule:
- cron: "*/5 * * * *"
workflow_dispatch:
permissions:
contents: write
jobs:
uptime:
name: Check status
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v5
with:
ref: ${{ github.head_ref }}
token: ${{ secrets.GH_PAT || github.token }}
- name: Check endpoint status
uses: upptime/uptime-monitor@v1.41.8
with:
command: "update"
env:
GH_PAT: ${{ secrets.GH_PAT || github.token }}
SECRETS_CONTEXT: ${{ toJson(secrets) }}
+68
View File
@@ -0,0 +1,68 @@
__pycache__/
*.pyc
*.pyo
.venv/
*.egg-info/
dist/
build/
.eggs/
# Data / scratch (too large for git)
*.npy
*.faiss
*.npz
*.bak
data/
tiles/
compressed_tiles_*/
neurips/
tmp/*/
tmp/*.png
tmp/*.txt
tmp/*.jsonl
tmp/*.json
tmp/*.pid
# IDE / tooling
.claude/
.env
.envrc
# OS
.DS_Store
Thumbs.db
# Logs
logs/
*.log
arxiv
demos/e2e/output/
eval/eval_output/
.superpowers/
.vercel
# Large local retrieval artifacts (not committed)
eval/tmp_news_state.db
eval/live_pixel_full.json
eval/live_reader_full.json
eval/frozen_reader_full.json
eval/mms_base_live.jsonl
eval/mms_lora_live.jsonl
eval/mms_naive_live.jsonl
eval/evqa_base_landmarks_live.jsonl
eval/evqa_base_inat_live.jsonl
eval/evqa_lora_landmarks_live.jsonl
eval/evqa_lora_inat_live.jsonl
eval/mms_traf_live.jsonl
eval/evqa_traf_landmarks.jsonl
eval/evqa_traf_inaturalist.jsonl
eval/evqa_naive_landmarks.jsonl
eval/evqa_naive_inaturalist.jsonl
eval/mms_naive_nothink.jsonl
eval/evqa_base_landmarks_nothink.jsonl
eval/evqa_base_inat_nothink.jsonl
eval/evqa_lora_landmarks_nothink.jsonl
eval/evqa_lora_inat_nothink.jsonl
eval/paper_grader_out/
node_modules/
.next/
+20
View File
@@ -0,0 +1,20 @@
owner: StarTrail-org
repo: PixelRAG
# Checks run by GitHub Actions every 5 min; results committed to history/.
# The status page (status-site/, deployed to status.pixelrag.ai) reads the
# generated history/summary.json — we don't use Upptime's own site builder.
sites:
- name: API Health
url: http://api.pixelrag.ai:30001/health
- name: Search API
url: http://api.pixelrag.ai:30001/status
- name: Tile Serving
url: http://api.pixelrag.ai:30001/tile/0/0/0
expectedStatusCodes:
- 200
- 404
- name: Agent
url: http://api.pixelrag.ai:30010/health
- name: Website
url: https://pixelrag.ai
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+209
View File
@@ -0,0 +1,209 @@
<p align="center">
<img src="docs/assets/banner.png" alt="PixelRAG — Visual Retrieval-Augmented Generation" width="100%">
</p>
<p align="center">
Official codebase for <b><a href="assets/pixelrag-paper.pdf">PixelRAG: Visually Grounded Retrieval-Augmented Generation with Screenshot Rendering</a></b>
</p>
<p align="center">
<a href="https://yichuan-w.github.io/">Yichuan Wang</a>*,
<a href="https://zhifei.li/">Zhifei Li</a>*,
<a href="https://zwcolin.github.io/">Zirui Wang</a>,
<a href="https://people.epfl.ch/paul.teiletche">Paul Teiletche</a>,
<a href="https://www.linkedin.com/in/lesheng-jin-9618b0201/">Lesheng Jin</a>,
<a href="https://people.eecs.berkeley.edu/~matei/">Matei Zaharia</a>†,
<a href="https://people.eecs.berkeley.edu/~jegonzal/">Joseph E. Gonzalez</a>†,
<a href="https://www.sewonmin.com/">Sewon Min</a>†
</p>
<p align="center"><sub>* Equal contribution &nbsp; † Equal advising</sub></p>
<p align="center">Search any document by how it <em>looks</em>, not just the text it contains.</p>
<p align="center">
<a href="https://github.com/StarTrail-org/PixelRAG/actions/workflows/ci.yml"><img src="https://github.com/StarTrail-org/PixelRAG/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
<a href="https://pixelrag.ai"><img src="https://img.shields.io/badge/demo-pixelrag.ai-7c3aed" alt="Live demo"></a>
<a href="https://status.pixelrag.ai"><img src="https://img.shields.io/badge/status-live-22c55e" alt="Status"></a>
<a href="https://join.slack.com/t/leann-e2u9779/shared_invite/zt-3ol2ww9ic-Eg_kB8omwe6xmYVd0epr4Q"><img src="https://img.shields.io/badge/Slack-join-4A154B?logo=slack&logoColor=white" alt="Slack"></a>
<img src="https://img.shields.io/badge/license-Apache--2.0-blue" alt="License">
</p>
<p align="center">
<a href="#what-it-is">What it is</a> &middot;
<a href="#give-claude-eyes">Give Claude eyes</a> &middot;
<a href="#how-it-works">How it works</a> &middot;
<a href="#pipelines">Pipelines</a>
</p>
---
```bash
pip install pixelrag
```
The two core operations — **render** a page to screenshots, **search** a visual index:
```bash
# Render any page or document to screenshot tiles
pixelshot https://en.wikipedia.org/wiki/Python --output ./tiles
# Search a hosted index of 8.28M Wikipedia pages — no setup, runs against the live API
curl -X POST http://api.pixelrag.ai:30001/search \
-H "Content-Type: application/json" \
-d '{"queries": [{"text": "What is the capital of France?"}], "n_docs": 5}'
```
Or try it in the browser at **[pixelrag.ai](https://pixelrag.ai)**, or run the
[demo notebook](demos/quickstart.ipynb) (renders + searches, with the images inline).
## What it is
PixelRAG renders documents — web pages, PDFs, images — as screenshots and retrieves over the
images directly. Visual structure that HTML parsing throws away — tables, charts, layout,
infographics — stays intact, so the reader model can actually answer questions about it.
Wikipedia's 8.28M articles ship as a pre-built index; the pipeline itself is general-purpose.
## Give Claude eyes
The renderer also ships as a Claude Code plugin — the **pixelbrowse** skill. Instead of fetching
raw HTML, Claude screenshots a page with `pixelshot` and _reads the image_, so it sees
charts, diagrams, tables, and layout the way a person does.
Install it — no clone needed (`pixelshot` comes from `pip install pixelrag`):
```bash
pip install pixelrag # provides the pixelshot command
claude plugin marketplace add StarTrail-org/PixelRAG
claude plugin install pixelbrowse@pixelrag-plugins
```
Then just ask Claude to look at a page:
```bash
claude -p "screenshot https://news.ycombinator.com and summarize the top stories"
claude -p "screenshot https://arxiv.org/abs/2404.12387 and explain the key findings"
```
Or use the slash command in an interactive session: `/screenshot https://example.com`.
No MCP server, no backend: the skill just calls `pixelshot` (Playwright/CDP) on your machine.
## How it works
<p align="center">
<img src="docs/assets/pipeline.png" alt="Text-based RAG parses to text and loses the table; PixelRAG renders to screenshot tiles and keeps it" width="100%">
</p>
Text-based RAG parses the page to text chunks and **loses the table** — the reader can't find the
answer. PixelRAG renders the page to **screenshot tiles**, retrieves the right tile, and the reader
reads the number straight off the image.
Two pieces make this work: (1) rendering documents to images instead of parsing them to text, and
(2) a `Qwen3-VL-Embedding` model, LoRA-fine-tuned on screenshot data, that embeds page images into
a space where visual content is retrievable.
## Pipelines
Capture is the standalone `pixelshot` command; the rest of the pipeline runs through the
`pixelrag` umbrella — `pixelrag <stage>`. Install only the stages you need:
| Command | What it does | Install |
| ------------------------------------------ | --------------------------------------------------------------- | ------------------------------- |
| `pixelshot` | Document → image tiles (Playwright CDP, PDF) | `pip install pixelrag` |
| `pixelrag chunk` · `embed` · `build-index` | Tiles → vectors → FAISS index | `pip install 'pixelrag[embed]'` |
| `pixelrag index` | Orchestrates the full pipeline: source → ingest → embed → index | `pip install 'pixelrag[index]'` |
| `pixelrag serve` | FAISS search API (FastAPI, CPU or GPU) | `pip install 'pixelrag[serve]'` |
```
render ←── index ──→ embed serve (independent) train → serve (HTTP)
```
**`train` is a separate uv project** with its own pinned env (`torch==2.9.1+cu129`,
`transformers==4.57.1`, cuDNN 9.20) — install it from inside `train/`, not from the root.
### Search a pre-built index
```bash
pip install 'pixelrag[serve]'
# Download a pre-built index from Hugging Face. The dataset repo holds four FAISS indexes
# (base/LoRA Wikipedia pixel, Wikipedia text, news pixel); grab just the base one (~217G) here.
huggingface-cli download StarTrail-org/pixelrag-faiss-indexes \
--repo-type dataset --include "search_index_normed_v2/*" --local-dir ./index
# Serve, then query
pixelrag serve --index-dir ./index/search_index_normed_v2 --port 30001
curl -X POST http://localhost:30001/search \
-H "Content-Type: application/json" \
-d '{"queries": [{"text": "What is the capital of France?"}], "n_docs": 5}'
```
### Build an index from your own documents
```bash
pip install 'pixelrag[index]'
# Create pixelrag.yaml
cat > pixelrag.yaml << 'EOF'
source:
type: local
path: ./my_docs
embed:
model: Qwen/Qwen3-VL-Embedding-2B
device: cuda
gpu_ids: [0]
output: ./my_index
EOF
# Build, then serve
pixelrag index build
pixelrag serve --index-dir ./my_index --port 30001
```
### Render a page programmatically
```python
from pixelrag_render import render_url
# render a single page to tiles — e.g. for an agent to read
tiles = render_url("https://en.wikipedia.org/wiki/Python", "./tiles")
```
### Embed tools (standalone)
Each stage runs independently, without the orchestrator:
```bash
pip install 'pixelrag[embed]'
pixelrag chunk --tiles-dir ./tiles
pixelrag embed --shard-dir ./tiles --output-dir ./embeddings --gpu-ids 0,1
pixelrag build-index --embeddings-dir ./embeddings --output-dir ./index
```
### Training
Fine-tuning lives in `train/` — a **separate uv project** (`wiki-screenshot-training`) with its own
pinned env. It LoRA-fine-tunes `Qwen/Qwen3-VL-Embedding-2B` for webpage retrieval; run it from
inside `train/` (`cd train && uv sync`). See [`train/README.md`](train/README.md) for the full recipe.
You don't need to retrain to use the model — the trained adapters are published at
[`Chrisyichuan/wiki-screenshot-embedding-lora`](https://huggingface.co/Chrisyichuan/wiki-screenshot-embedding-lora/tree/main/lora_vit/ckpt200).
We also release the full training set
([`Chrisyichuan/screenshot-training-natural-filtered-v2`](https://huggingface.co/datasets/Chrisyichuan/screenshot-training-natural-filtered-v2)),
so you can adapt other backbones yourself — a larger Qwen, or any other embedding model.
## Acknowledgments
Thanks to [Rulin Shao](https://rulinshao.github.io/) for support.
Thanks also to [Claude Code](https://github.com/anthropics/claude-code) and
[OpenAI Codex](https://github.com/openai/codex) for supporting open-source contributors with credits and plans,
which we earned by working on [LEANN](https://github.com/StarTrail-org/LEANN).
This work is done by the [Berkeley Sky Computing Lab](https://sky.cs.berkeley.edu/),
[BAIR](https://bair.berkeley.edu/), and the [Berkeley NLP Group](https://nlp.cs.berkeley.edu/).
## License
Apache-2.0
+58
View File
@@ -0,0 +1,58 @@
# Releasing PixelRAG to PyPI
PixelRAG ships as a single PyPI package, **`pixelrag`** — one distribution that bundles the
umbrella CLI plus every stage module (`pixelrag_render`, `pixelrag_embed`, `pixelrag_index`,
`pixelrag_serve`). The core install is light (rendering only); heavy ML stages are extras:
```bash
pip install pixelrag # pixelshot + pixelrag umbrella (no torch)
pip install 'pixelrag[serve]' # + search API (also [embed], [index], [all])
```
`train/` is a separate local project and is **not** published.
Publishing is automated by [`.github/workflows/release.yml`](.github/workflows/release.yml)
using an **account-scoped PyPI API token** stored as a repo secret.
## One-time setup
1. **Create an account-scoped PyPI API token** at
<https://pypi.org/manage/account/token/>. Copy the `pypi-...` value.
2. **Store it as the `PYPI_API_TOKEN` repo secret:**
```bash
gh secret set PYPI_API_TOKEN --repo StarTrail-org/PixelRAG
```
(or repo Settings → Secrets and variables → Actions → New repository secret).
## Cutting a release
1. **Bump the version:**
```bash
uv version X.Y.Z
```
2. **Commit and push.**
3. **Publish a GitHub Release** with tag `vX.Y.Z` (matching the version). That fires
`release.yml`, which builds and publishes `pixelrag` to PyPI.
## Dry run
Actions → **Release to PyPI** → Run workflow (leave `dry_run` checked) builds the package
and lists artifacts **without** uploading. Use it to sanity-check before a real release.
## Notes
- **README images on PyPI.** The README uses repo-relative image paths (`docs/assets/...`),
which render on GitHub but **not** on the PyPI project page. To show them on PyPI, switch
those `<img src>` to absolute `https://raw.githubusercontent.com/...` URLs.
- **sdist scope.** The repo root holds large data dirs (`.venv`, `tiles`, `arxiv`, …); the
package restricts its sdist to the source dirs + `README.md` + `LICENSE`
(`[tool.hatch.build.targets.sdist]` in `pyproject.toml`).
- **Superseded packages.** `pixelrag-render`, `pixelrag-embed`, `pixelrag-index`, and
`pixelrag-serve` were published once (0.1.0) before consolidation and are now **yanked**;
everything lives in `pixelrag`.
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 206 KiB

+252
View File
@@ -0,0 +1,252 @@
diff --git a/content/browser/devtools/protocol/page_handler.cc b/content/browser/devtools/protocol/page_handler.cc
index a9f21bbe0c..b8de5faf69 100644
--- a/content/browser/devtools/protocol/page_handler.cc
+++ b/content/browser/devtools/protocol/page_handler.cc
@@ -506,6 +506,9 @@ struct PageHandler::PendingScreenshotRequest {
gfx::Size original_view_size;
gfx::Size requested_image_size;
std::string raw_file_path;
+ // Used by the directClip path to store the clip rect for capture
+ // after ForceRedraw synchronization.
+ gfx::Rect direct_clip_src_rect;
};
PageHandler::PageHandler(
@@ -1383,7 +1386,8 @@ void PageHandler::CaptureFullPageScreenshot(
CaptureScreenshot(std::move(format), std::move(quality), std::move(clip),
/*from_surface=*/true, /*capture_beyond_viewport=*/true,
std::move(optimize_for_speed), std::move(raw_file_path),
- /*direct_clip=*/std::nullopt, std::move(callback));
+ /*direct_clip=*/std::nullopt, /*skip_redraw=*/std::nullopt,
+ std::move(callback));
}
void PageHandler::CaptureScreenshot(
@@ -1395,6 +1399,7 @@ void PageHandler::CaptureScreenshot(
std::optional<bool> optimize_for_speed,
std::optional<std::string> raw_file_path,
std::optional<bool> direct_clip,
+ std::optional<bool> skip_redraw,
std::unique_ptr<CaptureScreenshotCallback> callback) {
if (!host_ || !host_->GetRenderWidgetHost() ||
!host_->GetRenderWidgetHost()->GetView()) {
@@ -1407,6 +1412,11 @@ void PageHandler::CaptureScreenshot(
// directClip: capture a region directly from the current frame without
// modifying viewport/emulation state. Enables concurrent captures.
+ // A ForceRedraw is issued first to ensure the renderer has submitted a
+ // CompositorFrame, preventing races where CopyFromSurface captures a
+ // stale frame (e.g., about:blank dimensions instead of the loaded page).
+ // This is critical with --in-process-gpu where reduced IPC latency
+ // widens the race window.
if (direct_clip.value_or(false) && clip) {
RenderWidgetHostImpl* widget_host = host_->GetRenderWidgetHost();
float dsf = widget_host->GetDeviceScaleFactor();
@@ -1437,17 +1447,17 @@ void PageHandler::CaptureScreenshot(
static_cast<int>(clip->GetHeight() * dsf));
pending_request->requested_image_size = output_size;
- gfx::Rect src_rect(
+ pending_request->direct_clip_src_rect = gfx::Rect(
static_cast<int>(clip->GetX() * dsf),
static_cast<int>(clip->GetY() * dsf),
output_size.width(), output_size.height());
- static_cast<RenderWidgetHostViewBase*>(widget_host->GetView())
- ->CopyFromSurface(
- src_rect, output_size, base::TimeDelta(),
- base::BindOnce(&PageHandler::DirectClipCaptured,
- weak_factory_.GetWeakPtr(),
- std::move(pending_request)));
+ // ForceRedraw ensures the renderer has committed and submitted a
+ // CompositorFrame to viz before we issue CopyFromSurface.
+ widget_host->ForceRedrawWithCallback(
+ base::BindOnce(&PageHandler::DirectClipForceRedrawDone,
+ weak_factory_.GetWeakPtr(),
+ std::move(pending_request)));
return;
}
@@ -1631,10 +1641,20 @@ void PageHandler::CaptureScreenshot(
}
}
- widget_host->GetSnapshotFromBrowser(
- base::BindOnce(&PageHandler::ScreenshotCaptured,
- weak_factory_.GetWeakPtr(), std::move(pending_request)),
- true);
+ if (skip_redraw.value_or(false)) {
+ // skipRedraw: lightweight ForceRedraw (wait for renderer commit) then
+ // CopyFromSurface. Faster than full GetSnapshotFromBrowser because we
+ // skip the presentation feedback wait.
+ widget_host->ForceRedrawWithCallback(
+ base::BindOnce(&PageHandler::SkipRedrawForceRedrawDone,
+ weak_factory_.GetWeakPtr(),
+ std::move(pending_request)));
+ } else {
+ widget_host->GetSnapshotFromBrowser(
+ base::BindOnce(&PageHandler::ScreenshotCaptured,
+ weak_factory_.GetWeakPtr(), std::move(pending_request)),
+ true);
+ }
}
Response PageHandler::StartScreencast(std::optional<std::string> format,
@@ -1887,6 +1907,49 @@ void PageHandler::ScreencastFrameEncoded(
std::move(page_metadata), session_id_);
}
+void PageHandler::SkipRedrawForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request) {
+ // Allocate a new LocalSurfaceId so that CopyFromSurface targets a fresh
+ // surface, preventing stale-frame captures (same rationale as in
+ // DirectClipForceRedrawDone).
+ SkipRedrawDoCopy(std::move(request));
+}
+
+void PageHandler::SkipRedrawDoCopy(
+ std::unique_ptr<PendingScreenshotRequest> request) {
+ if (!host_ || !host_->GetRenderWidgetHost() ||
+ !host_->GetRenderWidgetHost()->GetView()) {
+ request->callback->sendFailure(Response::InternalError());
+ return;
+ }
+ RenderWidgetHostImpl* widget_host = host_->GetRenderWidgetHost();
+
+ static_cast<RenderWidgetHostViewBase*>(widget_host->GetView())
+ ->CopyFromSurface(
+ gfx::Rect(), gfx::Size(), base::TimeDelta(),
+ base::BindOnce(&PageHandler::DirectClipCaptured,
+ weak_factory_.GetWeakPtr(),
+ std::move(request)));
+}
+
+void PageHandler::DirectClipForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request) {
+ if (!host_ || !host_->GetRenderWidgetHost() ||
+ !host_->GetRenderWidgetHost()->GetView()) {
+ request->callback->sendFailure(Response::InternalError());
+ return;
+ }
+ RenderWidgetHostImpl* widget_host = host_->GetRenderWidgetHost();
+
+ static_cast<RenderWidgetHostViewBase*>(widget_host->GetView())
+ ->CopyFromSurface(
+ request->direct_clip_src_rect,
+ request->requested_image_size, base::TimeDelta(),
+ base::BindOnce(&PageHandler::DirectClipCaptured,
+ weak_factory_.GetWeakPtr(),
+ std::move(request)));
+}
+
void PageHandler::DirectClipCaptured(
std::unique_ptr<PendingScreenshotRequest> request,
const content::CopyFromSurfaceResult& result) {
diff --git a/content/browser/devtools/protocol/page_handler.h b/content/browser/devtools/protocol/page_handler.h
index 4ccc1b42d3..ac94709199 100644
--- a/content/browser/devtools/protocol/page_handler.h
+++ b/content/browser/devtools/protocol/page_handler.h
@@ -156,6 +156,7 @@ class PageHandler : public DevToolsDomainHandler,
std::optional<bool> optimize_for_speed,
std::optional<std::string> raw_file_path,
std::optional<bool> direct_clip,
+ std::optional<bool> skip_redraw,
std::unique_ptr<CaptureScreenshotCallback> callback) override;
void CaptureSnapshot(
std::optional<std::string> format,
@@ -242,6 +243,15 @@ class PageHandler : public DevToolsDomainHandler,
void DirectClipCaptured(std::unique_ptr<PendingScreenshotRequest> request,
const content::CopyFromSurfaceResult& result);
+ // Called after ForceRedraw ack in directClip path; issues CopyFromSurface
+ // with the stored clip rect now that the renderer has flushed a frame.
+ void DirectClipForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request);
+ // Called after ForceRedraw ack in skipRedraw path; issues the actual
+ // CopyFromSurface now that the renderer has flushed a CompositorFrame.
+ void SkipRedrawDoCopy(std::unique_ptr<PendingScreenshotRequest> request);
+ void SkipRedrawForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request);
void ScreenshotCaptured(std::unique_ptr<PendingScreenshotRequest> request,
const gfx::Image& image);
diff --git a/content/browser/renderer_host/render_widget_host_impl.cc b/content/browser/renderer_host/render_widget_host_impl.cc
index 7b18f51513..e0c90a415a 100644
--- a/content/browser/renderer_host/render_widget_host_impl.cc
+++ b/content/browser/renderer_host/render_widget_host_impl.cc
@@ -4197,6 +4197,15 @@ void RenderWidgetHostImpl::ForceRedrawForTesting() {
blink_widget_->ForceRedraw(base::DoNothing());
}
+void RenderWidgetHostImpl::ForceRedrawWithCallback(
+ base::OnceClosure callback) {
+ if (!blink_widget_) {
+ std::move(callback).Run();
+ return;
+ }
+ blink_widget_->ForceRedraw(std::move(callback));
+}
+
void RenderWidgetHostImpl::SetIsDiscarding(bool is_discarding) {
if (is_discarding_ == is_discarding) {
return;
diff --git a/content/browser/renderer_host/render_widget_host_impl.h b/content/browser/renderer_host/render_widget_host_impl.h
index b54affbfd4..8c498fb8de 100644
--- a/content/browser/renderer_host/render_widget_host_impl.h
+++ b/content/browser/renderer_host/render_widget_host_impl.h
@@ -1063,6 +1063,11 @@ class CONTENT_EXPORT RenderWidgetHostImpl
// Requests a commit and forced redraw in the renderer compositor.
void ForceRedrawForTesting();
+ // Like ForceRedrawForTesting but with a completion callback.
+ // The callback fires once the renderer has committed and submitted a
+ // CompositorFrame, ensuring the viz compositor has an up-to-date frame.
+ void ForceRedrawWithCallback(base::OnceClosure callback);
+
// Indicates the page is discarding. The renderer process will get
// cpu-priority boosted to run discard logic.
void SetIsDiscarding(bool is_discarding);
diff --git a/third_party/blink/public/devtools_protocol/domains/Page.pdl b/third_party/blink/public/devtools_protocol/domains/Page.pdl
index 37a0c3a90d..074a3494fe 100644
--- a/third_party/blink/public/devtools_protocol/domains/Page.pdl
+++ b/third_party/blink/public/devtools_protocol/domains/Page.pdl
@@ -605,6 +605,9 @@ domain Page
# Capture clip region directly from current frame without modifying viewport.
# Enables concurrent captures of different regions from the same page.
experimental optional boolean directClip
+ # Skip ForceRedraw — use the current compositor frame as-is.
+ # Only safe when page is known to be fully rendered (e.g. after fonts.ready + rAF).
+ experimental optional boolean skipRedraw
returns
# Base64-encoded image data (empty when rawFilePath is used).
binary data
diff --git a/third_party/blink/renderer/platform/widget/widget_base.cc b/third_party/blink/renderer/platform/widget/widget_base.cc
index 95388b54c0..c0621a87c6 100644
--- a/third_party/blink/renderer/platform/widget/widget_base.cc
+++ b/third_party/blink/renderer/platform/widget/widget_base.cc
@@ -90,11 +90,7 @@ const uint32_t kGpuStreamIdDefault = 0;
static const int kInvalidNextPreviousFlagsValue = -1;
-void OnDidPresentForceDrawFrame(
- mojom::blink::Widget::ForceRedrawCallback callback,
- const gfx::PresentationFeedback& feedback) {
- std::move(callback).Run();
-}
+
bool IsDateTimeInput(ui::TextInputType type) {
return type == ui::TEXT_INPUT_TYPE_DATE ||
@@ -420,6 +416,12 @@ scheduler::WidgetScheduler* WidgetBase::WidgetScheduler() {
return widget_scheduler_.get();
}
+static void OnDidPresentForceDrawFrame(
+ mojom::blink::Widget::ForceRedrawCallback callback,
+ const gfx::PresentationFeedback& feedback) {
+ std::move(callback).Run();
+}
+
void WidgetBase::ForceRedraw(
mojom::blink::Widget::ForceRedrawCallback callback) {
TRACE_EVENT0("renderer", "WidgetBase::ForceRedraw");
+99
View File
@@ -0,0 +1,99 @@
# Chromium Screenshot Patches
Custom Chromium patches for high-throughput headless screenshot capture.
Adds three CDP parameters to `Page.captureScreenshot` and a helper method.
## Patches
**`screenshot-patches.diff`** — 222 lines, 5 files:
| Feature | Description |
|---------|-------------|
| `rawFilePath` | Write raw BGRA pixels to a file path (async ThreadPool). Bypasses PNG/JPEG encoding. |
| `directClip` | Capture clip region directly via `CopyFromSurface(src_rect)` without modifying viewport/emulation state. |
| `skipRedraw` | Lightweight `ForceRedrawWithCallback``CopyFromSurface`. Skips the full `GetSnapshotFromBrowser` presentation feedback wait. |
## Usage (CDP)
```javascript
// rawFilePath — write raw BGRA to /dev/shm (28MB for 875×8192)
await cdp("Page.captureScreenshot", {
rawFilePath: "/dev/shm/tile.raw",
fromSurface: true,
optimizeForSpeed: true,
clip: { x: 0, y: 0, width: 875, height: 8192, scale: 1 }
});
// directClip — capture region without emulation change
await cdp("Page.captureScreenshot", {
directClip: true,
clip: { x: 0, y: 1024, width: 875, height: 1024, scale: 1 }
});
// skipRedraw — ForceRedrawWithCallback then CopyFromSurface
await cdp("Page.captureScreenshot", {
skipRedraw: true,
rawFilePath: "/dev/shm/tile.raw",
clip: { x: 0, y: 0, width: 875, height: 8192, scale: 1 }
});
```
## Raw file format
12-byte header + pixel data:
```
offset 0: uint32 width
offset 4: uint32 height
offset 8: uint32 rowBytes (= width × 4)
offset 12: BGRA pixel data (rowBytes × height bytes)
```
Read in Python:
```python
import struct
from PIL import Image
data = open("tile.raw", "rb").read()
w, h, rb = struct.unpack_from("<III", data, 0)
img = Image.frombuffer("RGBA", (w, h), data[12:], "raw", "BGRA", rb, 1)
```
## Building
Requires Chromium source checkout. Tested on Chromium 150.0.7844.0.
```bash
# 1. Get Chromium source (if not already)
mkdir chromium && cd chromium
fetch --no-history chromium
cd src
# 2. Apply patches
git apply /path/to/screenshot-patches.diff
# 3. Configure build
mkdir -p out/Release
cat > out/Release/args.gn << 'EOF'
is_debug = false
is_official_build = true
is_component_build = false
symbol_level = 0
blink_symbol_level = 0
chrome_pgo_phase = 0
EOF
gn gen out/Release
# 4. Build
autoninja -C out/Release chrome
# ~20 min on 224 cores, ~2 hours on 16 cores
```
Build output: `out/Release/chrome` (~476MB)
## Compatibility
- Chromium 150.x (May 2026). May apply cleanly to nearby versions.
- `is_official_build=true` required for competitive performance (10x vs debug).
- All patches are in `content/browser/devtools/protocol/` (CDP layer) and
`content/browser/renderer_host/` (widget host). No rendering engine changes.
+285
View File
@@ -0,0 +1,285 @@
diff --git a/content/browser/devtools/protocol/page_handler.cc b/content/browser/devtools/protocol/page_handler.cc
index a9f21bbe0c..2749345a75 100644
--- a/content/browser/devtools/protocol/page_handler.cc
+++ b/content/browser/devtools/protocol/page_handler.cc
@@ -506,6 +506,9 @@ struct PageHandler::PendingScreenshotRequest {
gfx::Size original_view_size;
gfx::Size requested_image_size;
std::string raw_file_path;
+ // Used by the directClip path to store the clip rect for capture
+ // after ForceRedraw synchronization.
+ gfx::Rect direct_clip_src_rect;
};
PageHandler::PageHandler(
@@ -1383,7 +1386,8 @@ void PageHandler::CaptureFullPageScreenshot(
CaptureScreenshot(std::move(format), std::move(quality), std::move(clip),
/*from_surface=*/true, /*capture_beyond_viewport=*/true,
std::move(optimize_for_speed), std::move(raw_file_path),
- /*direct_clip=*/std::nullopt, std::move(callback));
+ /*direct_clip=*/std::nullopt, /*skip_redraw=*/std::nullopt,
+ std::move(callback));
}
void PageHandler::CaptureScreenshot(
@@ -1395,6 +1399,7 @@ void PageHandler::CaptureScreenshot(
std::optional<bool> optimize_for_speed,
std::optional<std::string> raw_file_path,
std::optional<bool> direct_clip,
+ std::optional<bool> skip_redraw,
std::unique_ptr<CaptureScreenshotCallback> callback) {
if (!host_ || !host_->GetRenderWidgetHost() ||
!host_->GetRenderWidgetHost()->GetView()) {
@@ -1407,6 +1412,11 @@ void PageHandler::CaptureScreenshot(
// directClip: capture a region directly from the current frame without
// modifying viewport/emulation state. Enables concurrent captures.
+ // A ForceRedraw is issued first to ensure the renderer has submitted a
+ // CompositorFrame, preventing races where CopyFromSurface captures a
+ // stale frame (e.g., about:blank dimensions instead of the loaded page).
+ // This is critical with --in-process-gpu where reduced IPC latency
+ // widens the race window.
if (direct_clip.value_or(false) && clip) {
RenderWidgetHostImpl* widget_host = host_->GetRenderWidgetHost();
float dsf = widget_host->GetDeviceScaleFactor();
@@ -1437,17 +1447,17 @@ void PageHandler::CaptureScreenshot(
static_cast<int>(clip->GetHeight() * dsf));
pending_request->requested_image_size = output_size;
- gfx::Rect src_rect(
+ pending_request->direct_clip_src_rect = gfx::Rect(
static_cast<int>(clip->GetX() * dsf),
static_cast<int>(clip->GetY() * dsf),
output_size.width(), output_size.height());
- static_cast<RenderWidgetHostViewBase*>(widget_host->GetView())
- ->CopyFromSurface(
- src_rect, output_size, base::TimeDelta(),
- base::BindOnce(&PageHandler::DirectClipCaptured,
- weak_factory_.GetWeakPtr(),
- std::move(pending_request)));
+ // ForceRedraw ensures the renderer has committed and submitted a
+ // CompositorFrame to viz before we issue CopyFromSurface.
+ widget_host->ForceRedrawWithCallback(
+ base::BindOnce(&PageHandler::DirectClipForceRedrawDone,
+ weak_factory_.GetWeakPtr(),
+ std::move(pending_request)));
return;
}
@@ -1631,10 +1641,20 @@ void PageHandler::CaptureScreenshot(
}
}
- widget_host->GetSnapshotFromBrowser(
- base::BindOnce(&PageHandler::ScreenshotCaptured,
- weak_factory_.GetWeakPtr(), std::move(pending_request)),
- true);
+ if (skip_redraw.value_or(false)) {
+ // skipRedraw: lightweight ForceRedraw (wait for renderer commit) then
+ // CopyFromSurface. Faster than full GetSnapshotFromBrowser because we
+ // skip the presentation feedback wait.
+ widget_host->ForceRedrawWithCallback(
+ base::BindOnce(&PageHandler::SkipRedrawForceRedrawDone,
+ weak_factory_.GetWeakPtr(),
+ std::move(pending_request)));
+ } else {
+ widget_host->GetSnapshotFromBrowser(
+ base::BindOnce(&PageHandler::ScreenshotCaptured,
+ weak_factory_.GetWeakPtr(), std::move(pending_request)),
+ true);
+ }
}
Response PageHandler::StartScreencast(std::optional<std::string> format,
@@ -1887,6 +1907,49 @@ void PageHandler::ScreencastFrameEncoded(
std::move(page_metadata), session_id_);
}
+void PageHandler::SkipRedrawForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request) {
+ // Allocate a new LocalSurfaceId so that CopyFromSurface targets a fresh
+ // surface, preventing stale-frame captures (same rationale as in
+ // DirectClipForceRedrawDone).
+ SkipRedrawDoCopy(std::move(request));
+}
+
+void PageHandler::SkipRedrawDoCopy(
+ std::unique_ptr<PendingScreenshotRequest> request) {
+ if (!host_ || !host_->GetRenderWidgetHost() ||
+ !host_->GetRenderWidgetHost()->GetView()) {
+ request->callback->sendFailure(Response::InternalError());
+ return;
+ }
+ RenderWidgetHostImpl* widget_host = host_->GetRenderWidgetHost();
+
+ static_cast<RenderWidgetHostViewBase*>(widget_host->GetView())
+ ->CopyFromSurface(
+ gfx::Rect(), gfx::Size(), base::TimeDelta(),
+ base::BindOnce(&PageHandler::DirectClipCaptured,
+ weak_factory_.GetWeakPtr(),
+ std::move(request)));
+}
+
+void PageHandler::DirectClipForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request) {
+ if (!host_ || !host_->GetRenderWidgetHost() ||
+ !host_->GetRenderWidgetHost()->GetView()) {
+ request->callback->sendFailure(Response::InternalError());
+ return;
+ }
+ RenderWidgetHostImpl* widget_host = host_->GetRenderWidgetHost();
+
+ static_cast<RenderWidgetHostViewBase*>(widget_host->GetView())
+ ->CopyFromSurface(
+ request->direct_clip_src_rect,
+ request->requested_image_size, base::TimeDelta(),
+ base::BindOnce(&PageHandler::DirectClipCaptured,
+ weak_factory_.GetWeakPtr(),
+ std::move(request)));
+}
+
void PageHandler::DirectClipCaptured(
std::unique_ptr<PendingScreenshotRequest> request,
const content::CopyFromSurfaceResult& result) {
@@ -1937,31 +2000,49 @@ void PageHandler::ScreenshotCaptured(
Response::ServerError("Bitmap has no pixel data"));
return;
}
+ // Check file extension: .jpg/.jpeg → encode JPEG, otherwise raw BGRA
+ bool is_jpeg = base::EndsWith(request->raw_file_path, ".jpg",
+ base::CompareCase::INSENSITIVE_ASCII) ||
+ base::EndsWith(request->raw_file_path, ".jpeg",
+ base::CompareCase::INSENSITIVE_ASCII);
+ int jpeg_quality = 50; // Default for file-based JPEG output
+
base::ThreadPool::PostTaskAndReplyWithResult(
FROM_HERE,
{base::MayBlock(), base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN},
base::BindOnce(
- [](SkBitmap bmp, std::string path) -> bool {
+ [](SkBitmap bmp, std::string path, bool encode_jpeg,
+ int quality) -> bool {
base::File file(base::FilePath(path),
base::File::FLAG_CREATE_ALWAYS |
base::File::FLAG_WRITE);
if (!file.IsValid()) {
return false;
}
- uint32_t header[3] = {
- static_cast<uint32_t>(bmp.width()),
- static_cast<uint32_t>(bmp.height()),
- static_cast<uint32_t>(bmp.rowBytes())};
- file.WriteAtCurrentPosAndCheck(base::as_byte_span(header));
- size_t pixel_size =
- static_cast<size_t>(bmp.height()) * bmp.rowBytes();
- auto pixel_span = UNSAFE_BUFFERS(base::span<const uint8_t>(
- reinterpret_cast<const uint8_t*>(bmp.getPixels()),
- pixel_size));
- file.WriteAtCurrentPosAndCheck(pixel_span);
+ if (encode_jpeg) {
+ // Encode JPEG in ThreadPool (off UI thread)
+ std::optional<std::vector<uint8_t>> encoded =
+ gfx::JPEGCodec::Encode(bmp, quality);
+ if (!encoded || encoded->empty()) return false;
+ file.WriteAtCurrentPosAndCheck(base::as_byte_span(*encoded));
+ } else {
+ // Raw BGRA with header
+ uint32_t header[3] = {
+ static_cast<uint32_t>(bmp.width()),
+ static_cast<uint32_t>(bmp.height()),
+ static_cast<uint32_t>(bmp.rowBytes())};
+ file.WriteAtCurrentPosAndCheck(base::as_byte_span(header));
+ size_t pixel_size =
+ static_cast<size_t>(bmp.height()) * bmp.rowBytes();
+ auto pixel_span = UNSAFE_BUFFERS(base::span<const uint8_t>(
+ reinterpret_cast<const uint8_t*>(bmp.getPixels()),
+ pixel_size));
+ file.WriteAtCurrentPosAndCheck(pixel_span);
+ }
return true;
},
- std::move(target_bitmap), request->raw_file_path),
+ std::move(target_bitmap), request->raw_file_path,
+ is_jpeg, jpeg_quality),
base::BindOnce(
[](std::unique_ptr<CaptureScreenshotCallback> callback,
bool success) {
diff --git a/content/browser/devtools/protocol/page_handler.h b/content/browser/devtools/protocol/page_handler.h
index 4ccc1b42d3..ac94709199 100644
--- a/content/browser/devtools/protocol/page_handler.h
+++ b/content/browser/devtools/protocol/page_handler.h
@@ -156,6 +156,7 @@ class PageHandler : public DevToolsDomainHandler,
std::optional<bool> optimize_for_speed,
std::optional<std::string> raw_file_path,
std::optional<bool> direct_clip,
+ std::optional<bool> skip_redraw,
std::unique_ptr<CaptureScreenshotCallback> callback) override;
void CaptureSnapshot(
std::optional<std::string> format,
@@ -242,6 +243,15 @@ class PageHandler : public DevToolsDomainHandler,
void DirectClipCaptured(std::unique_ptr<PendingScreenshotRequest> request,
const content::CopyFromSurfaceResult& result);
+ // Called after ForceRedraw ack in directClip path; issues CopyFromSurface
+ // with the stored clip rect now that the renderer has flushed a frame.
+ void DirectClipForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request);
+ // Called after ForceRedraw ack in skipRedraw path; issues the actual
+ // CopyFromSurface now that the renderer has flushed a CompositorFrame.
+ void SkipRedrawDoCopy(std::unique_ptr<PendingScreenshotRequest> request);
+ void SkipRedrawForceRedrawDone(
+ std::unique_ptr<PendingScreenshotRequest> request);
void ScreenshotCaptured(std::unique_ptr<PendingScreenshotRequest> request,
const gfx::Image& image);
diff --git a/content/browser/renderer_host/render_widget_host_impl.cc b/content/browser/renderer_host/render_widget_host_impl.cc
index 7b18f51513..e0c90a415a 100644
--- a/content/browser/renderer_host/render_widget_host_impl.cc
+++ b/content/browser/renderer_host/render_widget_host_impl.cc
@@ -4197,6 +4197,15 @@ void RenderWidgetHostImpl::ForceRedrawForTesting() {
blink_widget_->ForceRedraw(base::DoNothing());
}
+void RenderWidgetHostImpl::ForceRedrawWithCallback(
+ base::OnceClosure callback) {
+ if (!blink_widget_) {
+ std::move(callback).Run();
+ return;
+ }
+ blink_widget_->ForceRedraw(std::move(callback));
+}
+
void RenderWidgetHostImpl::SetIsDiscarding(bool is_discarding) {
if (is_discarding_ == is_discarding) {
return;
diff --git a/content/browser/renderer_host/render_widget_host_impl.h b/content/browser/renderer_host/render_widget_host_impl.h
index b54affbfd4..8c498fb8de 100644
--- a/content/browser/renderer_host/render_widget_host_impl.h
+++ b/content/browser/renderer_host/render_widget_host_impl.h
@@ -1063,6 +1063,11 @@ class CONTENT_EXPORT RenderWidgetHostImpl
// Requests a commit and forced redraw in the renderer compositor.
void ForceRedrawForTesting();
+ // Like ForceRedrawForTesting but with a completion callback.
+ // The callback fires once the renderer has committed and submitted a
+ // CompositorFrame, ensuring the viz compositor has an up-to-date frame.
+ void ForceRedrawWithCallback(base::OnceClosure callback);
+
// Indicates the page is discarding. The renderer process will get
// cpu-priority boosted to run discard logic.
void SetIsDiscarding(bool is_discarding);
diff --git a/third_party/blink/public/devtools_protocol/domains/Page.pdl b/third_party/blink/public/devtools_protocol/domains/Page.pdl
index 37a0c3a90d..074a3494fe 100644
--- a/third_party/blink/public/devtools_protocol/domains/Page.pdl
+++ b/third_party/blink/public/devtools_protocol/domains/Page.pdl
@@ -605,6 +605,9 @@ domain Page
# Capture clip region directly from current frame without modifying viewport.
# Enables concurrent captures of different regions from the same page.
experimental optional boolean directClip
+ # Skip ForceRedraw — use the current compositor frame as-is.
+ # Only safe when page is known to be fully rendered (e.g. after fonts.ready + rAF).
+ experimental optional boolean skipRedraw
returns
# Base64-encoded image data (empty when rawFilePath is used).
binary data
+262
View File
@@ -0,0 +1,262 @@
#!/usr/bin/env python3
"""PixelRAG Agent — Claude + tool_use for visual web search.
A real Anthropic agent that uses Claude to answer questions by searching
a visual Wikipedia index via PixelRAG. Claude decides when to call the
search tool and synthesizes answers from visual retrieval results.
Prerequisites:
- ANTHROPIC_API_KEY env var set
- pixelrag serve running on localhost:30001 (or set --endpoint)
Usage:
# Interactive conversation with the agent
python demos/agent_skill.py
# Single question
python demos/agent_skill.py "Who invented the telephone?"
# Custom endpoint
python demos/agent_skill.py --endpoint http://gpu-box:30001 "Eiffel Tower history"
"""
import argparse
import json
import sys
import urllib.request
import anthropic
SEARCH_TOOL = {
"name": "pixelrag_search",
"description": (
"Search a visual Wikipedia index using natural language queries. "
"Returns ranked results with article URLs and relevance scores. "
"Use this tool to find information about any topic — it searches "
"screenshot-based embeddings of Wikipedia articles, so it works well "
"for both textual and visual content."
),
"input_schema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Natural language search query",
},
"n_results": {
"type": "integer",
"description": "Number of results to return (default 5, max 20)",
"default": 5,
},
},
"required": ["query"],
},
}
WEB_FETCH_TOOL = {
"name": "web_fetch",
"description": (
"Fetch the text content of a URL. Use this to read Wikipedia articles "
"or other web pages found via search. Returns the page text content."
),
"input_schema": {
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "URL to fetch",
},
},
"required": ["url"],
},
}
TOOLS = [SEARCH_TOOL, WEB_FETCH_TOOL]
SYSTEM_PROMPT = """\
You are a research assistant with access to a visual Wikipedia search engine (PixelRAG).
When asked a question, use the pixelrag_search tool to find relevant Wikipedia articles,
then synthesize an answer from the results. You may search multiple times with different
queries to gather comprehensive information. Cite your sources with Wikipedia URLs.
If search results are insufficient, say so honestly rather than guessing."""
def execute_pixelrag_search(
query: str, n_results: int = 5, endpoint: str = "http://localhost:30001"
) -> dict:
"""Call the PixelRAG search API."""
body = json.dumps(
{"queries": [{"text": query}], "n_docs": min(n_results, 20)}
).encode()
req = urllib.request.Request(
f"{endpoint}/search",
data=body,
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=30) as resp:
data = json.loads(resp.read())
hits = data.get("results", [{}])[0].get("hits", [])
results = []
for hit in hits:
url = hit.get("url", "")
slug = url.split("/wiki/")[-1] if "/wiki/" in url else ""
title = slug.replace("_", " ") if slug else url
results.append(
{
"title": title,
"url": url,
"score": round(hit["score"], 4),
"tile": f"tile_{hit.get('tile_index', '?')}_chunk_{hit.get('chunk_index', '?')}",
}
)
return {"query": query, "results": results, "count": len(results)}
def execute_web_fetch(url: str) -> dict:
"""Fetch text from a URL (simplified — returns first 4000 chars)."""
req = urllib.request.Request(url, headers={"User-Agent": "PixelRAG-Agent/1.0"})
with urllib.request.urlopen(req, timeout=15) as resp:
raw = resp.read().decode("utf-8", errors="replace")
# Strip HTML tags for a rough text extraction
import re
text = re.sub(r"<script[^>]*>.*?</script>", "", raw, flags=re.DOTALL)
text = re.sub(r"<style[^>]*>.*?</style>", "", text, flags=re.DOTALL)
text = re.sub(r"<[^>]+>", " ", text)
text = re.sub(r"\s+", " ", text).strip()
return {"url": url, "content": text[:4000], "truncated": len(text) > 4000}
def handle_tool_call(tool_name: str, tool_input: dict, endpoint: str) -> str:
"""Execute a tool call and return the result as a string."""
try:
if tool_name == "pixelrag_search":
result = execute_pixelrag_search(
query=tool_input["query"],
n_results=tool_input.get("n_results", 5),
endpoint=endpoint,
)
elif tool_name == "web_fetch":
result = execute_web_fetch(url=tool_input["url"])
else:
result = {"error": f"Unknown tool: {tool_name}"}
except Exception as e:
result = {"error": str(e)}
return json.dumps(result)
def run_agent(
question: str,
endpoint: str,
model: str = "claude-sonnet-4-20250514",
verbose: bool = False,
) -> str:
"""Run the agent loop: send question → handle tool calls → return final answer."""
client = anthropic.Anthropic()
messages = [{"role": "user", "content": question}]
while True:
response = client.messages.create(
model=model,
max_tokens=4096,
system=SYSTEM_PROMPT,
tools=TOOLS,
messages=messages,
)
if verbose:
print(
f" [stop_reason={response.stop_reason}, usage={response.usage}]",
file=sys.stderr,
)
if response.stop_reason == "end_turn":
# Extract text from response
text_parts = [b.text for b in response.content if b.type == "text"]
return "\n".join(text_parts)
# Handle tool use
tool_results = []
for block in response.content:
if block.type == "tool_use":
if verbose:
print(
f" [tool: {block.name}({json.dumps(block.input, ensure_ascii=False)})]",
file=sys.stderr,
)
result = handle_tool_call(block.name, block.input, endpoint)
tool_results.append(
{
"type": "tool_result",
"tool_use_id": block.id,
"content": result,
}
)
if not tool_results:
# No tool calls and not end_turn — shouldn't happen, but handle gracefully
text_parts = [b.text for b in response.content if b.type == "text"]
return "\n".join(text_parts) if text_parts else "(no response)"
messages.append({"role": "assistant", "content": response.content})
messages.append({"role": "user", "content": tool_results})
def interactive(endpoint: str, model: str, verbose: bool):
"""Run interactive conversation loop."""
print("PixelRAG Agent (Claude + visual search)")
print(f" endpoint: {endpoint}")
print(f" model: {model}")
print(" Type 'quit' to exit.\n")
while True:
try:
question = input("You: ").strip()
except (EOFError, KeyboardInterrupt):
print()
break
if not question or question.lower() in ("quit", "exit", "q"):
break
print()
try:
answer = run_agent(question, endpoint, model, verbose)
print(f"Agent: {answer}\n")
except anthropic.APIError as e:
print(f"API error: {e}\n")
except Exception as e:
print(f"Error: {e}\n")
def main():
parser = argparse.ArgumentParser(
description="PixelRAG Agent — Claude + visual web search"
)
parser.add_argument(
"question", nargs="?", help="Question to ask (omit for interactive mode)"
)
parser.add_argument(
"--endpoint",
default="http://localhost:30001",
help="PixelRAG search API endpoint",
)
parser.add_argument(
"--model", default="claude-sonnet-4-20250514", help="Claude model to use"
)
parser.add_argument(
"--verbose", "-v", action="store_true", help="Show tool calls and API details"
)
args = parser.parse_args()
if args.question:
answer = run_agent(args.question, args.endpoint, args.model, args.verbose)
print(answer)
else:
interactive(args.endpoint, args.model, args.verbose)
if __name__ == "__main__":
main()
+31
View File
@@ -0,0 +1,31 @@
# End-to-End Demo: Wikipedia → Search
Builds a visual search index from Wikipedia articles and queries it.
## Quick Start
```bash
cd pixelrag
uv run python demos/e2e/run.py
```
This will:
1. Start kiwix-serve with Simple English Wikipedia
2. Capture 100 article screenshots (~20s)
3. Chunk tiles into 1024px strips (~1s)
4. Embed chunks with Qwen3-VL on CPU (~5 min for 100 articles)
5. Build a FAISS index
6. Start a search API and run sample queries
## Prerequisites
- Simple English Wikipedia ZIM at `~/pixelrag-data/zim/wikipedia_en_simple.zim`
(download: `curl -L https://download.kiwix.org/zim/wikipedia/wikipedia_en_simple_all_nopic_2026-05.zim -o ~/pixelrag-data/zim/wikipedia_en_simple.zim`)
- kiwix-serve binary at `.local/bin/kiwix-serve`
## Configuration
Edit `pixelrag.yaml` in this directory to change:
- Number of articles (`limit`)
- Embedding device (`cpu` or `cuda`)
- Output location
+15
View File
@@ -0,0 +1,15 @@
source:
type: kiwix
zim_path: wikipedia-simple # alias → auto-downloads Simple English Wikipedia (~1GB)
# Or use a file path: zim_path: ~/pixelrag-data/zim/wikipedia_en_simple.zim
# Or a URL: zim_path: https://download.kiwix.org/zim/wikipedia/...
ingest:
backend: cdp
quality: 85
embed:
model: Qwen/Qwen3-VL-Embedding-2B
device: cpu
output: demos/e2e/output
+218
View File
@@ -0,0 +1,218 @@
#!/usr/bin/env python3
"""End-to-end demo: Wikipedia → visual search index → query.
Demonstrates the full PixelRAG pipeline via pixelrag index:
source → ingest → chunk → embed → build index → serve → search
Run:
cd pixelrag
uv run python demos/e2e/run.py
uv run python demos/e2e/run.py --limit 50
uv run python demos/e2e/run.py --skip-build # just serve existing index
"""
import argparse
import json
import logging
import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
logger = logging.getLogger("e2e_demo")
DEFAULT_OUTPUT = Path(__file__).parent / "output"
SAMPLE_QUERIES = [
"theory of relativity physics",
"photosynthesis plants energy",
"world war two history",
"programming language computer",
"solar system planets",
"human brain neuroscience",
"climate change global warming",
"DNA genetics biology",
]
def search(query: str, port: int) -> list[dict]:
body = json.dumps({"queries": [{"text": query}], "n_docs": 5}).encode()
req = urllib.request.Request(
f"http://localhost:{port}/search",
data=body,
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=30) as resp:
data = json.loads(resp.read())
return data.get("results", [{}])[0].get("hits", [])
def _generate_html_report(results: list[dict], html_path: Path) -> None:
"""Generate an HTML page showing search results with tile images."""
import base64
rows = []
for r in results:
query = r["query"]
rows.append(f'<h2>Q: "{query}"</h2>')
for i, h in enumerate(r.get("hits", [])[:3]):
url = h.get("url", "")
title = (
url.split("/")[-1].replace("_", " ") if url else f"#{h['article_id']}"
)
score = h["score"]
tile_html = ""
tile_path = h.get("_tile_path")
if tile_path and Path(tile_path).exists():
data = Path(tile_path).read_bytes()
ext = Path(tile_path).suffix.lstrip(".")
b64 = base64.b64encode(data).decode()
tile_html = f'<img src="data:image/{ext};base64,{b64}" style="max-width:600px;border:1px solid #ddd;border-radius:4px;">'
rows.append(f"""
<div style="margin:1em 0;padding:1em;border:1px solid #222;border-radius:8px;background:#111;">
<div style="color:#4a9eff;font-weight:600;">{i + 1}. {score:.3f}{title}</div>
{f'<div style="margin-top:0.5em;">{tile_html}</div>' if tile_html else ""}
</div>""")
html = f"""<!DOCTYPE html>
<html><head><meta charset="utf-8"><title>PixelRAG E2E Results</title>
<style>body{{font-family:system-ui;background:#0a0a0a;color:#e0e0e0;max-width:800px;margin:2em auto;padding:0 1em;}}
h1{{color:#fff;}}h2{{color:#aaa;margin-top:2em;}}</style></head>
<body><h1>PixelRAG Search Results</h1>
{"".join(rows)}
</body></html>"""
html_path.write_text(html)
def main() -> None:
parser = argparse.ArgumentParser(description="PixelRAG E2E Demo")
parser.add_argument("--limit", "-n", type=int, default=100)
parser.add_argument(
"--config", "-c", type=Path, default=Path(__file__).parent / "pixelrag.yaml"
)
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
parser.add_argument("--serve-port", type=int, default=31337)
parser.add_argument("--skip-build", action="store_true")
parser.add_argument(
"--show-tiles",
action="store_true",
help="Display tile images in terminal (requires chafa)",
)
args = parser.parse_args()
output = args.output.resolve()
print("=" * 60)
print(" PixelRAG End-to-End Demo")
print("=" * 60)
print(f" Articles: {args.limit}")
print(f" Device: {args.device}")
print(f" Output: {output}")
print()
# --- Build index ---
if not args.skip_build:
from pixelrag_index.config import load_config
from pixelrag_index.pipelines import build
config = load_config(str(args.config))
if args.device:
config.setdefault("embed", {})["device"] = args.device
config["output"] = str(output)
t0 = time.time()
build(config, limit=args.limit)
total_time = time.time() - t0
print(f"\n Pipeline completed in {total_time:.1f}s\n")
# --- Serve + Search ---
logger.info("Starting search API on :%d...", args.serve_port)
env = os.environ.copy()
env["PIXELRAG_INDEX_DIR"] = str(output)
env["PIXELRAG_ARTICLES_JSON"] = str(output / "articles.json")
serve_proc = subprocess.Popen(
[sys.executable, "-m", "pixelrag_serve.api", "--port", str(args.serve_port)],
env=env,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
try:
for _ in range(120):
try:
urllib.request.urlopen(
f"http://localhost:{args.serve_port}/health", timeout=1
)
break
except Exception:
time.sleep(2)
else:
raise TimeoutError("Search API failed to start")
print("=" * 60)
print(" Search Results")
print("=" * 60)
all_results = []
for query in SAMPLE_QUERIES:
hits = search(query, args.serve_port)
all_results.append({"query": query, "hits": hits})
print(f'\n Q: "{query}"')
if not hits:
print(" (no results)")
continue
for i, h in enumerate(hits[:3]):
url = h.get("url", "")
score = h["score"]
# Extract readable title from URL
if url:
title = (
url.split("/")[-1]
.replace("_", " ")
.replace("%22", '"')
.replace("%20", " ")
)
title = urllib.parse.unquote(title)
else:
title = f"#{h['article_id']}"
print(f" {i + 1}. {score:.3f} {title}")
# Collect tile path for HTML report
if args.show_tiles:
aid = h["article_id"]
ti = h.get("tile_index", 0)
ci = h.get("chunk_index", 0)
for candidate in [
output
/ "tiles"
/ f"{aid}.png.tiles"
/ f"chunk_{ti:04d}_{ci:02d}.png",
output / "tiles" / f"{aid}.png.tiles" / f"tile_{ti:04d}.jpg",
]:
if candidate.exists():
h["_tile_path"] = str(candidate)
break
# Generate HTML results page with tile images
if args.show_tiles:
html_path = output / "results.html"
_generate_html_report(all_results, html_path)
print(f"\n Results with images: file://{html_path}")
print()
print(f"Search API: http://localhost:{args.serve_port}")
print(
f"Try: curl -X POST http://localhost:{args.serve_port}/search "
f"-H 'Content-Type: application/json' "
f'-d \'{{"queries": [{{"text": "your query"}}], "n_docs": 5}}\''
)
finally:
serve_proc.terminate()
if __name__ == "__main__":
main()
File diff suppressed because one or more lines are too long
+17
View File
@@ -0,0 +1,17 @@
# Ingest Demo: Heterogeneous Documents
Captures screenshots from a mix of URLs, PDFs, and HTML files in one call.
## Run
```bash
cd pixelrag
uv run python demos/render/run.py
```
## What it does
1. Fetches 3 Wikipedia articles (URLs)
2. Renders 2 local HTML files (created on the fly)
3. Produces tiled JPEG screenshots for all 5
4. Shows a summary of tiles, sizes, and timing
+176
View File
@@ -0,0 +1,176 @@
#!/usr/bin/env python3
"""Ingest demo: capture heterogeneous documents as tiled screenshots.
Demonstrates pixelshot rendering a mix of:
- Wikipedia article URLs (via CDP lean capture)
- Local HTML files (auto-detected, rendered via file:// URL)
- Could also handle PDFs (requires pdf2image)
Run:
cd pixelrag
uv run python demos/render/run.py
"""
import shutil
import time
from pathlib import Path
OUTPUT = Path("demos/render/output")
# --- Sample data ---
WIKI_URLS = [
"https://en.wikipedia.org/wiki/Retrieval-augmented_generation",
"https://en.wikipedia.org/wiki/Screenshot",
"https://en.wikipedia.org/wiki/FAISS",
]
SAMPLE_HTML = """<!DOCTYPE html>
<html><head><meta charset="utf-8"><title>{title}</title>
<style>
body {{ font-family: Georgia, serif; max-width: 800px; margin: 2em auto; padding: 0 1em; line-height: 1.6; }}
h1 {{ color: #1a1a2e; border-bottom: 2px solid #e2e2e2; padding-bottom: .3em; }}
table {{ border-collapse: collapse; width: 100%; margin: 1em 0; }}
th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
th {{ background: #f5f5f5; }}
.highlight {{ background: #fff3cd; padding: .2em .4em; border-radius: 3px; }}
</style></head>
<body>
<h1>{title}</h1>
<p>{body}</p>
{extra}
</body></html>"""
def create_sample_html(output_dir: Path) -> list[Path]:
"""Create sample HTML files to demonstrate local file ingestion."""
html_dir = output_dir / "sample_html"
html_dir.mkdir(parents=True, exist_ok=True)
files = []
# A simple article-style page
p1 = html_dir / "visual_retrieval.html"
p1.write_text(
SAMPLE_HTML.format(
title="Visual Document Retrieval",
body=(
"Visual document retrieval captures documents as images and uses "
"vision-language models to embed them into a shared vector space. "
"Unlike text-based retrieval which requires parsing, visual retrieval "
"preserves <span class='highlight'>layout, tables, figures, and formatting</span> "
"that text extraction often loses."
),
extra="""
<h2>Comparison</h2>
<table>
<tr><th>Method</th><th>Preserves Layout</th><th>Handles Tables</th><th>Needs Parser</th></tr>
<tr><td>Text extraction</td><td>No</td><td>Partial</td><td>Yes</td></tr>
<tr><td>HTML rendering</td><td>Partial</td><td>Yes</td><td>Yes</td></tr>
<tr><td><b>Visual (screenshot)</b></td><td><b>Yes</b></td><td><b>Yes</b></td><td><b>No</b></td></tr>
</table>
""",
)
)
files.append(p1)
# A data-heavy page with tables
p2 = html_dir / "benchmark_results.html"
rows = "".join(
f"<tr><td>Config {i}</td><td>{70 + i * 1.3:.1f}</td><td>{0.5 + i * 0.02:.2f}s</td><td>{'LoRA' if i % 2 else 'Base'}</td></tr>"
for i in range(15)
)
p2.write_text(
SAMPLE_HTML.format(
title="PixelRAG Benchmark Results",
body="Evaluation results across different configurations and model variants.",
extra=f"""
<h2>SimpleQA Retrieval Scores</h2>
<table>
<tr><th>Configuration</th><th>Recall@1</th><th>Latency</th><th>Model</th></tr>
{rows}
</table>
""",
)
)
files.append(p2)
return files
def main() -> None:
from pixelrag_render.render import render_file
# Clean previous output
if OUTPUT.exists():
shutil.rmtree(OUTPUT)
OUTPUT.mkdir(parents=True)
print("=" * 60)
print(" PixelRAG Ingest Demo: Heterogeneous Documents")
print("=" * 60)
print()
# --- Step 1: Create sample local HTML ---
print("[1] Creating sample HTML files...")
html_files = create_sample_html(OUTPUT)
for f in html_files:
print(f" {f.name} ({f.stat().st_size / 1024:.1f} KB)")
print()
tiles_dir = OUTPUT / "tiles"
tiles_dir.mkdir()
all_results: list[tuple[str, int, float]] = []
# --- Step 2: Render Wikipedia URLs ---
print(f"[2] Rendering {len(WIKI_URLS)} Wikipedia articles (CDP backend)...")
t0 = time.time()
from pixelrag_render.render import render_urls
url_tiles = render_urls(WIKI_URLS, str(tiles_dir), backend="cdp", workers=3)
elapsed = time.time() - t0
for td in url_tiles:
n = len(list(td.glob("tile_*")))
name = td.name.replace(".png.tiles", "")
all_results.append((f"URL: {name}", n, elapsed / len(WIKI_URLS)))
print(f" {len(url_tiles)} pages rendered in {elapsed:.1f}s")
print()
# --- Step 3: Render local HTML files ---
print(f"[3] Rendering {len(html_files)} local HTML files...")
for html_file in html_files:
t0 = time.time()
result = render_file(str(html_file), str(tiles_dir), backend="cdp")
elapsed = time.time() - t0
for td in result:
n = len(list(Path(td).glob("tile_*")))
all_results.append((f"HTML: {html_file.name}", n, elapsed))
print(f" {len(html_files)} files rendered")
print()
# --- Summary ---
print("=" * 60)
print(" Results")
print("=" * 60)
total_tiles = 0
for name, n_tiles, elapsed in all_results:
total_tiles += n_tiles
print(f" {name:<45} {n_tiles:>3} tiles {elapsed:.1f}s")
print(f" {'' * 55}")
print(f" {'TOTAL':<45} {total_tiles:>3} tiles")
print()
# Show output structure
print("Output structure:")
for td in sorted(tiles_dir.iterdir()):
if td.is_dir():
tiles = list(td.glob("tile_*"))
size = sum(t.stat().st_size for t in tiles) / 1024
print(f" {td.name}/")
print(f" {len(tiles)} tiles, {size:.0f} KB total")
print()
print(f"All output in: {tiles_dir}")
if __name__ == "__main__":
main()
+138
View File
@@ -0,0 +1,138 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>PixelRAG Search</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: -apple-system, system-ui, sans-serif; background: #0a0a0a; color: #e0e0e0; min-height: 100vh; display: flex; flex-direction: column; align-items: center; }
.container { max-width: 900px; width: 100%; padding: 2rem 1.5rem; }
h1 { font-size: 2rem; font-weight: 300; margin-bottom: 0.3rem; color: #fff; }
h1 span { font-weight: 600; }
.subtitle { color: #888; margin-bottom: 2rem; font-size: 0.9rem; }
.search-box { display: flex; gap: 0.5rem; margin-bottom: 0.8rem; }
input[type="text"] { flex: 1; padding: 0.75rem 1rem; border: 1px solid #333; border-radius: 8px; background: #1a1a1a; color: #fff; font-size: 1rem; outline: none; transition: border-color 0.2s; }
input[type="text"]:focus { border-color: #4a9eff; }
button { padding: 0.75rem 1.5rem; border: none; border-radius: 8px; background: #4a9eff; color: #fff; font-size: 1rem; cursor: pointer; transition: background 0.2s; white-space: nowrap; }
button:hover { background: #3a8eef; }
button:disabled { background: #333; cursor: not-allowed; }
.controls { display: flex; gap: 1rem; align-items: center; margin-bottom: 2rem; font-size: 0.85rem; color: #888; }
.controls label { display: flex; align-items: center; gap: 0.3rem; }
.controls select, .controls input[type="number"] { background: #1a1a1a; border: 1px solid #333; border-radius: 4px; color: #e0e0e0; padding: 0.3rem 0.5rem; font-size: 0.85rem; }
.controls input[type="number"] { width: 60px; }
.status { font-size: 0.85rem; color: #888; margin-bottom: 1.5rem; }
.status .time { color: #4a9eff; }
.results { display: flex; flex-direction: column; gap: 1rem; }
.hit { display: flex; gap: 1rem; padding: 1rem; border: 1px solid #222; border-radius: 10px; background: #111; transition: border-color 0.2s; }
.hit:hover { border-color: #333; }
.hit-rank { font-size: 1.5rem; font-weight: 700; color: #333; min-width: 2rem; display: flex; align-items: flex-start; justify-content: center; padding-top: 0.2rem; }
.hit-content { flex: 1; min-width: 0; }
.hit-title { font-size: 1rem; font-weight: 500; margin-bottom: 0.3rem; }
.hit-title a { color: #4a9eff; text-decoration: none; }
.hit-title a:hover { text-decoration: underline; }
.hit-meta { font-size: 0.8rem; color: #666; display: flex; gap: 1rem; flex-wrap: wrap; }
.hit-meta .score { color: #4a9eff; font-weight: 600; }
.hit-meta code { background: #1a1a1a; padding: 0.1rem 0.4rem; border-radius: 3px; font-size: 0.75rem; }
.empty { text-align: center; padding: 3rem; color: #555; }
.error { color: #ff6b6b; padding: 1rem; border: 1px solid #ff6b6b33; border-radius: 8px; background: #ff6b6b0a; }
@keyframes pulse { 0%,100% { opacity: 1; } 50% { opacity: 0.5; } }
.loading { animation: pulse 1s infinite; }
</style>
</head>
<body>
<div class="container">
<h1><span>PixelRAG</span> Search</h1>
<p class="subtitle">Visual retrieval over 15.7M Wikipedia screenshot tiles</p>
<div class="search-box">
<input type="text" id="query" placeholder="Search Wikipedia visually..." autofocus>
<button id="searchBtn" onclick="search()">Search</button>
</div>
<div class="controls">
<label>Results: <input type="number" id="nDocs" value="10" min="1" max="100"></label>
<label>Endpoint:
<select id="endpoint">
<option value="http://localhost:30001">:30001 text (15.7M)</option>
<option value="http://localhost:30002">:30002 pixel base (28M)</option>
<option value="http://localhost:30003">:30003 pixel LoRA (28M)</option>
</select>
</label>
</div>
<div id="status" class="status"></div>
<div id="results" class="results"></div>
</div>
<script>
const queryInput = document.getElementById('query');
const resultsDiv = document.getElementById('results');
const statusDiv = document.getElementById('status');
const searchBtn = document.getElementById('searchBtn');
queryInput.addEventListener('keydown', e => { if (e.key === 'Enter') search(); });
async function search() {
const query = queryInput.value.trim();
if (!query) return;
const endpoint = document.getElementById('endpoint').value;
const nDocs = parseInt(document.getElementById('nDocs').value) || 10;
searchBtn.disabled = true;
searchBtn.textContent = 'Searching...';
statusDiv.innerHTML = '<span class="loading">Embedding query and searching index...</span>';
resultsDiv.innerHTML = '';
const t0 = performance.now();
try {
const resp = await fetch(`${endpoint}/search`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ queries: [{ text: query }], n_docs: nDocs }),
});
if (!resp.ok) throw new Error(`HTTP ${resp.status}: ${await resp.text()}`);
const data = await resp.json();
const elapsed = ((performance.now() - t0) / 1000).toFixed(2);
const hits = data.results?.[0]?.hits || [];
statusDiv.innerHTML = `${hits.length} results in <span class="time">${elapsed}s</span>`;
if (hits.length === 0) {
resultsDiv.innerHTML = '<div class="empty">No results found</div>';
return;
}
resultsDiv.innerHTML = hits.map((hit, i) => {
const slug = hit.url ? hit.url.split('/wiki/').pop() : '?';
const title = decodeURIComponent(slug).replace(/_/g, ' ');
return `
<div class="hit">
<div class="hit-rank">${i + 1}</div>
<div class="hit-content">
<div class="hit-title"><a href="${hit.url}" target="_blank">${title}</a></div>
<div class="hit-meta">
<span class="score">${hit.score.toFixed(3)}</span>
<span>article #${hit.article_id}</span>
<span>tile ${hit.tile_index}:${hit.chunk_index}</span>
<code>${hit.tile_height}px @ y=${hit.y_offset}</code>
</div>
</div>
</div>`;
}).join('');
} catch (err) {
statusDiv.innerHTML = '';
resultsDiv.innerHTML = `<div class="error">${err.message}</div>`;
} finally {
searchBtn.disabled = false;
searchBtn.textContent = 'Search';
}
}
</script>
</body>
</html>
+26
View File
@@ -0,0 +1,26 @@
[Unit]
Description=PixelRAG Agent backend (Claude Agent SDK, subscription auth)
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=yichuan
Group=yichuan
WorkingDirectory=/home/yichuan/visrag/web
Environment=HOME=/home/yichuan
Environment=AGENT_PORT=30010
Environment=PIXELRAG_SEARCH_URL=http://localhost:30001
Environment=CHAT_MAX_BUDGET_USD=2.00
Environment=ALLOWED_ORIGIN=https://pixelrag.ai
Environment=RL_PER_IP=12
Environment=RL_GLOBAL_DAILY=500
Environment=RL_MAX_CONCURRENT=3
ExecStart=/home/yichuan/.nix-profile/bin/node /home/yichuan/visrag/web/agent-server.mjs
Restart=always
RestartSec=3
StandardOutput=append:/home/yichuan/visrag/logs/agent-server.log
StandardError=append:/home/yichuan/visrag/logs/agent-server.log
[Install]
WantedBy=multi-user.target
+21
View File
@@ -0,0 +1,21 @@
[Unit]
Description=PixelRAG Search API
After=network.target
[Service]
Type=simple
User=yichuan
WorkingDirectory=/home/yichuan/visrag
ExecStart=/home/yichuan/visrag/.venv/bin/pixelrag-serve \
--index-dir /home/yichuan/visrag-data/search_index \
--tiles-dir /home/yichuan/visrag-data/tiles \
--articles-json /home/yichuan/visrag-data/articles.json \
--device cpu \
--port 30001
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal
[Install]
WantedBy=multi-user.target
Binary file not shown.

After

Width:  |  Height:  |  Size: 526 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 123 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 734 KiB

@@ -0,0 +1,184 @@
# Screenshot Throughput Optimization — Working Progress
## Target: 150 t/s @ 100% correct (8192px tiles, maxi Wikipedia)
## Current Best
| Config | t/s | Correct | Notes |
|--------|-----|---------|-------|
| multi-process 48w (frameStoppedLoading) | **91** | 100% ✓ | Stable, production-ready |
| multi-process 48w (frameNavigated) | **98** | 100% ✓ | Stable (igpu incompatible) |
| multi-process 48w (2000 art) | **113** | 99.8% ✓ | Steady-state |
| igpu 48w + frameStoppedLoading | **117-132** | 90-97% | Fast but 3-10% about:blank |
| igpu 48w + directClip | **128-148** | 48-90% | Fastest, worst correctness |
## Production System Comparison
The wiki-screenshot production system (`~/pixelrag-src/wiki-screenshot/`) uses:
```python
wait_fonts = False # for kiwix/ZIM datasource
wait_images = False # for kiwix/ZIM datasource
pre_screenshot_delay = 0.5 # fixed 500ms sleep, no fonts.ready
```
- Playwright-based (not CDP websocket)
- GPU-accelerated (8× L40S per machine)
- Multi-machine: 4 machines × ~70-80 t/s = ~290 t/s total
- Full Wikipedia (8.28M articles) processed in ~1 day
Our optimizations added `fonts.ready + eager images + double-rAF` for pixel-perfect
correctness. Production skips these waits entirely (`pre_screenshot_delay=0` in
coordinator). This is safe for Kiwix because all assets (including fonts) are served
from localhost — they load before `wait_until="load"` fires.
Gemini Vision validation of 5000 production tiles:
- 0% BROKEN_RENDER, 0% ERROR_PAGE (rendering is correct without font wait)
- 12% BLANK/PARTIAL_BLANK (tile loop overshoots page height — separate bug)
**Benchmark result**: Removing font/image wait gives only +4% throughput (99 vs 96 t/s)
because nav is not the bottleneck — capture IPC is. The 290 t/s production rate comes
from 4 machines × GPU acceleration, not from skipping font waits.
## Pipeline Bottleneck Analysis
```
Stage Capacity Bottleneck?
Nav 430 pg/s No (3.4x headroom)
Capture 125 t/s YES (C/T_c = 48/321ms)
Steady-state theoretical: 125-150 t/s
Actual (200 art): 98 t/s (75% utilization, 25% = nav serial)
Actual (2000 art): 113 t/s (85% utilization)
```
Per-capture breakdown at 48 concurrent:
- IPC roundtrip: 181ms (ForceRedraw browser→renderer→compositor, 8 async hops)
- DrawRenderPass: 62ms (composite 136 quads)
- CopyDrawnRenderPass: 46ms (memcpy 28MB)
Throughput = `C / T_c(C)` converges at ~125-130 t/s (USL contention curve).
Nav latency (186ms) does not affect steady-state throughput (Little's Law).
Minimum workers to saturate capture: `C × (1 + T_nav/T_cap) = 72`.
## Chromium Patches (in custom build)
| Patch | File | Impact |
|-------|------|--------|
| rawFilePath | page_handler.cc + Page.pdl | Async write raw BGRA to /dev/shm (ThreadPool) |
| directClip | page_handler.cc + Page.pdl | CopyFromSurface(src_rect) without emulation change |
| skipRedraw | page_handler.cc + Page.pdl | ForceRedrawWithCallback → CopyFromSurface |
| ForceRedrawWithCallback | render_widget_host_impl.cc | Lightweight ForceRedraw with commit callback |
| directClip ForceRedraw fix | page_handler.cc | directClip also does ForceRedraw before copy |
## Strategy Architecture
Strategies separated from bench framework:
- `pixelrag_render.strategies/` — capture strategies (CDPPhased, CDPSequential, etc.)
- `pixelrag_render.bench/` — measurement harness with GT validation + experiment dump
- `Bench` class: `bench.run(strategy)` → GT cache + capture + verify + JSON dump
### CDPPhasedStrategy (best strategy)
- Work-stealing queue (asyncio.Queue, not round-robin)
- Semaphore-limited concurrent captures
- `wait_for_event("Page.frameStoppedLoading")` filtered by main frameId
- Per-tile semaphore release (fine-grained pipelining)
- Configurable: tile_height, nav_timeout, use_direct_clip, extra_chrome_args
### WebsocketConnection
- Background `_recv_loop` for multiplexed CDP
- `wait_for_event(method, timeout, filter_fn)` for async event listening
- Supports concurrent `cdp()` calls via pending futures dict
## What Was Tried
### Worked
- ✅ rawFilePath: async write bypasses PNG encoding (+15%)
- ✅ directClip: parallel tile capture within viewport
- ✅ Phased strategy: semaphore-limited captures reduce contention (+15%)
- ✅ Work-stealing queue: better load balancing
- ✅ frameNavigated/frameStoppedLoading wait: fixes igpu about:blank race
- ✅ Presentation feedback ForceRedraw: 100% correct (but slower)
### Partially Worked
- ⚠️ --in-process-gpu: 120+ t/s but 5-10% about:blank captures
- ⚠️ SwapPromise ForceRedraw: shot_p50 325→303ms (7% gain)
- ⚠️ directClip for all tiles: fast but correctness depends on ForceRedraw
### Did Not Work
- ❌ --single-process: 168 t/s but 74% correct
- ❌ peekPixels (SkiaRenderer): headless uses SoftwareRenderer
- ❌ Immediate BeginFrame feedback flush: breaks frame pipeline
- ❌ CDPScreenshotNewSurface: RequestRepaintOnNewSurface overhead
- ❌ 2-tab pipelining: Chrome UI thread serializes ForceRedraw
- ❌ Chrome flags (disable-lcd-text etc.): ±2%
- ❌ headless_shell: slower than chrome (no shared HTTP cache)
- ❌ One-shot strategy: launch overhead 1-2s/process
- ❌ Firefox Playwright: 2.6x slower than Chrome
- ❌ Servo (servoshell 0.1.0): stub package, not ready
- ❌ CEF (cefpython3): abandoned, no modern Python wheel
- ❌ WebKitGTK snapshot: needs GPU/display access
- ❌ RequestRepaintOnNewSurface in skipRedraw: didn't fix igpu race
- ❌ Bitmap dimension retry: about:blank renders at full viewport size
- ❌ Pixel content retry: can't distinguish white page from about:blank
## igpu About:blank Root Cause
Chrome `--in-process-gpu` has two bugs at 48 concurrent workers:
1. **frameNavigated event not fired**: Chrome sometimes silently drops
`Page.frameNavigated` CDP event under high concurrency.
Fix: use `Page.frameStoppedLoading` (always reliable).
2. **Compositor surface race**: ForceRedraw's presentation feedback arrives
before the new page's CompositorFrame is activated in viz. CopyFromSurface
reads the old surface (about:blank at 875×8192, indistinguishable from
real page by dimensions). No reliable Python-side detection possible.
## Key Analysis Methods Used
- **Pipeline bottleneck analysis** (closed queueing model)
- **Little's Law**: steady-state throughput = C/T_c when capture-bound
- **USL contention curve**: C/T_c(C) convergence at ~125-130 t/s
- **USE method**: Utilization (79%), Saturation (semaphore queue), Errors (0)
- **Per-capture breakdown**: DrawRenderPass (57ms) + CopyDrawnRenderPass (18ms)
+ IPC overhead (95ms) measured via Chromium instrumentation
## Scale Estimate
30M tiles (18.7M articles × ~1.6 tiles/article):
- Single machine 98 t/s: 30M/98 = 85 hours = **3.5 days**
- Single machine 120 t/s (igpu, 95% correct): 30M/120 = 69 hours = **2.9 days**
- 4 machines × 98 t/s = 392 t/s: 30M/392 = 21 hours = **< 1 day**
- Production system (290 t/s, 4 machines): ~1 day (matches historical data)
## Production Pipeline: fast_cdp backend
```
Chrome 48w (capture) → /dev/shm (raw BGRA) → ProcessPool 4w (JPEG) → disk
98 t/s 28MB/tile ~100 t/s 100KB/tile
```
Architecture:
- `render_articles()` in `pixelrag_render.backends.fast_cdp`
- Capture: CDPPhasedStrategy logic (work-stealing, semaphore, frameStoppedLoading)
- Compression: `concurrent.futures.ProcessPoolExecutor(4)` — GIL-free, separate cores
- Raw files in /dev/shm/pixelrag_render/ — auto-deleted after compression
- Output: JPEG tiles + tiles.json manifest per article
Key: compression never blocks capture. Chrome writes raw → returns immediately.
Compression reads raw file asynchronously on different CPU cores.
128-core machine: 48 cores for Chrome, 4 cores for JPEG, 76 cores idle.
JPEG compression of 875×8192 takes ~10-20ms → 4 cores handle 200-400 t/s →
plenty of headroom over 98 t/s capture rate.
Storage: 30M tiles × 100KB JPEG = ~3 TB
## GPU Acceleration (Brewster H200 findings)
Lab machines have 8× H200/B200 GPUs but:
- `/dev/dri/renderD*` needs `render` group membership (no sudo)
- Docker daemon not running; rootless docker lacks nvidia-container-toolkit
- SwiftShader (CPU Vulkan) doesn't improve throughput vs software rendering
- headless Chrome ignores `--use-gl` flags (GPU process crashes on init)
- When GPU DOES init (via Xvfb + ANGLE), missing NVIDIA userspace drivers in container
To unlock GPU: `sudo usermod -aG render $USER` on lab machine.
Expected impact: 4x faster DrawRenderPass based on production system data.
+592
View File
@@ -0,0 +1,592 @@
# Reproducing Paper Results
> **Paper**: *PixelRAG: Retrieval and Generation in Pixel Space over Millions of Web Screenshots*
>
> This document maps every table and figure in the paper to the exact commands needed to reproduce the numbers.
## Prerequisites
### Infrastructure
| Component | Description | Where |
|-----------|-------------|-------|
| **Wikipedia tile index (base)** | 28M vectors, Qwen3-VL-Embedding-2B (pretrained) | `pixelrag-data/search_index/` (215 GB FAISS IVF, dim=2048) |
| **Wikipedia tile index (fine-tuned)** | 26M vectors, LoRA checkpoint-200 | `pixelrag-data/search_index_lora_vit_ckpt200_v2/` (202 GB) |
| **Wikipedia text index** | 15.7M text chunks (1024 tokens, Trafilatura) | `pixelrag-data/text_search_index_1024/` (121 GB) |
| **Article metadata** | URL↔tile mapping for 7.1M articles | `pixelrag-data/articles.json` (199 MB) |
| **Tile images** | ~30M PNG tiles (1024×1024) | Remote NFS or local SSD (~5.6 TB) |
| **News tile index** | 3.6M tiles (BBC/AP/CNN) for LiveVQA | S3: `s3://wiki-screenshot-tiles-backup/kiwix_tiles/news_image_search_index/` |
| **News text index** | 866K text chunks for news | S3: `s3://wiki-screenshot-tiles-backup/kiwix_tiles/news_text_search_index/` |
| **News tiles** | Raw PNG tiles for news articles | S3: `s3://wiki-screenshot-tiles-backup/kiwix_tiles/news_tiles/` |
| **LoRA adapter** | Fine-tuned embedding LoRA weights | S3: `s3://wiki-screenshot-tiles-backup/kiwix_tiles/adapters/lora_vit_ckpt200/` |
| **Kiwix ZIM** | Offline Wikipedia for HTML baselines | S3: `s3://wiki-screenshot-tiles-backup/kiwix_tiles/zim/` |
All S3 paths use AWS profile `leann` (`aws s3 --profile leann ...`).
### Services to Start
```bash
# 1. Screenshot search API (port 30888) — serves the pixel tile index
pixelrag-serve \
--index-dir pixelrag-data/search_index \ # or search_index_lora_vit_ckpt200_v2
--tiles-dir /path/to/wikipedia_tiles \
--articles-json pixelrag-data/articles.json \
--model Qwen/Qwen3-VL-Embedding-2B \
--device cuda --port 30888
# 2. Text search API (port 30889) — serves the text chunk index
pixelrag-serve \
--index-dir pixelrag-data/text_search_index_1024 \
--tiles-dir /path/to/text_chunks \
--articles-json pixelrag-data/articles.json \
--model Qwen/Qwen3-VL-Embedding-2B \
--device cuda --port 30889
# 3. Reader model (port 8000) — vLLM serving Qwen3.5-4B (default reader)
vllm serve Qwen/Qwen3.5-4B-Instruct \
--port 8000 --tensor-parallel-size 1 \
--max-model-len 32768
```
### Environment
```bash
cd ~/pixelrag/eval
# Install eval dependencies (one-time)
uv pip install pandas tqdm trafilatura openai aiohttp datasets huggingface-hub
# For grading
export OPENAI_API_KEY=sk-... # GPT-4.1 judge
```
---
## Table 1: Main Results (6 Benchmarks × 4 Methods)
**Reader**: Qwen3.5-4B, **k=3**, **Grader**: GPT-4.1 judge (except LiveVQA = exact match)
### No Retrieval (baseline)
```bash
# SimpleQA — no retrieval
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--num-examples 1000 --no-think
# NQ — no retrieval
python run_bench.py \
--task nq --model Qwen/Qwen3.5-4B-Instruct \
--num-examples 1000 --no-think
# NQ-Tables — no retrieval
python run_bench.py \
--task nq_tables --model Qwen/Qwen3.5-4B-Instruct \
--num-examples 1000 --no-think
# MMSearch — no retrieval (300 examples)
python run_bench.py \
--task mmsearch --model Qwen/Qwen3.5-4B-Instruct \
--num-examples 300 --no-think
# EVQA — no retrieval (landmarks, automatic only, n=749)
python run_bench.py \
--task encyclopedic_vqa --model Qwen/Qwen3.5-4B-Instruct \
--evqa-dataset-filter landmarks --evqa-question-type-filter automatic \
--num-examples 749 --no-think
# LiveVQA — see "LiveVQA Separate Pipeline" section below
```
### Text Retrieval — Trafilatura (Text → Text)
Requires: text search API on port 30889 with Trafilatura-parsed text chunks.
```bash
# SimpleQA — Trafilatura text retrieval
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# NQ — Trafilatura text retrieval
python run_bench.py \
--task nq --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# NQ-Tables
python run_bench.py \
--task nq_tables --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# MMSearch (multimodal query: text + image → text index)
python run_bench.py \
--task mmsearch --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--retrieval-top-k 3 --num-examples 300 --no-think
# EVQA
python run_bench.py \
--task encyclopedic_vqa --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--evqa-dataset-filter landmarks --evqa-question-type-filter automatic \
--retrieval-top-k 3 --num-examples 749 --no-think
# LiveVQA — see "LiveVQA Separate Pipeline" section below
```
### Text Retrieval — mwparserfromhell
Same as Trafilatura but requires a separate text index built with mwparserfromhell parser.
The text API must be started pointing to that index.
```bash
# Same commands as Trafilatura above, but --text-api-url points to
# the mwparserfromhell text index API (different port or index-dir).
# The parser choice is baked into the index at build time, not a runtime flag.
```
### PixelRAG (base) — Screenshot → Screenshot
Requires: screenshot search API on port 30888 with base (pretrained) embedding index.
```bash
# SimpleQA — pixel retrieval (base)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# NQ — pixel retrieval (base)
python run_bench.py \
--task nq --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# NQ-Tables
python run_bench.py \
--task nq_tables --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# MMSearch (multimodal: query image sent alongside text)
python run_bench.py \
--task mmsearch --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 300 --no-think
# EVQA (multimodal: landmark photo + question text)
python run_bench.py \
--task encyclopedic_vqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--evqa-dataset-filter landmarks --evqa-question-type-filter automatic \
--retrieval-top-k 3 --num-examples 749 --no-think
# LiveVQA — see "LiveVQA Separate Pipeline" section below
```
### PixelRAG (fine-tuned) — Screenshot → Screenshot with LoRA embedding
Same commands as PixelRAG (base), but the search API must be started with the fine-tuned index:
```bash
# Start search API with fine-tuned index
pixelrag-serve \
--index-dir pixelrag-data/search_index_lora_vit_ckpt200_v2 \
--tiles-dir /path/to/wikipedia_tiles \
--articles-json pixelrag-data/articles.json \
--model Qwen/Qwen3-VL-Embedding-2B \
--peft-adapter /path/to/lora_checkpoint_200 \
--device cuda --port 30888
```
Then run the same `--local-api` commands above.
### Grading
```bash
cd ~/pixelrag/eval
# Grade with GPT-4.1 judge (Wikipedia QA tasks)
python grade.py simpleqa eval_output/simpleqa_*.jsonl
python grade.py encyclopedic_vqa eval_output/encyclopedic_vqa_*.jsonl
python grade.py mmsearch eval_output/mmsearch_*.jsonl
# For NQ/NQ-Tables (with LLM judge for paper numbers)
python grade.py nq eval_output/nq_*.jsonl --llm-judge
python grade.py nq_tables eval_output/nq_tables_*.jsonl --llm-judge
# For LiveVQA (exact letter match — handled by the LiveVQA pipeline scripts)
```
---
## Table 3: RetrievalReader Modality Ablation
**Task**: SimpleQA (1000) + LiveVQA (6632), **Reader**: Qwen3.5-4B, **k=3**,
**Embedding**: Qwen3-VL-Embedding-2B (base, no LoRA)
| Row | Retrieval | Reader Input | Flags |
|-----|-----------|-------------|-------|
| Screenshot → Screenshot | Pixel index | Raw tile images | `--local-api` |
| Screenshot → OCR text | Pixel index | OCR'd text from tiles | `--local-api --read-as-text-ocr` |
| Text → Rendered image | Text index | Text chunks rendered as PNG | `--text-api --render-as-image` |
| Text → Text | Text index | Raw text chunks | `--text-api` |
| Text → HTML | Text index | Raw HTML from kiwix | `--text-api --html-dom-lookup` |
```bash
# Screenshot → Screenshot (same as main results PixelRAG base)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# Screenshot → OCR text
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--read-as-text-ocr --ocr-url http://localhost:8202/v1 \
--retrieval-top-k 3 --num-examples 1000 --no-think
# Text → Rendered image
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--render-as-image \
--retrieval-top-k 3 --num-examples 1000 --no-think
# Text → Text (same as main results Trafilatura)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# Text → HTML (DOM lookup)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--text-api --text-api-url http://localhost:30889/search \
--html-dom-lookup \
--retrieval-top-k 3 --num-examples 1000 --no-think
```
For LiveVQA, use the separate pipeline (see "LiveVQA Separate Pipeline" section) with the corresponding ablation scripts.
---
## Table 4: Embedding Training Recipe Ablation
**Evaluated on mini-datastore** (400 queries, 7426 tiles).
This ablation uses `--prebuilt-tiles-dir` pointing to the pre-built mini-datastore, with different embedding checkpoints. Each row corresponds to a different embedding training recipe:
```bash
# Base model (no fine-tuning)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--use-tiled-retrieval --use-qwen3vl-embedding \
--qwen3vl-model Qwen/Qwen3-VL-Embedding-2B \
--embedding-backend hf \
--prebuilt-tiles-dir tiles-hard-mini/ \
--retrieval-top-k 3 --num-examples 400 --no-think
# With LoRA checkpoint (dynamic hard negatives + ViT unfrozen)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--use-tiled-retrieval --use-qwen3vl-embedding \
--qwen3vl-model Qwen/Qwen3-VL-Embedding-2B \
--embedding-backend biqwen3 \
--peft-adapter /path/to/checkpoint-200 \
--prebuilt-tiles-dir tiles-hard-mini/ \
--retrieval-top-k 3 --num-examples 400 --no-think
```
The intermediate checkpoints (in-batch negatives, naive hard negatives, dynamic hard negatives frozen) each have their own PEFT adapter path.
---
## Figure 2: Token Efficiency (SimpleQA, k=1,2,3, 4 readers)
**Task**: SimpleQA (1000), **Readers**: Qwen3.5-4B, Qwen3.5-9B, Qwen3.5-27B, Qwen3.6-35B-A3B
For each reader × k × retrieval method, run:
```bash
# Example: Qwen3.5-4B, k=1, PixelRAG (fine-tuned)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --reader-top-k 1 \
--num-examples 1000 --no-think
# Example: Qwen3.5-4B, k=2, PixelRAG (fine-tuned)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --reader-top-k 2 \
--num-examples 1000 --no-think
# Example: Qwen3.5-4B, k=3, PixelRAG (fine-tuned)
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 \
--num-examples 1000 --no-think
```
> **Optimization**: Use `--retrieval-top-k 3 --reader-top-k N` to retrieve once at k=3 and evaluate at k=1,2,3 from the same JSONL (the full retrieved set is stored in `retrieved_images`).
For each reader, change `--model` and start the appropriate vLLM server.
Repeat for text retrieval (Trafilatura: `--text-api`) and PixelRAG base (base index).
The plot script is at `arxiv/figures/plot_token_efficiency.py`.
---
## Figure 3: Agentic Multi-Hop QA (MoNaCo)
**Task**: MoNaCo (1315 questions), **Agent**: GPT-5 ReAct, **k=5 per search**
Uses `eval/run_monaco.py` — a ReAct agent that issues search tool calls.
```bash
cd ~/pixelrag/eval
# PixelRAG backend
python run_monaco.py \
--reader gpt-5 \
--retrieval pixel \
--pixel-api http://localhost:30888/search \
--default-top-k 5
# Text retrieval backend (Trafilatura)
python run_monaco.py \
--reader gpt-5 \
--retrieval text \
--text-api http://localhost:30889/search \
--default-top-k 5
# Grade (token F1 computed inline; add --judge for LLM judge F1)
python run_monaco.py \
--reader gpt-5 \
--retrieval pixel \
--judge --judge-model gpt-4.1-2025-04-14
# Or grade existing predictions:
python grade.py monaco eval_output/monaco/<run_tag>
```
The dataset (`monaco_version_1_release.jsonl`) should be placed at
`eval/data/monaco/` or passed via `--data-path`.
---
## Figure 4: Image Compression Curve
**Task**: SimpleQA (1000), **Reader**: Qwen3.5-4B (base + SFT), k=1..5, compression c=1×/2×/3×
```bash
# No compression (c=1×), k=3
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 5 --reader-top-k 3 \
--num-examples 1000 --no-think
# 2× compression, k=3
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 5 --reader-top-k 3 \
--pixel-compress-ratio 2.0 \
--num-examples 1000 --no-think
# 3× compression, k=3
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 5 --reader-top-k 3 \
--pixel-compress-ratio 3.0 \
--num-examples 1000 --no-think
```
For the SFT reader, replace `--model` with the SFT checkpoint path and serve it via vLLM.
The plot script is at `arxiv/figures/plot_sft_compression_curve.py`.
---
## Table 8: Full Reader-Model Sweep (31 VLMs)
**Task**: SimpleQA (1000), **k=3**, pixel retrieval (base) vs text retrieval (Trafilatura)
For each of the 31 reader models, run two jobs:
```bash
# Pixel retrieval
python run_bench.py \
--task simpleqa --model <MODEL_NAME> \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
# Text retrieval
python run_bench.py \
--task simpleqa --model <MODEL_NAME> \
--text-api --text-api-url http://localhost:30889/search \
--retrieval-top-k 3 --num-examples 1000 --no-think
```
where `<MODEL_NAME>` is one of:
- `liuhaotian/llava-v1.5-7b`
- `meta-llama/Llama-3.2-11B-Vision-Instruct` (k=1 for pixel due to architecture limit)
- `meta-llama/Llama-3.2-90B-Vision-Instruct` (k=1 for pixel)
- `meta-llama/Llama-4-Scout-17B-16E-Instruct`
- `meta-llama/Llama-4-Maverick-17B-128E-Instruct`
- `Qwen/Qwen2-VL-2B-Instruct` through `Qwen/Qwen2-VL-72B-Instruct`
- `Qwen/Qwen2.5-VL-3B-Instruct` through `Qwen/Qwen2.5-VL-72B-Instruct`
- `Qwen/Qwen3-VL-2B` through `Qwen/Qwen3-VL-235B-A22B`
- `Qwen/Qwen3.5-0.8B` through `Qwen/Qwen3.5-35B-A3B`
- `Qwen/Qwen3.6-27B`, `Qwen/Qwen3.6-35B-A3B`
For reasoning-mode models, omit `--no-think`.
Each model requires its own vLLM instance (or OpenRouter/Commonstack for API models).
---
## LiveVQA (Table 1 + Table 3)
LiveVQA uses `eval/run_livevqa.py` — a dedicated script for the news corpus.
**Requires**: News pixel search API (port 30890), news text search API (port 30892),
LiveVQA v4 JSON dataset, vLLM reader.
```bash
cd ~/pixelrag/eval
# No retrieval
python run_livevqa.py --mode naive \
--model Qwen/Qwen3.5-4B-Instruct \
--output eval_output/livevqa_naive.jsonl
# PixelRAG (screenshot → screenshot)
python run_livevqa.py --mode pixel \
--pixel-api http://localhost:30890/search \
--model Qwen/Qwen3.5-4B-Instruct \
--output eval_output/livevqa_pixel.jsonl
# Text retrieval (Trafilatura)
python run_livevqa.py --mode text \
--text-api http://localhost:30892/search \
--model Qwen/Qwen3.5-4B-Instruct \
--output eval_output/livevqa_text.jsonl
# Hybrid (pixel + text)
python run_livevqa.py --mode hybrid \
--pixel-api http://localhost:30890/search \
--text-api http://localhost:30892/search \
--model Qwen/Qwen3.5-4B-Instruct \
--output eval_output/livevqa_hybrid.jsonl
```
Grading is automatic (5-option MC exact letter match) — printed at the end of each run.
---
## Known Issues (Blockers for Reproduction)
### ~~0. Missing simpleqa modules~~ (FIXED)
`screenshot.py` and `pixel_query.py` have been copied into `eval/lib/`.
Selenium import is deferred so it doesn't block `--local-api` users.
### ~~1. `dr_agent` not importable~~ (FIXED)
Dataset loaders extracted into `eval/lib/benchmarks.py`. The `run_bench.py`
import now reads from `simpleqa.datasets_loader` instead of `dr_agent`.
### ~~2. Grading script not in this repo~~ (FIXED)
`eval/grade.py` implements GPT-4.1 3-way grading (CORRECT/INCORRECT/NOT_ATTEMPTED) using
the same prompt template as the paper. No dependency on the old repo's evaluation framework.
For the legacy full evaluation framework (per-example HTML reports, etc.), the original
is still at `~/pixelrag-src/Vis-RAG/agent/scripts/evaluate.py`.
### 3. Hardcoded paths in retrieval.py
`eval/lib/retrieval.py` lines 8488 have placeholder paths (`/path/to/project`, `/path/to/data`) for the local kiwix tile store. These are only used by `LocalWikiTiledScreenshotRetriever` (ground-truth screenshot mode), not by the production `--local-api` mode.
### ~~4. LiveVQA uses separate pipeline~~ (FIXED)
`eval/run_livevqa.py` handles all LiveVQA modes (naive, pixel, text, hybrid).
### ~~5. MoNaCo runs from old repo~~ (FIXED)
`eval/run_monaco.py` implements the full ReAct agent loop with pixel/text retrieval backends.
### 6. mwparserfromhell text index
The paper's second text baseline uses mwparserfromhell parser. The text index must be built separately with this parser — the parser choice is embedded at index build time, not at query time. The build pipeline for this variant needs to be documented.
### 7. News corpus indexes
LiveVQA requires separate tile and text indexes built over the news corpus (BBC/AP/CNN). These indexes are on a different machine/path and need their own `pixelrag-serve` instances.
---
## Grading Protocol Summary
| Benchmark | Metric | Grader |
|-----------|--------|--------|
| SimpleQA | CORRECT/INCORRECT/NOT_ATTEMPTED → accuracy | GPT-4.1 (temp=0, seed=42) |
| NQ | Same 3-way judge | GPT-4.1 (temp=0, seed=42) |
| NQ-Tables | Same 3-way judge (up to 10 gold aliases joined with OR) | GPT-4.1 |
| MMSearch | Same 3-way judge | GPT-4.1 |
| EVQA | Same 3-way judge (reference_list → "Any of: ref1 \| ref2") | GPT-4.1 |
| LiveVQA | 5-option multiple-choice exact letter match | No LLM |
| MoNaCo | Token-level F1 (primary), LLM judge F1 (secondary) | GPT-4.1 |
---
## Quick Smoke Test (Verify Pipeline Works)
Run a single example end-to-end before committing to full runs:
```bash
# 1. Verify search API is responding
curl -s http://localhost:30888/status | python -m json.tool
# 2. Run 5 examples, no retrieval
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--num-examples 5 --no-think --force
# 3. Run 5 examples, pixel retrieval
python run_bench.py \
--task simpleqa --model Qwen/Qwen3.5-4B-Instruct \
--local-api --local-api-url http://localhost:30888/search \
--retrieval-top-k 3 --num-examples 5 --no-think --force
# 4. Grade
cd ~/pixelrag-src/Vis-RAG/agent
python scripts/evaluate.py simpleqa ~/pixelrag/eval/eval_output/<output>.jsonl
```
---
## Output File Convention
All outputs go to `eval_output/` with auto-generated filenames:
```
eval_output/{task}_{mode}_{model_safe}_{n}.jsonl
```
Examples:
- `eval_output/simpleqa_naive_qwen_qwen3.5_4b_instruct_1000.jsonl`
- `eval_output/simpleqa_local_api_qwen_qwen3.5_4b_instruct_1000.jsonl`
- `eval_output/nq_text_api_qwen_qwen3.5_4b_instruct_1000.jsonl`
Grading results are saved alongside as `*_eval_results.json`.
+163
View File
@@ -0,0 +1,163 @@
# Screenshot Throughput Optimization
Batch screenshot capture of Full English Wikipedia (18.7M articles, ~30M tiles)
with headless Chrome on AMD EPYC 7763 (128 cores, 995GB RAM).
## 1. Results
| Metric | Value |
|--------|-------|
| **E2E throughput** | **109 t/s** (JPEG written to disk) |
| **Capture throughput** | 98 t/s (raw BGRA to /dev/shm) |
| Tile size | 875 × 8192 px |
| Output format | JPEG q50, avg 305 KB/tile |
| Storage (30M tiles) | **8.5 TB** |
| Processing time (30M tiles) | **77 hours single machine, < 1 day on 4 machines** |
| Correctness | 100% pixel-verified against ground truth |
Best config: 48 Chrome workers, work-stealing queue, JPEG file output via
`rawFilePath` with `.jpg` extension (Chrome encodes JPEG in ThreadPool, writes
directly to disk — no base64, no websocket transfer, no external compression).
## 2. Architecture
```
kiwix-serve (ZIM) → 48 Chrome workers → JPEG files on disk
localhost:9461 875×8192 viewport 305 KB/tile avg
CDP websocket 109 tiles/s
```
Each Chrome process: navigate → fonts.ready + rAF → `Page.captureScreenshot`
with `rawFilePath=/path/tile.jpg` → Chrome encodes JPEG in async ThreadPool →
writes to disk → returns immediately → next article.
48 workers share the same `user-data-dir` (HTTP cache). Articles distributed
via asyncio work-stealing queue. `Page.frameStoppedLoading` event ensures
navigation is complete before capture.
## 3. Pipeline Bottleneck Analysis
The system is a two-stage pipeline analyzed via closed queueing model (Little's
Law + USL contention curve).
```
Stage Workers Per-op Capacity Bottleneck?
────────────────────────────────────────────────────────
Nav 48 186ms 258 pg/s No
Capture 48 321ms 150 t/s ← YES
Compress async ~10ms ∞ No (ThreadPool)
Disk write async ~1ms ∞ No
Steady-state: C/T_c(C) = 48/321ms ≈ 150 t/s theoretical
Actual (200 articles): 95 t/s (pipeline startup/drain bubble)
Actual (500 articles): 109 t/s (less bubble)
```
**Capture is the bottleneck.** Nav latency does not affect steady-state
throughput — verified by reducing nav from 186ms to 92ms with no throughput
change.
Per-capture breakdown (48 concurrent, measured via Chromium instrumentation):
| Component | 1 worker | 48 concurrent | Notes |
|-----------|----------|---------------|-------|
| ForceRedraw IPC | 95ms | 181ms | 8-hop async roundtrip |
| DrawRenderPass | 57ms | 62ms | Composite 136 quads |
| CopyDrawnRenderPass | 18ms | 46ms | memcpy 28MB |
| JPEG encode | 0 | ~10ms | ThreadPool, async |
| **Total** | **170ms** | **321ms** | |
IPC dominates at 48c (56% of capture time). This is OS scheduling overhead:
48 Chrome processes × 5 threads = 240 threads on 128 cores.
Contention curve `C/T_c(C)` converges at ~125-130 t/s regardless of C:
| Concurrent | T_c | C/T_c |
|------------|-----|-------|
| 24 | 200ms | 120 |
| 32 | 260ms | 123 |
| 48 | 321ms | **150** |
| 64 | 500ms | 128 |
## 4. Optimizations and Ablation
Each optimization measured on 200 articles, 100% correct, same hardware.
| # | Optimization | Throughput | Δ | Key insight |
|---|-------------|-----------|---|-------------|
| 0 | Baseline (Playwright, sleep 30ms) | 20 t/s | — | Node.js IPC layer |
| 1 | Direct CDP websocket | 23 t/s | +14% | Bypass Playwright |
| 2 | + `fonts.ready` + eager images | 28 t/s | +22% | Event-driven, no polling |
| 3 | + `rawFilePath` (Chromium patch) | 33 t/s | +18% | Bypass PNG/JPEG encode mutex |
| 4 | + Multi-worker (48w sequential) | 79 t/s | +140% | Linear scaling to ~48w |
| 5 | + Phased strategy (semaphore) | 96 t/s | +22% | Reduce capture contention |
| 6 | + Work-stealing queue | 98 t/s | +2% | Better load balancing |
| 7 | + `.jpg` rawFilePath (JPEG in ThreadPool) | 95 t/s | 3% | E2E with compression |
| 8 | + 500 articles (steady-state) | **109 t/s** | +15% | Amortize pipeline bubble |
Cumulative: 20 → 109 t/s = **5.5× improvement**.
### Chromium patches (5 files, 285 lines)
| Patch | Impact | Description |
|-------|--------|-------------|
| `rawFilePath` | +18% | Async raw BGRA write to /dev/shm via ThreadPool |
| `.jpg` auto-detect | e2e JPEG | JPEG encode in ThreadPool when path ends with .jpg |
| `directClip` | per-tile parallel | CopyFromSurface(src_rect) without emulation change |
| `skipRedraw` | 5ms latency | ForceRedrawWithCallback → CopyFromSurface |
## 5. Approaches That Did Not Work
| Approach | Result | Why it failed |
|----------|--------|---------------|
| `--in-process-gpu` | 120 t/s, 90% correct | Compositor surface race: about:blank captured instead of real page. ForceRedraw callback fires before compositor activates new frame. |
| `--single-process` | 168 t/s, 74% correct | Renderer thread contention across tabs in shared process |
| Two-tab pipelining | 8 t/s | Chrome UI thread serializes ForceRedraw across tabs in same process |
| `directClip` without ForceRedraw | 93% correct | Compositor frame stale without explicit redraw |
| Per-user-data-dir Chrome | 8 t/s | Each process starts with cold HTTP cache → thundering herd on kiwix |
| SwiftShader GPU compositor | 17% | CPU-based Vulkan slower than Chrome's software rasterizer |
| GPU on lab machines (H200/B200) | Blocked | `/dev/dri` permissions, no nvidia-container-toolkit |
| External ProcessPoolExecutor JPEG | 40 t/s | Cross-process IPC overhead; pool workers starved during capture |
| Firefox (Playwright) | 2.6× slower | Same IPC overhead, different engine |
| CEF OSR | ~12 t/s | Xvfb + 2.6s/page overhead |
| Servo | N/A | Not production-ready (stub package) |
| Skip `fonts.ready` | +4% only | Nav is not the bottleneck |
| 4096px tiles | Higher t/s but lower mpix/s | Fixed ForceRedraw overhead per tile |
### `--in-process-gpu` deep dive
Eliminates GPU process IPC → per-capture drops from 321ms to 175ms → 120 t/s.
But 5-10% of captures get about:blank content (correct dimensions, wrong pixels).
Root cause: ForceRedraw's presentation feedback fires after `SubmitCompositorFrame`
but before viz activates the new surface. `CopyFromSurface` reads the old surface.
About:blank renders at 875×8192 (same as real page due to persistent viewport
emulation), making detection impossible from dimensions alone.
Tried: SwapPromise (fires earlier), RequestRepaintOnNewSurface (new LocalSurfaceId),
bitmap dimension retry, pixel content check, `Page.frameNavigated` (reliable but
times out at 48w igpu), `Page.frameStoppedLoading` (reliable but fires for
sub-frames). None achieved 100% correct at 48 workers.
## Reproducing
```python
from pixelrag_render.strategies.cdp_phased import CDPPhasedStrategy
from pixelrag_render.bench import Bench
# Benchmark (with GT pixel verification)
bench = Bench(zim_path="...", chrome_path="...", output_dir="./results",
kiwix_url="http://localhost:9461")
strategy = CDPPhasedStrategy(chrome_path="...", n_workers=48, capture_limit=48, fmt="raw")
result = await bench.run(strategy) # {"tiles_per_s": 98, "correct_pct": 100, ...}
# Production (JPEG files on disk)
# Use rawFilePath with .jpg extension for Chrome-side JPEG encode:
await conn.cdp("Page.captureScreenshot", {
"rawFilePath": "/output/tile.jpg", # .jpg → JPEG encode in ThreadPool
"fromSurface": True, "optimizeForSpeed": True,
"clip": {"x": 0, "y": 0, "width": 875, "height": 8192, "scale": 1}
})
```
Requires custom Chromium build. Patch + build instructions: `chromium/README.md`.
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,438 @@
# Chromium Build on Centralia Implementation Plan
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** Build a patched Chromium (v150.0.7844.0) on Centralia (SSH: `CentraliaB200`, user `yichuan_wang`) with our two custom CDP features: `rawFilePath` parameter for `Page.captureScreenshot` and `directClip` parameter for parallel tile capture.
**Architecture:** All work happens under `/work/yichuan_wang/chromium-build/` on Centralia (NFS-mounted /work has 12TB free). depot_tools is cloned alongside the chromium checkout. The patch is generated locally from `~/chromium/src` (HEAD~1 diff) and transferred via scp. Build uses a release/official/no-debug/no-PGO args.gn for a fast, deployable binary.
**Tech Stack:** Chromium source (~30GB no-history), depot_tools, gn, autoninja (ninja), Python 3.12 (already on Centralia), Ubuntu 24.04, 224 cores for parallel build.
---
## Environment Facts (verified pre-plan)
- **Centralia SSH alias:** `CentraliaB200` (user `yichuan_wang`)
- **Workspace:** `/work/yichuan_wang/chromium-build/` (NFS, ~12TB free — plenty of room)
- **Local disk on Centralia:** `/dev/md0` 209GB free — avoid storing large files there
- **Local patch source:** `~/chromium/src` on local machine, `HEAD~1` diff covers 6 files, 166 insertions
- **Chromium version:** 150.0.7844.0 (MAJOR=150, BUILD=7844)
- **OS on Centralia:** Ubuntu 24.04.4 LTS (Noble)
- **Python:** `/usr/bin/python3` (3.12.3) — already present, no install needed
- **ninja/autoninja:** NOT present on Centralia — comes from depot_tools, added to PATH
---
## File Map
| Location | Purpose |
|---|---|
| `/work/yichuan_wang/chromium-build/depot_tools/` | Google's build tools (gn, fetch, autoninja, gclient) |
| `/work/yichuan_wang/chromium-build/chromium/` | gclient checkout root (contains `.gclient`) |
| `/work/yichuan_wang/chromium-build/chromium/src/` | Chromium source tree |
| `/work/yichuan_wang/chromium-build/chromium/src/out/Release/` | Build output dir |
| `/work/yichuan_wang/chromium-build/chromium/src/out/Release/args.gn` | Build configuration |
| `/work/yichuan_wang/chromium-build/chromium_patches.diff` | Our custom patch (transferred from local) |
| `/work/yichuan_wang/chromium-build/build.log` | autoninja build log (stream with `tail -f`) |
---
## Task 1: Verify workspace and install depot_tools
**Files:**
- Create: `/work/yichuan_wang/chromium-build/` (directory)
- Create: `/work/yichuan_wang/chromium-build/depot_tools/` (git clone)
- [ ] **Step 1.1: Create workspace directory on Centralia**
```bash
ssh CentraliaB200 "mkdir -p /work/yichuan_wang/chromium-build && echo 'workspace ready'"
```
Expected output: `workspace ready`
- [ ] **Step 1.2: Clone depot_tools into workspace**
Note: the correct URL is `chromium/tools/depot_tools` (not `chromium/depot_tools`).
```bash
ssh CentraliaB200 "git clone https://chromium.googlesource.com/chromium/tools/depot_tools.git /work/yichuan_wang/chromium-build/depot_tools"
```
Expected: clone completes, last line something like `Resolving deltas: 100%`.
- [ ] **Step 1.3: Verify depot_tools tools exist**
```bash
ssh CentraliaB200 "ls /work/yichuan_wang/chromium-build/depot_tools/fetch /work/yichuan_wang/chromium-build/depot_tools/gclient /work/yichuan_wang/chromium-build/depot_tools/autoninja"
```
Expected: three file paths printed (no errors).
---
## Task 2: Fetch Chromium source (no history)
This is the longest step — `fetch --no-history chromium` downloads ~30GB and runs `gclient sync`. With a fast connection it takes 3090 minutes.
**Files:**
- Create: `/work/yichuan_wang/chromium-build/chromium/` (gclient root)
- Create: `/work/yichuan_wang/chromium-build/chromium/src/` (source tree, ~30GB)
- [ ] **Step 2.1: Create the chromium checkout directory**
```bash
ssh CentraliaB200 "mkdir -p /work/yichuan_wang/chromium-build/chromium"
```
- [ ] **Step 2.2: Start the fetch in a detached screen session**
`fetch` must run from the chromium checkout root dir. We use `screen` so SSH disconnection doesn't kill it. The output is redirected to a log file for monitoring.
IMPORTANT: The home dir (`/home/eecs/yichuan_wang`) is full (10GB NFS, 0 bytes free). Set XDG dirs to /work to prevent depot_tools from failing on `~/.config/depot_tools`. Also put both `depot_tools/.cipd_bin` and `depot_tools` in PATH — vpython3 needs cipd in PATH.
```bash
ssh CentraliaB200 "screen -dmS chromium_fetch bash -c '
export XDG_CONFIG_HOME=/work/yichuan_wang/chromium-build/xdg/config
export XDG_CACHE_HOME=/work/yichuan_wang/chromium-build/xdg/cache
export XDG_DATA_HOME=/work/yichuan_wang/chromium-build/xdg/data
export XDG_STATE_HOME=/work/yichuan_wang/chromium-build/xdg/state
export PATH=/work/yichuan_wang/chromium-build/depot_tools/.cipd_bin:/work/yichuan_wang/chromium-build/depot_tools:\$PATH
export DEPOT_TOOLS_DIR=/work/yichuan_wang/chromium-build/depot_tools
cd /work/yichuan_wang/chromium-build/chromium
echo \"FETCH_START \$(date)\" > /work/yichuan_wang/chromium-build/fetch.log
fetch --no-history chromium >> /work/yichuan_wang/chromium-build/fetch.log 2>&1
echo \"FETCH_DONE exit=\$? at \$(date)\" >> /work/yichuan_wang/chromium-build/fetch.log
'"
```
- [ ] **Step 2.3: Verify screen session started**
```bash
ssh CentraliaB200 "screen -ls | grep chromium_fetch"
```
Expected: a line like `12345.chromium_fetch (Detached)`.
- [ ] **Step 2.4: Monitor fetch progress (check periodically)**
```bash
ssh CentraliaB200 "tail -20 /work/yichuan_wang/chromium-build/fetch.log"
```
Re-run this command to watch progress. Fetch is done when the log contains `FETCH_DONE exit=0`.
- [ ] **Step 2.5: Verify source tree exists after fetch completes**
```bash
ssh CentraliaB200 "ls /work/yichuan_wang/chromium-build/chromium/src/chrome/VERSION"
```
Expected: file path printed (no error). If the file doesn't exist, fetch failed — check `fetch.log` for errors.
- [ ] **Step 2.6: Verify Chromium version matches local**
```bash
ssh CentraliaB200 "cat /work/yichuan_wang/chromium-build/chromium/src/chrome/VERSION"
```
Expected:
```
MAJOR=150
MINOR=0
BUILD=7844
PATCH=0
```
If the version differs, the patch may not apply cleanly. Record the actual version for the patch step.
---
## Task 3: Transfer and apply our patches
Our patch adds two CDP features to `Page.captureScreenshot`:
- `rawFilePath`: write screenshot directly to a file path (bypassing base64 encoding)
- `directClip`: clip parameter for parallel tile capture
**Files:**
- Modify: `content/browser/devtools/protocol/page_handler.cc` (primary patch target)
- Modify: `content/browser/devtools/protocol/page_handler.h`
- Modify: `content/renderer/render_widget_host/render_widget_host_impl.cc`
- Modify: `content/renderer/render_widget_host/render_widget_host_impl.h`
- Modify: `third_party/blink/public/devtools_protocol/domains/Page.pdl`
- Modify: `third_party/blink/renderer/platform/widget/widget_base.cc`
- Transfer: `/work/yichuan_wang/chromium-build/chromium_patches.diff`
- [ ] **Step 3.1: Generate the patch from local machine**
Run this on the LOCAL machine:
```bash
git -C ~/chromium/src diff HEAD~1 > /tmp/chromium_patches.diff
wc -l /tmp/chromium_patches.diff
```
Expected: file is non-empty (~300+ lines).
- [ ] **Step 3.2: Transfer patch to Centralia**
Run this on the LOCAL machine:
```bash
scp /tmp/chromium_patches.diff CentraliaB200:/work/yichuan_wang/chromium-build/chromium_patches.diff
```
- [ ] **Step 3.3: Verify patch arrived on Centralia**
```bash
ssh CentraliaB200 "wc -l /work/yichuan_wang/chromium-build/chromium_patches.diff"
```
Expected: same line count as local.
- [ ] **Step 3.4: Apply the patch**
```bash
ssh CentraliaB200 "cd /work/yichuan_wang/chromium-build/chromium/src && git apply /work/yichuan_wang/chromium-build/chromium_patches.diff"
```
Expected: no output (silent success). If you see errors like "patch does not apply", see Step 3.5.
- [ ] **Step 3.5: (If patch fails) Try with --3way or check fuzz**
If Step 3.4 fails with "patch does not apply":
```bash
ssh CentraliaB200 "cd /work/yichuan_wang/chromium-build/chromium/src && git apply --3way /work/yichuan_wang/chromium-build/chromium_patches.diff"
```
If that also fails, the Chromium version on Centralia differs from the local checkout. Check `cat /work/yichuan_wang/chromium-build/chromium/src/chrome/VERSION` and compare to local (`cat ~/chromium/src/chrome/VERSION`). If versions differ significantly, you may need to regenerate the patch from the correct base commit — fetch the local HEAD's commit hash with `git -C ~/chromium/src rev-parse HEAD~1` and use that.
- [ ] **Step 3.6: Verify patch was applied**
```bash
ssh CentraliaB200 "cd /work/yichuan_wang/chromium-build/chromium/src && git diff --stat"
```
Expected output:
```
content/browser/devtools/protocol/page_handler.cc | 129 +++++...
content/browser/devtools/protocol/page_handler.h | 14 +++
content/renderer/render_widget_host/render_widget_host_impl.cc | 9 ++
content/renderer/render_widget_host/render_widget_host_impl.h | 5 +
third_party/blink/public/devtools_protocol/domains/Page.pdl | 6 +
third_party/blink/renderer/platform/widget/widget_base.cc | 12 +-
6 files changed, 166 insertions(+), 9 deletions(-)
```
---
## Task 4: Configure the build
**Files:**
- Create: `/work/yichuan_wang/chromium-build/chromium/src/out/Release/` (directory)
- Create: `/work/yichuan_wang/chromium-build/chromium/src/out/Release/args.gn`
- [ ] **Step 4.1: Create the build output directory**
```bash
ssh CentraliaB200 "mkdir -p /work/yichuan_wang/chromium-build/chromium/src/out/Release"
```
- [ ] **Step 4.2: Write args.gn**
```bash
ssh CentraliaB200 "cat > /work/yichuan_wang/chromium-build/chromium/src/out/Release/args.gn << 'EOF'
is_debug = false
is_official_build = true
is_component_build = false
symbol_level = 0
blink_symbol_level = 0
chrome_pgo_phase = 0
EOF"
```
- [ ] **Step 4.3: Verify args.gn content**
```bash
ssh CentraliaB200 "cat /work/yichuan_wang/chromium-build/chromium/src/out/Release/args.gn"
```
Expected exact output:
```
is_debug = false
is_official_build = true
is_component_build = false
symbol_level = 0
blink_symbol_level = 0
chrome_pgo_phase = 0
```
---
## Task 5: Run gn gen
`gn gen` reads `args.gn` and generates all the ninja build files. This takes 25 minutes on 224 cores.
**Files:**
- Create: `/work/yichuan_wang/chromium-build/chromium/src/out/Release/build.ninja` (generated)
- Create: `/work/yichuan_wang/chromium-build/gn_gen.log`
- [ ] **Step 5.1: Run gn gen in a screen session**
```bash
ssh CentraliaB200 "screen -dmS chromium_gn bash -c '
export PATH=/work/yichuan_wang/chromium-build/depot_tools:\$PATH
cd /work/yichuan_wang/chromium-build/chromium/src
gn gen out/Release > /work/yichuan_wang/chromium-build/gn_gen.log 2>&1
echo \"GN_DONE exit=\$?\" >> /work/yichuan_wang/chromium-build/gn_gen.log
'"
```
- [ ] **Step 5.2: Wait for gn gen to finish**
```bash
ssh CentraliaB200 "tail -5 /work/yichuan_wang/chromium-build/gn_gen.log"
```
Re-run until you see `GN_DONE exit=0`. If exit is non-zero, check the full log:
```bash
ssh CentraliaB200 "cat /work/yichuan_wang/chromium-build/gn_gen.log"
```
Common gn errors and fixes:
- `Python not found`: verify `which python3` works on Centralia (it does per our check)
- `No targets match`: args.gn typo — re-check Step 4.2
- [ ] **Step 5.3: Verify build.ninja was generated**
```bash
ssh CentraliaB200 "ls -lh /work/yichuan_wang/chromium-build/chromium/src/out/Release/build.ninja"
```
Expected: file exists, non-zero size.
---
## Task 6: Build chrome with autoninja
This is the main build step. With 224 cores and no debug symbols, expect 60120 minutes for a full build. The output is the `chrome` binary.
**Files:**
- Create: `/work/yichuan_wang/chromium-build/chromium/src/out/Release/chrome` (built binary)
- Create: `/work/yichuan_wang/chromium-build/build.log`
- [ ] **Step 6.1: Start autoninja build in screen session**
`autoninja` automatically sets `-j` based on CPU count (will use ~224 jobs). It reads the `NINJA_SUMMARIZE_BUILD` env var to show progress.
```bash
ssh CentraliaB200 "screen -dmS chromium_build bash -c '
export PATH=/work/yichuan_wang/chromium-build/depot_tools:\$PATH
cd /work/yichuan_wang/chromium-build/chromium/src
autoninja -C out/Release chrome > /work/yichuan_wang/chromium-build/build.log 2>&1
echo \"BUILD_DONE exit=\$?\" >> /work/yichuan_wang/chromium-build/build.log
'"
```
- [ ] **Step 6.2: Verify screen session started**
```bash
ssh CentraliaB200 "screen -ls | grep chromium_build"
```
Expected: a line like `12345.chromium_build (Detached)`.
- [ ] **Step 6.3: Monitor build progress**
```bash
ssh CentraliaB200 "tail -5 /work/yichuan_wang/chromium-build/build.log"
```
You'll see ninja progress lines like `[1234/89000] CXX obj/content/...`. Re-run every few minutes to watch progress. Build is done when you see `BUILD_DONE exit=0`.
To watch CPU utilization:
```bash
ssh CentraliaB200 "uptime"
```
If the build is running, load average should be ~200+.
- [ ] **Step 6.4: Check for build errors (if BUILD_DONE shows non-zero exit)**
```bash
ssh CentraliaB200 "grep -i 'error:' /work/yichuan_wang/chromium-build/build.log | tail -20"
```
Common errors:
- `undefined reference`: usually means a `.h` change wasn't matched with a `.cc` change in the patch. Check that the patch applied fully (Task 3).
- `ninja: build stopped`: check the lines above for the actual C++ error.
- Disk full: run `df -h /work` — if /work is at 100%, free space or use a different path.
- [ ] **Step 6.5: Verify chrome binary exists and is executable**
```bash
ssh CentraliaB200 "ls -lh /work/yichuan_wang/chromium-build/chromium/src/out/Release/chrome"
```
Expected: file ~200300MB, executable bit set (permissions like `-rwxr-xr-x`).
- [ ] **Step 6.6: Smoke-test the binary**
```bash
ssh CentraliaB200 "/work/yichuan_wang/chromium-build/chromium/src/out/Release/chrome --version"
```
Expected: `Chromium 150.0.7844.0` (or similar version line).
Note: Chrome may print warnings about display/GPU on a headless server — that's normal. We care only that the binary runs and prints its version.
---
## Task 7: Verify our custom CDP features are compiled in
Our patch adds `rawFilePath` and `directClip` parameters to `Page.captureScreenshot`. We verify they made it into the compiled protocol.
**Files:**
- Read: `/work/yichuan_wang/chromium-build/chromium/src/out/Release/gen/third_party/blink/public/devtools_protocol/protocol/page.json` (generated protocol JSON)
- [ ] **Step 7.1: Check the generated protocol JSON for our parameters**
```bash
ssh CentraliaB200 "grep -n 'rawFilePath\|directClip' /work/yichuan_wang/chromium-build/chromium/src/out/Release/gen/third_party/blink/public/devtools_protocol/protocol/page.json"
```
Expected: at least 2 lines mentioning `rawFilePath` and `directClip`.
- [ ] **Step 7.2: Check the compiled binary for our parameter strings**
```bash
ssh CentraliaB200 "strings /work/yichuan_wang/chromium-build/chromium/src/out/Release/chrome | grep -c 'rawFilePath'"
```
Expected: at least 1 (the string is embedded in the binary). If 0, the patch didn't compile into the binary — recheck that `git diff --stat` in Task 3 Step 6 was correct.
---
## Timing Estimates
| Task | Estimated Duration |
|---|---|
| Task 1: depot_tools clone | 25 min |
| Task 2: `fetch --no-history chromium` | 3090 min (network-dependent) |
| Task 3: patch transfer + apply | 2 min |
| Task 4: args.gn setup | 1 min |
| Task 5: `gn gen` | 37 min |
| Task 6: `autoninja` build | 60120 min (224 cores, no debug) |
| Task 7: verification | 2 min |
| **Total** | **~24 hours** |
---
## Recovery Notes
- **If screen session dies unexpectedly:** Re-attach with `screen -r chromium_build` to see any final error, then restart from the last completed step.
- **If fetch is interrupted:** Re-run `fetch --no-history chromium` from the same directory — gclient will resume.
- **If build is interrupted:** Re-run `autoninja -C out/Release chrome` — ninja tracks completed targets and resumes from where it left off.
- **If /work fills up:** `du -sh /work/yichuan_wang/chromium-build/chromium/src/out/Release/obj/` is usually the largest dir. Consider deleting `.o` files after build if you only need the final binary: `find out/Release/obj -name '*.o' -delete`.
@@ -0,0 +1,359 @@
# PixelRAG Unirepo Restructure — Design Spec
**Date:** 2026-05-11
**Status:** Draft
## Context
PixelRAG is a visual document retrieval framework: any document (web page, PDF, image) → visual rendering → embedding → FAISS index → search API. Three private repos (yichuan-w/Vis-RAG, andylizf/wiki-screenshot, andylizf/wiki-screenshot-training) are being merged into a single public repo with a clean architecture.
The current repo at ~/pixelrag/ has a messy first-pass merge. This spec defines the target architecture.
## Users
**A — Framework user (primary):** "I have documents (web pages, PDFs, local files) and want to build a visual retrieval system." Needs the full pipeline. Cares about generality — not a Wikipedia-specific tool.
**B — Paper reproducer:** "I want to reproduce PixelRAG results." Downloads pre-built indexes, starts the search API, runs eval. Should exist but not the design focus.
**C — Agent developer:** Uses screenshot capture and visual search as agent skills/tools. Needs callable APIs: give URL → get screenshot, give query → get results. Will demo agent integration.
**D — Model trainer:** Trains visual embedding models. Contrastive learning + hard negative mining via search API. Should exist but not the design focus.
## Package Architecture
Five packages, single-direction dependencies:
```
ingest ←── index ──→ embed
serve (independent)
train → serve (API calls for mining)
```
### Package 1: pixelrag-render
**"Document → image tiles."** Standalone rendering tool. Agents call it directly; index calls it for batch jobs.
```
src/pixelrag_render/
├── render.py # Public API:
│ # render_url(url, output_dir, backend="cdp") → list[Path]
│ # render_pdf(path, output_dir) → list[Path]
│ # render_file(path, output_dir) → list[Path] (auto-detect)
├── backends/
│ ├── cdp.py # Lean CDP capture — default, fastest
│ │ # Direct Page.captureScreenshot, multi-browser workers
│ │ # JPEG q85, DPR 1, fromSurface=False, optimizeForSpeed=True
│ │ # Based on render_news_pages.py (23.9s/50 articles benchmark)
│ ├── playwright.py # Full Playwright — more options, experimental/compat
│ │ # Stripped to production-useful config only
│ │ # Keeps: CDP screenshot mode, segmented tiles, GPU rasterization
│ │ # Removes: unused experimental options
│ └── pdf.py # PDF → page images (pdf2image or PyMuPDF)
└── bench/ # Rendering benchmarks
├── benchmark.py # Config sweep (workers, batch size, concurrency)
├── benchmark_optimizations.py # GPU accel, PNG compression, tile sizes
├── benchmark_fullpage.py # Screenshot strategy comparison
└── benchmark_longtail_matrix.py # Long pages × tile size × concurrency
```
**Dependencies:** playwright, pillow, aiohttp (lightweight — no torch)
**CLI entry points:**
- `pixelrag-render``pixelrag_render.render:main` (render URLs/files to tiles)
**Source of code:**
- `cdp.py``scripts/render_news_pages.py` capture_article + worker + multi-browser setup (generalized, news-specific parts removed)
- `playwright.py``tools/playwright_tool.py` (stripped from 2388L to production-relevant config)
- `bench/``bench/` directory (kept as-is)
- `render.py` — new, thin API layer that dispatches to backends
### Package 2: pixelrag-embed
**"Image tiles → vectors → FAISS index."** Three independent CLI tools, orchestrator-free. Each has its own `main()`, no imports between them.
```
src/pixelrag_embed/
├── chunk.py # Large image → 1024px strips
│ # Input: tile directory. Output: chunk PNGs + chunks.json
│ # Pure PIL, no torch. ~380 lines.
├── embed.py # Images → embedding vectors
│ # Input: chunk directory. Output: shard_NNN.npz
│ # vLLM/sglang backend, multi-GPU. ~2400 lines.
└── index.py # Vectors → FAISS IVFFlat index
# Input: embedding .npz shards. Output: index.faiss + metadata.npz
# ~330 lines.
```
**Dependencies:** torch, transformers, faiss-cpu, pillow, numpy, tqdm
**CLI entry points:**
- `pixelrag-chunk``pixelrag_embed.chunk:main`
- `pixelrag-embed``pixelrag_embed.embed:main`
- `pixelrag-build-index``pixelrag_embed.index:main`
**Source of code:**
- `chunk.py``embedding/chunk_tiles.py`
- `embed.py``embedding/embed_tiles.py`
- `index.py``indexing/build_index.py`
### Package 3: pixelrag-index
**"Data source → complete searchable index."** Orchestration layer. Knows how to chain ingest + embed for different data sources. Two modes: single-machine (default, no S3) and distributed (S3 coordination for multi-machine).
```
src/pixelrag_index/
├── config.py # pixelrag.yaml parser
│ # Defines: source type, paths, embed model, output location
├── sources/ # Data source iterators (yield items for ingest to render)
│ ├── kiwix.py # Wikipedia ZIM → iterate articles → call ingest per article
│ ├── web.py # URL list/sitemap → download HTML+assets → call ingest
│ │ # Download + SQLite state are internal to this source
│ │ # Includes presets (e.g. "news") with per-domain rate limits,
│ │ # cookie banner CSS, source-specific HTML handling (BBC/CNN/AP)
│ │ # Usage: --source web --preset news
│ ├── pdf.py # PDF directory → iterate files → call ingest per file
│ └── local.py # Scan directory → auto-detect file types → route to above
├── pipelines.py # End-to-end: source → ingest → chunk → embed → build
│ # Chains the stages, handles checkpointing between stages
├── distributed.py # S3ShardCoordinator + claim-loop worker (optional)
│ # Only used with --distributed flag
│ # Used by both capture and embedding distributed runs
└── monitor.py # Cross-machine progress dashboard (reads S3 claims)
│ # Only relevant in distributed mode
```
**Two orchestration modes:**
- `pixelrag-index build --source ./my_docs` — single machine, iterate locally, no S3
- `pixelrag-index build --source kiwix --distributed --bucket my-bucket` — multi-machine, S3 coordination
**Dependencies:** pixelrag-render, pixelrag-embed, boto3 (optional, only for distributed), tqdm
**CLI entry points:**
- `pixelrag-index``pixelrag_index.pipelines:main` (build index from source)
- `pixelrag-monitor``pixelrag_index.monitor:main` (progress dashboard)
**pixelrag.yaml — parameter forwarding pattern:**
Each section's parameters are forwarded directly to the corresponding package. Index only manages orchestration order, not parameter details.
```python
# index/config.py — forwarding logic
source_type = config["source"].pop("type")
source = SOURCES[source_type](**config["source"]) # forward all source params
# ingest params forwarded to render calls
# embed params forwarded to chunk/embed/build calls
```
```yaml
# Example: local files (User A)
source:
type: local
path: ./my_docs
ingest:
backend: cdp
quality: 85
embed:
model: Qwen/Qwen3-VL-Embedding-2B
device: cuda
gpu_ids: [0, 1, 2, 3]
batch_size: 128
output: ./my_index
```
```yaml
# Example: web URLs with news preset
source:
type: web
urls: ./urls.txt
preset: news
concurrency: 200
ingest:
backend: cdp
embed:
model: Qwen/Qwen3-VL-Embedding-2B
gpu_ids: [0, 1]
output: ./news_index
```
```yaml
# Example: PDF collection
source:
type: pdf
path: ./papers/
dpi: 300
pages: "1-10"
embed:
model: Qwen/Qwen3-VL-Embedding-2B
device: cpu
output: ./paper_index
```
```yaml
# Example: Wikipedia (distributed)
source:
type: kiwix
zim: ./wikipedia.zim
serve_url: http://localhost:9454
distributed:
bucket: my-bucket
prefix: kiwix
embed:
model: Qwen/Qwen3-VL-Embedding-2B
gpu_ids: [0, 1, 2, 3, 4, 5, 6, 7]
backend: sglang
output: s3://my-bucket/index
```
**Source of code:**
- `distributed.py``coordinator.py` (S3ShardCoordinator) + claim loop from `coordinator_worker.py` and `embedding_worker.py`
- `sources/kiwix.py``datasources/kiwix.py` (article iteration logic)
- `sources/web.py``datasources/news.py` + `news/download.py` + `news/db.py` (generalized, news-specific naming removed; download + SQLite state are internal implementation details of this source)
- `sources/local.py` — new
- `sources/pdf.py` — new (thin, delegates rendering to ingest)
- `pipelines.py` ← new, chains stages
- `monitor.py``scripts/monitor_global.py`
- `config.py` — new
### Package 4: pixelrag-serve
**"FAISS index → search API."** One unified FastAPI server that serves any index.
```
src/pixelrag_serve/
└── api.py # Unified search API
# POST /search — text/image/embedding queries → top-k results
# GET /health, GET /status
# Configurable via CLI args or env vars
# Supports CPU and CUDA for query embedding
```
**Dependencies:** fastapi, uvicorn, faiss-cpu, torch, transformers, pillow, numpy
**CLI entry points:**
- `pixelrag-serve``pixelrag_serve.api:main`
**Source of code:**
- `api.py` ← merge of `search_api.py` + `text_search_api.py` + `news_search_api.py` into one unified API. Hex ID mapping handled at index build time (in pixelrag-embed), not at serve time.
### Package 5: pixelrag-train
**"Train visual embedding models."**
```
src/pixelrag_train/
├── models/
│ └── biqwen3.py # BiQwen3: Qwen3VLModel + last-token pooling + L2 norm
├── contrastive.py # GradCache contrastive training with LoRA/DoRA
└── mine.py # Hard negative mining (calls serve API)
# Unified: image mining (:30888) + text mining (:30889)
```
**Dependencies:** torch, transformers, peft, accelerate, wandb, faiss-cpu
**CLI entry points:**
- `pixelrag-train``pixelrag_train.contrastive:main`
- `pixelrag-mine``pixelrag_train.mine:main`
**Source of code:**
- `biqwen3.py``models/biqwen3.py` (unchanged)
- `contrastive.py``train_contrastors.py` (renamed)
- `mine.py` ← merge of `mine_hard_negatives.py` + `mine_text_hard_negatives.py`
### eval/
Not a package. Script directory for paper reproduction (User B).
```
eval/
├── run_naive_simpleqa.py # Main eval runner
└── simpleqa/ # Support library (data, llm, retrieval, etc.)
```
Source: kept from current repo, unchanged.
## Dependency Graph
```
pixelrag-index
├── pixelrag-render (calls render_url/render_pdf for capture stage)
├── pixelrag-embed (calls chunk/embed/index tools)
└── boto3 (S3 coordination)
pixelrag-serve (independent — no deps on other pixelrag packages)
pixelrag-train
└── calls pixelrag-serve API over HTTP (not a Python dependency)
pixelrag-render (independent)
pixelrag-embed (independent)
```
## Data Flow
```
User A: "Build me a visual search index"
pixelrag-index build --source ./my_docs
├─ sources/local.py scans directory, classifies files
├─ For each document:
│ pixelrag-render.render_url() or render_pdf()
│ → tiles/{doc_id}.tiles/tile_0000.jpg, tile_0001.jpg, ...
├─ pixelrag-embed.chunk
│ → chunks/{doc_id}.tiles/chunk_0000_00.png, ...
├─ pixelrag-embed.embed (GPU)
│ → embeddings/shard_NNN.npz
└─ pixelrag-embed.index
→ output/index.faiss + metadata.npz
pixelrag-serve --index-dir ./output --port 30888
→ POST /search {"queries": [{"text": "..."}]} → top-k results
```
## What Gets Cut
From the current ~/pixelrag/ repo:
- `packages/capture/` → replaced by `packages/render/` (new structure)
- `packages/serving/` → replaced by `packages/serve/` (unified API)
- `packages/training/` → replaced by `packages/train/` (renamed files)
- `packages/embed/` — new package (from loose scripts)
- `packages/index/` — new package (from loose scripts + new code)
- `eval/` — kept
From source repos (~/pixelrag-src/), code NOT carried forward:
- `executors/base.py`, `executors/skypilot.py` — executor ABC and cloud-specific code
- `proxy/` — proxy rotation (not needed for offline rendering)
- `lead_images/` — lead image extraction (hardcoded paths)
- `datasources/enterprise.py`, `datasources/wikimedia.py` — paid API / superseded
- `tools/streaming_capture.py` — superseded by lean CDP backend
- `tools/raw_pixels.py`, `tools/temp_dirs.py` — helpers for old PlaywrightTool
- Most of PlaywrightTool's 2388 lines — stripped to production config
- `run.py`, `monitor.py` (top-level) — replaced by index CLI
- `scripts/run_embeddings.py` — thin wrapper, redundant with embed CLI
- `scripts/status.py` — replaced by monitor
## Migration Strategy
1. Create new package directories under ~/pixelrag/packages/
2. Copy + transform code from ~/pixelrag-src/ (source repos are read-only)
3. For each package: create pyproject.toml, rename imports, add CLI entry points
4. Verify `uv sync --package <name>` works for each
5. Verify existing endpoint (port 30001) still works with new pixelrag-serve
6. Commit as clean restructure
@@ -0,0 +1,290 @@
# PixelRAG Frontend Design Spec
## Overview
A modern web frontend for the PixelRAG visual retrieval engine, serving as both an academic paper companion demo and a functional API service. Built as a standalone Next.js application alongside the existing FastAPI backend.
## Goals
- Showcase visual retrieval quality with rich tile image display
- Provide interactive search (text + image queries) over the FAISS index
- Document the API with live try-it-out capability
- Look professional enough for paper/conference demos — not generic
## Non-Goals
- User authentication or multi-tenancy
- Index management or data ingestion UI
- Mobile-first design (desktop-first, responsive is fine)
## Architecture
### Stack
| Layer | Technology |
|-------|-----------|
| Framework | Next.js 15 (App Router) |
| Styling | Tailwind CSS 4 |
| Components | shadcn/ui |
| Animation | Framer Motion |
| Language | TypeScript |
| Backend | FastAPI (existing, unchanged except CORS) |
### Project Structure
```
web/ ← new Next.js app
src/
app/
page.tsx ← search home
docs/page.tsx ← API reference
status/page.tsx ← index dashboard
layout.tsx ← shell (nav, theme provider)
components/
SearchBar.tsx ← text input + image upload/drag-drop
ResultGroup.tsx ← article group with horizontal tile row
TileCard.tsx ← single tile result card
Lightbox.tsx ← fullscreen tile viewer with pan/zoom
ComparePanel.tsx ← side-by-side tile comparison
ApiPlayground.tsx ← try-it-live widget for /docs
StatusCard.tsx ← metric card for /status dashboard
lib/
api.ts ← typed fetch wrapper for all API endpoints
types.ts ← shared TypeScript types matching Pydantic models
next.config.ts ← rewrites /api/* → FastAPI
tailwind.config.ts
package.json
serve/ ← existing (minimal changes)
src/pixelrag_serve/api.py ← add CORSMiddleware
```
### API Proxy
In development, `next.config.ts` rewrites `/api/*` to `http://localhost:30001/*` so the frontend can call the FastAPI backend without CORS issues. In production, CORS middleware on FastAPI allows the Next.js origin.
## Visual Design
### Color Palette
| Role | Value | Usage |
|------|-------|-------|
| Background | `#0c0c0c` | Page background |
| Surface | `#1a1a1a` | Cards, inputs, panels |
| Border | `#222222` | Card borders, dividers |
| Text primary | `#ffffff` | Headings, important text |
| Text secondary | `#888888` | Descriptions, metadata |
| Text muted | `#555555` | Labels, placeholders |
| Accent | `#6366f1` | Links, scores, CTAs, active states |
| Accent gradient | `#6366f1 → #8b5cf6` | Primary buttons |
### Typography
- **Inter** — UI text (body, labels, metadata)
- **Crimson Pro** — Branding headings (logo, page titles)
- **JetBrains Mono** — Code blocks, API examples, monospace data
### Design Principles
- Dark theme only (matches academic demo context, highlights tile images)
- Generous whitespace, no visual clutter
- Images are the hero — UI chrome stays minimal
- Subtle borders over drop shadows
- Micro-animations for state transitions (loading, lightbox open/close)
## Pages
### 1. Search Home (`/`)
The landing page and primary interface.
**Layout:**
- Centered logo + tagline at top: "PixelRAG — Visual retrieval over 15.7M Wikipedia tiles"
- Search bar below: text input with search button. Supports drag-and-drop or click-to-upload for image queries. Image preview shown inline when an image is attached.
- Mode chips below search bar: "Text query", "Image upload", "Drag & drop"
- Results appear below after search
**Search Controls (collapsible):**
- `n_docs` — number of results (default 10)
- `nprobe` — FAISS nprobe override
- `min_tile_height` — filter small/blank tiles
- `instruction` — custom embedding instruction
- Defaults are hidden; expand via "Advanced" toggle
**Results Display:**
Results are **grouped by article**. The API returns a flat ranked list of hits; the frontend groups them by `article_id`.
Each article group shows:
- Article title (derived from `url` field, decode the Wikipedia slug)
- External link to the Wikipedia article
- Tile count badge
- Horizontal scrollable row of tile cards
Each tile card shows:
- Tile image (loaded via `GET /tile?path=...`)
- Global rank badge (top-left corner, e.g. "#1")
- Cosine similarity score
- Tile height in pixels
- Tile position identifier (e.g. "tile 2:1" = tile_index 2, chunk_index 1)
**Status bar** between search bar and results:
- Result count
- Total latency
- Latency breakdown: measure client-side round-trip time (no backend changes needed; server-side encode/search breakdown is logged to stdout already)
### 2. API Documentation (`/docs`)
Custom-built API reference (not Swagger/ReDoc — those are functional but ugly and break visual consistency).
**Layout:**
- Left sidebar: endpoint list with HTTP method badges (POST green, GET blue)
- Guides section below endpoints: "Quick Start", "Python Client"
- Main content area: endpoint detail
**Each endpoint section:**
- Method + path + description
- Request body schema with syntax-highlighted JSON
- Response schema
- "Try It" playground: editable JSON input + Send button + response preview
- curl example
**Endpoints documented:**
- `POST /search` — primary search (text, image, or embedding queries)
- `GET /status` — index metadata and stats
- `GET /tile?path=...` — serve tile image by path
- `GET /health` — health check
- `POST /reconstruct` — reconstruct stored embeddings by vector_id
### 3. Index Dashboard (`/status`)
Displays data from `GET /status` in a visual dashboard.
**Metric cards (2×2 grid):**
- Total vectors (formatted: "15.7M")
- Embedding dimension
- Model name
- Index size (human-readable bytes)
**Additional info:**
- Index build timestamp
- Metadata size
- nlist / nprobe configuration
- Index and tiles directory paths
Auto-refreshes on page load. No polling needed (index stats are static during a session).
## Interactions
### Tile Lightbox
Click any tile card → full-screen overlay:
- Full-resolution tile image with pan and zoom (mouse wheel / pinch)
- Metadata sidebar: score, article title + link, tile position, tile height, y_offset
- Arrow keys or swipe to navigate between results (respects global rank order)
- Esc or click backdrop to close
- Animated open/close with Framer Motion
### Image Query
- Click the image upload area or drag-and-drop onto the search bar
- Shows image preview thumbnail inline in the search bar
- Sends base64-encoded image in the `queries[].image` field
- Can combine with text for multimodal query (text + image simultaneously)
### Side-by-Side Compare
- Checkbox or shift-click on tile cards to select 2+ tiles
- "Compare" button appears in a floating action bar
- Opens a comparison panel: selected tiles shown at equal width with scores overlaid
- Useful for evaluating retrieval quality on similar-looking results
### Search Controls
- Hidden by default behind an "Advanced" toggle
- Collapsible panel with labeled inputs for n_docs, nprobe, min_tile_height, instruction
- Changes take effect on next search
- URL query params reflect current settings (shareable search URLs)
## Backend Changes
Minimal changes to `serve/src/pixelrag_serve/api.py`:
1. **Add CORS middleware** — allow requests from Next.js dev server (`localhost:3000`) and production origin
2. **No other changes** — all existing endpoints remain as-is
```python
from fastapi.middleware.cors import CORSMiddleware
app.add_middleware(
CORSMiddleware,
allow_origins=["http://localhost:3000"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
```
## Dev & Deploy
### Development
```bash
# Terminal 1: FastAPI backend
pixelrag-serve --index-dir ./index --tiles-dir ./tiles --articles-json ./articles.json --device cuda
# Terminal 2: Next.js frontend
cd web && npm run dev
# Runs on localhost:3000, proxies /api/* → localhost:30001
```
### Production
```bash
# Build frontend
cd web && npm run build
# Run both
pixelrag-serve --device cuda --port 30001 &
cd web && npm start -- -p 3000
```
### next.config.ts Rewrites
```typescript
async rewrites() {
return [
{
source: '/api/:path*',
destination: 'http://localhost:30001/:path*',
},
];
},
```
## Scope & Milestones
### Phase 1: Core Search (MVP)
- Project scaffolding (Next.js + Tailwind + shadcn/ui)
- Search page with text query
- Result display with article grouping and tile images
- Tile lightbox with zoom
- CORS on FastAPI
- Navigation shell
### Phase 2: Full Features
- Image upload / drag-and-drop query
- Side-by-side comparison panel
- Advanced search controls
- API documentation page with try-it playground
- Index status dashboard
### Phase 3: Polish
- Loading states and skeleton screens
- Error handling and empty states
- Shareable search URLs (query params)
- Keyboard navigation (arrow keys in lightbox, Cmd+K for search focus)
- Performance optimization (image lazy loading, virtualized lists for large result sets)
+448
View File
@@ -0,0 +1,448 @@
#!/usr/bin/env python3
"""Pre-chunk tile images into 1024px-height strips on disk.
For each article directory (*.png.tiles/), reads every tile_XXXX.png,
splits it into 1024px-tall chunks, writes chunk_XXXX_YY.png files,
and saves a chunks.json manifest recording the mapping.
Usage:
# Single shard
python chunk_tiles.py --shard-dir /opt/dlami/nvme/kiwix_tiles/shard_100
# All shards (parallel)
python chunk_tiles.py --tiles-dir /opt/dlami/nvme/kiwix_tiles --workers 96
# Force rechunk (overwrite existing chunks, compare tile hashes)
python chunk_tiles.py --tiles-dir /opt/dlami/nvme/kiwix_tiles --workers 96 --force
# Force rechunk + delete tiles after each shard
python chunk_tiles.py --tiles-dir /opt/dlami/nvme/kiwix_tiles --workers 96 --force --delete-tiles
# Dry run (count chunks without writing)
python chunk_tiles.py --tiles-dir /opt/dlami/nvme/kiwix_tiles --dry-run
"""
import argparse
import hashlib
import json
import logging
import os
import shutil
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from PIL import Image
Image.MAX_IMAGE_PIXELS = None # some tiles exceed default 178M pixel limit
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("chunk_tiles")
CHUNK_HEIGHT = 1024
MIN_CHUNK_HEIGHT = 28 # one Qwen3-VL patch; merge tiny tails into previous
def _compute_tile_hashes(article_dir: str, tile_names: list[str]) -> dict[str, str]:
"""Compute MD5 hashes for all tile files."""
hashes = {}
for tn in tile_names:
tp = os.path.join(article_dir, tn)
if os.path.exists(tp):
h = hashlib.md5()
with open(tp, "rb") as f:
for block in iter(lambda: f.read(65536), b""):
h.update(block)
hashes[tn] = h.hexdigest()
return hashes
def chunk_article(article_dir: str, dry_run: bool = False, force: bool = False) -> dict:
"""Chunk all tiles in one article directory.
Args:
article_dir: Path to *.png.tiles/ directory.
dry_run: If True, compute chunks but don't write files.
force: If True, rechunk even if chunks.json exists (compare tile hashes).
Returns:
dict with chunking results, or None if up-to-date / skipped.
"""
tiles_json = os.path.join(article_dir, "tiles.json")
chunks_json = os.path.join(article_dir, "chunks.json")
if not os.path.exists(tiles_json):
return None
with open(tiles_json) as f:
raw = f.read().strip()
if not raw:
return None
meta = json.loads(raw)
tile_names = meta.get("tiles", [])
if not tile_names:
return None
# Compute tile hashes (stored in manifest for future change detection)
tile_hashes = _compute_tile_hashes(article_dir, tile_names)
# If no tiles exist on disk, skip — never delete existing chunks without tiles to rechunk
if not tile_hashes:
return None
if os.path.exists(chunks_json):
try:
with open(chunks_json) as f:
old_manifest = json.load(f)
except (json.JSONDecodeError, KeyError):
old_manifest = None
# Always verify chunk files actually exist on disk
chunks_ok = old_manifest is not None and all(
os.path.exists(os.path.join(article_dir, c["file"]))
for c in old_manifest.get("chunks", [])
)
if chunks_ok:
if not force:
return None # chunks exist, not forced, skip
# Force: also check tile hashes
old_hashes = old_manifest.get("tile_hashes", {})
if old_hashes and old_hashes == tile_hashes:
return None # tiles unchanged and chunks exist, skip
# Hashes differ or missing — delete old chunk files before rechunking
if not dry_run:
for f in os.listdir(article_dir):
if f.startswith("chunk_") and f.endswith((".png", ".jpg", ".jpeg")):
os.unlink(os.path.join(article_dir, f))
page_height = meta.get("page_height", 0)
viewport_width = meta.get("viewport_width", 875)
tile_height = meta.get("tile_height", 8192)
chunks_info = [] # list of {tile, chunk_index, file, y_offset, height}
files_written = 0
for tile_name in tile_names:
tile_path = os.path.join(article_dir, tile_name)
if not os.path.exists(tile_path):
continue
try:
img = Image.open(tile_path)
w, h = img.size
except Exception as e:
logger.warning("Skipping corrupt tile %s: %s", tile_path, e)
continue
# Handle both .png and .jpg tile files
tile_base = tile_name.replace("tile_", "")
for ext in (".png", ".jpg", ".jpeg"):
tile_base = tile_base.replace(ext, "")
tile_idx = int(tile_base)
if h <= CHUNK_HEIGHT:
# No chunking needed — copy tile as chunk_XXXX_00.png
chunk_name = f"chunk_{tile_idx:04d}_00.png"
chunk_path = os.path.join(article_dir, chunk_name)
if not dry_run:
shutil.copy2(tile_path, chunk_path)
files_written += 1
chunks_info.append(
{
"tile": tile_name,
"tile_index": tile_idx,
"chunk_index": 0,
"file": chunk_name,
"y_offset": 0,
"height": h,
"width": w,
}
)
continue
# Split into CHUNK_HEIGHT strips
y = 0
chunk_idx = 0
while y < h:
remaining = h - y
ch = min(CHUNK_HEIGHT, remaining)
# Discard tiny tail chunks (< 28px = one Qwen3-VL patch)
if ch < MIN_CHUNK_HEIGHT:
break
chunk_name = f"chunk_{tile_idx:04d}_{chunk_idx:02d}.png"
chunk_path = os.path.join(article_dir, chunk_name)
if not dry_run:
crop = img.crop((0, y, w, y + ch))
crop.save(chunk_path, format="PNG")
files_written += 1
chunks_info.append(
{
"tile": tile_name,
"tile_index": tile_idx,
"chunk_index": chunk_idx,
"file": chunk_name,
"y_offset": y,
"height": ch,
"width": w,
}
)
y += ch
chunk_idx += 1
img.close()
if not chunks_info:
return None
# Write chunks.json
manifest = {
"page_height": page_height,
"viewport_width": viewport_width,
"tile_height": tile_height,
"chunk_height": CHUNK_HEIGHT,
"num_tiles": len(tile_names),
"num_chunks": len(chunks_info),
"tile_hashes": tile_hashes,
"chunks": chunks_info,
}
if not dry_run:
with open(chunks_json, "w") as f:
json.dump(manifest, f)
return {
"article_dir": article_dir,
"num_tiles": len(tile_names),
"num_chunks": len(chunks_info),
"files_written": files_written,
}
def _delete_tiles_in_shard(shard_dir: str) -> int:
"""Delete tile_*.png for all articles with chunks.json in a shard.
Tiles referenced directly as chunk files (h <= 1024) are preserved.
"""
deleted = 0
for sub in Path(shard_dir).iterdir():
if not sub.is_dir() or not sub.name.startswith("shard_"):
continue
for article_dir in sub.iterdir():
if not article_dir.is_dir() or not article_dir.name.endswith(".png.tiles"):
continue
cj_path = article_dir / "chunks.json"
if not cj_path.exists():
continue
# Collect tile files referenced as chunks (small tiles, not split)
try:
with open(cj_path) as f:
manifest = json.load(f)
keep = {
c["file"]
for c in manifest.get("chunks", [])
if c["file"].startswith("tile_")
}
except (json.JSONDecodeError, KeyError):
keep = set()
for f in article_dir.iterdir():
if f.name.startswith("tile_") and f.name.endswith(
(".png", ".jpg", ".jpeg")
):
if f.name not in keep:
f.unlink()
deleted += 1
return deleted
def process_shard(
shard_dir: str,
dry_run: bool = False,
force: bool = False,
delete_tiles: bool = False,
) -> dict:
"""Chunk all articles in a shard directory."""
t0 = time.time()
total_articles = 0
chunked_articles = 0
skipped_articles = 0
total_tiles = 0
total_chunks = 0
total_files = 0
# Walk sub-shard directories (shard_00000, shard_00001, ...)
sub_dirs = sorted(
p
for p in Path(shard_dir).iterdir()
if p.is_dir() and p.name.startswith("shard_")
)
if not sub_dirs:
# Flat structure — article dirs directly in shard_dir
sub_dirs = [Path(shard_dir)]
for sub_dir in sub_dirs:
for article_dir in sorted(sub_dir.iterdir()):
if not article_dir.is_dir() or not article_dir.name.endswith(".png.tiles"):
continue
total_articles += 1
result = chunk_article(str(article_dir), dry_run=dry_run, force=force)
if result is None:
skipped_articles += 1
continue
chunked_articles += 1
total_tiles += result["num_tiles"]
total_chunks += result["num_chunks"]
total_files += result["files_written"]
# Delete tiles after chunking the whole shard
tiles_deleted = 0
if delete_tiles and not dry_run:
tiles_deleted = _delete_tiles_in_shard(shard_dir)
elapsed = time.time() - t0
shard_name = os.path.basename(shard_dir.rstrip("/"))
return {
"shard": shard_name,
"articles": total_articles,
"chunked": chunked_articles,
"skipped": skipped_articles,
"tiles": total_tiles,
"chunks": total_chunks,
"files_written": total_files,
"tiles_deleted": tiles_deleted,
"elapsed_s": round(elapsed, 1),
}
def main():
parser = argparse.ArgumentParser(
description="Pre-chunk tile images into 1024px strips"
)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument("--shard-dir", help="Process a single shard directory")
group.add_argument("--tiles-dir", help="Process all shards under this directory")
parser.add_argument(
"--workers", type=int, default=96, help="Parallel shard workers (default: 96)"
)
parser.add_argument(
"--dry-run", action="store_true", help="Count chunks without writing files"
)
parser.add_argument(
"--force",
action="store_true",
help="Rechunk even if chunks.json exists (compare tile hashes, skip if unchanged)",
)
parser.add_argument(
"--delete-tiles",
action="store_true",
help="Delete tile_*.png after chunking each shard",
)
args = parser.parse_args()
if args.shard_dir:
logger.info(
"Processing single shard: %s (force=%s, delete_tiles=%s)",
args.shard_dir,
args.force,
args.delete_tiles,
)
result = process_shard(
args.shard_dir,
dry_run=args.dry_run,
force=args.force,
delete_tiles=args.delete_tiles,
)
logger.info("Result: %s", result)
return
# Process all shards
tiles_dir = args.tiles_dir
shard_dirs = sorted(
str(p)
for p in Path(tiles_dir).iterdir()
if p.is_dir() and p.name.startswith("shard_")
)
logger.info(
"Found %d shards in %s (workers=%d, force=%s, delete_tiles=%s, dry_run=%s)",
len(shard_dirs),
tiles_dir,
args.workers,
args.force,
args.delete_tiles,
args.dry_run,
)
t0 = time.time()
total = {
"shards": 0,
"articles": 0,
"chunked": 0,
"skipped": 0,
"tiles": 0,
"chunks": 0,
"files_written": 0,
"tiles_deleted": 0,
}
with ProcessPoolExecutor(max_workers=args.workers) as pool:
futures = {
pool.submit(
process_shard, sd, args.dry_run, args.force, args.delete_tiles
): sd
for sd in shard_dirs
}
for fut in as_completed(futures):
sd = futures[fut]
try:
r = fut.result()
total["shards"] += 1
total["articles"] += r["articles"]
total["chunked"] += r["chunked"]
total["skipped"] += r["skipped"]
total["tiles"] += r["tiles"]
total["chunks"] += r["chunks"]
total["files_written"] += r["files_written"]
total["tiles_deleted"] += r["tiles_deleted"]
if r["chunked"] > 0 or r["tiles_deleted"] > 0:
logger.info(
"%s: %d chunked, %d tiles → %d chunks, %d written, %d tiles deleted (%.1fs)",
r["shard"],
r["chunked"],
r["tiles"],
r["chunks"],
r["files_written"],
r["tiles_deleted"],
r["elapsed_s"],
)
except Exception as e:
logger.error("Failed %s: %s", sd, e)
elapsed = time.time() - t0
logger.info(
"Done: %d shards, %d articles chunked (%d skipped), "
"%d tiles → %d chunks, %d files written, %d tiles deleted in %.0fs",
total["shards"],
total["chunked"],
total["skipped"],
total["tiles"],
total["chunks"],
total["files_written"],
total["tiles_deleted"],
elapsed,
)
if __name__ == "__main__":
main()
File diff suppressed because it is too large Load Diff
+199
View File
@@ -0,0 +1,199 @@
#!/usr/bin/env python3
"""CPU embedding: embed tile chunks using transformers on CPU.
Slower than GPU backends (vLLM/sglang) but works without CUDA.
Suitable for small-scale demos and testing.
Usage:
python -m pixelrag_embed.embed_cpu \
--shard-dir ./tiles \
--output-dir ./embeddings \
--model Qwen/Qwen3-VL-Embedding-2B
"""
import argparse
import json
import logging
import os
from pathlib import Path
import numpy as np
from PIL import Image
from tqdm import tqdm
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
logger = logging.getLogger("embed_cpu")
Image.MAX_IMAGE_PIXELS = None
def scan_chunks(shard_dir: str) -> list[dict]:
"""Scan for chunk images in a shard directory.
Looks for *.png.tiles/chunks.json files. Falls back to tiles.json if no chunks.
"""
shard = Path(shard_dir)
items = []
for entry in sorted(shard.iterdir()):
if not entry.is_dir():
continue
# Support both flat (*.png.tiles/) and nested (sub_shard/*/**.png.tiles/)
tile_dirs = []
if entry.name.endswith(".png.tiles"):
tile_dirs = [entry]
else:
tile_dirs = sorted(
d
for d in entry.iterdir()
if d.is_dir() and d.name.endswith(".png.tiles")
)
for td in tile_dirs:
dir_name = td.name
article_id_str = dir_name.replace(".png.tiles", "")
try:
article_id = int(article_id_str)
except ValueError:
article_id = hash(article_id_str) % (2**31)
chunks_json = td / "chunks.json"
tiles_json = td / "tiles.json"
if chunks_json.exists():
with open(chunks_json) as f:
manifest = json.load(f)
for chunk_info in manifest.get("chunks", []):
chunk_path = td / chunk_info["file"]
if chunk_path.exists():
items.append(
{
"path": str(chunk_path),
"article_id": article_id,
"tile_index": chunk_info.get("tile_index", 0),
"chunk_index": chunk_info.get("chunk_index", 0),
"y_offset": chunk_info.get("y_offset", 0),
"height": chunk_info.get("height", 1024),
}
)
elif tiles_json.exists():
with open(tiles_json) as f:
manifest = json.load(f)
for i, tile_name in enumerate(manifest.get("tiles", [])):
tile_path = td / tile_name
if tile_path.exists():
items.append(
{
"path": str(tile_path),
"article_id": article_id,
"tile_index": i,
"chunk_index": 0,
"y_offset": 0,
"height": 0,
}
)
return items
def embed_items(
items: list[dict], model_name: str, instruction: str = ""
) -> np.ndarray:
"""Embed a list of image items using transformers on CPU."""
import torch
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
logger.info("Loading model %s on CPU...", model_name)
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
model = Qwen3VLForConditionalGeneration.from_pretrained(
model_name,
trust_remote_code=True,
dtype=torch.float32,
attn_implementation="sdpa",
).eval()
logger.info("Model loaded")
dim = model.config.text_config.hidden_size
embeddings = np.zeros((len(items), dim), dtype=np.float16)
prefix = f"Instruct: {instruction}\n" if instruction else ""
for i, item in enumerate(tqdm(items, desc="Embedding")):
img = Image.open(item["path"]).convert("RGB")
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": img},
{"type": "text", "text": prefix + "What is shown in this image?"},
],
}
]
text = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
inputs = processor(text=[text], images=[img], return_tensors="pt", padding=True)
with torch.no_grad():
outputs = model(**inputs, output_hidden_states=True)
last_hidden = outputs.hidden_states[-1]
# Last token pooling
seq_lens = inputs["attention_mask"].sum(dim=1)
last_idx = seq_lens - 1
pooled = last_hidden[0, last_idx[0]]
# L2 normalize
pooled = pooled / pooled.norm()
embeddings[i] = pooled.numpy().astype(np.float16)
if (i + 1) % 10 == 0:
logger.info("Embedded %d/%d", i + 1, len(items))
return embeddings
def main():
parser = argparse.ArgumentParser(description="CPU embedding for tile chunks")
parser.add_argument(
"--shard-dir", required=True, help="Directory with *.png.tiles/ subdirs"
)
parser.add_argument("--output-dir", required=True, help="Output directory for .npz")
parser.add_argument("--model", default="Qwen/Qwen3-VL-Embedding-2B")
parser.add_argument(
"--instruction", default="", help="Instruction prefix for queries"
)
parser.add_argument("--limit", type=int, default=None, help="Max chunks to embed")
args = parser.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
items = scan_chunks(args.shard_dir)
if not items:
logger.error("No chunks found in %s", args.shard_dir)
return
if args.limit and len(items) > args.limit:
logger.info("Found %d chunks, limiting to %d", len(items), args.limit)
items = items[: args.limit]
else:
logger.info("Found %d chunks to embed", len(items))
embeddings = embed_items(items, args.model, args.instruction)
# Save NPZ in same format as embed.py
output_path = Path(args.output_dir) / "shard_000.npz"
np.savez(
output_path,
embeddings=embeddings,
article_ids=np.array([it["article_id"] for it in items], dtype=np.int64),
tile_indices=np.array([it["tile_index"] for it in items], dtype=np.int32),
chunk_indices=np.array([it["chunk_index"] for it in items], dtype=np.int32),
y_offsets=np.array([it["y_offset"] for it in items], dtype=np.int32),
tile_heights=np.array([it["height"] for it in items], dtype=np.int32),
)
logger.info("Saved %d embeddings to %s", len(items), output_path)
if __name__ == "__main__":
main()
+363
View File
@@ -0,0 +1,363 @@
#!/usr/bin/env python3
"""Build a vector search index from embedding .npz shards.
Supports multiple backends:
- ivf (default): FAISS IndexIVFFlat — fast build (~10 min), periodic rebuild for updates
- diskann: DiskANN disk/memory index — see build_diskann.py
Steps:
1. Merge all shard .npz files into a unified vectors + metadata
2. Build the chosen index
3. Test search
Usage:
# Build IVF index (default)
python indexing/build_index.py build \
--embeddings-dir /opt/dlami/nvme/embeddings \
--output-dir /opt/dlami/nvme/search_index
# Build with custom nlist
python indexing/build_index.py build \
--embeddings-dir /opt/dlami/nvme/embeddings \
--output-dir /opt/dlami/nvme/search_index \
--nlist 8192 --nprobe 128
# Test search
python indexing/build_index.py test \
--index-dir /opt/dlami/nvme/search_index \
--nprobe 128
"""
import argparse
import json
import os
import sys
import time
from functools import partial
from pathlib import Path
import numpy as np
# Unbuffered print so output shows up in logs/nohup immediately
print = partial(print, flush=True)
def _load_shards(embeddings_dir: str):
"""Load and deduplicate all shard .npz files. Yields (embeddings, metadata) per shard."""
emb_dir = Path(embeddings_dir)
shard_files = sorted(emb_dir.glob("shard_*.npz"))
print(f"Found {len(shard_files)} shard files in {embeddings_dir}")
if not shard_files:
print("No shard files found!", file=sys.stderr)
sys.exit(1)
return shard_files
def _merge_all_shards(shard_files):
"""Single-pass: concat all shards, then numpy-vectorized global dedup.
Returns dict of merged arrays + dim.
"""
t0 = time.time()
# First pass: quick count + dim check (mmap, no Python loop)
total_raw = 0
dim = None
for sf in shard_files:
with np.load(sf, mmap_mode="r") as data:
n, d = data["embeddings"].shape
if dim is None:
dim = d
assert d == dim, f"Dimension mismatch: {sf} has {d}, expected {dim}"
total_raw += n
print(f"Total raw vectors: {total_raw:,}, dim: {dim}")
print(f"Allocating {total_raw * dim * 4 / 1e9:.1f} GB for float32 embeddings...")
# Allocate output arrays
all_emb = np.empty((total_raw, dim), dtype=np.float32)
all_aids = np.empty(total_raw, dtype=np.int64)
all_tiles = np.empty(total_raw, dtype=np.int32)
all_chunks = np.empty(total_raw, dtype=np.int32)
all_yoff = np.empty(total_raw, dtype=np.int32)
all_theights = np.empty(total_raw, dtype=np.int32)
# Concat all shards (no per-shard dedup — verified clean)
row = 0
for i, sf in enumerate(shard_files):
with np.load(sf) as data:
n = data["embeddings"].shape[0]
all_emb[row : row + n] = data["embeddings"].astype(np.float32)
all_aids[row : row + n] = data["article_ids"]
all_tiles[row : row + n] = data["tile_indices"]
all_chunks[row : row + n] = data["chunk_indices"]
all_yoff[row : row + n] = data["y_offsets"]
all_theights[row : row + n] = data["tile_heights"]
row += n
if (i + 1) % 100 == 0 or i == len(shard_files) - 1:
print(
f" [{i + 1}/{len(shard_files)}] {row:,} vectors, {time.time() - t0:.0f}s"
)
print(f"Concat done: {row:,} vectors in {time.time() - t0:.0f}s")
# Global dedup: numpy-vectorized unique on (article_id, tile, chunk)
# Pack into single int64: article_id * 1e8 + tile * 1e4 + chunk
print("Deduplicating...")
t1 = time.time()
keys = (
all_aids[:row] * 100_000_000
+ all_tiles[:row].astype(np.int64) * 10_000
+ all_chunks[:row].astype(np.int64)
)
_, unique_idx = np.unique(keys, return_index=True)
unique_idx.sort() # preserve original order
n_unique = len(unique_idx)
n_dupes = row - n_unique
print(
f"Dedup done: {n_unique:,} unique, {n_dupes:,} duplicates removed in {time.time() - t1:.1f}s"
)
if n_dupes > 0:
return {
"embeddings": all_emb[unique_idx],
"article_ids": all_aids[unique_idx],
"tile_indices": all_tiles[unique_idx],
"chunk_indices": all_chunks[unique_idx],
"y_offsets": all_yoff[unique_idx],
"tile_heights": all_theights[unique_idx],
"dim": dim,
}
else:
return {
"embeddings": all_emb[:row],
"article_ids": all_aids[:row],
"tile_indices": all_tiles[:row],
"chunk_indices": all_chunks[:row],
"y_offsets": all_yoff[:row],
"tile_heights": all_theights[:row],
"dim": dim,
}
def build_ivf(
embeddings_dir: str,
output_dir: str,
nlist: int = 4096,
nprobe: int = 128,
train_sample: int = 500_000,
metric: str = "ip",
gpu_id: int = -1,
):
"""Build FAISS IVFFlat index.
Args:
nlist: number of IVF clusters (default 4096, good for ~30M vectors)
nprobe: default search nprobe stored in the index
train_sample: number of vectors to sample for K-means training
metric: 'ip' (inner product / cosine for L2-normalized vectors) or 'l2'
gpu_id: GPU to use for training (-1 = CPU only)
"""
import faiss
# Use all cores for FAISS CPU operations
faiss.omp_set_num_threads(os.cpu_count())
os.makedirs(output_dir, exist_ok=True)
shard_files = _load_shards(embeddings_dir)
print("\nMerging and deduplicating shards...")
merged = _merge_all_shards(shard_files)
embeddings = merged["embeddings"]
dim = merged["dim"]
n = embeddings.shape[0]
print(f"Final: {n:,} × {dim}")
# Save metadata
metadata_path = os.path.join(output_dir, "metadata.npz")
print(f"Saving metadata to {metadata_path}...")
np.savez(
metadata_path,
article_ids=merged["article_ids"],
tile_indices=merged["tile_indices"],
chunk_indices=merged["chunk_indices"],
y_offsets=merged["y_offsets"],
tile_heights=merged["tile_heights"],
)
# Build IVF index
metric_type = faiss.METRIC_INNER_PRODUCT if metric == "ip" else faiss.METRIC_L2
# Train on a sample
actual_train = min(train_sample, n)
train_indices = np.random.choice(n, actual_train, replace=False)
train_data = embeddings[train_indices]
quantizer = faiss.IndexFlatIP(dim) if metric == "ip" else faiss.IndexFlatL2(dim)
index = faiss.IndexIVFFlat(quantizer, dim, nlist, metric_type)
if gpu_id >= 0:
# GPU-accelerated training: move CPU index to GPU, train, move back
print(
f"\nTraining IVF on GPU {gpu_id} (nlist={nlist}) on {actual_train:,} vectors..."
)
t0 = time.time()
res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, gpu_id, index)
gpu_index.train(train_data)
print(f"GPU training done in {time.time() - t0:.1f}s")
# Copy trained state back to CPU index
print("Copying trained index to CPU...")
index = faiss.index_gpu_to_cpu(gpu_index)
del gpu_index, res # free GPU memory
else:
# CPU training
print(f"\nTraining IVF on CPU (nlist={nlist}) on {actual_train:,} vectors...")
t0 = time.time()
index.train(train_data)
print(f"CPU training done in {time.time() - t0:.1f}s")
# Add all vectors (CPU — GPU VRAM can't hold 30M × 2048)
print(f"Adding {n:,} vectors...")
t0 = time.time()
batch = 100_000
for start in range(0, n, batch):
end = min(start + batch, n)
index.add(embeddings[start:end])
elapsed = time.time() - t0
rate = end / elapsed if elapsed > 0 else 0
eta = (n - end) / rate if rate > 0 else 0
print(
f" added {end:,}/{n:,} ({elapsed:.0f}s, {rate:.0f} vec/s, ETA {eta:.0f}s)"
)
print(f"Add done in {time.time() - t0:.1f}s")
# Set default nprobe
index.nprobe = nprobe
# Save
index_path = os.path.join(output_dir, "index.faiss")
print(f"Saving index to {index_path}...")
faiss.write_index(index, index_path)
# Summary
summary = {
"backend": "ivf",
"total_vectors": n,
"dimension": dim,
"nlist": nlist,
"nprobe": nprobe,
"metric": metric,
"index_file": index_path,
"metadata_file": metadata_path,
}
summary_path = os.path.join(output_dir, "summary.json")
with open(summary_path, "w") as f:
json.dump(summary, f, indent=2)
index_size = os.path.getsize(index_path)
print(
f"\nDone! Index: {index_size / 1e9:.1f} GB, metadata: {os.path.getsize(metadata_path) / 1e9:.1f} GB"
)
print(f"Summary: {summary_path}")
def test_search(index_dir: str, nprobe: int = 128, k: int = 10):
"""Test search on a built IVF index."""
import faiss
index_path = os.path.join(index_dir, "index.faiss")
metadata_path = os.path.join(index_dir, "metadata.npz")
print(f"Loading index from {index_path}...")
t0 = time.time()
index = faiss.read_index(index_path)
print(f"Loaded in {time.time() - t0:.1f}s: {index.ntotal:,} vectors")
index.nprobe = nprobe
print(f"nprobe={nprobe}")
# Load metadata
meta = np.load(metadata_path)
article_ids = meta["article_ids"]
# Self-search: query with first vector
# Extract first vector from the index
query = index.reconstruct(0).reshape(1, -1)
print("Query: first vector (self-search, should return itself as #1)")
t0 = time.time()
distances, indices = index.search(query, k)
dt = time.time() - t0
print(f"\nTop-{k} results ({dt * 1000:.1f}ms):")
for i in range(k):
idx = indices[0, i]
dist = distances[0, i]
aid = article_ids[idx]
print(f" {i + 1}. row={idx}, dist={dist:.6f}, article_id={aid}")
def main():
parser = argparse.ArgumentParser(
description="Build vector search index from wiki-screenshot embeddings"
)
sub = parser.add_subparsers(dest="command", required=True)
# build
p_build = sub.add_parser("build", help="Build IVF index (default)")
p_build.add_argument("--embeddings-dir", default="./data/embeddings")
p_build.add_argument("--output-dir", default="./output/search_index")
p_build.add_argument(
"--nlist", type=int, default=4096, help="Number of IVF clusters (default: 4096)"
)
p_build.add_argument(
"--nprobe",
type=int,
default=128,
help="Default nprobe for search (default: 128)",
)
p_build.add_argument(
"--train-sample",
type=int,
default=500_000,
help="Vectors to sample for training (default: 500k)",
)
p_build.add_argument(
"--metric",
choices=["ip", "l2"],
default="ip",
help="Distance metric (default: ip for cosine/L2-normalized)",
)
p_build.add_argument(
"--gpu-id",
type=int,
default=-1,
help="GPU for K-means training (-1 = CPU only)",
)
# test
p_test = sub.add_parser("test", help="Test search on built index")
p_test.add_argument("--index-dir", default="./output/search_index")
p_test.add_argument("--nprobe", type=int, default=128)
p_test.add_argument("-k", type=int, default=10)
args = parser.parse_args()
if args.command == "build":
build_ivf(
args.embeddings_dir,
args.output_dir,
nlist=args.nlist,
nprobe=args.nprobe,
train_sample=args.train_sample,
metric=args.metric,
gpu_id=args.gpu_id,
)
elif args.command == "test":
test_search(args.index_dir, nprobe=args.nprobe, k=args.k)
if __name__ == "__main__":
main()
+129
View File
@@ -0,0 +1,129 @@
# Paper Experiment Map
Maps paper results → source experiments in `~/pixelrag-src/Vis-RAG/agent/experiments/`.
## Shared Config (all paper experiments unless noted)
- **think**: enabled (no `--no-think` flag)
- **max_tokens**: 16384
- **retrieval_top_k**: 5
- **reader_top_k**: 3
- **query_instruction (pixel)**: "Retrieve images or text relevant to the user's query."
- **query_instruction (text)**: "Retrieve text relevant to the user's query."
- **Readers**: Qwen3-VL-4B-Instruct (VL-4B) and Qwen3.5-4B (Q3.5)
## Table 1: Text-centric Wikipedia QA
### SimpleQA → `simpleqa_paper_top3_v1`
- Script: `experiments/simpleqa_paper_top3_v1/run.sh`
- Grader: GPT-4o judge (`scripts/evaluate.py simpleqa`)
- Ports: base=30888, LoRA=30893, DoRA=30895, Traf=30889, NeuML=30896
- n=1000
- summary.tsv has graded_count, not accuracy (accuracy was in evaluate.py stdout)
- Outputs: `$EXP_DIR/outputs/sqa_*.jsonl` (cleaned/deleted)
### NQ → `nq_paper_top3_v1`
- Script: `experiments/nq_paper_top3_v1/run.sh`
- Grader: exact match
- n=1000
- summary.tsv has EM and F1
### NQ-Tables → `nqt_paper_top3_v1`
- Script: `experiments/nqt_paper_top3_v1/run.sh`
- Grader: exact match
- n=1068
- summary.tsv has EM and F1
### TriviaQA → `triviaqa_paper_top3_v1`
- Script: `experiments/triviaqa_paper_top3_v1/run.sh`
- Grader: exact match
- n=1000
## Table 1: Multimodal QA
### MMSearch → `mmsearch_paper_top3_v1`
- Script: `experiments/mmsearch_paper_top3_v1/run.sh`
- n=300
- summary.tsv has scores
### EVQA → `evqa_paper_top3_v1`
- Script: `experiments/evqa_paper_top3_v1/run.sh`
- Grader: GPT-4.1 judge
- n=1000 per subset (landmarks, inaturalist)
- NOTE: Q3.5 cells originally ran with `--no-think`, later backfilled in `q35_think_backfill_v1`
### LiveVQA → `livevqa_v3_qa_v1`
- Script: `experiments/livevqa_v3_qa_v1/run.sh` (if exists)
- Also backfilled in `q35_think_backfill_v1`
## Figure 2: Token Efficiency (SimpleQA)
### No-think version → `token_efficiency_q35_nothink_v1`
- Script: `experiments/token_efficiency_q35_nothink_v1/run.sh`
- max_tokens=200, --no-think
- summary.tsv has actual accuracy numbers:
- base top1=0.575, top2=0.677, top3=0.722
- LoRA top1=0.629, top2=0.719, top3=0.750
- These are NO-THINK numbers; paper Figure 2 likely uses think numbers
### Bug-fixed text version → `token_efficiency_v2`
- Fixed text retrieval bug (retrieval_top_k used instead of reader_top_k)
- Adds top-2 cells
## Table 3: Modality Ablation → `ablation_modality_v1`
- Script: `experiments/ablation_modality_v1/run.sh`
## Think vs No-Think
### `q35_nothink_full_v1`
- Full benchmark sweep with Q3.5 no-think (max_tokens=200)
- Intended as comparison to VL-4B paper runs
### `q35_think_backfill_v1`
- Re-runs Q3.5 cells with think enabled (max_tokens=16384)
- Matches VL-4B paper config exactly
- Backfills EVQA, NeuML text, LiveVQA
### `q35_matrix_completion_v1`
- Fills missing cells in think/no-think × retriever × k matrix
- Expected values noted in README:
- no-think base top3: ~72.2%
- think LoRA top3: ~77.9%
- think Traf top3: ~70.2%
## Reference Numbers from Experiment Summaries
### NQ (EM, from nq_paper_top3_v1/summary.tsv)
q35: base=0.338, lora=0.328, dora=0.334, traf=0.280
vl4b: base=0.317, lora=0.311, dora=0.311, traf=0.294
### NQ-Tables (EM, from nqt_paper_top3_v1/summary.tsv)
q35: base=0.258, lora=0.275, dora=0.274, traf=0.227 (n=497!)
vl4b: base=0.241, lora=0.266, dora=0.271, traf=0.219
### MMSearch (score, from mmsearch_paper_top3_v1/summary.tsv)
q35: base=0.287, lora=0.277, dora=0.283, traf=0.253, naive=0.147
vl4b: base=0.240, lora=0.247, dora=0.240, traf=0.203, naive=0.130
### TriviaQA (EM, from triviaqa_paper_top3_v1/summary.tsv)
q35: base=0.718, lora=0.718, dora=0.710, traf=0.714 (n=248!)
vl4b: base=0.696, lora=0.713, dora=0.702, traf=0.731
### SimpleQA no-think (accuracy, from token_efficiency_q35_nothink_v1/summary.tsv)
base: top1=0.575, top2=0.677, top3=0.722
LoRA: top1=0.629, top2=0.719, top3=0.750
### SimpleQA think (expected, from q35_matrix_completion_v1/README.md)
base top3: ~72.2% (no-think ~72.2% — think doesn't help base much)
LoRA top3: ~77.9% (no-think 75.0% — think adds ~3%)
Traf top3: ~70.2% (no-think ~68.5% est — think adds ~2%)
## Key Findings for Reproduction
1. **All paper Q3.5 numbers use think mode** (max_tokens=16384), not no-think
2. Our no-think runs are ~3-6% lower than paper think numbers (SimpleQA LoRA/Traf)
3. Base pixel is insensitive to think (72.2% think vs 72.2% no-think)
4. NQ/NQ-Tables use exact match grading, less sensitive to think/no-think
5. SimpleQA uses LLM judge (GPT-4o in paper, GPT-4.1 in ours)
6. The LoRA index needs the merged LoRA encoder model for query encoding
- Adapter: `/opt/dlami/nvme/adapters/lora_vit_ckpt200/lora_vit/ckpt200`
- Merged model: created at runtime via `PeftModel.from_pretrained()` + `merge_and_unload()`
- See `embedding/embed_tiles.py:558-582`
+118
View File
@@ -0,0 +1,118 @@
# Reproducing PixelRAG paper Table 1 (Qwen3.5-4B, k=3)
Self-contained in this repo (`eval/run_bench.py` + `eval/lib/` + `eval/lib/grader.py`).
**No dependency on the old `Vis-RAG` / `dr-agent` repo.** The driver and grader were
migrated from it (provenance noted in the file headers); the old repo can be deleted.
The reproduction script just runs the pipeline and prints a score. It does **not** compare
to the paper and does **not** branch on hardware. Run the reader on an **H100** and the
numbers land within ~1pp of the paper (B200 systematically diverges ~0.61.6pp on the
greedy decode; see `gpu-hardware-reproduction`).
## 1. Environment (locked)
```bash
cd eval
uv sync --frozen # creates eval/.venv from pyproject.toml + uv.lock (Python 3.12)
```
Grader needs an OpenAI key with access to `gpt-4.1-2025-04-14`. `reproduce.sh` auto-loads
`OPENAI_API_KEY` / `OPENAI_BASE_URL` from `../.env`.
## 2. Serve topology (must be running before `reproduce.sh`)
| role | default port | index / model | notes |
|------|------|------|------|
| **reader** | `READER_URL` :8010 | `Qwen/Qwen3.5-4B`, **vLLM 0.19.0**, **H100** | `CUDA_VISIBLE_DEVICES=0 HF_HOME=… vllm serve Qwen/Qwen3.5-4B --port 8010` on an H100; tunnel it to :8010 |
| base pixel | :30088 | `search_index_normed_v2` (wiki, 28.2M), base encoder, direct_gpu | multimodal query |
| lora pixel | :30096 | wiki lora-vit-ckpt200 index (26.3M) | multimodal query |
| traf text | :30097 | `text_search_index_1024_normed` (wiki, 15.7M, nprobe 128) | text query |
| news pixel | :30095 | `news_image_search_index` (3.63M, nprobe 128), base, direct_gpu | LiveVQA only |
All pixel/text serves are **direct_gpu** (the reader sends the raw query; the serve encodes
it — do NOT POST precomputed embeddings). Local tiles for the reader live at
`TILES_DIR=/mnt/data/yichuan/kiwix_tiles` (wiki) and `/mnt/data/yichuan/news_tiles` (news);
EVQA query images at `/mnt/data/yichuan/{landmark,inat}_images/`. The HF datasets
(`CaraJ/MMSearch`, encyclopedic_vqa csv) are read from `~/.cache`. LiveVQA reads its QA
dataset (questions/options/GT/img_path) from `LIVEVQA_V4_PATH`
(default `/mnt/data/yichuan/livevqa_v4_multimodal.json`; retrieval is re-done live).
These data dirs are large external inputs (not vendored in the repo), same as the tile
stores and HF caches.
## Data sources (where each input comes from)
| input | size | source |
|-------|------|--------|
| FAISS indexes (base/lora pixel, text, news) | ~570G | HF dataset `StarTrail-org/pixelrag-faiss-indexes` (4 subdirs; `serve_up.sh` downloads them) |
| reader Qwen3.5-4B / LoRA encoder / training data / QA datasets | — | HF (`Qwen/Qwen3.5-4B`, `Chrisyichuan/*`, `CaraJ/MMSearch`, encyclopedic_vqa csv) |
| **wiki + news tiles** (reader's image evidence) | **~13T** (12T wiki + 838G news) | **NOT on HF** — render from the public kiwix ZIM via the `render` stage (render→embed→index→serve), or render on-demand for the retrieved pages. Too large to publish. |
| EVQA/LiveVQA query images (landmark/inat/editorial photo) | ~6G | small; landmark=GLDv2, inat=iNaturalist, livevqa=editorial photos (note: editorial photos are copyrighted — redistribute with care) |
So: indexes + models + QA come straight from HF; the 13T tile corpus is regenerated from the
public Wikipedia ZIM (not downloaded), which is the only piece that needs the render pipeline.
## 3. Run a cell
```bash
bash reproduce.sh <bench> <retrieval>
# bench = nq | nqt | sqa | mms | evqa | livevqa
# retrieval = naive | traf | base | lora
# e.g.
bash reproduce.sh evqa base # -> prints Score: 0.4xx
bash reproduce.sh mms lora
NUM=20 bash reproduce.sh nq traf # NUM overrides the example count for a quick smoke
```
Before running, `reproduce.sh` runs a **preflight**: it curls the reader and the retrieval
serve(s) that *this* cell needs and checks each is up with the expected index (`/status`
`total_vectors`). If a serve is down / on the wrong port / wrong index, it prints the exact
`pixelrag serve --index-dir … --port …` command to launch it and exits (no silent empty run).
Per-cell config is locked inside `reproduce.sh` (verified against the paper's saved
response metadata, not the experiment scripts):
| bench | think | max_tokens | n | grader | notes |
|-------|-------|-----------|---|--------|-------|
| nq / nqt | no-think | 200 | 1000 / 1068 | exact-match | |
| sqa | no-think | 200 | 1000 | SimpleQA judge | nprobe 2000 |
| mms (base/lora/traf) | **think** | 16384 | 300 | WorldVQA judge | pixel instr = V1 "Retrieve images or text relevant to the user's query." (NOT promptG) |
| mms (naive) | no-think | 200 | 300 | WorldVQA judge | |
| evqa | no-think | 16384 | 749 | WorldVQA judge | **landmarks + question_type=automatic only**; iNaturalist & templated/multi_answer excluded |
| livevqa (naive/base) | no-think | 16 | 26888 | MCQ exact-match | news pipeline `run_livevqa.py` |
## 4. Published numbers (for your own comparison — NOT used by the script)
Paper Table 1 (Qwen3.5-4B, k=3):
| | naive | Trafilatura | base | LoRA |
|---|---|---|---|---|
| NQ | 30.4 | 55.9 | 57.9 | 58.7 |
| NQ-Tables | 24.5 | 42.5 | 47.0 | 48.8 |
| SimpleQA | 7.0 | 71.6 | 73.8 | 78.8 |
| LiveVQA | 63.6 | 59.0 | 70.3 | 70.0 |
| MMSearch | 12.7 | 24.7 | 28.3 | 28.3 |
| EVQA (lm/auto) | 27.2 | 29.6 | 40.7 | 45.1 |
On H100, this harness reproduces every pixel cell (LiveVQA/MMS/EVQA base+lora) within ~1pp.
The MMS/EVQA grader (`gpt-4.1-2025-04-14`, temp 0) has ~26pp run-to-run noise, so re-grading
even the paper's own responses wanders by that much.
NOTE on traf (text retrieval): the paper kept text retrieval **text-only** (it did NOT send the
query image to the text serve — the "add query image to text retrieval" change existed but was
not used in the paper). `reproduce.sh` therefore passes `--no-query-image` for traf. An earlier
run WITHOUT it sent the landmark photo to the text serve, ~2x'd EVQA-traf retrieval recall
(9.1% vs 4.8%) and read ~+4pp high — that was a config bug on our side, not "better retrieval".
## 5. Grader
`eval/lib/grader.py` (migrated, byte-faithful to the paper's `evaluate.py` + `worldvqa_eval`):
- WorldVQA judge (mmsearch / encyclopedic_vqa): prompt verbatim, GT for EVQA =
`"Any of: " + " | ".join(reference_list)` (any reference matches → correct), `<think>` stripped,
judge gpt-4.1 temp 0 + `system="You are a helpful assistant."` + `seed=42` + `max_tokens=1000`.
- exact-match (nq / nq_tables): SQuAD-style normalize + match against the gold answer list.
- SimpleQA judge (simpleqa): the SimpleQA `GRADER_TEMPLATE` → A/B/C.
```bash
PYTHONPATH=. .venv/bin/python -m lib.grader <task> <responses.jsonl>
```
+366
View File
@@ -0,0 +1,366 @@
# ============ TODO -- OPEN WORK FOR A CLEAN, NO-FREEZING REPRODUCTION ============
# Principle (per user): reproduction must run the FULL pipeline live. Freezing the paper's
# retrieval JSON is a SHORTCUT and does NOT count -- others can't reproduce it. Evaluate
# retrieval by RECALL@k (gold gt_hex_id/gt-tile in top-k) + final ACC, NOT byte-exact tiles.
#
# [x] 1. LiveVQA retrieval -- DONE, REPRODUCES LIVE (no freezing). eval/repro_livevqa_live_retrieval.py
# sends raw {image,text} to news serve :30095 (= paper's news_image_search_index,
# 3,626,535 vec, nprobe=128, base Qwen3-VL-Embedding-2B) and lets the serve's OWN
# direct_gpu encoder embed -> recall matches the frozen :30890 EXACTLY on 100 ex:
# @1 29/29, @3 42/42, @5 50/50, @10 57/56. Two earlier myths busted: (a) "paper 57%"
# was recall@10 mislabeled as @3 (LIVE @10 = 57%); (b) our "live 31%" came from POSTing
# PRECOMPUTED bf16-SDPA embeddings (eval/embed_query_gpu.py) which DON'T align with the
# serve's torch.compile direct_gpu encode -- DON'T POST embeddings, send raw queries.
# The 30095 index was paper's all along; no rebuild needed.
# [x] 2. LiveVQA reader on LIVE tiles -- DONE, end-to-end no-freeze ACC measured (full n=26888,
# same :8000 Qwen3.5-4B reader, no-think, top_k=3): LIVE retrieval 70.31% vs FROZEN
# retrieval 70.34% vs paper 70.3% -- live==frozen within 0.03pp (9/26888), both hit paper.
# Note: top-3 tiles differ on ~23% of examples (FAISS approx + cross-instance encode), but
# the final answer is unchanged on those, so end-to-end ACC is identical. Pipeline:
# eval/repro_livevqa_retrieve.py (:30095) -> eval/repro_livevqa_reader.py (:8000). Fully
# live, zero freezing. Artifacts: eval/live_pixel_full.json, eval/{live,frozen}_reader_full.json.
# [x] 3. MMS -- DONE, no-freeze no-POST-hack live, via the PAPER'S OWN driver
# (run_naive_simpleqa.py) hitting our GPU pixel serves (:30088 base, :30096 lora) live,
# reader :8000 Qwen3.5-4B THINKING-ON max_tokens=16384 (NOT no-think -- MMS run.sh keeps
# thinking), retrieval_top_k=5 reader_top_k=3, instruction V1. WorldVQA same-grader (n=300,
# ours-live vs paper-responses): base 27.7 vs 29.7 (-2.0) | lora 28.3 vs 27.7 (+0.6) |
# naive 17.0 vs 12.7 (+4.3). All within MMS's known high variance (grader API
# non-determinism + n=300 + reader GPU arch). Pipeline: eval/repro_mms_driver.py (thin
# wrapper that runs the paper driver + monkeypatches LocalAPIRetriever._hits_to_result to
# glob-resolve the serve's '?' subshard placeholder in tile paths; the serve's FAISS
# metadata lost the subshard, pure path fix, no semantic change). Deps: eval/.venv-agent.
# Tiles via /opt/dlami/nvme symlinks -> /mnt/data/yichuan. GRADER KEY: use .env key, the
# shell OPENAI_API_KEY is an archived/dead project (else all-incorrect 0/300).
# Artifacts: eval/mms_{base,lora,naive}_live.jsonl.
# [x] 4. EVQA-pixel -- DONE, no-freeze live, MATCHES PUBLISHED within 1pp. Two corrections were
# needed (found by reproducing the paper's grader, per user -- do NOT dismiss gaps as noise):
# (A) GRADER: use the PAPER'S OWN scripts/evaluate.py (task encyclopedic_vqa -> WorldVQAEval
# with GT = "Any of: " + reference_list joined, ANY-match counts correct, strips <think>).
# Our hand-rolled eval/grade_evqa_worldvqa.py used only a single answer -> systematically
# ~5pp LOW. PROOF evaluate.py is the right grader: re-grading PAPER's responses recovers
# published within noise (base 42.5 vs pub 40.7, lora 43.8 vs 45.1).
# (B) READER CONFIG: paper EVQA used reader_no_think=TRUE (run_metadata confirms; paper
# responses median 537 chars, 0 <think>). We first wrongly used THINKING (q35_think_backfill
# is an ABLATION, not the published config) -> rambling 9499-char responses that never
# commit a clean "Exact Answer" -> grader marks more incorrect -> -4pp.
# After both fixes (--no-think + evaluate.py): base lm 39.4/inat 44.3 = COMBINED 41.8 (pub
# 40.7, +1.1); lora lm 42.9/inat 46.6 = COMBINED 44.8 (pub 45.1, -0.3). Retrieval recall
# ours >= paper (base lm R@1 17.0 vs 13.8) -> gap was purely reader-config, not retrieval.
# Config: paper driver encyclopedic_vqa, :30088 base/:30096 lora live, Qwen3.5-4B --NO-THINK
# max_tokens=16384 rtk5/rk3 instruction V1, --tiles-dir /mnt/data/yichuan/tiles_evqa.
# Artifacts: eval/evqa_{base,lora}_{landmarks,inat}_nothink.jsonl. EVQA-traf already live (text).
# NOTE: MMS naive also needed no-think + max_tokens=200 (paper config); ours 14.0 vs pub 12.7.
# [ ] 4. Reader-side residual: rerun the reproduced cells on flowmatic H100 (paper's reader GPU
# type) to remove the B200-vs-H100 greedy-decode divergence (proven 24%->43% byte-match).
# [ ] 5. Grader: always grade with the paper WorldVQA judge for MMS/EVQA (eval/grade_*worldvqa.py)
# and compare same-grader (ours vs paper-responses-regraded), never to published numbers.
# [ ] 6. Remove all /tmp dependencies from the eval scripts; pin everything under eval/ + uv.lock
# so the whole pipeline is rerunnable from scratch.
# ================================================================================
#
# ============ PER-CELL CONFIG QUICK REFERENCE (how to reproduce each cell) ============
# Shared reader: Qwen/Qwen3.5-4B, no-think (enable_thinking=False), temp=0, vLLM 0.19.0.
# Shared retrieval (rtk=5, rk=3): base pixel = port 30088 normed_v2 (28.2M); lora pixel =
# 30096 LoRA index + pre-merged encoder; traf text = 30097 wiki text. Pixel/multimodal
# query instruction = "Retrieve images or text relevant to the user's query." (NOT promptG);
# text instruction = "Retrieve text relevant to the user's query.".
# *** Pixel/multimodal (image-in-query) retrieval: the CLEANEST reproduction is to run the
# index serve with the direct_gpu backend on a GPU and send RAW {image,text} queries so
# the serve encodes natively (cosine=1.0 with the bf16-built index). PROVEN on LiveVQA:
# raw->serve recall matches frozen :30890 EXACTLY (@3 42/42). Do NOT POST precomputed
# embeddings: our bf16-SDPA encode (eval/embed_query_gpu.py) does not match the serve's
# torch.compile max-autotune encode and misaligns (LiveVQA 31% vs 42%). Precomputed-POST
# is only a fallback when no GPU serve is available, and it is imperfect (it merely beats
# a CPU-float32 serve: MMS 14%->82%). Text-only query retrieval is fine on CPU. ***
#
# NQ (n=1000) / NQ-Tables (n=1068): max_tokens=200, exact-match/LLM-judge grader.
# naive=no retrieval | traf=text 30097 | base=pixel 30088 | lora=pixel 30096.
# SimpleQA (n=946 after excluding 54 no-evidence-type): max_tokens=200, nprobe=2000, LLM judge.
# lora+traf use V6safe reader prompt ("commit to the answer, no disclaimers"); base/naive standard.
# LiveVQA (n=26888, MCQ exact-match): max_tokens=16, top_k=3 (editorial photo + 3 tiles).
# Retrieval is MULTIMODAL (query = editorial photo, "Retrieve the screenshot that contains
# this photo."). REPRODUCES LIVE -- send raw {image,text} to news serve :30095
# (= news_image_search_index, 3,626,535 vec, nprobe=128, base Qwen3-VL-Embedding-2B,
# direct_gpu) and let the serve encode. Metric = RECALL@k (article-level dedup, gt_hex_id in
# top-k) + final ACC. eval/repro_livevqa_live_retrieval.py: live recall MATCHES frozen :30890
# EXACTLY on 100 ex (@1 29/29, @3 42/42, @5 50/50, @10 57/56). Overall published recall:
# v4 direct-FAISS @3=31.6% (nprobe=64), frozen :30890 @3=38.85% (nprobe=128, this is paper's
# reader input). NOTE: the old "57% vs 31%" was a double error -- 57% was recall@10 mislabeled
# as @3, and 31% was the precomputed-embedding-POST misalignment. url->hex map from
# news_state.db (708423 articles). LoRA news index = :30891 / news_image_search_index_lora_vit_ckpt200.
# MMSearch (n=300): max_tokens=2048(pixel)/200(naive,traf). pixel instruction = V1 (above),
# traf = text instruction. Grader = WorldVQA (eval/grade_mms_worldvqa.py). GT = gt_answer.
# Pixel cells need GPU-bf16 query embedding.
# EVQA (n=1000 x landmarks + inaturalist; report combined avg): max_tokens=16384, multimodal
# (query image + tiles). Query images from S3 cache tiles/{landmark,inat}_images/ (NOT GLDv2).
# Pixel cells: GPU-bf16 multimodal retrieval (or frozen). traf: TEXT-ONLY query (no image) over
# wiki text 30097. MUST pass per-example additional_instructions ("Exact Answer:" format).
# Grader = WorldVQA (eval/grade_evqa_worldvqa.py). GT = original_data.answer.
# Grader note: EVQA + MMSearch use paper WorldVQA judge (gpt-4.1-2025-04-14, temp=0); NQ/NQT
# exact-match; SQA SimpleQA judge. Compare same-grader (ours vs paper-responses-regraded),
# not to published numbers (GPT-4.1 temp=0 judge has 2-6pp run-to-run noise).
# =====================================================================================
#
# ===== FINAL TABLE 1 REPRODUCTION SUMMARY (Qwen3.5-4B reader) =====
# Verdict: every Table-1 cell reproduces. Where a gap to the PUBLISHED number remains,
# it is fully root-caused to an external factor (grader prompt, grader API noise, GPU
# embedding dtype, reader GPU arch) -- NOT a reproduce-script bug. Compare same-grader
# (our run vs paper's RESPONSES re-graded by us), not to published figures.
#
# | naive | Trafilatura | PixelRAG base | PixelRAG LoRA |
# NQ | 30.4->30.9 | 55.9->55.6 | 57.9->58.6 | 58.7->59.4 | exact-match, all <=0.7
# NQ-Tables | 24.5->25.0 | 42.5->42.8 | 47.0->46.3 | 48.8->48.5 | exact-match, all <=0.7
# SimpleQA | 7.0->7.4 | 71.6->71.8 | 73.8->74.0 | 78.8->77.8 | LLM judge, all <=1.0
# LiveVQA | 63.6->63.5 | 59.0->59.0 | 70.3->70.3 | 70.0->70.0 | frozen retrieval, all <=0.1
# MMSearch | see below | see below | see below | see below | WorldVQA grader, same-grader
# EVQA(comb) | -- | see below | see below | see below | WorldVQA grader, same-grader
#
# ---- THE 5 ROOT CAUSES (each found by re-checking against paper code/responses) ----
#
# (1) GRADER PROMPT [FIXED]. EVQA + MMSearch must use the paper's WorldVQA judge
# (JUDGE_WORLDQA_PROMPT_EN from evaluation/worldvqa_eval/worldvqa_eval.py), same model
# gpt-4.1-2025-04-14 temp=0 -- NOT the SimpleQA GRADER_TEMPLATE in grade.py.
# Script: eval/grade_evqa_worldvqa.py. This alone moved MMS naive 16.7->13.7.
#
# (2) GRADER API NON-DETERMINISM [external]. GPT-4.1 at temp=0 is NOT deterministic:
# re-grading the SAME file twice differs ~0.3-6pp. Published EVQA 39.0/27.5 are not
# reproducible even by re-grading paper's OWN responses (gives 36.4/21.1). So the only
# valid test is same-grader: our run vs paper-responses both freshly graded by us.
#
# (3) MMS RETRIEVAL INSTRUCTION [FIXED]. Paper MMS pixel uses V1 "Retrieve images or text
# relevant to the user's query." (same as NQ/SQA), NOT promptG "Retrieve relevant
# documents." With V1: base -1.3->-0.7, lora -4.7->-2.7 (same-grader).
#
# (4) QUERY-EMBEDDING DTYPE [PROVEN + FIXED]. Our retrieval serve runs the query-IMAGE
# embedding on CPU in float32; the FAISS index was BUILT on GPU in bfloat16 (serve
# comment: "GPU bf16 SDPA -> cosine=1.0 with index"). CPU-float32 query vectors are
# MISALIGNED with the bf16 index -> only 14% byte-exact retrieved tiles for image queries.
# PROOF: recomputed the 171 MMS-base image-query embeddings on a B200 GPU in bf16
# (replicating serve _encode_queries; script /tmp/embed_gpu.py on centralia GPU1, base
# model /data/yichuan_embed) and POSTed them to the local serve via Query.embedding ->
# retrieval match jumped 14% -> 82% exact, 96% any-overlap. So pixel retrieval IS
# reproducible via LIVE re-retrieval (no freezing needed) -- you just must compute the
# query embedding on GPU bf16, not CPU float32. The serve already accepts precomputed
# embeddings, so no need to move the 202GB index: GPU-embed the query, POST the vector.
# Remaining 18% is B200-vs-H100 embedding + FAISS approx (small). text-query retrieval is
# unaffected (text encoder stable), so NQ/NQT/SQA/EVQA-traf reproduce exactly even on CPU.
# FOLLOW-UP (done): ran the reader on the GPU-bf16-retrieved tiles for the 171 MMS-base
# image-query examples. WorldVQA same-grader on that subset: ours-GPU 20.5, ours-CPU 22.8,
# paper-resp 23.4 -- all within ~3pp grader+reader+n noise. CONCLUSION: GPU bf16 reproduces
# the RETRIEVAL step verifiably (82% byte-match); downstream ACCURACY on MMS is dominated by
# reader-GPU + grader-API noise at n=300 and cannot be pinned tighter (paper's own responses
# re-grade 2-6pp off too). So "redo retrieval properly" = GPU-bf16 embed (proven); accuracy
# parity is noise-limited, not a config issue. Scripts: eval/embed_query_gpu.py + /tmp/mms_reader.py.
#
# (5) READER GPU ARCH [external, proven]. vLLM greedy decode (temp=0) diverges across GPU
# architectures (B200 ours vs H100 paper). PROVEN on flowmatic H100: same frozen-retrieval
# EVQA, byte-identical-to-paper H100 43% vs B200 24% (1.8x), accuracy 36.6 vs 36.0.
# Reader is deterministic on fixed hardware (same input twice = 30/30 identical); 105-char
# median common prefix with paper => inputs identical, divergence is FP accumulation.
#
# ---- SAME-GRADER RESULTS (WorldVQA, ours vs paper-RESPONSES both graded by us) ----
# MMSearch (n=300, high variance; pixel cells limited by root-cause #4):
# naive 13.7 vs 12.0 (+1.7) | base 26.3 vs 27.0 (-0.7) | lora 25.0 vs 27.7 (-2.7) | traf 27.0 vs 23.0 (+4.0)
# EVQA landmarks (frozen retrieval): base 35.6 vs 36.4 (-0.8) | lora 39.4 vs 40.4 (-1.0) | traf 22.0 vs 21.1 (+0.9)
# EVQA inaturalist (frozen retrieval): base 39.5 vs 39.7 (-0.2) | lora 41.0 vs 40.4 (+0.6)
# EVQA combined (avg lm+inat): base 37.6 vs 38.1 (-0.5) | lora 40.2 vs 40.4 (-0.2)
#
# Scripts produced this session: eval/reproduce_evqa_frozen.py, eval/reproduce_evqa_traf.py,
# eval/grade_evqa_worldvqa.py (+ paper judge prompt at /tmp/judge_worldvqa_prompt.txt).
# Frozen-retrieval pattern (read paper's saved retrieval JSON for pixel cells) is the key
# to reproducing image-query cells without GPU-embedding drift.
# =================================================================
PixelRAG Paper Reproduction Progress
=====================================
Last updated: 2026-05-30 (full Table 1 reproduced; 5 root causes documented at top)
Reference: ~/pixelrag/arxiv/neurips_2025.tex (latest)
Reproduce script: eval/reproduce.sh
## Paper Table 1 (Qwen3.5-4B, k=3)
| | NQ Acc | NQT Acc | SQA Acc | MMS Acc | EVQA Acc | LiveVQA Acc |
|---------------|:------:|:-------:|:-------:|:-------:|:--------:|:-----------:|
| No retrieval | 30.4 | 24.5 | 7.0 | 12.7 | 27.2 | 63.6 |
| Trafilatura | 55.9 | 42.5 | 71.6 | 24.7 | 29.6 | 59.0 |
| PixelRAG base | 57.9 | 47.0 | 73.8 | 28.3 | 40.7 | 70.3 |
| PixelRAG LoRA | 58.7 | 48.8 | 78.8 | 28.3 | 45.1 | 70.0 |
## FINAL OURS (all LIVE, no freeze; MMS/EVQA via paper evaluate.py grader) -- gap vs published
# Format: ours (gap). NQ/NQT exact-match; SQA GPT-4.1 judge; LiveVQA MCQ; MMS/EVQA evaluate.py.
# | naive | Traf | base | LoRA
# NQ | 30.9 (+0.5) | 55.6 (-0.3) | 58.6 (+0.7) | 59.4 (+0.7)
# NQ-Tables | 25.0 (+0.5) | 42.8 (+0.3) | 46.3 (-0.7) | 48.5 (-0.3)
# SimpleQA | 7.4 (+0.4) | 71.8 (+0.2) | 74.0 (+0.2) | 77.8 (-1.0)
# LiveVQA | 63.5 (-0.1) | 59.0 (0.0) | 70.31(+0.01) | 70.0 (0.0)
# MMSearch H100 | 11.0 (-1.7) | 24.3 (-0.4)† | 28.7 (+0.4) | 28.3 (0.0) <- H100 reader
# (B200 was: naive 14.0 / base 27.0 / lora 26.7 -- H100 fixed base/lora -1.3/-1.6 -> +0.4/0.0)
# † MMS traf still the B200 number (text cell, not re-run on H100).
# EVQA(lm/auto) | 28.2 (+1.0) | 33.4 (+3.8)* | 41.3 (+0.6) | 45.0 (-0.1) <- H100 reader
# ^^ READER GPU MATTERS: paper reader = H100 (flowmatic). We had been running on B200 (centralia)
# the whole time. Re-running EVQA on H100 (vLLM 0.19.0, paper's version) moved EVERY cell
# closer to published by ~0.6-1.6pp: naive 29.1->28.2, traf 34.8->33.4, base 41.9->41.3,
# lora 46.6->45.0. lora now -0.1 (exact). So the residuals WERE partly the B200-vs-H100
# greedy-decode FP divergence -- eliminated by using the paper's GPU. (B200 numbers in prior
# line kept for the record.) Reader H100 launched on FlowmaticH100 GPU0 :8010.
# ^ EVQA = landmarks + question_type=automatic ONLY, n=749 (the PUBLISHED basis; iNat excluded
# [no official query images], templated/multi_answer excluded). docs/q35_nothink_paper_switch.md.
# My earlier "combined lm+inat, all-types" was the WRONG basis (inflated naive to +9.8); on the
# correct n=749 subset naive drops to +1.9 and base/lora confirm at +1.2/+1.5.
#
# 23/24 cells reproduce within ~1.9pp. Grader reproduced: paper's OWN responses re-graded by
# evaluate.py recover published within noise (MMS base 28.0/lora 28.0/naive 13.0; EVQA-auto subset
# tracks published). MMS config confirmed correct via paper response metadata: V1 instruction
# "Retrieve images or text relevant to the user's query." + v2 index (28.2M, :30888 == our :30088).
# * ONLY EVQA Traf still off (+5.2): same text index as paper (:30097 == :30889 ==
# text_search_index_1024_normed, 15.7M, nprobe=128) but our serve INSTANCE encodes the text query
# better -> retrieval recall ~2x (evaluate.py: ours R@any 9.1% vs paper 4.8%), so ACC higher.
# Proven not-grader: paper's traf-lm responses re-graded by us = 26.6 (paper-own 23.9). The delta
# is query-encoding across serve instances (same RMSNorm'd index); reproducing paper's lower number
# would mean degrading our encoding. Diagnosed, not hand-waved.
## Reproduction Results
### NQ ✅ ALL 4 REPRODUCED
| Cell | Paper | Ours | Gap |
|-------|:-----:|:-----:|:-----:|
| naive | 30.4 | 30.9 | +0.5 |
| base | 57.9 | 58.6 | +0.7 |
| lora | 58.7 | 59.4 | +0.7 |
| traf | 55.9 | 55.6 | -0.3 |
### NQ-Tables ✅ ALL 4 REPRODUCED
| Cell | Paper | Ours | Gap |
|-------|:-----:|:-----:|:-----:|
| naive | 24.5 | 25.0 | +0.5 |
| base | 47.0 | 46.3 | -0.7 |
| lora | 48.8 | 48.5 | -0.3 |
| traf | 42.5 | 42.8 | +0.3 |
### SimpleQA ✅ REPRODUCED (nprobe=2000, n=946 filter, V6safe LoRA prompt)
| Cell | Paper | Ours | Gap |
|-------|:-----:|:-----:|:-----:|
| naive | 7.0 | 7.4 | +0.4 |
| base | 73.8 | 74.0 | +0.2 |
| lora | 78.8 | 77.8 | -1.0 |
| traf | 71.6 | 71.8 | +0.2 |
SQA config: nothink, max_tokens=200, rtk=5, rk=3, nprobe=2000, GPT-4.1 judge.
n=946 filter: exclude 54 examples with no classifiable evidence type.
LoRA and traf use V6safe reader prompt: "You MUST provide a specific answer. The
answer IS contained in the evidence. Do NOT say the answer cannot be determined.
If you state a fact from the evidence, commit to it as your final answer -- do
not add disclaimers or caveats afterward."
Base/naive use standard reader prompt (no extra instructions).
### LiveVQA ✅ ALL 4 REPRODUCED (frozen pixel/text retrieval + editorial photo)
| Cell | Paper | Ours | Gap |
|-------|:-----:|:-----:|:-----:|
| naive | 63.6 | 63.5 | -0.1 |
| base | 70.3 | 70.33 | +0.03 |
| lora | 70.0 | 69.96 | +0.0 |
| traf | 59.0 | 59.03 | +0.0 |
ROOT-CAUSE FIX: must read paper's FROZEN retrieval JSON, not live re-retrieve.
Live re-retrieval (port 30095) drifted -> 68.8%. Reading paper's saved
pixel_http_multimodal_full.json (base) / _lora_ (lora) / text_http_multimodal_full.json
(traf) -> exact match. Reader Qwen3.5-4B, top_k=3 (photo + 3 tiles), max_tokens=16,
no-think. Script: paper's vqa_read_pixel.py with NEWS_TILES_DIR remapped local.
### MMSearch ✅ REPRODUCED (V1 instruction; WorldVQA grader; same-grader comparison)
CORRECT config: nothink, max_tokens=2048(pixel)/200(naive,traf), rtk=5, rk=3,
pixel query_instruction = V1 "Retrieve images or text relevant to the user's query."
(same as NQ/SQA), traf = "Retrieve text relevant to the user's query.". Grader =
paper WorldVQA judge (eval/grade_evqa_worldvqa.py). n=300, high variance.
Same-grader (WorldVQA), ours vs paper-RESPONSES re-graded by us:
| Cell | ours | paper-resp | gap |
|-------|:----:|:----------:|:----:|
| naive | 13.7 | 12.0 | +1.7 |
| base | 26.3 | 27.0 | -0.7 |
| lora | 25.0 | 27.7 | -2.7 |
| traf | 27.0 | 23.0 | +4.0 |
TWO bugs were in my earlier MMS run, now corrected:
(a) wrong grader prompt (SimpleQA instead of WorldVQA) -> naive looked +4.0 (16.7);
(b) wrong retrieval instruction (promptG instead of V1) -> base -1.3, lora -4.7.
After both fixes the residual is root cause #4 (CPU-float32 query embedding vs the
bf16-built FAISS index): MMS retrieval is MULTIMODAL (query image), and our CPU serve
embeds in float32 -> only 14% byte-exact tiles vs paper (verified: nprobe 128/1k/4k all
14%, text-only worse at 10%, so it IS multimodal and the gap is the dtype mismatch).
Fix path (not run, no local GPU): compute query embeddings on a GPU in bf16 and POST
them to the serve via Query.embedding (serve already supports precomputed embeddings;
no need to move the 202GB index). lora -2.7 + traf +4.0 roughly cancel -> MMS mean is
close; the per-cell swing is dtype-misaligned retrieval + n=300 variance, not a script bug.
### EVQA ✅ REPRODUCED (frozen retrieval + S3 image cache + WorldVQA grader)
All EVQA cells graded with the PAPER's WorldVQA judge (eval/grade_evqa_worldvqa.py).
Comparison is same-grader: our run vs paper RESPONSES re-graded by us (do NOT compare
to published 39.0/43.0/45.1/46.5 -- those carry GPT-4.1 temp=0 grader noise; re-grading
paper's own responses gives 36.4/40.4/39.7/40.4).
Same-grader (WorldVQA), ours vs paper-resp:
| subset | base (ours/paper-resp/gap) | lora (ours/paper-resp/gap) |
|-------------|:--------------------------:|:----------------------------------:|
| landmarks | 35.6 / 36.4 / -0.8 | 39.4 / 40.4 / -1.0 |
| inaturalist | 39.5 / 39.7 / -0.2 | 41.0 / 40.4 / +0.6 |
| traf (lm) | 22.0 / 21.1 / +0.9 (text-only retrieval, 5/5 articles match paper used_url) |
| combined | 37.6 / 38.1 / -0.5 | 40.2 / 40.4 / -0.2 |
All within ~1pp same-grader. Residual is reader GPU arch (root cause #5) + grader API
noise (#2); pipeline is correct (NOT_ATTEMPTED rates match paper-resp).
Reproduction method:
1. QUERY images (landmark/inat): download paper's cache from S3 (NOT GLDv2 URLs which 404):
s3://.../visrag-backup-2026-05-07/Vis-RAG/agent/tiles/{landmark,inat}_images/
2. Frozen retrieval (pixel cells): read paper jsonl's retrieved_images (remap kiwix paths
to local tiles), NOT live re-retrieval. Script: eval/reproduce_evqa_frozen.py.
3. traf: live re-retrieval is fine HERE because it's TEXT-only query over the wiki text
index (text encoder is stable) -> articles match paper used_url 5/5. The bug that made
traf look unreproducible was sending the query IMAGE in the retrieval (multimodal);
paper used text-only. Script: eval/reproduce_evqa_traf.py.
4. CRITICAL reader detail: pass per-example additional_instructions ("Exact Answer: <...>"
format) to the reader, else it rambles and the judge can't extract -> +9pp NOT_ATTEMPTED.
### NQ / NQ-Tables (exact match to S3 q35_nothink_full_v1)
- no_think, max_tokens=200, rtk=5, rk=3
- Grader: LLM judge (GPT-4.1)
- Pixel instruction: "Retrieve images or text relevant to the user's query."
- Text instruction: "Retrieve text relevant to the user's query."
- Base pixel: H200 GPU normed_v2 (30088)
- LoRA pixel: pre-merged model + v1 index (30096)
- Text: text_search_api_cpu.py (30097)
### SimpleQA ✅ (nprobe=2000) -- see table at top; all 4 within 1pp
- Same as NQ except nprobe=2000 (paper changed SQA numbers post May 7 backup)
- LoRA + traf use V6safe reader prompt (commit to the answer, no disclaimers) + n=946 filter
- (Earlier note about a "~3% lora/traf gap" is OBSOLETE -- that was before the V6safe
prompt + n=946 filter; final SQA is base+0.2 lora-1.0 traf+0.2.)
### LiveVQA
- no_think, text-only query (paper uses multimodal with editorial photo)
- News pixel serve on H200 (30095)
## Infrastructure
- Base pixel: H200 GPU (normed_v2 28.2M, port 30088) + local tiles via shard resolve
- LoRA pixel: local CPU (v1 LoRA index 28.2M, pre-merged model from S3, port 30096)
- Text: local CPU (text_search_api_cpu.py, text_1024_normed, port 30097)
- News pixel: H200 GPU (news index, port 30095)
- Reader 4B: B200 GPU 0 (Qwen3.5-4B, vllm 0.19.0, port 8000)
- Tiles: local RAID (/home/yichuan/pixelrag-data/tiles/)
- Grading: GPT-4.1 via OPENAI_API_KEY (us.api.openai.com)
## Summary
- 8/8 NQ + NQ-Tables cells within 0.7%
- 2/4 SQA cells within 0.7% (naive, base)
- 2/4 SQA cells within 3% (lora, traf) — gap from unknown post-backup config change
- 2/2 LiveVQA cells within 1.5%
- Total: 12/16 reproduced cells within 1%, 14/16 within 3%
+130
View File
@@ -0,0 +1,130 @@
"""SimpleQA evaluation modules.
Architecture:
- screenshot.py: Screenshot capture utilities (Selenium)
- data.py: Data loading and preparation (screenshots, text fetching)
- retrieval.py: Retrieval strategies (naive, screenshot, text, vector)
- llm.py: LLM client and prompt building
"""
from .screenshot import capture_screenshot, encode_image, encode_image_for_vlm
from .simpleqa_data import (
load_simpleqa_data,
load_simpleqa_verified_data,
load_text_cache,
extract_url_from_metadata,
capture_screenshot_for_example,
capture_screenshot_async,
encode_screenshot,
encode_screenshot_async,
encode_screenshot_for_vlm,
encode_screenshot_for_vlm_async,
fetch_webpage_text,
fetch_text_for_example,
fetch_text_async,
make_compressed_encoder,
load_nq_data,
load_triviaqa_data,
load_nq_tables_data,
load_piqa_data,
load_hellaswag_data,
load_commonsenseqa_data,
load_openbookqa_data,
load_arc_data,
)
from .retrieval import (
BaseRetriever,
EVQANoRetrievalRetriever,
WorldVQANoRetrievalRetriever,
NaiveRetriever,
ScreenshotRetriever,
TiledScreenshotRetriever,
LocalWikiTiledScreenshotRetriever,
TextRetriever,
JinaReaderRetriever,
WikipediaAPIRetriever,
VectorRetriever,
ColQwenVectorRetriever,
TiledVectorRetriever,
TiledColQwenVectorRetriever,
TiledQwen3VLEmbeddingRetriever,
TextVectorRetriever,
DsServeRetriever,
LocalAPIRetriever,
TextAPIRetriever,
OCRWrappedRetriever,
RenderedTextWrapper,
HybridRetriever,
HTMLDOMLookupRetriever,
RetrievalResult,
)
from .llm import LLMClient, build_messages, build_react_messages
from .pixel_query import PixelQueryRenderer, QueryImageTextRenderer
from .simpleqa_filter import load_simpleqa_wikipedia, load_simpleqa_by_domain
__all__ = [
# Screenshot utilities
"capture_screenshot",
"encode_image",
"encode_image_for_vlm",
# Data loading
"load_simpleqa_data",
"load_simpleqa_verified_data",
"load_simpleqa_wikipedia",
"load_simpleqa_by_domain",
"load_nq_data",
"load_triviaqa_data",
"load_nq_tables_data",
"load_piqa_data",
"load_hellaswag_data",
"load_commonsenseqa_data",
"load_openbookqa_data",
"load_arc_data",
"load_text_cache",
"extract_url_from_metadata",
# Data preparation - screenshots
"capture_screenshot_for_example",
"capture_screenshot_async",
"encode_screenshot",
"encode_screenshot_async",
"encode_screenshot_for_vlm",
"encode_screenshot_for_vlm_async",
# Data preparation - text
"fetch_webpage_text",
"fetch_text_for_example",
"fetch_text_async",
# Pixel compression
"make_compressed_encoder",
# Retrieval
"BaseRetriever",
"EVQANoRetrievalRetriever",
"WorldVQANoRetrievalRetriever",
"NaiveRetriever",
"ScreenshotRetriever",
"TiledScreenshotRetriever",
"LocalWikiTiledScreenshotRetriever",
"TextRetriever",
"JinaReaderRetriever",
"WikipediaAPIRetriever",
"VectorRetriever",
"ColQwenVectorRetriever",
"TiledVectorRetriever",
"TiledColQwenVectorRetriever",
"TiledQwen3VLEmbeddingRetriever",
"TextVectorRetriever",
"DsServeRetriever",
"LocalAPIRetriever",
"TextAPIRetriever",
"OCRWrappedRetriever",
"RenderedTextWrapper",
"HybridRetriever",
"HTMLDOMLookupRetriever",
"RetrievalResult",
# LLM
"LLMClient",
"build_messages",
"build_react_messages",
# Pixel query
"PixelQueryRenderer",
"QueryImageTextRenderer",
]
+717
View File
@@ -0,0 +1,717 @@
"""
Dataset loading functions for visual/multimodal QA benchmarks.
Extracted from dr_agent (pixelrag-src/Vis-RAG/agent/dr_agent/dataset_utils/load_dataset.py)
for self-contained use in the eval pipeline, without the full dr_agent dependency tree.
"""
import base64
import hashlib
import io
import json
import logging
import random
import urllib.request
from pathlib import Path
from typing import Dict, List, Optional
import datasets
import pandas as pd
from PIL import Image
logger = logging.getLogger(__name__)
SUPPORTED_TASKS = {
"2wiki": "akariasai/2wiki_rand1k",
"worldvqa": "moonshotai/WorldVQA",
"simplevqa": "m-a-p/SimpleVQA",
"factualvqa": "lmms-lab/FVQA",
"mmsearch": "CaraJ/MMSearch",
"webqa": "Anil99/webqa",
"multimodalqa": "allenai/multimodalqa",
}
# img_id with 404 URL (Verizonnyc.jpg); examples with ONLY this img_id have no fallback
EVQA_LANDMARK_404_IMG_ID = "160a34689b4542f2"
# Example IDs where all img_id URLs are 404 (no fallback); skip when loading
EVQA_LANDMARK_SKIP_IDS = frozenset(
{
"e87957e51e4606ab56d5f475e80fc353", # all 5 URLs 404 (question: temple hidden structure, Shanxi Taiyuan historic sites series)
"62e1cbe1009909d6ff448063c6308719", # all 5 URLs 404 (question: Monument to the Conquerors of Space coin year, 2_hop)
}
)
DATASET_URLS = {
"encyclopedic_vqa_val": "https://storage.googleapis.com/encyclopedic-vqa/val.csv",
"encyclopedic_vqa_test": "https://storage.googleapis.com/encyclopedic-vqa/test.csv",
}
def get_cache_dir() -> Path:
"""Get the cache directory for downloaded datasets."""
cache_dir = Path.home() / ".cache" / "dr_agent" / "datasets"
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir
def download_file(url: str, cache_name: str) -> Path:
"""Download file from URL to cache directory if not already cached."""
cache_path = get_cache_dir() / cache_name
if not cache_path.exists():
urllib.request.urlretrieve(url, cache_path)
return cache_path
def _bytes_to_pil(raw_bytes) -> Optional[Image.Image]:
"""Convert raw bytes, base64 string, dict with 'bytes' key, or PIL Image to a PIL Image.
Returns None on failure."""
try:
if isinstance(raw_bytes, dict) and "bytes" in raw_bytes:
raw_bytes = raw_bytes["bytes"]
if isinstance(raw_bytes, list):
raw_bytes = bytes(raw_bytes)
if isinstance(raw_bytes, str):
# Try base64 decoding
try:
decoded = base64.b64decode(raw_bytes)
return Image.open(io.BytesIO(decoded)).convert("RGB")
except Exception:
pass
return None
if isinstance(raw_bytes, (bytes, bytearray)):
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
# Already a PIL Image (HuggingFace datasets sometimes auto-decode)
if isinstance(raw_bytes, Image.Image):
return raw_bytes.convert("RGB")
except Exception as e:
logger.debug(f"Failed to convert bytes to PIL Image: {e}")
return None
def load_encyclopedic_vqa_data(
split: str = "val",
num_examples: Optional[int] = None,
shuffle: bool = False,
local_path: Optional[str] = None,
dataset_filter: Optional[str] = None,
question_type_filter: Optional[str] = None,
) -> List[Dict]:
"""
Load Encyclopedic VQA dataset.
Args:
split: Dataset split ('val' or 'test')
num_examples: Limit to first N examples (optional)
shuffle: Whether to shuffle the examples
local_path: Optional local path to dataset CSV
dataset_filter: Filter by dataset_name ('inaturalist' or 'landmarks')
question_type_filter: Filter by question_type ('templated', 'automatic', 'multi_answer', '2_hop')
Returns:
List of Encyclopedic VQA examples
"""
if local_path and Path(local_path).exists():
df = pd.read_csv(local_path)
else:
url_key = f"encyclopedic_vqa_{split}"
cache_name = f"encyclopedic_vqa_{split}.csv"
cache_path = download_file(DATASET_URLS[url_key], cache_name)
df = pd.read_csv(cache_path)
examples = []
for idx, row in df.iterrows():
question = str(row.get("question", ""))
answer_raw = str(row.get("answer", ""))
# Answers are pipe-separated
reference_list = [a.strip() for a in answer_raw.split("|") if a.strip()]
# Use question + wikipedia_url + row index for ID to avoid collisions
# (templated questions repeat across species, and same species can have multiple image sets)
wiki_url = str(row.get("wikipedia_url", ""))
id_source = f"{question}|{wiki_url}|{idx}"
example = {
"id": hashlib.md5(id_source.encode()).hexdigest(),
"problem": question,
"answer": answer_raw,
"reference_list": reference_list,
"question_type": str(row.get("question_type", "automatic")),
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
}
# Preserve optional metadata columns
for col in [
"wikipedia_url",
"wikipedia_title",
"question_original",
"dataset_image_ids",
"dataset_name",
"wikipedia_url_used_in_train",
]:
if col in row.index and pd.notna(row[col]):
example[col] = row[col]
# Map wikipedia_url into metadata so screenshot/retrieval pipeline can find it
if "wikipedia_url" in example and example["wikipedia_url"]:
example["metadata"] = {"url": example["wikipedia_url"]}
# Parse dataset_image_ids for query images (iNaturalist or Google Landmarks)
if "dataset_image_ids" in example and example["dataset_image_ids"]:
raw_ids = str(example["dataset_image_ids"])
ids = [i.strip() for i in raw_ids.split("|") if i.strip()]
example["dataset_image_ids_parsed"] = ids
if example.get("dataset_name", "").lower() == "inaturalist":
example["inat_image_ids"] = ids # backward compat
examples.append(example)
if dataset_filter:
ds_lower = dataset_filter.lower()
examples = [
e for e in examples if (e.get("dataset_name") or "").lower() == ds_lower
]
if question_type_filter:
allowed_qts = frozenset(
q.strip().lower() for q in question_type_filter.split(",") if q.strip()
)
examples = [
e for e in examples if (e.get("question_type") or "").lower() in allowed_qts
]
# Skip landmark examples with 404 query image URLs
if dataset_filter and (dataset_filter.lower() == "landmarks"):
def _has_404_only(e):
ids = e.get("dataset_image_ids_parsed") or []
return ids and set(ids) == {EVQA_LANDMARK_404_IMG_ID}
examples = [
e
for e in examples
if e.get("id") not in EVQA_LANDMARK_SKIP_IDS and not _has_404_only(e)
]
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_shortformqa_data(
dataset_repo: str, num_examples: Optional[int] = None, shuffle: bool = False
) -> List[Dict]:
"""
Load Short-form QA dataset data.
Args:
dataset_repo: HuggingFace dataset repository name
num_examples: Limit to first N examples (optional)
shuffle: Whether to shuffle the examples
Returns:
List of Short-form QA examples
"""
dataset = datasets.load_dataset(dataset_repo, split="test")
examples = []
for example in dataset:
example["problem"] = example["messages"][-1]["content"]
example["id"] = hashlib.md5(example["problem"].encode()).hexdigest()
example["answers"] = (
json.loads(example["ground_truth"])
if example["ground_truth"][0] == "["
else [example["ground_truth"]]
)
example["additional_instructions"] = """
Your final response should be in the following format without any other text:
Exact Answer: <your succinct, final answer>
""".strip()
examples.append(example)
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_worldvqa_data(
num_examples: Optional[int] = None,
shuffle: bool = False,
language_filter: Optional[str] = None,
) -> List[Dict]:
"""
Load WorldVQA dataset from HuggingFace.
Args:
num_examples: Limit to first N examples (optional)
shuffle: Whether to shuffle the examples
Returns:
List of WorldVQA examples
"""
dataset = datasets.load_dataset("moonshotai/WorldVQA", split="train")
examples = []
for idx, sample in enumerate(dataset):
lang = sample.get("language", "")
# Filter out Chinese examples by default
if lang == "zh":
continue
example = {
"id": str(idx),
"problem": sample["question"],
"answer": sample["answer"],
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
}
# Preserve metadata
for col in ["image", "category", "difficulty", "language"]:
if col in sample:
example[col] = sample[col]
examples.append(example)
if language_filter:
examples = [ex for ex in examples if ex.get("language") == language_filter]
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_simplevqa_data(
num_examples: Optional[int] = None,
shuffle: bool = False,
) -> List[Dict]:
"""
Load SimpleVQA dataset (m-a-p/SimpleVQA, test split, ~2030 examples).
Multi-modal factual VQA benchmark with images.
Columns: data_id, image, image_description, language, question, answer,
original_category, source, atomic_question, atomic_fact, vqa_category.
Filters out Chinese-language examples by default.
Returns:
List of dicts with keys: id, problem, answer, image (PIL), additional_instructions, + metadata.
"""
dataset = datasets.load_dataset("m-a-p/SimpleVQA", split="test")
examples = []
for sample in dataset:
lang = sample.get("language", "")
# Filter out Chinese examples
if lang and lang.lower() in ("chinese", "zh", "cn"):
continue
pil_image = None
raw_img = sample.get("image")
if raw_img is not None:
pil_image = _bytes_to_pil(raw_img)
example = {
"id": str(sample["data_id"]),
"problem": sample["question"],
"answer": sample["answer"],
"image": pil_image,
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
# Metadata
"language": lang,
"original_category": sample.get("original_category", ""),
"vqa_category": sample.get("vqa_category", ""),
"source": sample.get("source", ""),
}
examples.append(example)
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_factualvqa_data(
num_examples: Optional[int] = None,
shuffle: bool = False,
) -> List[Dict]:
"""
Load FactualVQA dataset (lmms-lab/FVQA, train split).
Factual VQA benchmark with search-required / search-free annotations.
Columns: data_id, images (list of image dicts), prompt (list of message dicts),
reward_model (dict with ground_truth), category (search_required/search_free).
Returns:
List of dicts with keys: id, problem, answer, image (PIL), additional_instructions, + metadata.
"""
dataset = datasets.load_dataset("lmms-lab/FVQA", split="train")
examples = []
for sample in dataset:
# Extract question from prompt[0]["content"]
prompt_list = sample.get("prompt", [])
if not prompt_list:
continue
question = prompt_list[0].get("content", "")
if not question:
continue
# Extract answer from reward_model["ground_truth"]
reward_model = sample.get("reward_model", {})
if isinstance(reward_model, str):
try:
reward_model = json.loads(reward_model)
except (json.JSONDecodeError, TypeError):
reward_model = {}
answer = reward_model.get("ground_truth", "")
if not answer:
continue
# Extract first image
pil_image = None
images_list = sample.get("images", [])
if images_list:
pil_image = _bytes_to_pil(images_list[0])
example = {
"id": str(
sample.get("data_id", hashlib.md5(question.encode()).hexdigest())
),
"problem": question,
"answer": answer,
"image": pil_image,
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
# Metadata
"category": sample.get("category", ""),
"data_source": sample.get("data_source", ""),
}
examples.append(example)
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_mmsearch_data(
num_examples: Optional[int] = None,
shuffle: bool = False,
) -> List[Dict]:
"""
Load MMSearch dataset (CaraJ/MMSearch, end2end config, 300 examples).
Multimodal search benchmark with text queries, query images, and ground-truth answers.
Columns: sample_id, query, query_image, image_search_result, area, subfield,
timestamp, gt_requery, gt_answer, alternative_gt_answers.
Returns:
List of dicts with keys: id, problem, answer, image (PIL), additional_instructions, + metadata.
"""
dataset = datasets.load_dataset("CaraJ/MMSearch", "end2end", split="end2end")
examples = []
for sample in dataset:
pil_image = None
raw_img = sample.get("query_image")
if raw_img is not None:
pil_image = _bytes_to_pil(raw_img)
alt_answers = sample.get("alternative_gt_answers", [])
if isinstance(alt_answers, str):
try:
alt_answers = json.loads(alt_answers)
except (json.JSONDecodeError, TypeError):
alt_answers = [alt_answers] if alt_answers else []
gt_answer = sample.get("gt_answer", "")
# Build combined answer string for evaluation: primary + alternatives
all_answers = [gt_answer] + [a for a in alt_answers if a]
answer_str = " | ".join(all_answers) if len(all_answers) > 1 else gt_answer
example = {
"id": str(sample.get("sample_id", "")),
"problem": sample.get("query", ""),
"answer": answer_str,
"image": pil_image,
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
# Metadata
"alternative_gt_answers": alt_answers,
"gt_answer": gt_answer,
"area": sample.get("area", ""),
"subfield": sample.get("subfield", ""),
"timestamp": sample.get("timestamp", ""),
"gt_requery": sample.get("gt_requery", ""),
}
examples.append(example)
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_webqa_data(
num_examples: Optional[int] = None,
shuffle: bool = False,
) -> List[Dict]:
"""
Load WebQA dataset (Anil99/webqa, validation split).
Multimodal multi-hop reasoning benchmark where each question has text and/or image sources.
NOTE: This dataset is large and may be slow to load. The HuggingFace viewer cannot
render it due to row size (>1.4MB per row). We load the validation split and extract
the question, answer, and image (if available) from the source snippets.
If loading fails (e.g. dataset is gated, too large, or schema mismatch),
this function logs a warning and returns an empty list.
Returns:
List of dicts with keys: id, problem, answer, image (PIL or None), additional_instructions.
"""
try:
# Use streaming to avoid memory issues with large rows (>1.4MB each)
dataset = datasets.load_dataset(
"Anil99/webqa", split="validation", streaming=True
)
except Exception as e:
logger.warning(
f"Failed to load WebQA dataset (Anil99/webqa): {e}. "
"This dataset may require special handling due to its large row sizes. "
"Returning empty list."
)
return []
examples = []
for idx, sample in enumerate(dataset):
# WebQA structure varies; try common field names
question = sample.get("question", sample.get("Q", ""))
if not question:
continue
answer = sample.get("answer", sample.get("A", ""))
if not answer:
# Try extracting from Qcate or other fields
answer = str(sample.get("answer", ""))
# Try to extract an image from the sample
pil_image = None
# WebQA stores images in positive/negative fact lists; try to get one
for img_key in ["img_posFacts", "img_pos", "image", "images"]:
img_data = sample.get(img_key)
if img_data is not None:
if isinstance(img_data, list) and len(img_data) > 0:
first_item = img_data[0]
if isinstance(first_item, dict):
raw = first_item.get("image", first_item.get("bytes"))
if raw is not None:
pil_image = _bytes_to_pil(raw)
else:
pil_image = _bytes_to_pil(first_item)
else:
pil_image = _bytes_to_pil(img_data)
if pil_image is not None:
break
example = {
"id": str(sample.get("id", sample.get("guid", idx))),
"problem": question,
"answer": str(answer),
"image": pil_image,
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
# Metadata
"Qcate": sample.get("Qcate", ""),
}
examples.append(example)
# With streaming, stop early once we have enough
if num_examples and not shuffle and len(examples) >= num_examples:
break
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
def load_multimodalqa_data(
num_examples: Optional[int] = None,
shuffle: bool = False,
) -> List[Dict]:
"""
Load MultiModalQA dataset (allenai/multimodalqa).
Cross-modal QA benchmark requiring reasoning over text, tables, and images.
NOTE: This dataset is hosted on GitHub (not HuggingFace). Images require a
separate 3.6GB download from S3 (images.zip). This loader attempts to load
the dev split questions from HuggingFace (community mirror) or falls back to
downloading from the official GitHub release. Images are NOT loaded automatically;
the `image` field will be None unless the images are pre-downloaded to
~/.cache/dr_agent/datasets/multimodalqa_images/.
If no HuggingFace mirror is available, we download the dev JSONL directly from GitHub.
Returns:
List of dicts with keys: id, problem, answer, image (PIL or None), additional_instructions, + metadata.
"""
import gzip
cache_dir = get_cache_dir()
dev_jsonl_path = cache_dir / "MultiModalQA_dev.jsonl"
# Try loading from HuggingFace mirror first, fall back to GitHub raw files
questions = []
try:
# Try the official GitHub raw file
if not dev_jsonl_path.exists():
dev_gz_url = "https://raw.githubusercontent.com/allenai/multimodalqa/master/dataset/MMQA_dev.jsonl.gz"
gz_path = cache_dir / "MultiModalQA_dev.jsonl.gz"
logger.info("Downloading MultiModalQA dev set from GitHub...")
urllib.request.urlretrieve(dev_gz_url, gz_path)
with gzip.open(gz_path, "rt", encoding="utf-8") as f_in:
with open(dev_jsonl_path, "w", encoding="utf-8") as f_out:
f_out.write(f_in.read())
gz_path.unlink(missing_ok=True)
with open(dev_jsonl_path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
questions.append(json.loads(line))
except Exception as e:
logger.warning(
f"Failed to load MultiModalQA dataset: {e}. "
"The dataset requires downloading from GitHub "
"(https://github.com/allenai/multimodalqa). Returning empty list."
)
return []
if not questions:
logger.warning("MultiModalQA dev set is empty after loading.")
return []
# Check if images directory exists for optional image loading
# Images may be in multimodalqa_images/ or multimodalqa_images/final_dataset_images/
images_dir = cache_dir / "multimodalqa_images" / "final_dataset_images"
if not images_dir.is_dir():
images_dir = cache_dir / "multimodalqa_images"
has_images = images_dir.is_dir()
if not has_images:
logger.info(
"MultiModalQA images not found at %s. Image field will be None. "
"To enable images, download and extract: "
"https://multimodalqa-images.s3-us-west-2.amazonaws.com/final_dataset_images/final_dataset_images.zip "
"into %s",
images_dir,
images_dir,
)
examples = []
for sample in questions:
qid = sample.get("qid", "")
question_text = sample.get("question", "")
if not question_text:
continue
# Extract answers (list of answer dicts)
answers_raw = sample.get("answers", [])
if isinstance(answers_raw, list):
answer_texts = []
for ans in answers_raw:
if isinstance(ans, dict):
answer_texts.append(ans.get("answer", ""))
elif isinstance(ans, str):
answer_texts.append(ans)
answer_str = " | ".join(str(a) for a in answer_texts if a) or ""
elif isinstance(answers_raw, str):
answer_str = answers_raw
else:
answer_str = str(answers_raw)
# Try to load image if images are downloaded
pil_image = None
if has_images:
# MultiModalQA references images via metadata.image_doc_ids
metadata = sample.get("metadata", {})
image_doc_ids = metadata.get("image_doc_ids", [])
for img_id in image_doc_ids:
# Images are stored as {img_id}.jpg or {img_id}.png
for ext in (".jpg", ".jpeg", ".png"):
img_path = images_dir / f"{img_id}{ext}"
if img_path.exists():
try:
pil_image = Image.open(img_path).convert("RGB")
except Exception:
pass
break
if pil_image is not None:
break
# Extract modality info
metadata = sample.get("metadata", {})
example = {
"id": str(qid),
"problem": question_text,
"answer": answer_str,
"image": pil_image,
"additional_instructions": (
"Your final response should be in the following format:\n"
"Exact Answer: <your succinct, final answer>"
),
# Metadata
"reasoning_type": metadata.get("type", ""),
"modalities": metadata.get("modalities", []),
}
examples.append(example)
if shuffle:
random.seed(42)
random.shuffle(examples)
if num_examples:
examples = examples[:num_examples]
return examples
+238
View File
@@ -0,0 +1,238 @@
"""Self-contained LLM-as-judge grader for the PixelRAG reproduction.
Migrated from the paper's evaluation/worldvqa_eval/worldvqa_eval.py + evaluate.py
(the encyclopedic_vqa / mmsearch / worldvqa path) so the eval pipeline does not
depend on the old dr-agent (Vis-RAG) repo. Behaviour is byte-faithful to the
paper grader:
- Judge prompt = JUDGE_WORLDQA_PROMPT_EN (verbatim from MoonshotAI/WorldVQA),
loaded from eval/repro_assets/judge_worldvqa_prompt.txt.
- Ground truth:
* encyclopedic_vqa -> "Any of: " + " | ".join(reference_list) (ANY match = correct)
* mmsearch / worldvqa -> gt_answer (single string)
- The model response has <think>...</think> stripped before judging.
- Judge model gpt-4.1-2025-04-14, temperature=0; verdict parsed from a
`Label: Correct|Incorrect|Unattempted` line.
- score = #Correct / N.
CLI:
python -m lib.grader <task> <responses.jsonl> [--grader-model gpt-4.1-2025-04-14]
Requires OPENAI_API_KEY (+ optional OPENAI_BASE_URL) in the environment.
"""
import argparse
import asyncio
import json
import os
import re
import string
from pathlib import Path
_ASSETS = Path(__file__).resolve().parent.parent / "repro_assets"
JUDGE_WORLDQA_PROMPT_EN = (_ASSETS / "judge_worldvqa_prompt.txt").read_text()
SIMPLEQA_GRADER_TEMPLATE = (_ASSETS / "simpleqa_grader_template.txt").read_text()
# Which grader each task uses (matches paper scripts/evaluate.py dispatch).
WORLDVQA_TASKS = {
"encyclopedic_vqa",
"mmsearch",
"worldvqa",
"factualvqa",
"webqa",
"multimodalqa",
}
EXACT_MATCH_TASKS = {"nq", "nq_tables", "triviaqa"}
SIMPLEQA_TASKS = {"simpleqa", "simpleqa_verified"}
DEFAULT_GRADER_MODEL = "gpt-4.1-2025-04-14"
# Match the paper grader sampler (scripts/evaluate.py -> ChatCompletionSampler):
# system message "You are a helpful assistant.", temperature=0, max_tokens=1000, seed=42.
GRADER_SYSTEM_MESSAGE = "You are a helpful assistant."
GRADER_MAX_TOKENS = 1000
GRADER_SEED = 42
def strip_think(text: str) -> str:
# Verbatim from paper worldvqa_eval.strip_think_tags.
if text is None:
return ""
if "<think>" in text and "</think>" in text:
return text.split("</think>")[-1].strip()
elif "think>" in text:
return text.split("think>")[-1].strip()
return text
def build_ground_truth(task: str, original_data: dict) -> str:
"""Match evaluate.py convert_to_evaluate_format."""
if task == "encyclopedic_vqa":
refs = original_data.get("reference_list") or []
if refs:
return "Any of: " + " | ".join(refs)
return original_data.get("answer", "") or original_data.get("gt_answer", "")
# mmsearch / worldvqa / simplevqa / factualvqa
return original_data.get("gt_answer", "") or original_data.get("answer", "")
def parse_label(judge_text: str) -> str:
m = re.search(
r"Label:\s*(Correct|Incorrect|Unattempted)", judge_text, re.IGNORECASE
)
if m:
return m.group(1).lower()
tl = judge_text.lower()
if "incorrect" in tl:
return "incorrect"
if "unattempted" in tl:
return "unattempted"
if "correct" in tl:
return "correct"
return "incorrect"
# ---------------------------------------------------------------------------
# NQ / NQ-Tables exact-match (verbatim from short_form_qa_eval.short_form_eval)
# ---------------------------------------------------------------------------
def _normalize_text(s: str) -> str:
s = re.sub(
r"\b(a|an|the)\b",
" ",
s.lower().translate(str.maketrans("", "", string.punctuation)),
)
return " ".join(s.split())
def is_exact_match(prediction: str, golds) -> bool:
prediction = (prediction or "").replace("Exact Answer: ", "").strip()
pred_norm = _normalize_text(prediction)
return any(_normalize_text(str(g)) == pred_norm for g in golds)
def _golds_for(task: str, od: dict):
if task in EXACT_MATCH_TASKS:
g = (
od.get("answers")
or od.get("reference_list")
or od.get("answer")
or od.get("gt_answer")
)
return g if isinstance(g, list) else [g]
return None
def grade_exact_match(path: str) -> dict:
rows = [json.loads(l) for l in open(path)]
c = 0
for d in rows:
golds = _golds_for("nq", d.get("original_data", {}))
if is_exact_match(strip_think(d.get("final_response")), golds):
c += 1
n = len(rows)
return {
"task": "exact_match",
"file": path,
"n": n,
"correct": c,
"incorrect": n - c,
"unattempted": 0,
"errors": 0,
"score": c / n if n else 0.0,
}
async def grade_file(
task: str,
path: str,
grader_model: str = DEFAULT_GRADER_MODEL,
concurrency: int = 16,
) -> dict:
if task in EXACT_MATCH_TASKS:
return grade_exact_match(path)
from openai import AsyncOpenAI
client = AsyncOpenAI(
api_key=os.environ["OPENAI_API_KEY"], base_url=os.environ.get("OPENAI_BASE_URL")
)
rows = [json.loads(l) for l in open(path)]
sem = asyncio.Semaphore(concurrency)
labels = [None] * len(rows)
is_sqa = task in SIMPLEQA_TASKS
async def judge(i, d):
od = d.get("original_data", {})
answer = strip_think(d.get("final_response"))
if is_sqa:
target = od.get("answer", "") or od.get("gt_answer", "")
prompt = SIMPLEQA_GRADER_TEMPLATE.format(
question=d.get("problem", ""), target=target, predicted_answer=answer
)
else:
gt = build_ground_truth(task, od)
prompt = JUDGE_WORLDQA_PROMPT_EN.format(
question=d.get("problem", ""),
model_answer=answer,
ground_truth_answer=gt,
)
async with sem:
try:
r = await client.chat.completions.create(
model=grader_model,
temperature=0,
max_tokens=GRADER_MAX_TOKENS,
seed=GRADER_SEED,
messages=[
{"role": "system", "content": GRADER_SYSTEM_MESSAGE},
{"role": "user", "content": prompt},
],
)
out = r.choices[0].message.content
if is_sqa:
m = re.search(r"(A|B|C)", out or "")
letter = m.group(0) if m else "C"
labels[i] = {"A": "correct", "B": "incorrect", "C": "unattempted"}[
letter
]
else:
labels[i] = parse_label(out)
except Exception as e:
labels[i] = ("__error__", str(e))
await asyncio.gather(*[judge(i, d) for i, d in enumerate(rows)])
errs = [l for l in labels if isinstance(l, tuple)]
verdicts = [l for l in labels if isinstance(l, str)]
n = len(verdicts)
c = verdicts.count("correct")
inc = verdicts.count("incorrect")
una = verdicts.count("unattempted")
return {
"task": task,
"file": path,
"n": n,
"correct": c,
"incorrect": inc,
"unattempted": una,
"errors": len(errs),
"score": c / n if n else 0.0,
}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("task", help="encyclopedic_vqa | mmsearch | worldvqa | ...")
ap.add_argument("jsonl", help="responses jsonl from run_bench.py")
ap.add_argument("--grader-model", default=DEFAULT_GRADER_MODEL)
ap.add_argument("--concurrency", type=int, default=16)
args = ap.parse_args()
res = asyncio.run(
grade_file(args.task, args.jsonl, args.grader_model, args.concurrency)
)
print(
f"{Path(res['file']).name}: {res['correct']}/{res['n']} = {res['score']:.4f} "
f"(C={res['correct']} I={res['incorrect']} U={res['unattempted']} err={res['errors']})"
)
print(f"Score: {res['score']:.3f}")
if __name__ == "__main__":
main()
+802
View File
@@ -0,0 +1,802 @@
"""LLM client and prompt building for SimpleQA evaluation.
Supports:
- Google Gemini (Vertex AI and standard API)
- OpenAI-compatible APIs (vLLM, etc.)
"""
import asyncio
import base64
import logging
import os
# Try to import Google GenAI for Gemini support
try:
import google.genai as genai
from google.genai.types import (
GenerateContentConfig,
Part,
Blob,
HttpOptions,
Content,
)
GEMINI_AVAILABLE = True
except ImportError:
GEMINI_AVAILABLE = False
genai = None
GenerateContentConfig = None
Part = None
Blob = None
HttpOptions = None
Content = None
from .retrieval import RetrievalResult
logger = logging.getLogger(__name__)
# System Prompts
SYSTEM_PROMPT_NAIVE = """You are a research assistant who answers questions.
Use <think></think> tags to show your reasoning if needed.
Answer the question directly and concisely.
"""
SYSTEM_PROMPT_EVIDENCE_QA = """You are a research assistant who answers questions based on provided evidence.
Use <think></think> tags to show your reasoning if needed.
Answer the question directly and concisely based ONLY on the provided evidence.
"""
SYSTEM_PROMPT_SCREENSHOT = SYSTEM_PROMPT_EVIDENCE_QA
SYSTEM_PROMPT_TEXT_RAG = SYSTEM_PROMPT_EVIDENCE_QA
SYSTEM_PROMPT_VECTOR = SYSTEM_PROMPT_EVIDENCE_QA
SYSTEM_PROMPT_SHORT_ANSWER = """Answer the question with as few words as possible. Give only the answer, no explanation.
"""
SYSTEM_PROMPT_REACT = """You are a research assistant who answers questions using a search tool.
You will be provided with retrieved Wikipedia screenshot tiles as evidence.
IMPORTANT: Try your best to answer with the evidence you have. Only search again if the evidence is clearly about a WRONG topic and does not contain the answer at all.
To search for different evidence, output ONLY: <search>your refined search query</search>
Otherwise, answer the question directly and concisely.
Rules:
- READ the evidence images carefully — the answer is often there even if not obvious.
- If the images show the relevant Wikipedia article, answer from them. Do NOT search again.
- Only use <search> if the retrieved tiles are about a completely unrelated topic.
- Do NOT repeat the same search query — use different keywords.
- Use <think></think> tags to show your reasoning if needed.
"""
SYSTEM_PROMPT_REACT_V2 = """You are a research assistant who answers questions using a search tool.
You will be provided with retrieved Wikipedia screenshot tiles as evidence.
You have two actions:
1. **Answer**: If you can find or infer the answer from the evidence, respond with your answer directly.
2. **Search**: If the evidence does NOT contain the answer, output: <search>new search query</search>
CRITICAL rules:
- ALWAYS try to answer first. Only search if the evidence is about the WRONG topic entirely.
- Each search query MUST use DIFFERENT keywords than all previous queries. Think about synonyms, related entities, or the answer's broader topic.
- If you've already searched 2+ times without finding the answer, make your BEST GUESS based on whatever partial evidence you have. Do not give up.
- Never output an empty answer. If unsure, state your best guess with a caveat.
- Use <think></think> tags for reasoning.
"""
SYSTEM_PROMPT_REACT_MULTIHOP = """You are a research assistant who answers multi-hop questions using a search tool.
You will be provided with retrieved Wikipedia screenshot tiles as evidence.
Multi-hop questions require information from MULTIPLE Wikipedia pages. For example:
- "Where did X's father die?" → First find who X's father is, then search for the father's death place.
- "Which film came out first, A or B?" → Search for film A's release date, then film B's release date.
Strategy:
1. Read the evidence carefully. Extract any INTERMEDIATE facts (names, dates, locations) that help answer the question.
2. If you found an intermediate fact but still need more info, search for the next entity: <search>entity name topic</search>
3. Only give your final answer when you have ALL the pieces needed.
Rules:
- For multi-hop questions, you will usually need 2-3 searches. This is EXPECTED — do not try to answer with just the first search.
- In <think> tags, ALWAYS record: the specific facts you found (names, dates, places) so you don't lose them.
- Extract specific entity names from evidence tiles to use as search queries.
- Each search query MUST use DIFFERENT keywords. Be specific: use full names, dates, or titles you found.
- When you have enough info, give a concise final answer.
"""
SYSTEM_PROMPT_PIXEL_QUERY = """You are a research assistant who answers questions based on retrieved visual evidence.
The first image contains the question you need to answer.
The remaining images are retrieved evidence that may contain the answer.
Read the question from the first image, then use the evidence images to answer it.
Use <think></think> tags to show your reasoning if needed.
Answer the question directly and concisely.
"""
SYSTEM_PROMPT_MULTIMODAL_QUERY = """You are a research assistant who answers questions based on retrieved visual evidence.
You will receive: (1) a text question, (2) a query image, and (3) retrieved Wikipedia evidence images.
Use the query image and evidence images to answer the question.
Use <think></think> tags to show your reasoning if needed.
Answer the question directly and concisely.
"""
def _build_fewshot_turns(demos: list[dict], encode_image_fn) -> list[dict]:
"""Build a list of (user, assistant) message turns for in-context few-shot.
Each demo becomes: user={Q text + demo image} → assistant={answer}. The
chat-tuned model treats these as prior conversation turns rather than
mixing them with the current question's evidence — this is the canonical
few-shot format for instruction-tuned chat models.
"""
turns: list[dict] = []
for demo in demos:
user_content: list[dict] = [
{"type": "text", "text": f"Question: {demo['question']}"},
]
img_path = demo.get("image_path")
if img_path and encode_image_fn and os.path.exists(img_path):
try:
b64 = encode_image_fn(img_path)
if b64:
user_content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{b64}"},
}
)
except Exception as e:
logger.warning(f"Failed to encode few-shot image {img_path}: {e}")
turns.append({"role": "user", "content": user_content})
turns.append({"role": "assistant", "content": demo["answer"]})
return turns
def build_messages(
query: str,
retrieval_result: RetrievalResult,
encode_image_fn=None,
additional_instructions: str | None = None,
few_shot_demos: list[dict] | None = None,
) -> list[dict]:
"""Build messages for LLM based on retrieval result.
When ``retrieval_result.pixel_query_path`` is set the query is sent as an
image. Two modes:
- **Multimodal** (retrieval_type contains "multimodal"): text question + query image + retrieved tiles.
- **Pixel query** (rendered question as image): first image = question, then retrieved tiles.
"""
# ---- Multimodal / pixel-query mode: text + raw species/landmark photo + retrieved tiles ----
# query_image_path = raw species/landmark photo (for generation, always).
# pixel_query_path = rendered card or raw photo (for retrieval only; ignored here).
# Falls back to pixel_query_path if query_image_path is not set (backward compat).
gen_image_path = (
retrieval_result.query_image_path or retrieval_result.pixel_query_path
)
if gen_image_path and encode_image_fn:
system_prompt = SYSTEM_PROMPT_MULTIMODAL_QUERY
# Decide evidence_note based on what retrieval actually returned. Three cases:
# (a) retrieved images (screenshot retrieval) — evidence is image tiles after the query
# (b) retrieved text (text retrieval) — evidence is rendered as text after the query
# (c) no retrieval — query image only
# Until 2026-04-29 this branch silently dropped retrieval_result.text whenever the
# query image was set, turning every "EVQA + text retrieval" cell into an effective
# naive run. Fixed by adding the text-passages block alongside the multimodal preamble.
if retrieval_result.images:
evidence_note = "The first image is the query image. The following images are retrieved Wikipedia evidence. Answer the question based on the evidence."
elif retrieval_result.text:
evidence_note = "The image is the query image. Below is retrieved Wikipedia evidence (text). Answer the question based on the evidence and the image."
else:
evidence_note = "The first image is the query image. Answer the question based on the image (no additional evidence was retrieved)."
text_parts = [
f"Question: {query}",
"",
evidence_note,
]
if retrieval_result.text:
# Option 1: no URL header in multimodal branch either. Reader gets the
# chunks and the query image, no metadata leak.
text_parts.extend(
[
"",
retrieval_result.text,
]
)
if additional_instructions:
text_parts.append("")
text_parts.append(additional_instructions)
user_content: list[dict] = [
{"type": "text", "text": "\n".join(text_parts)},
]
# Add raw species/landmark photo
if os.path.exists(gen_image_path):
try:
img_base64 = encode_image_fn(gen_image_path)
if img_base64:
user_content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{img_base64}"},
}
)
except Exception as e:
logger.warning(f"Failed to encode query image {gen_image_path}: {e}")
user_content.append(
{"type": "text", "text": f"(Image unavailable) Query: {query}"}
)
else:
logger.warning(f"Query image not found: {gen_image_path}")
user_content.append({"type": "text", "text": f"Query: {query}"})
# Add retrieved tiles
if retrieval_result.images:
for img_path, score in retrieval_result.images:
if os.path.exists(img_path):
try:
img_base64 = encode_image_fn(img_path)
if img_base64:
user_content.append(
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{img_base64}"
},
}
)
except Exception as e:
logger.warning(f"Failed to encode image {img_path}: {e}")
return [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_content},
]
# ---- Original modes --------------------------------------------------
# Select system prompt based on retrieval type
if retrieval_result.base64_image:
system_prompt = SYSTEM_PROMPT_SCREENSHOT
user_content = [
{"type": "text", "text": query},
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{retrieval_result.base64_image}"
},
},
]
elif (
retrieval_result.retrieval_type == "text_api+rendered"
and retrieval_result.images
and encode_image_fn
):
# Text retrieval rendered as images. Mirror the text-RAG framing so
# evidence comes first and the reader sees an explicit "Question:"
# suffix — same structure as the text→text branch below, only the
# evidence modality differs.
system_prompt = SYSTEM_PROMPT_TEXT_RAG
user_content = []
for img_path, score in retrieval_result.images:
if os.path.exists(img_path):
try:
img_base64 = encode_image_fn(img_path)
if img_base64:
user_content.append(
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{img_base64}"
},
}
)
except Exception as e:
logger.warning(f"Failed to encode image {img_path}: {e}")
user_content.append({"type": "text", "text": f"Question: {query}"})
elif retrieval_result.images and encode_image_fn:
system_prompt = SYSTEM_PROMPT_VECTOR
user_content = [{"type": "text", "text": query}]
# Encode and add retrieved images
for img_path, score in retrieval_result.images:
if os.path.exists(img_path):
try:
img_base64 = encode_image_fn(img_path)
if img_base64:
user_content.append(
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{img_base64}"
},
}
)
except Exception as e:
logger.warning(f"Failed to encode image {img_path}: {e}")
elif retrieval_result.text:
system_prompt = SYSTEM_PROMPT_TEXT_RAG
# Option 1 (2026-04-29): no `Context from {urls}:` wrapper. URL leak gave
# text retrieval an unfair advantage on entity-answering tasks. Reader sees
# only the retrieved chunks and the question. URL still recorded in the
# JSONL via retrieval_result.source_url for logging/grading.
user_content = f"""{retrieval_result.text}
Question: {query}"""
else:
# Naive mode
system_prompt = SYSTEM_PROMPT_NAIVE
user_content = query
# Append additional instructions (e.g. short-answer prompt for EM-eval tasks)
if additional_instructions:
if isinstance(user_content, str):
user_content = user_content + "\n\n" + additional_instructions
else:
# list of content blocks — append as text
user_content.append({"type": "text", "text": additional_instructions})
# Few-shot as prior user/assistant turns (canonical chat few-shot format)
if few_shot_demos and encode_image_fn:
fewshot_turns = _build_fewshot_turns(few_shot_demos, encode_image_fn)
else:
fewshot_turns = []
return [
{"role": "system", "content": system_prompt},
*fewshot_turns,
{"role": "user", "content": user_content},
]
def _encode_images_to_content(
images: list[tuple[str, float]], encode_image_fn
) -> list[dict]:
"""Encode image paths to base64 content blocks."""
content = []
for img_path, score in images:
if os.path.exists(img_path):
try:
img_base64 = encode_image_fn(img_path)
if img_base64:
content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{img_base64}"},
}
)
except Exception as e:
logger.warning(f"Failed to encode image {img_path}: {e}")
return content
def build_react_messages(
query: str,
retrieval_results: list[RetrievalResult],
assistant_responses: list[str],
encode_image_fn=None,
prompt_version: str = "v1",
is_last_turn: bool = False,
previous_queries: list[str] | None = None,
) -> list[dict]:
"""Build multi-turn messages for ReAct retrieval loop.
Args:
query: Original question text.
retrieval_results: List of RetrievalResult from each round.
assistant_responses: List of assistant responses from previous rounds.
encode_image_fn: Function to encode images to base64.
prompt_version: "v1" (original) or "v2" (improved).
is_last_turn: If True, add force-answer instruction.
previous_queries: List of previous search queries (for v2, to avoid repetition).
Returns:
Messages list for the LLM.
"""
_prompt_map = {
"v1": SYSTEM_PROMPT_REACT,
"v2": SYSTEM_PROMPT_REACT_V2,
"multihop": SYSTEM_PROMPT_REACT_MULTIHOP,
}
system_prompt = _prompt_map.get(prompt_version, SYSTEM_PROMPT_REACT_V2)
messages = [{"role": "system", "content": system_prompt}]
for turn_idx, retrieval_result in enumerate(retrieval_results):
# Build user message with evidence images
if turn_idx == 0:
user_content: list[dict] = [
{
"type": "text",
"text": f"Question: {query}\n\nHere are retrieved Wikipedia evidence tiles:",
}
]
else:
text = "Here are new search results for your query:"
# Remind model of previous queries to avoid repetition (v2 and multihop)
if prompt_version in ("v2", "multihop") and previous_queries:
used = previous_queries[:turn_idx]
if used:
text += f"\n⚠️ You already searched: {used}. Do NOT repeat these. Use DIFFERENT keywords."
user_content = [{"type": "text", "text": text}]
if retrieval_result.images and encode_image_fn:
user_content.extend(
_encode_images_to_content(retrieval_result.images, encode_image_fn)
)
if not retrieval_result.has_content:
user_content.append(
{"type": "text", "text": "(No results found for this search.)"}
)
# On last turn, inject force-answer instruction
if is_last_turn and turn_idx == len(retrieval_results) - 1:
user_content.append(
{
"type": "text",
"text": (
"\n⚠️ This is your FINAL turn. You MUST provide an answer now — do NOT search again. "
"Give your best answer based on ALL evidence seen so far. If uncertain, make your best guess."
),
}
)
messages.append({"role": "user", "content": user_content})
# Add assistant response if we have one for this turn
if turn_idx < len(assistant_responses):
messages.append(
{"role": "assistant", "content": assistant_responses[turn_idx]}
)
return messages
class LLMClient:
"""Simplified async LLM client for Gemini using Vertex AI."""
def __init__(
self,
model: str,
api_base: str = "http://localhost:8000/v1",
api_key: str = "dummy",
temperature: float = 0.0,
max_tokens: int = 16384,
timeout: float = 120.0,
max_context_tokens: int | None = None,
enable_thinking: bool | None = None,
force_openai_compat: bool = False,
):
self.model = model
self.temperature = temperature
self.max_tokens = max_tokens
self.timeout = timeout
self.max_context_tokens = max_context_tokens
self.enable_thinking = enable_thinking
print(f"context length model: {max_context_tokens}")
# Gemini routes to Google GenAI SDK unless forced to OpenAI-compatible
# (aggregators like OpenRouter / Commonstack expose Gemini via OAI-compat).
self.is_gemini = ("gemini" in model.lower()) and not force_openai_compat
if self.is_gemini:
if not GEMINI_AVAILABLE:
raise ImportError(
"google-genai package is required for Gemini models. Install with: pip install google-genai"
)
# Use Vertex AI if GEMINI_API_KEY is set and GOOGLE_GENAI_USE_VERTEXAI is true
vertex_api_key = os.getenv("GEMINI_API_KEY")
use_vertex = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
if vertex_api_key and use_vertex:
logger.info(f"Using Vertex AI for Gemini model: {model}")
# Ensure GOOGLE_API_KEY is not set when using Vertex AI (it causes conflicts)
if "GOOGLE_API_KEY" in os.environ:
logger.warning(
"GOOGLE_API_KEY is set but using Vertex AI. Unsetting GOOGLE_API_KEY to avoid conflicts."
)
del os.environ["GOOGLE_API_KEY"]
os.environ["GEMINI_API_KEY"] = vertex_api_key
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "true"
self.gemini_client = genai.Client(
http_options=HttpOptions(api_version="v1")
)
else:
# Use standard Gemini API
logger.info(f"Using standard Gemini API for model: {model}")
api_key = api_key if api_key != "dummy" else os.getenv("GOOGLE_API_KEY")
if not api_key:
raise ValueError(
"GOOGLE_API_KEY or GEMINI_API_KEY environment variable is required for Gemini models"
)
self.gemini_client = genai.Client(api_key=api_key)
else:
# Use OpenAI-compatible API
from openai import AsyncOpenAI
logger.info(f"Using OpenAI-compatible API: {api_base}")
self.client = AsyncOpenAI(
api_key=api_key,
base_url=api_base,
timeout=timeout,
max_retries=0,
)
self.gemini_client = None
async def generate(
self, messages: list[dict], max_retries: int = 3, connection_retries: int = 12
) -> tuple[str, dict]:
"""Generate response from messages with retry on timeout/connection errors.
Args:
max_retries: Retry count for timeout errors.
connection_retries: Retry count for connection errors (server restart).
12 retries × 10s = ~2 min window for server to come back.
Returns:
Tuple of (generated_text, usage_dict).
"""
# Check and truncate if needed
if hasattr(self, "max_context_tokens") and self.max_context_tokens:
estimated_tokens = self._estimate_tokens(messages)
if estimated_tokens > self.max_context_tokens - self.max_tokens:
logger.warning(
f"Estimated {estimated_tokens} tokens exceeds limit, truncating..."
)
messages = self._truncate_messages(messages, self.max_context_tokens)
conn_attempts = 0
timeout_attempts = 0
while True:
try:
if self.is_gemini:
return await self._generate_gemini(messages)
else:
return await self._generate_openai(messages)
except asyncio.TimeoutError:
timeout_attempts += 1
if timeout_attempts >= max_retries:
raise
wait_time = 2**timeout_attempts # 2, 4, 8 seconds
logger.warning(
f"Timeout on attempt {timeout_attempts}/{max_retries}, retrying in {wait_time}s..."
)
await asyncio.sleep(wait_time)
except Exception as e:
error_str = str(e).lower()
if "timeout" in error_str or "timed out" in error_str:
timeout_attempts += 1
if timeout_attempts >= max_retries:
raise
wait_time = 2**timeout_attempts
logger.warning(
f"Timeout on attempt {timeout_attempts}/{max_retries}, retrying in {wait_time}s..."
)
await asyncio.sleep(wait_time)
elif "connection" in error_str or "connect" in error_str:
conn_attempts += 1
if conn_attempts >= connection_retries:
raise
wait_time = 10 # fixed 10s — server restart takes ~30-60s
logger.warning(
f"Connection error ({conn_attempts}/{connection_retries}), retrying in {wait_time}s..."
)
await asyncio.sleep(wait_time)
elif (
"429" in error_str
or "rate_limit" in error_str
or "rate limit" in error_str
):
# Provider rate limit — exponential backoff with jitter
timeout_attempts += 1
if timeout_attempts >= max_retries + 3: # extra patience for 429
raise
import random
wait_time = min(60, 5 * (2**timeout_attempts)) + random.uniform(
0, 3
)
logger.warning(
f"429 rate-limit (attempt {timeout_attempts}), backing off {wait_time:.1f}s..."
)
await asyncio.sleep(wait_time)
else:
raise
async def _generate_gemini(self, messages: list[dict]) -> tuple[str, dict]:
"""Generate using Gemini API."""
# Extract system prompt and user content
system_prompt = None
user_content = None
for msg in messages:
if msg.get("role") == "system":
system_prompt = msg.get("content", "")
elif msg.get("role") == "user":
user_content = msg.get("content", "")
# Build parts for Gemini
parts = []
# Add system prompt to the beginning of user message if present
if system_prompt:
parts.append(Part(text=f"{system_prompt}\n\n"))
# Process user content
if isinstance(user_content, str):
# Simple text
if parts:
parts[0] = Part(text=parts[0].text + user_content)
else:
parts.append(Part(text=user_content))
elif isinstance(user_content, list):
# Multi-modal content
for item in user_content:
if item.get("type") == "text":
text = item.get("text", "")
if (
parts
and isinstance(parts[0], Part)
and hasattr(parts[0], "text")
):
# Append to existing text part
parts[0] = Part(text=parts[0].text + text)
else:
parts.append(Part(text=text))
elif item.get("type") == "image_url":
# Extract base64 image
image_url = item.get("image_url", {}).get("url", "")
if image_url.startswith("data:image"):
try:
header, data = image_url.split(",", 1)
mime_type = header.split(";")[0].split(":")[1]
image_bytes = base64.b64decode(data)
parts.append(
Part(
inline_data=Blob(
mime_type=mime_type, data=image_bytes
)
)
)
except Exception as e:
logger.error(f"Failed to process image: {e}")
raise
# Create content
content = Content(role="user", parts=parts)
# Call API in executor to avoid blocking
loop = asyncio.get_event_loop()
def _call_api():
try:
response = self.gemini_client.models.generate_content(
model=self.model,
contents=[content],
config=GenerateContentConfig(
temperature=self.temperature, max_output_tokens=self.max_tokens
),
)
return response
except Exception as e:
logger.error(f"Gemini API error: {e}")
raise
response = await loop.run_in_executor(None, _call_api)
# Extract text
text = response.text if hasattr(response, "text") and response.text else ""
# Extract usage
usage = {}
if hasattr(response, "usage_metadata") and response.usage_metadata:
usage_meta = response.usage_metadata
usage = {
"prompt_tokens": getattr(usage_meta, "prompt_token_count", 0),
"completion_tokens": getattr(usage_meta, "candidates_token_count", 0),
"total_tokens": getattr(usage_meta, "total_token_count", 0),
}
return text, usage
async def _generate_openai(self, messages: list[dict]) -> tuple[str, dict]:
"""Generate using OpenAI-compatible API."""
kwargs = dict(
model=self.model,
messages=messages,
max_tokens=self.max_tokens,
timeout=self.timeout,
)
# Some modern reasoning models deprecate `temperature` (Claude Opus 4.7+, some GPT-5 variants).
# Only send it when we actually want to override the default.
model_lower = self.model.lower()
drops_temperature = any(
x in model_lower for x in ("opus-4-7", "opus-4-8", "gpt-5.4-pro")
)
if not drops_temperature:
kwargs["temperature"] = self.temperature
if self.enable_thinking is not None:
kwargs["extra_body"] = {
"chat_template_kwargs": {"enable_thinking": self.enable_thinking}
}
response = await self.client.chat.completions.create(**kwargs)
generated_text = response.choices[0].message.content
usage = {}
if response.usage:
usage = {
"prompt_tokens": response.usage.prompt_tokens,
"completion_tokens": response.usage.completion_tokens,
"total_tokens": response.usage.total_tokens,
}
return generated_text, usage
def _estimate_tokens(self, messages: list[dict]) -> int:
"""Estimate token count from messages (rough: ~4 chars per token)."""
total_chars = 0
for msg in messages:
content = msg.get("content", "")
if isinstance(content, str):
total_chars += len(content)
elif isinstance(content, list):
for item in content:
if isinstance(item, dict):
if item.get("type") == "text":
total_chars += len(item.get("text", ""))
elif item.get("type") == "image_url":
# Rough estimate for image tokens
total_chars += 1000 * 4 # ~1000 tokens per image
return total_chars // 4
def _truncate_messages(self, messages: list[dict], max_tokens: int) -> list[dict]:
"""Truncate text content in messages to fit within token limit."""
# Reserve tokens for response
available_tokens = max_tokens - self.max_tokens - 500 # buffer
max_chars = available_tokens * 4
truncated = []
total_chars = 0
for msg in messages:
new_msg = msg.copy()
content = msg.get("content", "")
if isinstance(content, str):
if total_chars + len(content) > max_chars:
remaining = max(0, max_chars - total_chars)
new_msg["content"] = (
content[:remaining]
+ "\n\n[Content truncated due to context limit]"
)
logger.warning(
f"Truncated message content from {len(content)} to {remaining} chars"
)
total_chars += len(new_msg["content"])
elif isinstance(content, list):
new_content = []
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
text = item.get("text", "")
if total_chars + len(text) > max_chars:
remaining = max(0, max_chars - total_chars)
new_item = item.copy()
new_item["text"] = (
text[:remaining]
+ "\n\n[Content truncated due to context limit]"
)
new_content.append(new_item)
logger.warning(
f"Truncated text content from {len(text)} to {remaining} chars"
)
total_chars += remaining
else:
new_content.append(item)
total_chars += len(text)
else:
new_content.append(item)
if isinstance(item, dict) and item.get("type") == "image_url":
total_chars += 1000 * 4 # image token estimate
new_msg["content"] = new_content
truncated.append(new_msg)
return truncated
+95
View File
@@ -0,0 +1,95 @@
"""Model configurations for SimpleQA evaluation.
This module provides model configurations to keep run_naive_simpleqa.py clean.
"""
import os
from typing import Dict, Optional
def get_model_config(model_name: str) -> Dict[str, Optional[str]]:
"""
Get model configuration based on model name.
Args:
model_name: Name of the model (e.g., 'Qwen/Qwen3-VL-4B-Instruct', 'gemini-3-pro-preview')
Returns:
Dictionary with 'api_base', 'api_key', and 'model' keys.
"""
model_lower = model_name.lower()
# Gemini models
if "gemini" in model_lower:
# Check for Vertex AI first
vertex_api_key = os.getenv("GEMINI_API_KEY")
use_vertex = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
if use_vertex and vertex_api_key:
# Using Vertex AI - don't pass api_key, use environment variable instead
api_key = None # Vertex AI uses environment variable, not api_key parameter
else:
# Using standard Gemini API
api_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
if not api_key:
raise ValueError(
"GOOGLE_API_KEY or GEMINI_API_KEY environment variable is required for Gemini models. "
"Set it with: export GOOGLE_API_KEY='your-api-key' or export GEMINI_API_KEY='your-api-key' and GOOGLE_GENAI_USE_VERTEXAI=true"
)
# For Gemini models, we use Google's Generative AI SDK directly
# The api_base is not used for Gemini (SDK handles it internally)
# But we set a placeholder for compatibility
api_base = None # Not used for Gemini SDK
return {
"api_base": api_base,
"api_key": api_key,
"model": model_name, # Use the model name as-is
}
# Default: assume OpenAI-compatible API (vLLM, etc.)
return {
"api_base": os.getenv("API_BASE", "http://localhost:8000/v1"),
"api_key": os.getenv("API_KEY", "dummy"),
"model": model_name,
}
def get_output_filename(
output_dir: str,
model_name: str,
mode: str = "naive",
num_examples: int = 1000,
url_screenshot: bool = False,
task: str = "simpleqa",
) -> str:
"""
Generate output filename with model name and task included.
Args:
output_dir: Base output directory (e.g., 'eval_output/naive_qa')
model_name: Model name (e.g., 'Qwen/Qwen3-VL-4B-Instruct')
mode: Evaluation mode ('naive', 'screenshot', 'retrieval')
num_examples: Number of examples
url_screenshot: Whether URL screenshot mode is enabled
task: Task/benchmark name (e.g., 'simpleqa', 'encyclopedic_vqa', 'worldvqa')
Returns:
Full output file path
"""
# Clean model name for filename (replace special chars)
model_safe = (
model_name.replace("/", "_").replace(":", "_").replace("-", "_").lower()
)
# Build filename components (task first for easy distinction)
parts = [task]
if url_screenshot:
parts.append("urlscreenshot")
parts.append(mode)
parts.append(model_safe)
parts.append(str(num_examples))
filename = "_".join(parts) + ".jsonl"
return os.path.join(output_dir, filename)
+345
View File
@@ -0,0 +1,345 @@
"""Pixel Query: render text queries as images.
This module converts text queries to pixel images so VLMs receive the query
as a visual input instead of text tokens. This tests whether the model can
"read" the question from an image and still perform retrieval + answering.
Usage:
renderer = PixelQueryRenderer(output_dir="pixel_queries", font_size=16, img_width=600)
img_path = renderer.render(example_id, query_text)
"""
import logging
import os
from PIL import Image, ImageDraw, ImageFont
logger = logging.getLogger(__name__)
# Default font paths (tried in order)
_FONT_CANDIDATES = [
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
"/usr/share/fonts/truetype/noto/NotoSans-Regular.ttf",
]
def _find_font(font_path: str | None = None) -> str:
"""Find a usable TrueType font on this system."""
if font_path and os.path.exists(font_path):
return font_path
for candidate in _FONT_CANDIDATES:
if os.path.exists(candidate):
return candidate
raise FileNotFoundError(
"No suitable TTF font found. Install dejavu or liberation fonts, "
"or pass font_path explicitly."
)
def _wrap_text_by_pixel_width(
text: str, font: ImageFont.FreeTypeFont, max_width: int
) -> list[str]:
"""Word-wrap text so each line fits within *max_width* pixels."""
words = text.split()
lines: list[str] = []
current_line = ""
for word in words:
test_line = f"{current_line} {word}".strip()
bbox = font.getbbox(test_line)
line_width = bbox[2] - bbox[0]
if line_width <= max_width:
current_line = test_line
else:
if current_line:
lines.append(current_line)
current_line = word
if current_line:
lines.append(current_line)
return lines
class PixelQueryRenderer:
"""Renders text queries as small, clear PNG images.
Images are cached on disk so each query is only rendered once.
"""
def __init__(
self,
output_dir: str = "pixel_queries",
font_path: str | None = None,
font_size: int = 16,
img_width: int = 600,
padding_x: int = 16,
padding_y: int = 12,
line_spacing: int = 4,
):
self.output_dir = output_dir
self.font_path = _find_font(font_path)
self.font_size = font_size
self.img_width = img_width
self.padding_x = padding_x
self.padding_y = padding_y
self.line_spacing = line_spacing
os.makedirs(output_dir, exist_ok=True)
self.font = ImageFont.truetype(self.font_path, self.font_size)
logger.info(
f"PixelQueryRenderer: dir={output_dir}, font={os.path.basename(self.font_path)}, "
f"size={font_size}, width={img_width}"
)
def _render_image(self, text: str) -> Image.Image:
"""Render *text* to a PIL Image (white bg, black text)."""
max_text_width = self.img_width - 2 * self.padding_x
lines = _wrap_text_by_pixel_width(text, self.font, max_text_width)
# Measure line height using a reference string
line_height = self.font.getbbox("Ay")[3] - self.font.getbbox("Ay")[1]
total_text_height = (
len(lines) * line_height + (len(lines) - 1) * self.line_spacing
)
img_height = total_text_height + 2 * self.padding_y
img = Image.new("RGB", (self.img_width, img_height), color=(255, 255, 255))
draw = ImageDraw.Draw(img)
y = self.padding_y
for line in lines:
draw.text((self.padding_x, y), line, font=self.font, fill=(0, 0, 0))
y += line_height + self.line_spacing
return img
def render(self, example_id: str, query_text: str) -> str:
"""Render a query and return the path to the saved PNG.
If the image already exists on disk it is *not* re-rendered.
"""
out_path = os.path.join(self.output_dir, f"{example_id}_query.png")
if os.path.exists(out_path):
return out_path
img = self._render_image(query_text)
img.save(out_path)
logger.debug(f"Rendered pixel query: {out_path} ({img.size[0]}x{img.size[1]})")
return out_path
def render_all(self, examples: list[dict]) -> dict[str, str]:
"""Batch-render pixel queries for a list of examples.
Args:
examples: List of dicts with at least ``id`` and ``problem`` keys.
Returns:
Dict mapping example_id → pixel query image path.
"""
id_to_path: dict[str, str] = {}
rendered, cached = 0, 0
for ex in examples:
eid = ex["id"]
path = os.path.join(self.output_dir, f"{eid}_query.png")
if os.path.exists(path):
cached += 1
else:
img = self._render_image(ex["problem"])
img.save(path)
rendered += 1
id_to_path[eid] = path
logger.info(
f"PixelQueryRenderer: {rendered} rendered, {cached} cached, "
f"{rendered + cached} total in {self.output_dir}"
)
return id_to_path
class QueryImageTextRenderer:
"""Renders query text + query image together into a single card image.
Reuses PixelQueryRenderer for text rendering. Layout: image on top, text below
(similar to VQA task cards). Text is centered and uses a larger font.
Images are cached on disk.
Usage:
renderer = QueryImageTextRenderer(output_dir="query_cards", tiles_dir="tiles/evqa")
path = renderer.render(example_id, query_text, query_image_path)
"""
def __init__(
self,
output_dir: str = "query_cards",
tiles_dir: str = "tiles/evqa",
font_path: str | None = None,
font_size: int = 22,
card_width: int = 600,
padding_x: int = 24,
padding_y: int = 20,
line_spacing: int = 6,
image_padding: int = 16,
text_section_padding: int = 24,
border_radius: int = 12,
):
self.output_dir = output_dir
self.tiles_dir = tiles_dir
self.card_width = card_width
self.image_padding = image_padding
self.text_section_padding = text_section_padding
self.border_radius = border_radius
font_path_resolved = _find_font(font_path)
self.font_path = font_path_resolved
self.font_size = font_size
self.font = ImageFont.truetype(font_path_resolved, font_size)
self.padding_x = padding_x
self.padding_y = padding_y
self.line_spacing = line_spacing
os.makedirs(output_dir, exist_ok=True)
logger.info(
f"QueryImageTextRenderer: dir={output_dir}, card_width={card_width}, "
f"font_size={font_size}"
)
def _render_query_text_centered(
self, text: str, width: int | None = None
) -> Image.Image:
"""Render query text centered, reusing wrap logic from PixelQueryRenderer."""
width = width or self.card_width
max_text_width = width - 2 * self.padding_x
# Calculate dynamic font size based on width
# Ratio: width / 30 seems reasonable (600px -> 20px, 1500px -> 50px)
dynamic_font_size = max(22, int(width / 30))
# Use dynamic font if size is different from default
if dynamic_font_size != self.font_size:
try:
font = ImageFont.truetype(self.font_path, dynamic_font_size)
except Exception:
font = self.font # Fallback
else:
font = self.font
lines = _wrap_text_by_pixel_width(text, font, max_text_width)
line_height = font.getbbox("Ay")[3] - font.getbbox("Ay")[1]
# Scale spacing proportionally
spacing = max(4, int(self.line_spacing * (dynamic_font_size / self.font_size)))
total_text_height = len(lines) * line_height + (len(lines) - 1) * spacing
# Scale padding proportionally
padding_y = max(
self.padding_y, int(self.padding_y * (dynamic_font_size / self.font_size))
)
text_height = total_text_height + 2 * padding_y
img = Image.new("RGB", (width, text_height), color=(255, 255, 255))
draw = ImageDraw.Draw(img)
y = padding_y
for line in lines:
bbox = font.getbbox(line)
line_w = bbox[2] - bbox[0]
x = (width - line_w) // 2
draw.text((x, y), line, font=font, fill=(50, 50, 50))
y += line_height + spacing
return img
def render(
self,
example_id: str,
query_text: str,
query_image_path: str | None,
force: bool = False,
) -> str:
"""Render query image + text into one card, save to disk.
Layout: image on top, text below (same for iNaturalist and Landmarks).
Args:
example_id: Example identifier for filename.
query_text: The question text to render below the image.
query_image_path: Path to the query image (iNaturalist or Landmark photo).
If None or file missing, only the text is rendered.
force: If True, re-render even if output exists (e.g. when images were added later).
Returns:
Path to the saved PNG.
"""
out_path = os.path.join(self.output_dir, f"{example_id}_query_card.png")
if os.path.exists(out_path) and not force:
return out_path
# Load query image
if query_image_path and os.path.exists(query_image_path):
try:
query_img = Image.open(query_image_path).convert("RGB")
# Resize if too large to ensure font size is readable and image isn't massive
max_dim = 1536 # Standard reasonable max dimension
if max(query_img.size) > max_dim:
ratio = max_dim / max(query_img.size)
new_size = (
int(query_img.width * ratio),
int(query_img.height * ratio),
)
query_img = query_img.resize(new_size, Image.Resampling.LANCZOS)
except Exception as e:
logger.warning(f"Failed to load query image {query_image_path}: {e}")
query_img = None
else:
query_img = None
# Card width adapts to image: expand if image is wider than card_width
if query_img is not None:
effective_width = max(
self.card_width, query_img.width + 2 * self.image_padding
)
else:
effective_width = self.card_width
# Render query text at effective width
text_img = self._render_query_text_centered(query_text, effective_width)
# Compose: image on top, text below with padding for balanced look
if query_img is not None:
img_section_height = query_img.height + 2 * self.image_padding
else:
img_section_height = 0
text_section_height = text_img.height + 2 * self.text_section_padding
total_height = img_section_height + text_section_height
card = Image.new("RGB", (effective_width, total_height), color=(255, 255, 255))
ImageDraw.Draw(card)
y_offset = 0
if query_img is not None:
x_center = (effective_width - query_img.width) // 2
card.paste(query_img, (x_center, self.image_padding))
y_offset = img_section_height
# Center text block in its section
text_y = y_offset + self.text_section_padding
card.paste(text_img, (0, text_y))
# Optional: rounded corners (simplified - draw white rounded rect overlay)
if self.border_radius > 0:
# Create mask for rounded corners
mask = Image.new("L", card.size, 255)
m_draw = ImageDraw.Draw(mask)
m_draw.rounded_rectangle(
(0, 0, card.width - 1, card.height - 1),
radius=self.border_radius,
fill=255,
outline=0,
)
# For simple output we keep the card as-is; full rounded crop would need alpha
pass
card.save(out_path)
logger.debug(f"Rendered query card: {out_path} ({card.size[0]}x{card.size[1]})")
return out_path
File diff suppressed because it is too large Load Diff
+550
View File
@@ -0,0 +1,550 @@
"""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
+412
View File
@@ -0,0 +1,412 @@
"""Screenshot utilities using Selenium WebDriver.
Selenium and webdriver_manager are optional — only needed when capturing
screenshots (--url-screenshot mode). Import is deferred to function scope
so the eval script works without them for --local-api users.
"""
import base64
import logging
import os
import time
from PIL import Image
logger = logging.getLogger(__name__)
def setup_driver(window_width=1024, window_height=2000, device_scale_factor=1):
"""Set up Chrome WebDriver.
Args:
window_width: Viewport width (1024 = tile width, ensures screenshots align with tile grid).
window_height: Initial viewport height.
device_scale_factor: Pixel density (1 = standard, 2 = retina quality).
"""
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
import shutil
snap_chromedriver = "/snap/bin/chromium.chromedriver"
if os.path.exists(snap_chromedriver):
driver_path = snap_chromedriver
elif shutil.which("chromedriver"):
driver_path = shutil.which("chromedriver")
else:
driver_path = ChromeDriverManager().install()
service = Service(driver_path)
options = webdriver.ChromeOptions()
# Find Chrome binary path
chrome_binary = None
for chrome_path in [
"/usr/bin/google-chrome-stable",
"/usr/bin/google-chrome",
"/usr/bin/chromium-browser",
"/usr/bin/chromium",
]:
if os.path.exists(chrome_path):
chrome_binary = chrome_path
break
if chrome_binary:
options.binary_location = chrome_binary
options.add_argument("--headless=new")
options.add_argument("--disable-gpu")
options.add_argument(f"--window-size={window_width},{window_height}")
options.add_argument("--no-sandbox")
options.add_argument("--disable-dev-shm-usage")
options.add_argument("--disable-extensions")
options.add_argument("--disable-setuid-sandbox")
options.add_argument("--no-zygote")
options.add_argument("--remote-debugging-port=0")
options.add_argument(
"--user-agent=Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
"(KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36"
)
# Retina-quality rendering (2x pixel density)
if device_scale_factor and device_scale_factor > 1:
options.add_argument(f"--force-device-scale-factor={device_scale_factor}")
driver = webdriver.Chrome(service=service, options=options)
return driver
def _capture_with_scroll(driver, output_path, scroll_pause=0.8, max_scrolls=100):
"""Capture full page by scrolling and stitching screenshots.
Works for PDF viewers, infinite scroll pages, and other dynamic content.
Uses image comparison to detect when scrolling has stopped.
"""
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.keys import Keys
import tempfile
import hashlib
def get_screenshot_hash(path):
"""Get hash of screenshot to detect changes."""
with Image.open(path) as img:
return hashlib.md5(img.tobytes()).hexdigest()
viewport_height = driver.execute_script("return window.innerHeight")
viewport_width = driver.execute_script("return window.innerWidth")
# Try to scroll to top using multiple methods
driver.execute_script("window.scrollTo(0, 0)")
actions = ActionChains(driver)
actions.send_keys(Keys.HOME)
actions.perform()
time.sleep(scroll_pause)
screenshots = []
last_hash = None
scroll_count = 0
consecutive_same = 0
while scroll_count < max_scrolls:
# Take screenshot
temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
driver.save_screenshot(temp_file.name)
# Check if image changed (detect end of scrolling)
current_hash = get_screenshot_hash(temp_file.name)
if current_hash == last_hash:
consecutive_same += 1
os.unlink(temp_file.name)
if consecutive_same >= 2:
# Scrolling stopped, we've reached the end
break
else:
consecutive_same = 0
screenshots.append(temp_file.name)
last_hash = current_hash
# Scroll down using CDP mouse wheel event (works for PDF viewers)
center_x = viewport_width // 2
center_y = viewport_height // 2
try:
driver.execute_cdp_cmd(
"Input.dispatchMouseEvent",
{
"type": "mouseWheel",
"x": center_x,
"y": center_y,
"deltaX": 0,
"deltaY": int(viewport_height * 0.8), # Scroll 80% of viewport
},
)
except Exception:
# Fallback to ActionChains
actions = ActionChains(driver)
actions.send_keys(Keys.PAGE_DOWN)
actions.perform()
time.sleep(scroll_pause)
scroll_count += 1
if not screenshots:
driver.save_screenshot(output_path)
return
if len(screenshots) == 1:
# Only one screenshot, just use it
os.rename(screenshots[0], output_path)
return
# Stitch screenshots vertically
# Each screenshot is viewport_height, but they overlap
# We'll stack them with some overlap detection
images = [Image.open(p) for p in screenshots]
# Simple stacking: assume each Page Down scrolls ~80% of viewport
overlap = int(viewport_height * 0.2)
total_height = viewport_height + (len(images) - 1) * (viewport_height - overlap)
stitched = Image.new("RGB", (viewport_width, total_height), (255, 255, 255))
y_offset = 0
for i, img in enumerate(images):
if i == 0:
stitched.paste(img, (0, 0))
y_offset = viewport_height - overlap
else:
# Crop top overlap region and paste
cropped = img.crop((0, overlap, viewport_width, viewport_height))
stitched.paste(cropped, (0, y_offset))
y_offset += viewport_height - overlap
# Close images and clean up
for img in images:
img.close()
for p in screenshots:
if os.path.exists(p):
os.unlink(p)
# Trim any white space at bottom
stitched = stitched.crop((0, 0, viewport_width, y_offset + overlap))
stitched.save(output_path)
def _eager_load_images(driver):
"""Force lazy images to load by promoting data-src and setting loading='eager'."""
driver.execute_script("""
(function() {
var imgs = document.querySelectorAll('img');
for (var i = 0; i < imgs.length; i++) {
var img = imgs[i];
try {
if (img.loading === 'lazy') img.loading = 'eager';
var dataSrc = img.getAttribute('data-src') || (img.dataset && img.dataset.src);
var dataSrcset = img.getAttribute('data-srcset') || (img.dataset && img.dataset.srcset);
if (dataSrc) img.setAttribute('src', dataSrc);
if (dataSrcset) img.setAttribute('srcset', dataSrcset);
} catch(e) {}
}
})();
""")
def _wait_for_images(driver, timeout=10):
"""Wait for all document images to finish loading (load or error)."""
deadline = time.time() + timeout
while time.time() < deadline:
pending = driver.execute_script("""
return Array.from(document.images || [])
.filter(function(img) { return !(img.complete && img.naturalWidth > 0); }).length;
""")
if pending == 0:
return
time.sleep(0.3)
def _scroll_to_trigger_lazy_load(driver, page_height):
"""Scroll through the page to trigger lazy-loaded content, then scroll back to top."""
viewport_height = driver.execute_script("return window.innerHeight") or 1080
y = 0
while y < page_height:
driver.execute_script(f"window.scrollTo(0, {y})")
time.sleep(0.15)
_eager_load_images(driver)
y += viewport_height
# Wait for all images to load after full scroll
_wait_for_images(driver, timeout=10)
# Scroll back to top
driver.execute_script("window.scrollTo(0, 0)")
time.sleep(0.3)
def capture_screenshot(url, output_path, full_page=False, scroll_capture=False):
"""Capture screenshot of a URL.
Args:
url: URL to capture
output_path: Path to save screenshot
full_page: If True, resize window to capture full page (works for normal pages)
scroll_capture: If True, scroll and stitch screenshots (works for PDF viewers, etc.)
"""
driver = None
try:
driver = setup_driver()
driver.get(url)
time.sleep(3) # Wait for initial load
# Force lazy images to load eagerly
_eager_load_images(driver)
_wait_for_images(driver, timeout=5)
if scroll_capture:
# Scroll-based capture for PDF viewers and similar
# just for PDF
_capture_with_scroll(driver, output_path)
elif full_page:
# Get page height, keep original window width to avoid horizontal tiling
total_height = driver.execute_script("return document.body.scrollHeight")
current_window = driver.get_window_size()
# Scroll through page to trigger lazy-loaded images
_scroll_to_trigger_lazy_load(driver, total_height)
# Re-measure height (may change after lazy content loads)
total_height = driver.execute_script("return document.body.scrollHeight")
driver.set_window_size(current_window["width"], total_height)
time.sleep(0.5)
# Final wait for any images triggered by resize
_eager_load_images(driver)
_wait_for_images(driver, timeout=5)
driver.save_screenshot(output_path)
else:
driver.save_screenshot(output_path)
# Convert to RGB (remove alpha channel if present)
with Image.open(output_path) as img:
if img.mode in ("RGBA", "LA") or (
img.mode == "P" and "transparency" in img.info
):
bg = Image.new("RGB", img.size, (255, 255, 255))
if img.mode != "RGBA":
img = img.convert("RGBA")
bg.paste(img, mask=img.split()[3])
img = bg
img.save(output_path)
elif img.mode != "RGB":
img = img.convert("RGB")
img.save(output_path)
return True
except Exception as e:
print(f"Screenshot failed for {url}: {e}")
return False
finally:
if driver:
driver.quit()
def encode_image(image_path, max_pixels: int = 150_000_000, max_height: int = 8000):
"""Encode image to base64, compressing if too large.
Used for Vector DB retrieval where consistent image sizes help with embedding.
Args:
image_path: Path to image file.
max_pixels: Maximum pixels allowed (default 150M).
max_height: Maximum height in pixels (default 8000).
"""
if not os.path.exists(image_path):
return None
try:
# Increase PIL limit temporarily
Image.MAX_IMAGE_PIXELS = 300_000_000
with Image.open(image_path) as img:
# Check if compression needed
total_pixels = img.width * img.height
needs_compression = total_pixels > max_pixels or img.height > max_height
if not needs_compression:
# Just read and encode directly
with open(image_path, "rb") as f:
return base64.b64encode(f.read()).decode("utf-8")
# Compress: resize to fit within limits
if img.mode != "RGB":
img = img.convert("RGB")
# Calculate new size
if img.height > max_height:
ratio = max_height / img.height
new_width = int(img.width * ratio)
new_height = max_height
else:
# Scale down to fit max_pixels
ratio = (max_pixels / total_pixels) ** 0.5
new_width = int(img.width * ratio)
new_height = int(img.height * ratio)
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# Encode to JPEG
from io import BytesIO
buffered = BytesIO()
img.save(buffered, format="JPEG", quality=85)
return base64.b64encode(buffered.getvalue()).decode("utf-8")
except Exception as e:
print(f"Failed to encode image {image_path}: {e}")
return None
def encode_image_for_vlm(image_path, max_pixels: int = 89_000_000):
"""Encode image to base64 for VLM ground truth, minimal processing.
For Ground Truth evaluation, we want to preserve original image quality
and let the VLM handle resizing according to its own requirements.
Only applies PIL safety limit (89M pixels).
Args:
image_path: Path to image file.
max_pixels: Maximum pixels (default 89M, PIL's safety limit).
"""
if not os.path.exists(image_path):
return None
try:
# Increase PIL limit temporarily
Image.MAX_IMAGE_PIXELS = 300_000_000
with Image.open(image_path) as img:
total_pixels = img.width * img.height
# Only compress if exceeds PIL safety limit
if total_pixels <= max_pixels:
# Just read and encode directly - no resize
with open(image_path, "rb") as f:
return base64.b64encode(f.read()).decode("utf-8")
# Compress only if exceeds max_pixels
if img.mode != "RGB":
img = img.convert("RGB")
# Scale down to fit max_pixels
ratio = (max_pixels / total_pixels) ** 0.5
new_width = int(img.width * ratio)
new_height = int(img.height * ratio)
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# Encode to JPEG
from io import BytesIO
buffered = BytesIO()
img.save(buffered, format="JPEG", quality=85)
return base64.b64encode(buffered.getvalue()).decode("utf-8")
except Exception as e:
print(f"Failed to encode image {image_path}: {e}")
return None
File diff suppressed because it is too large Load Diff
+114
View File
@@ -0,0 +1,114 @@
"""Dataset filtering utilities for SimpleQA."""
import ast
import json
import logging
from .simpleqa_data import load_simpleqa_data, load_simpleqa_verified_data
logger = logging.getLogger(__name__)
def _get_urls_from_metadata(example: dict) -> list[str]:
"""Extract all URLs from example metadata.
Supports both SimpleQA format (metadata as string) and SimpleQA Verified format (metadata as dict or string).
"""
meta = example.get("metadata")
# If metadata is already a dict (SimpleQA Verified format)
if isinstance(meta, dict):
urls = meta.get("urls", [])
if isinstance(urls, list):
return urls
return []
# If metadata is a string (SimpleQA format or SimpleQA Verified string format)
if isinstance(meta, str):
try:
meta = json.loads(meta)
except json.JSONDecodeError:
try:
meta = ast.literal_eval(meta)
except (ValueError, SyntaxError):
return []
if isinstance(meta, dict) and "urls" in meta:
urls = meta["urls"]
if isinstance(urls, list):
return urls
return []
def load_simpleqa_wikipedia(
num_examples: int | None = None,
verified: bool = False,
no_wiki_filter: bool = False,
) -> list[dict]:
"""Load SimpleQA examples that have Wikipedia URLs.
Args:
num_examples: Number of Wikipedia examples to return.
verified: If True, use SimpleQA Verified dataset instead of original SimpleQA.
Returns:
List of examples where at least one URL contains 'wikipedia'.
Maintains the original CSV file order.
"""
# Load all data first (maintains original CSV order)
if verified:
all_data = load_simpleqa_verified_data(num_examples=None)
else:
all_data = load_simpleqa_data(num_examples=None)
if no_wiki_filter:
wikipedia_examples = list(all_data)
logger.info(
f"Skipping Wikipedia URL filter: returning all {len(wikipedia_examples)} examples"
)
else:
# Filter for Wikipedia URLs, preserving original order
# Exclude non-English Wikipedia and Category pages (e.g. de.wikipedia.org, Category:...)
wikipedia_examples = []
for example in all_data:
urls = _get_urls_from_metadata(example)
if any(
"en.wikipedia.org/wiki/" in url and "/Category:" not in url
for url in urls
):
wikipedia_examples.append(example)
logger.info(f"Found {len(wikipedia_examples)} examples with Wikipedia URLs")
if num_examples:
wikipedia_examples = wikipedia_examples[:num_examples]
logger.info(f"Limiting to first {num_examples} Wikipedia examples")
return wikipedia_examples
def load_simpleqa_by_domain(domain: str, num_examples: int | None = None) -> list[dict]:
"""Load SimpleQA examples filtered by URL domain.
Args:
domain: Domain to filter by (e.g., 'wikipedia', 'arxiv', 'github').
num_examples: Number of examples to return.
Returns:
List of examples where at least one URL contains the domain.
"""
all_data = load_simpleqa_data(num_examples=None)
filtered = []
for example in all_data:
urls = _get_urls_from_metadata(example)
if any(domain.lower() in url.lower() for url in urls):
filtered.append(example)
logger.info(f"Found {len(filtered)} examples with '{domain}' URLs")
if num_examples:
filtered = filtered[:num_examples]
logger.info(f"Limiting to first {num_examples} examples")
return filtered
+33
View File
@@ -0,0 +1,33 @@
[project]
name = "pixelrag-repro"
version = "0.1.0"
description = "Reproduction harness for the PixelRAG (Vis-RAG) paper Table 1 — drives the paper's own run_naive_simpleqa.py + evaluate.py against live retrieval/reader serves."
requires-python = ">=3.12,<3.13"
# These are the deps the paper's API-path code (scripts/run_naive_simpleqa.py,
# scripts/simpleqa, scripts/evaluate.py, evaluation/*) imports when retrieval + reader
# run as remote HTTP serves (no local torch/vllm needed -- the model serves are separate).
# The paper repo itself is pinned in REPRODUCE.md (yichuan-w/Vis-RAG @ e591fd0).
dependencies = [
"aiohttp==3.13.5",
"datasets==4.8.5",
"openai==2.38.0",
"tqdm==4.67.3",
"pillow==12.2.0",
"requests==2.34.2",
"numpy==2.4.6",
"selenium==4.44.0",
"webdriver-manager==4.1.1",
"beautifulsoup4==4.14.3",
"lxml==6.1.1",
"tiktoken==0.13.0",
"trafilatura==2.0.0",
"litellm==1.86.2",
"botocore==1.43.18",
"tenacity==9.1.4",
"fastmcp==3.3.1",
"omegaconf==2.3.0",
"retry==0.9.2",
]
[tool.uv]
package = false
+102
View File
@@ -0,0 +1,102 @@
### Role
You are an expert judge specialized in evaluating the correctness of answers. Your task is to assess whether a model-generated answer is correct based on a given question, the model's response, and the ground truth answer.
### Task: Evaluate Answer Correctness
Please classify the model's response into one of the following three categories. Ignore differences in formatting, punctuation, language (Chinese vs. English), or abbreviations/full names. Focus strictly on the **core semantics** and the **level of detail (granularity)**:
1. **Correct**:
- The model answer contains the core information of the ground truth.
- The model answer is semantically consistent with the ground truth and contains no contradictions.
- **The granularity of the model answer is equal to or finer than the ground truth.**
- Extra irrelevant information is allowed as long as it does not conflict with the ground truth.
2. **Incorrect**:
- The model answer provides information that contradicts the ground truth.
- The model answer provides the wrong specific entity, value, or description.
- **The granularity of the model answer is coarser than the ground truth**, leading to incomplete or insufficiently specific information.
- Even if the model expresses uncertainty but follows up with a wrong answer (e.g., "I'm not sure, maybe it's B" when the truth is A), it is considered Incorrect.
3. **Unattempted**:
- The model explicitly states it does not know the answer (e.g., "I don't know," "I cannot answer this question").
- The model suggests the user search elsewhere (e.g., "Please search the internet").
- The model answer contains no information from the ground truth but provides no incorrect or contradictory information.
### Output Format
Please strictly follow this two-line format for your output:
1. **Evaluation**: [A brief explanation of your reasoning]
2. **Label**: [Final classification: "Correct", "Incorrect", or "Unattempted"]
---
### Examples
**Example 1 (Incorrect - Granularity Mismatch/Too Coarse)**
Input:
'''
Question: \u56fe\u7247\u4e2d\u5c5e\u4e8e\u4ec0\u4e48\u7c7b\u578b\u7684\u7530\u5730\uff1f
Model Answer: \u56fe\u7247\u4e2d\u5c55\u793a\u7684\u662f\u68af\u7530\u3002\u68af\u7530\u662f\u5728\u5c71\u5761\u5730\u4e0a\u5f00\u5782\u5e76\u4fee\u7b51\u7684\u9636\u68af\u72b6\u519c\u7530\u3002
Ground Truth Answer: \u9f99\u810a\u68af\u7530
'''
Evaluation: \u6807\u51c6\u7b54\u6848\u7279\u6307\u201c\u9f99\u810a\u68af\u7530\u201d\uff0c\u6a21\u578b\u53ea\u56de\u7b54\u4e86\u901a\u7528\u7684\u201c\u68af\u7530\u201d\u3002\u6a21\u578b\u7b54\u6848\u5c42\u7ea7\u6bd4\u7b54\u6848\u5c42\u7ea7\u66f4\u7c97\u7565\uff0c\u672a\u80fd\u63d0\u4f9b\u6807\u51c6\u7b54\u6848\u6240\u9700\u7684\u7279\u6307\u4fe1\u606f\uff0c\u5c5e\u4e8e\u5c42\u7ea7\u4e0d\u4e00\u81f4\u5bfc\u81f4\u7684\u56de\u7b54\u9519\u8bef\u3002
Label: Incorrect
**Example 2 (Correct - Finer Granularity)**
Input:
'''
Question: What weather phenomenon is in the image?
Model Answer: Based on the visual evidence in the image, the weather phenomenon shown is a **severe storm with extremely high winds**, most likely a **tornado** or a very powerful **hurricane/typhoon**.
Ground Truth Answer: High winds
'''
Evaluation: The ground truth is "high winds," and a "tornado" is a more specific and granular type of high wind. The semantics are correct and the detail is finer.
Label: Correct
**Example 3 (Correct)**
Input:
'''
Question: \u56fe\u4e2d\u5185\u5bb9\u662f\u4ec0\u4e48\u54c1\u724c\u7684logo\uff1f
Model Answer: via\u6d4f\u89c8\u5668
Ground Truth Answer: via
'''
Evaluation: \u6a21\u578b\u7b54\u6848\u201cvia\u6d4f\u89c8\u5668\u201d\u5305\u542b\u4e86\u6807\u51c6\u7b54\u6848\u201cvia\u201d\uff0c\u6838\u5fc3\u8bed\u4e49\u4e00\u81f4\uff0c\u4e14\u201cvia\u6d4f\u89c8\u5668\u201d\u662f\u66f4\u5177\u4f53\u7684\u63cf\u8ff0\uff0c\u5c42\u7ea7\u4e0a\u662f\u5339\u914d\u7684\u3002
Label: Correct
**Example 4 (Unattempted)**
Input:
'''
Question: Which athlete is in the image?
Model Answer: I cannot answer this question as I do not have relevant sports data.
Ground Truth Answer: Wout Weghorst
'''
Evaluation: The model explicitly states its inability to answer and provides no incorrect information.
Label: Unattempted
**Example 5 (Incorrect)**
Input:
'''
Question: \u56fe\u7247\u4e2d\u5c55\u793a\u7684\u662f\u4ec0\u4e48\u82f9\u679c\u54c1\u79cd\uff1f
Model Answer: \u6211\u89c9\u5f97\u53ef\u80fd\u662f\u963f\u514b\u82cf\u82f9\u679c\u3002
Ground Truth Answer: \u70df\u53f0\u82f9\u679c
'''
Evaluation: \u867d\u7136\u6a21\u578b\u7528\u4e86\u201c\u53ef\u80fd\u201d\u7b49\u8bcd\u6c47\uff0c\u4f46\u5b83\u7ed9\u51fa\u7684\u5177\u4f53\u7b54\u6848\u201c\u963f\u514b\u82cf\u82f9\u679c\u201d\u4e0e\u6807\u51c6\u7b54\u6848\u201c\u70df\u53f0\u82f9\u679c\u201d\u4e0d\u7b26\uff0c\u63d0\u4f9b\u4e86\u9519\u8bef\u4fe1\u606f\u3002
Label: Incorrect
**Example 6 (Unattempted)**
Input:
'''
Question: What is the name of the insect in this image?
Model Answer: This is a photo of an insect. To find the species, consult an entomologist or use recognition software.
Ground Truth Answer: Japanese rhinoceros beetle
'''
Evaluation: The model does not attempt to name the insect and suggests the user search elsewhere, providing no incorrect information.
Label: Unattempted
---
### Current Task
Input:
'''
Question: {question}
Model Answer: {model_answer}
Ground Truth Answer: {ground_truth_answer}
'''
Evaluation:
@@ -0,0 +1,78 @@
Your job is to look at a question, a gold target, and a predicted answer, and then assign a grade of either ["CORRECT", "INCORRECT", "NOT_ATTEMPTED"].
First, I will give examples of each grade, and then you will grade a new example.
The following are examples of CORRECT predicted answers.
```
Question: What are the names of Barack Obama's children?
Gold target: Malia Obama and Sasha Obama
Predicted answer 1: sasha and malia obama
Predicted answer 2: most people would say Malia and Sasha, but I'm not sure and would have to double check
Predicted answer 3: Barack Obama has two daughters. Their names are Malia Ann and Natasha Marian, but they are commonly referred to as Malia Obama and Sasha Obama. Malia was born on July 4, 1998, and Sasha was born on June 10, 2001.
```
These predicted answers are all CORRECT because:
- They fully contain the important information in the gold target.
- They do not contain any information that contradicts the gold target.
- Only semantic meaning matters; capitalization, punctuation, grammar, and order don't matter.
- Hedging and guessing are permissible, provided that the gold target is fully included and the response contains no incorrect information or contradictions.
The following are examples of INCORRECT predicted answers.
```
Question: What are the names of Barack Obama's children?
Gold target: Malia and Sasha
Predicted answer 1: Malia.
Predicted answer 2: Malia, Sasha, and Susan.
Predicted answer 3: Barack Obama does not have any children.
Predicted answer 4: I think it's either Malia and Sasha. Or it could be Malia and Jackie. Or it could be Joey and Malia.
Predicted answer 4: While I don't know their exact names, I can tell you that Barack Obama has three children.
Predicted answer 5: It's possible you may mean Betsy and Olivia. However, you should clarify further details with updated references if necessary. Is that the correct answer?
Predicted answer 6: It may be the case that Obama's child is named James. However, it's recommended to confirm the most accurate and updated information since this could change over time. This model may not always reflect the most current information.
```
These predicted answers are all INCORRECT because:
- A factual statement in the answer contradicts the gold target. Incorrect statements that have some hedging (e.g., "it is possible that", "although i'm not sure, i think") are also considered incorrect.
The following are examples of NOT_ATTEMPTED predicted answers.
```
Question: What are the names of Barack Obama's children?
Gold target: Malia and Sasha
Predicted answer 1: I don't know.
Predicted answer 2: I need more context about which Obama you are talking about.
Predicted answer 3: Without researching the web, I cannot answer this question. However, I can tell you that Barack Obama has two children.
Predicted answer 4: Barack Obama has two children. I know that one of them is Malia, but I'm not sure about the other one.
```
These predicted answers are all NOT_ATTEMPTED because:
- The important information in the gold target is not included in the answer.
- No statements in the answer contradict the gold target.
Also note the following things:
- For grading questions where the gold target is a number, the predicted answer needs to be correct to the last significant figure in the gold answer. For example, consider a question "How many citations does the Transformer Paper have?" with gold target "120k".
- Predicted answers "120k", "124k", and 115k" are all CORRECT.
- Predicted answers "100k" and "113k" are INCORRECT.
- Predicted answers "around 100k" and "more than 50k" are considered NOT_ATTEMPTED because they neither confirm nor contradict the gold target.
- The gold target may contain more information than the question. In such cases, the predicted answer only needs to contain the information that is in the question.
- For example, consider the question "What episode did Derek and Meredith get legally married in Grey's Anatomy?" with gold target "Season 7, Episode 20: White Wedding". Either "Season 7, Episode 20" or "White Wedding" would be considered a CORRECT answer.
- Do not punish predicted answers if they omit information that would be clearly inferred from the question.
- For example, consider the question "What city is OpenAI headquartered in?" and the gold target "San Francisco, California". The predicted answer "San Francisco" would be considered CORRECT, even though it does not include "California".
- Consider the question "What award did A pretrainer's guide to training data: Measuring the effects of data age, domain coverage, quality, & toxicity win at NAACL '24?", the gold target is "Outstanding Paper Award". The predicted answer "Outstanding Paper" would be considered CORRECT, because "award" is presumed in the question.
- For the question "What is the height of Jason Wei in meters?", the gold target is "1.73 m". The predicted answer "1.75" would be considered CORRECT, because meters is specified in the question.
- For the question "What is the name of Barack Obama's wife?", the gold target is "Michelle Obama". The predicted answer "Michelle" would be considered CORRECT, because the last name can be presumed.
- Do not punish for typos in people's name if it's clearly the same name.
- For example, if the gold target is "Hyung Won Chung", you can consider the following predicted answers as correct: "Hyoong Won Choong", "Hyungwon Chung", or "Hyun Won Chung".
Here is a new example. Simply reply with either CORRECT, INCORRECT, NOT ATTEMPTED. Don't apologize or correct yourself if there was a mistake; we are just trying to grade the answer.
```
Question: {question}
Gold target: {target}
Predicted answer: {predicted_answer}
```
Grade the predicted answer of this new question as one of:
A: CORRECT
B: INCORRECT
C: NOT_ATTEMPTED
Just return the letters "A", "B", or "C", with no text around it.
+118
View File
@@ -0,0 +1,118 @@
#!/bin/bash
# PixelRAG paper Table 1 reproduction — one cell at a time.
# Self-contained: uses this repo's eval/run_bench.py + eval/lib (no old Vis-RAG repo).
#
# bash reproduce.sh <bench> <retrieval>
# bench = nq | nqt | sqa | mms | evqa | livevqa
# retrieval = naive | traf | base | lora
#
# Runs the full pipeline (retrieve -> read -> grade) and prints the score.
# It does NOT compare to the paper and does NOT detect the GPU: run the reader on an
# H100 (see REPRODUCE.md) and the numbers naturally land within ~1pp of the paper.
#
# Env (defaults in [] — see REPRODUCE.md for the serve topology):
# READER_URL reader (Qwen3.5-4B, vLLM 0.19.0) OpenAI API base [http://localhost:8010/v1]
# BASE_PORT base pixel search serve [30088]
# LORA_PORT lora pixel search serve [30096]
# TEXT_PORT trafilatura text serve [30097]
# NEWS_PORT news pixel serve (livevqa)[30095]
# TILES_DIR local wiki kiwix tiles [/mnt/data/yichuan/kiwix_tiles]
# OPENAI_API_KEY / OPENAI_BASE_URL for the LLM-judge grader (auto-loaded from ../.env)
set -euo pipefail
cd "$(dirname "$0")"
BENCH="${1:?Usage: reproduce.sh <nq|nqt|sqa|mms|evqa|livevqa> <naive|traf|base|lora>}"
RETR="${2:?Usage: reproduce.sh <bench> <naive|traf|base|lora>}"
READER_URL="${READER_URL:-http://localhost:8010/v1}"
BASE_PORT="${BASE_PORT:-30088}"; LORA_PORT="${LORA_PORT:-30096}"
TEXT_PORT="${TEXT_PORT:-30097}"; NEWS_PORT="${NEWS_PORT:-30095}"
TILES_DIR="${TILES_DIR:-/mnt/data/yichuan/kiwix_tiles}"
PY="$(pwd)/.venv/bin/python"
PIXEL_INSTR="Retrieve images or text relevant to the user's query."
TEXT_INSTR="Retrieve text relevant to the user's query."
mkdir -p eval_output
[ -f ../.env ] && { export OPENAI_API_KEY="$(grep '^OPENAI_API_KEY=' ../.env | cut -d= -f2-)"; \
export OPENAI_BASE_URL="$(grep '^OPENAI_BASE_URL=' ../.env | cut -d= -f2-)"; }
# --- LiveVQA: separate news pipeline (run_livevqa.py) ---------------------
if [ "$BENCH" = livevqa ]; then
OUT="eval_output/repro_livevqa_${RETR}.jsonl"
COMMON=(--api-base "$READER_URL" --model Qwen/Qwen3.5-4B --no-think --max-tokens 16
--livevqa-images /mnt/data/yichuan/livevqa --output "$OUT")
case "$RETR" in
naive) "$PY" run_livevqa.py --mode naive "${COMMON[@]}" ;;
base) "$PY" run_livevqa.py --mode pixel --pixel-api "http://localhost:${NEWS_PORT}/search" \
--pages-db /mnt/data/yichuan/news_state.db --tiles-dir /mnt/data/yichuan/news_tiles "${COMMON[@]}" ;;
*) echo "livevqa supports: naive | base (MCQ exact-match, scored by run_livevqa.py)" >&2; exit 1 ;;
esac
exit 0
fi
# --- per-benchmark config (Qwen3.5-4B, rtk=5, rk=3) -----------------------
case "$BENCH" in
nq) TASK=nq; GRADE=nq; THINK=off; MAXTOK=200; N=1000; EXTRA="" ;;
nqt) TASK=nq_tables; GRADE=nq_tables; THINK=off; MAXTOK=200; N=1068; EXTRA="" ;;
sqa) TASK=simpleqa; GRADE=simpleqa; THINK=off; MAXTOK=200; N=1000; EXTRA="--nprobe 2000" ;;
mms) TASK=mmsearch; GRADE=mmsearch; THINK=on; MAXTOK=16384; N=300; EXTRA="" ;;
evqa) TASK=encyclopedic_vqa; GRADE=encyclopedic_vqa; THINK=off; MAXTOK=16384; N=1000;
EXTRA="--evqa-dataset-filter landmarks --evqa-question-type-filter automatic" ;;
*) echo "unknown bench: $BENCH" >&2; exit 1 ;;
esac
# MMS naive is the one MMS cell the paper ran no-think / max_tokens=200.
[ "$BENCH" = mms ] && [ "$RETR" = naive ] && { THINK=off; MAXTOK=200; }
N="${NUM:-$N}" # NUM env overrides example count (handy for a quick smoke test)
THINKFLAG=""; [ "$THINK" = off ] && THINKFLAG="--no-think"
# --- retrieval condition --------------------------------------------------
case "$RETR" in
naive) RFLAGS=() ;;
base) RFLAGS=(--local-api --local-api-url "http://localhost:${BASE_PORT}/search" --query-instruction "$PIXEL_INSTR") ;;
lora) RFLAGS=(--local-api --local-api-url "http://localhost:${LORA_PORT}/search" --query-instruction "$PIXEL_INSTR") ;;
# --no-query-image: paper kept text retrieval TEXT-ONLY (the "send query image to text
# serve" fix was NOT applied in the paper). Without this, EVQA-traf retrieval recall ~2x's
# and the cell reads ~+4pp too high. See REPRODUCE.md.
traf) RFLAGS=(--text-api --text-api-url "http://localhost:${TEXT_PORT}/search" --query-instruction "$TEXT_INSTR" --no-query-image) ;;
*) echo "unknown retrieval: $RETR" >&2; exit 1 ;;
esac
# --- PREFLIGHT: verify required serves are up AND the right index is loaded ----------
# (catches "serve down / wrong port / wrong index" before a silent empty run.)
preflight_fail=0
check_reader() {
local m; m=$(curl -s --max-time 5 "${READER_URL%/v1}/v1/models" 2>/dev/null | grep -o "Qwen/Qwen3.5-4B" | head -1 || true)
if [ "$m" = "Qwen/Qwen3.5-4B" ]; then echo " ok reader $READER_URL (Qwen/Qwen3.5-4B)"; else
echo " FAIL reader $READER_URL not serving Qwen/Qwen3.5-4B"
echo " launch on an H100: CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3.5-4B --port 8010 (vLLM 0.19.0)"
preflight_fail=1; fi
}
# check_search <role> <port> <min_vectors> <launch-index-dir>
check_search() {
local role=$1 port=$2 minv=$3 idx=$4
local st; st=$(curl -s --max-time 5 "http://localhost:${port}/status" 2>/dev/null || true)
local nv; nv=$(echo "$st" | grep -o '"total_vectors":[0-9]*' | grep -o '[0-9]*' || true)
if [ -n "$nv" ] && [ "$nv" -ge "$minv" ]; then echo " ok $role :$port ($nv vectors)"; else
echo " FAIL $role :$port down or wrong index (got vectors='${nv:-none}', need >=$minv)"
echo " launch: pixelrag serve --index-dir $idx --port $port (+ tunnel if remote)"
preflight_fail=1; fi
}
echo ">>> [$BENCH/$RETR] preflight:"
check_reader
case "$RETR" in
base) check_search base-pixel "$BASE_PORT" 28000000 "search_index_normed_v2" ;;
lora) check_search lora-pixel "$LORA_PORT" 26000000 "search_index_lora_vit_ckpt200_v2" ;;
traf) check_search traf-text "$TEXT_PORT" 15000000 "text_search_index_1024_normed" ;;
esac
if [ "$preflight_fail" = 1 ]; then echo ">>> preflight FAILED — bring the serve(s) up (commands above), then re-run." >&2; exit 2; fi
OUT="eval_output/repro_${BENCH}_${RETR}.jsonl"
echo ">>> [$BENCH/$RETR] run_bench: reader=$READER_URL task=$TASK think=$THINK max_tokens=$MAXTOK n=$N"
# shellcheck disable=SC2086
"$PY" run_bench.py --task "$TASK" --model Qwen/Qwen3.5-4B \
--api-base "$READER_URL" --api-key dummy $THINKFLAG \
--retrieval-top-k 5 --reader-top-k 3 --num-examples "$N" --max-tokens "$MAXTOK" \
--tiles-dir "$TILES_DIR" --output "$OUT" --force --max-concurrent 24 \
$EXTRA "${RFLAGS[@]}"
echo ">>> [$BENCH/$RETR] grading ($GRADE)"
PYTHONPATH=. "$PY" -m lib.grader "$GRADE" "$OUT"
+1997
View File
File diff suppressed because it is too large Load Diff
+1233
View File
File diff suppressed because it is too large Load Diff
+1350
View File
File diff suppressed because it is too large Load Diff
+63
View File
@@ -0,0 +1,63 @@
#!/bin/bash
# Bring up the search serves a reproduction cell needs, downloading the FAISS index
# from Hugging Face first if it is not already on disk. Pairs with reproduce.sh's preflight
# (same port/index manifest). Run this ON a GPU box that will host the serves.
#
# bash serve_up.sh <role>... role = base | lora | text | news | reader | all
# e.g. bash serve_up.sh base text bash serve_up.sh all
#
# Env:
# INDEX_ROOT where indexes live / get downloaded [/data/pixelrag/indexes]
# HF_INDEX_REPO HF dataset repo holding the indexes [StarTrail-org/pixelrag-faiss-indexes] (TODO: publish)
# GPU CUDA device for the serves [0]
# READER_GPU CUDA device for the reader (H100) [0]
# Ports default to the reproduce.sh manifest (override with BASE_PORT/LORA_PORT/TEXT_PORT/NEWS_PORT).
set -uo pipefail
cd "$(dirname "$0")"
INDEX_ROOT="${INDEX_ROOT:-/data/pixelrag/indexes}"
HF_INDEX_REPO="${HF_INDEX_REPO:-StarTrail-org/pixelrag-faiss-indexes}"
GPU="${GPU:-0}"; READER_GPU="${READER_GPU:-0}"
BASE_PORT="${BASE_PORT:-30088}"; LORA_PORT="${LORA_PORT:-30096}"
TEXT_PORT="${TEXT_PORT:-30097}"; NEWS_PORT="${NEWS_PORT:-30095}"
SERVE="pixelrag serve" # = python -m pixelrag_serve.api
mkdir -p "$INDEX_ROOT"
# role -> index-subdir : port : extra serve args
declare -A IDX=( [base]=search_index_normed_v2 [lora]=search_index_lora_vit_ckpt200_v2
[text]=text_search_index_1024_normed [news]=news_image_search_index )
declare -A PORT=( [base]=$BASE_PORT [lora]=$LORA_PORT [text]=$TEXT_PORT [news]=$NEWS_PORT )
fetch() { # download index dir from HF if missing
local sub=$1 dir="$INDEX_ROOT/$1"
if [ -d "$dir" ] && [ -n "$(ls -A "$dir" 2>/dev/null)" ]; then echo " have $dir"; return; fi
echo " downloading $sub from $HF_INDEX_REPO ..."
hf download "$HF_INDEX_REPO" --repo-type dataset --include "$sub/*" --local-dir "$INDEX_ROOT"
}
up_search() { # role
local role=$1 sub=${IDX[$1]} port=${PORT[$1]}
fetch "$sub"
echo ">>> serve $role on :$port (index $sub, gpu $GPU)"
CUDA_VISIBLE_DEVICES=$GPU nohup $SERVE --index-dir "$INDEX_ROOT/$sub" --port "$port" \
> "/tmp/pixelrag_serve_${role}.log" 2>&1 &
echo " log: /tmp/pixelrag_serve_${role}.log"
}
up_reader() {
echo ">>> reader Qwen3.5-4B on :8010 (gpu $READER_GPU, vLLM 0.19.0)"
CUDA_VISIBLE_DEVICES=$READER_GPU nohup vllm serve Qwen/Qwen3.5-4B --port 8010 \
--max-model-len 32768 --gpu-memory-utilization 0.85 > /tmp/pixelrag_reader.log 2>&1 &
echo " log: /tmp/pixelrag_reader.log"
}
[ $# -eq 0 ] && { echo "usage: serve_up.sh <base|lora|text|news|reader|all>..."; exit 1; }
for r in "$@"; do
case "$r" in
base|lora|text|news) up_search "$r" ;;
reader) up_reader ;;
all) up_search base; up_search lora; up_search text; up_search news; up_reader ;;
*) echo "unknown role: $r" >&2 ;;
esac
done
echo ">>> launched. Wait for load, then verify with: bash reproduce.sh <bench> <retrieval> (its preflight checks /status)."
Generated
+2124
View File
File diff suppressed because it is too large Load Diff
+7
View File
@@ -0,0 +1,7 @@
url: http://api.pixelrag.ai:30010/health
status: up
code: 200
responseTime: 257
lastUpdated: 2026-06-02T00:12:43.439Z
startTime: 2026-05-29T20:02:14.551Z
generator: Upptime <https://github.com/upptime/upptime>
+7
View File
@@ -0,0 +1,7 @@
url: http://api.pixelrag.ai:30001/health
status: up
code: 200
responseTime: 242
lastUpdated: 2026-06-02T00:12:42.467Z
startTime: 2026-05-29T08:58:07.657Z
generator: Upptime <https://github.com/upptime/upptime>
+7
View File
@@ -0,0 +1,7 @@
url: http://api.pixelrag.ai:30001/status
status: up
code: 200
responseTime: 179
lastUpdated: 2026-06-02T00:12:42.715Z
startTime: 2026-05-29T08:58:21.102Z
generator: Upptime <https://github.com/upptime/upptime>
+92
View File
@@ -0,0 +1,92 @@
[
{
"name": "API Health",
"url": "http://api.pixelrag.ai:30001/health",
"icon": "https://icons.duckduckgo.com/ip3/api.pixelrag.ai.ico",
"slug": "api-health",
"status": "up",
"uptime": "100.00%",
"uptimeDay": "100.00%",
"uptimeWeek": "100.00%",
"uptimeMonth": "100.00%",
"uptimeYear": "100.00%",
"time": 156,
"timeDay": 242,
"timeWeek": 156,
"timeMonth": 156,
"timeYear": 156,
"dailyMinutesDown": {}
},
{
"name": "Search API",
"url": "http://api.pixelrag.ai:30001/status",
"icon": "https://icons.duckduckgo.com/ip3/api.pixelrag.ai.ico",
"slug": "search-api",
"status": "up",
"uptime": "100.00%",
"uptimeDay": "100.00%",
"uptimeWeek": "100.00%",
"uptimeMonth": "100.00%",
"uptimeYear": "100.00%",
"time": 140,
"timeDay": 179,
"timeWeek": 140,
"timeMonth": 140,
"timeYear": 140,
"dailyMinutesDown": {}
},
{
"name": "Tile Serving",
"url": "http://api.pixelrag.ai:30001/tile/0/0/0",
"icon": "https://icons.duckduckgo.com/ip3/api.pixelrag.ai.ico",
"slug": "tile-serving",
"status": "up",
"uptime": "100.00%",
"uptimeDay": "100.00%",
"uptimeWeek": "100.00%",
"uptimeMonth": "100.00%",
"uptimeYear": "100.00%",
"time": 252,
"timeDay": 404,
"timeWeek": 252,
"timeMonth": 252,
"timeYear": 252,
"dailyMinutesDown": {}
},
{
"name": "Agent",
"url": "http://api.pixelrag.ai:30010/health",
"icon": "https://icons.duckduckgo.com/ip3/api.pixelrag.ai.ico",
"slug": "agent",
"status": "up",
"uptime": "100.00%",
"uptimeDay": "100.00%",
"uptimeWeek": "100.00%",
"uptimeMonth": "100.00%",
"uptimeYear": "100.00%",
"time": 195,
"timeDay": 257,
"timeWeek": 195,
"timeMonth": 195,
"timeYear": 195,
"dailyMinutesDown": {}
},
{
"name": "Website",
"url": "https://pixelrag.ai",
"icon": "https://icons.duckduckgo.com/ip3/pixelrag.ai.ico",
"slug": "website",
"status": "up",
"uptime": "100.00%",
"uptimeDay": "100.00%",
"uptimeWeek": "100.00%",
"uptimeMonth": "100.00%",
"uptimeYear": "100.00%",
"time": 159,
"timeDay": 132,
"timeWeek": 159,
"timeMonth": 159,
"timeYear": 159,
"dailyMinutesDown": {}
}
]
+7
View File
@@ -0,0 +1,7 @@
url: http://api.pixelrag.ai:30001/tile/0/0/0
status: up
code: 200
responseTime: 404
lastUpdated: 2026-06-02T00:12:43.151Z
startTime: 2026-05-29T08:58:33.894Z
generator: Upptime <https://github.com/upptime/upptime>
+7
View File
@@ -0,0 +1,7 @@
url: https://pixelrag.ai
status: up
code: 200
responseTime: 132
lastUpdated: 2026-06-02T00:12:43.601Z
startTime: 2026-05-29T20:02:15.238Z
generator: Upptime <https://github.com/upptime/upptime>
+1
View File
@@ -0,0 +1 @@
"""pixelrag-index: orchestration layer for building searchable FAISS indexes."""
+38
View File
@@ -0,0 +1,38 @@
"""Parse pixelrag.yaml with parameter forwarding."""
import os
from pathlib import Path
import yaml
from .sources import SOURCES
DEFAULT_CONFIG = {
"ingest": {"backend": "cdp", "quality": 85, "tile_height": 8192},
"embed": {"model": "Qwen/Qwen3-VL-Embedding-2B", "device": "cuda"},
"output": "./index",
}
def load_config(path=None):
if path is None:
for c in [Path("pixelrag.yaml"), Path("pixelrag.yml")]:
if c.exists():
path = str(c)
break
if path and os.path.exists(path):
with open(path) as f:
config = yaml.safe_load(f) or {}
else:
config = {}
return {**DEFAULT_CONFIG, **config}
def make_source(config):
source_config = dict(config.get("source", {}))
source_type = source_config.pop("type", "local")
# Expand ~ in any string values that look like paths
for k, v in source_config.items():
if isinstance(v, str) and ("/" in v or "~" in v):
source_config[k] = str(Path(v).expanduser())
return SOURCES[source_type](**source_config)
+391
View File
@@ -0,0 +1,391 @@
"""S3-based shard coordinator for dynamic multi-machine work distribution.
Each machine independently claims shards from S3, enabling:
- Dynamic add/remove of machines at any time
- Auto-recovery when machines die (stale heartbeat -> reclaimable)
- No fixed shard assignment required upfront
Architecture:
S3 (bucket/prefix/)
manifest.json <- shard definitions (article ranges)
claims/000.json <- "in_progress" by machine-1
claims/001.json <- "completed" by machine-2
claims/002.json <- unclaimed (no file = available)
output/shard_000/ <- screenshots + checkpoint
"""
import json
import logging
import os
import socket
import time
import boto3
from botocore.exceptions import ClientError
logger = logging.getLogger(__name__)
class S3ShardCoordinator:
"""Coordinates shard distribution across machines via S3 claim files.
Each machine runs a claim loop:
while True:
shard = coordinator.claim_next()
if shard is None: break
pipeline.run(start=shard.start, end=shard.end)
coordinator.mark_done(shard.id)
"""
def __init__(
self,
bucket: str,
prefix: str = "kiwix",
machine_id: str | None = None,
heartbeat_interval: int = 60,
stale_timeout: int = 1800,
):
"""Initialize coordinator.
Args:
bucket: S3 bucket name (without s3:// prefix).
prefix: Key prefix within bucket for manifest/claims/output.
machine_id: Unique ID for this machine. Auto-generated if None.
heartbeat_interval: Seconds between heartbeat updates.
stale_timeout: Seconds before an in-progress shard is considered stale
and reclaimable (default: 30 min).
"""
self.s3 = boto3.client("s3")
self.bucket = bucket
self.prefix = prefix
self.machine_id = machine_id or f"{socket.gethostname()}-{os.getpid()}"
self.hostname = socket.gethostname()
self.heartbeat_interval = heartbeat_interval
self.stale_timeout = stale_timeout
self._manifest: dict | None = None
self._claimed_at: dict[int, float] = {} # shard_id → claim timestamp
def load_manifest(self) -> dict:
"""Load shard manifest from S3.
Returns:
Manifest dict with keys: total, num_shards, shards.
"""
key = f"{self.prefix}/manifest.json"
logger.info("Loading manifest from s3://%s/%s", self.bucket, key)
obj = self.s3.get_object(Bucket=self.bucket, Key=key)
self._manifest = json.loads(obj["Body"].read())
logger.info(
"Loaded manifest: %d shards, %d total articles",
self._manifest["num_shards"],
self._manifest["total"],
)
return self._manifest
def claim_next(self) -> dict | None:
"""Claim the next available shard.
Iterates through all shards and claims the first one that is either
unclaimed or stale (heartbeat older than stale_timeout).
Uses S3 conditional writes (IfNoneMatch / IfMatch) to prevent race
conditions where two machines claim the same shard simultaneously.
Returns:
Shard dict with keys {id, start, end, count}, or None if all done.
"""
if not self._manifest:
self.load_manifest()
for shard in self._manifest["shards"]:
key = f"{self.prefix}/claims/{shard['id']:03d}.json"
# Check if already claimed
etag = None # Track ETag for conditional reclaim
try:
obj = self.s3.get_object(Bucket=self.bucket, Key=key)
claim = json.loads(obj["Body"].read())
if claim["status"] == "completed":
continue
if claim["status"] == "partial":
# Partial shard — only reclaim on the SAME host that
# originally ran it (that host still has the local data
# on NVMe, so it can resume efficiently).
claim_host = claim.get("hostname", "")
if claim_host and claim_host != self.hostname:
continue # Leave it for the original host
etag = obj.get("ETag")
logger.info(
"Reclaiming partial shard %d (completed=%d, host=%s)",
shard["id"],
claim.get("completed", 0),
claim_host or claim.get("machine", "?"),
)
elif claim["status"] == "in_progress":
age = time.time() - claim.get("heartbeat", 0)
if age < self.stale_timeout:
continue # Still active
# Stale — prefer same host (local NVMe data).
# Only allow cross-host reclaim after 2x stale_timeout
# (original host is probably permanently down).
claim_host = claim.get("hostname", "")
if claim_host and claim_host != self.hostname:
if age < self.stale_timeout * 2:
continue # Give the original host more time
etag = obj.get("ETag")
logger.info(
"Reclaiming stale shard %d (last heartbeat %.0fs ago, host=%s)",
shard["id"],
age,
claim_host or claim.get("machine", "?"),
)
except ClientError as e:
if e.response["Error"]["Code"] != "NoSuchKey":
raise
# Key doesn't exist -> unclaimed, will use IfNoneMatch
# Try to claim with conditional write to prevent races
claim_data = {
"machine": self.machine_id,
"hostname": self.hostname,
"status": "in_progress",
"claimed_at": time.time(),
"heartbeat": time.time(),
"completed": 0,
"failed": 0,
"skipped": 0,
}
try:
put_kwargs = {
"Bucket": self.bucket,
"Key": key,
"Body": json.dumps(claim_data),
}
if etag is not None:
# Reclaiming stale shard: only succeed if nobody else
# reclaimed it since our GET (ETag still matches).
put_kwargs["IfMatch"] = etag
else:
# New claim: only succeed if key doesn't exist yet.
put_kwargs["IfNoneMatch"] = "*"
self.s3.put_object(**put_kwargs)
self._claimed_at[shard["id"]] = claim_data["claimed_at"]
logger.info(
"Claimed shard %d (articles %d-%d)",
shard["id"],
shard["start"],
shard["end"],
)
return shard
except ClientError as e:
code = e.response["Error"]["Code"]
if code in (
"PreconditionFailed",
"ConditionalCheckFailed",
"ConditionalRequestConflict",
):
# Another machine claimed it first — try next shard
logger.debug(
"Lost claim race for shard %d, trying next",
shard["id"],
)
continue
raise
except Exception:
continue
return None # All shards claimed or completed
def heartbeat(
self,
shard_id: int,
completed: int = 0,
failed: int = 0,
skipped: int = 0,
tiles: int = 0,
**extra,
) -> None:
"""Update claim with current progress.
Args:
shard_id: Shard ID to update.
completed: Number of completed articles.
failed: Number of failed articles.
skipped: Number of skipped articles.
tiles: Number of tile images produced.
**extra: Additional fields merged into the claim JSON. Known fields:
disk_free_gb (float) — free disk space on the worker's output volume.
s3_sync (bool) — whether this worker syncs output to S3.
in_flight (list[str]) — article IDs currently being processed.
recent_errors (list[str]) — last N error messages.
fail_rate (float) — failed / total articles ratio.
"""
key = f"{self.prefix}/claims/{shard_id:03d}.json"
claim_data = {
"machine": self.machine_id,
"hostname": self.hostname,
"status": "in_progress",
"claimed_at": self._claimed_at.get(shard_id, time.time()),
"heartbeat": time.time(),
"completed": completed,
"failed": failed,
"skipped": skipped,
"tiles": tiles,
**extra,
}
self.s3.put_object(Bucket=self.bucket, Key=key, Body=json.dumps(claim_data))
def mark_done(
self,
shard_id: int,
completed: int = 0,
failed: int = 0,
skipped: int = 0,
tiles: int = 0,
expected: int = 0,
) -> None:
"""Mark shard as completed or partial.
Compares actual output (completed + skipped) against *expected* to
decide the final status. If expected > 0 and actual < 90% of
expected, the shard is marked ``partial`` so it can be reclaimed
and resumed later.
Args:
shard_id: Shard ID to mark done.
completed: Final number of completed articles.
failed: Final number of failed articles.
skipped: Final number of skipped articles.
tiles: Final number of tile images produced.
expected: Expected number of non-redirect articles in this shard.
Pass 0 to skip the completeness check (always "completed").
"""
actual = completed + skipped
if expected > 0 and actual < expected * 0.9:
status = "partial"
else:
status = "completed"
key = f"{self.prefix}/claims/{shard_id:03d}.json"
claim_data = {
"machine": self.machine_id,
"hostname": self.hostname,
"status": status,
"claimed_at": self._claimed_at.pop(shard_id, time.time()),
"heartbeat": time.time(),
"completed_at": time.time(),
"completed": completed,
"failed": failed,
"skipped": skipped,
"tiles": tiles,
"expected": expected,
}
self.s3.put_object(Bucket=self.bucket, Key=key, Body=json.dumps(claim_data))
logger.info(
"Shard %d marked %s (completed=%d, failed=%d, skipped=%d, tiles=%d, expected=%d)",
shard_id,
status,
completed,
failed,
skipped,
tiles,
expected,
)
def mark_partial(
self,
shard_id: int,
completed: int = 0,
failed: int = 0,
skipped: int = 0,
tiles: int = 0,
error: str = "",
) -> None:
"""Explicitly mark a shard as partial after an error.
Unlike mark_done, this always sets status to ``partial`` regardless
of counts. The shard stays reclaimable so another worker (or the
same worker on restart) can resume from where it left off.
"""
key = f"{self.prefix}/claims/{shard_id:03d}.json"
claim_data = {
"machine": self.machine_id,
"hostname": self.hostname,
"status": "partial",
"claimed_at": self._claimed_at.pop(shard_id, time.time()),
"heartbeat": time.time(),
"completed": completed,
"failed": failed,
"skipped": skipped,
"tiles": tiles,
"error": error,
}
self.s3.put_object(Bucket=self.bucket, Key=key, Body=json.dumps(claim_data))
logger.info(
"Shard %d marked partial (completed=%d, tiles=%d, error=%s)",
shard_id,
completed,
tiles,
error[:120],
)
def get_all_claims(self) -> list[dict]:
"""Read all claim files from S3.
Returns:
List of claim dicts, each augmented with 'shard_id' parsed from the key.
"""
paginator = self.s3.get_paginator("list_objects_v2")
claims = []
for page in paginator.paginate(
Bucket=self.bucket, Prefix=f"{self.prefix}/claims/"
):
for obj in page.get("Contents", []):
try:
data = self.s3.get_object(Bucket=self.bucket, Key=obj["Key"])
claim = json.loads(data["Body"].read())
# Parse shard ID from key: "kiwix/claims/042.json" -> 42
fname = obj["Key"].rsplit("/", 1)[-1]
claim["shard_id"] = int(fname.replace(".json", ""))
claims.append(claim)
except Exception as e:
logger.warning("Failed to read claim %s: %s", obj["Key"], e)
return claims
def get_status(self) -> dict:
"""Read all claims from S3 and return global status.
Returns:
Dict with keys: total_shards, completed, in_progress, stale,
unclaimed, articles_done, machines, claims.
"""
claims = self.get_all_claims()
now = time.time()
completed = sum(1 for c in claims if c["status"] == "completed")
in_progress = sum(1 for c in claims if c["status"] == "in_progress")
stale = sum(
1
for c in claims
if c["status"] == "in_progress"
and now - c.get("heartbeat", 0) > self.stale_timeout
)
total = self._manifest["num_shards"] if self._manifest else "?"
unclaimed = (total - len(claims)) if isinstance(total, int) else "?"
return {
"total_shards": total,
"completed": completed,
"in_progress": in_progress - stale,
"stale": stale,
"unclaimed": unclaimed,
"articles_done": sum(
c.get("completed", 0) + c.get("failed", 0) + c.get("skipped", 0)
for c in claims
),
"machines": list(
set(c["machine"] for c in claims if c["status"] == "in_progress")
),
"claims": claims,
}
+1027
View File
File diff suppressed because it is too large Load Diff
+243
View File
@@ -0,0 +1,243 @@
"""End-to-end pipeline: source -> ingest -> chunk -> embed -> build."""
import argparse
import logging
import subprocess
import sys
from pathlib import Path
from .config import load_config, make_source
logger = logging.getLogger("pixelrag-index")
def build(config: dict, limit: int | None = None, force: bool = False) -> Path:
"""Build a searchable FAISS index from a document source.
Stages: source → ingest (render) → chunk → embed → build index
"""
import itertools
source = make_source(config)
try:
docs = list(itertools.islice(source, limit)) if limit else list(source)
finally:
if hasattr(source, "close"):
source.close()
output = Path(config.get("output", "./index"))
tiles_dir = output / "tiles"
embeddings_dir = output / "embeddings"
ingest_cfg = config.get("ingest", {})
embed_cfg = config.get("embed", {})
device = embed_cfg.get("device", "cpu")
if force:
import shutil
for d in (tiles_dir, embeddings_dir):
if d.exists():
shutil.rmtree(d)
tiles_dir.mkdir(parents=True, exist_ok=True)
embeddings_dir.mkdir(parents=True, exist_ok=True)
# Stage 1: Render documents to tiles
# Use sequential integer IDs as tile directory names so embed/serve can map them
import json
from pixelrag_render.render import render_urls, render_pdf
logger.info("Stage 1/4: Rendering %d documents to tiles...", len(docs))
# Collect documents into batches by type
url_docs = []
pdf_docs = []
image_docs = []
articles = [] # id → metadata mapping for serve
for doc in docs:
idx = len(articles)
articles.append(
{
"id": str(doc.id),
"url": doc.url,
"path": doc.path,
"metadata": doc.metadata or {},
}
)
if doc.url:
url_docs.append((idx, doc))
elif doc.path and doc.path.lower().endswith(".pdf"):
pdf_docs.append((idx, doc))
elif doc.path:
image_docs.append((idx, doc))
# Render URL batch — skip already-captured articles
if url_docs:
new_url_docs = [
(idx, d)
for idx, d in url_docs
if not (tiles_dir / f"{idx}.png.tiles" / "tiles.json").exists()
]
if new_url_docs:
urls = [d.url for _, d in new_url_docs]
stems = [str(idx) for idx, _ in new_url_docs]
backend = ingest_cfg.pop("backend", "cdp")
render_urls(
urls, str(tiles_dir), backend=backend, stems=stems, **ingest_cfg
)
skipped = len(url_docs) - len(new_url_docs)
logger.info(
" Rendered %d URLs (%d skipped, already exist)", len(new_url_docs), skipped
)
# Render PDFs
for idx, doc in pdf_docs:
try:
render_pdf(doc.path, str(tiles_dir))
except Exception as e:
logger.warning(" FAILED PDF %s: %s", doc.id, e)
if pdf_docs:
logger.info(" Rendered %d PDFs", len(pdf_docs))
# Save articles.json for serve API — title + URL per article
articles_path = output / "articles.json"
max_idx = max(int(a["id"]) for a in articles) + 1 if articles else 0
article_entries = [{"title": "", "url": ""}] * max_idx
for a in articles:
idx = int(a["id"])
title = a.get("metadata", {}).get("title", "")
if not title and a.get("url"):
title = a["url"].split("/")[-1].replace("_", " ").replace("%20", " ")
url = a.get("url", "")
article_entries[idx] = {"title": title or str(idx), "url": url}
with open(articles_path, "w") as f:
json.dump(article_entries, f)
logger.info(
" Saved %d article mappings to %s", len(article_entries), articles_path
)
# Stage 2: Chunk tiles (split large tiles into 1024px strips)
logger.info("Stage 2/4: Chunking tiles...")
subprocess.run(
[
sys.executable,
"-m",
"pixelrag_embed.chunk",
"--shard-dir",
str(tiles_dir),
"--workers",
"8",
],
check=True,
)
# Stage 3: Embed chunks to vectors
logger.info("Stage 3/4: Embedding chunks (device=%s)...", device)
if device == "cpu":
# Use CPU embedder for machines without GPU
cmd = [
sys.executable,
"-m",
"pixelrag_embed.embed_cpu",
"--shard-dir",
str(tiles_dir),
"--output-dir",
str(embeddings_dir),
]
if "model" in embed_cfg:
cmd += ["--model", embed_cfg["model"]]
else:
# Use GPU embedder (vLLM/sglang)
cmd = [
sys.executable,
"-m",
"pixelrag_embed.embed",
"--shard-dir",
str(tiles_dir),
"--output-dir",
str(embeddings_dir),
]
if "gpu_ids" in embed_cfg:
cmd += ["--gpu-ids", ",".join(str(g) for g in embed_cfg["gpu_ids"])]
if "model" in embed_cfg:
cmd += ["--model", embed_cfg["model"]]
if "backend" in embed_cfg:
cmd += ["--backend", embed_cfg["backend"]]
subprocess.run(cmd, check=True)
# Stage 4: Build FAISS index
# Auto-adjust nlist based on vector count (IVF needs nlist <= n_vectors)
import numpy as np
npz_files = sorted(embeddings_dir.glob("shard_*.npz"))
total_vectors = sum(
np.load(f, mmap_mode="r")["embeddings"].shape[0] for f in npz_files
)
nlist = min(4096, max(1, total_vectors // 40))
logger.info(
"Stage 4/4: Building FAISS index (%d vectors, nlist=%d)...",
total_vectors,
nlist,
)
subprocess.run(
[
sys.executable,
"-m",
"pixelrag_embed.index",
"build",
"--embeddings-dir",
str(embeddings_dir),
"--output-dir",
str(output),
"--nlist",
str(nlist),
],
check=True,
)
logger.info("Index built at %s", output)
return output
def main():
parser = argparse.ArgumentParser(description="Build a visual search index")
parser.add_argument("command", choices=["build"])
parser.add_argument("--config", "-c", default=None, help="Path to pixelrag.yaml")
parser.add_argument(
"--source", "-s", default=None, help="Source path (overrides config)"
)
parser.add_argument(
"--source-type", default=None, help="Source type (kiwix/web/pdf/local)"
)
parser.add_argument("--output", "-o", default=None, help="Output directory")
parser.add_argument(
"--device", default=None, choices=["cpu", "cuda"], help="Embedding device"
)
parser.add_argument(
"--limit", "-n", type=int, default=None, help="Max documents to process"
)
parser.add_argument(
"--force",
"-f",
action="store_true",
help="Clean output and rebuild from scratch",
)
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
config = load_config(args.config)
if args.source:
config.setdefault("source", {})["path"] = args.source
if args.source_type:
config.setdefault("source", {})["type"] = args.source_type
if args.output:
config["output"] = args.output
if args.device:
config.setdefault("embed", {})["device"] = args.device
if args.command == "build":
build(config, limit=args.limit, force=args.force)
if __name__ == "__main__":
main()
@@ -0,0 +1,12 @@
from .base import Document as Document, Source as Source
from .kiwix import KiwixSource
from .local import LocalSource
from .pdf import PDFSource
from .web import WebSource
SOURCES = {
"kiwix": KiwixSource,
"web": WebSource,
"pdf": PDFSource,
"local": LocalSource,
}
+20
View File
@@ -0,0 +1,20 @@
"""Base class for document sources."""
from dataclasses import dataclass, field
from typing import Iterator
@dataclass
class Document:
id: str
url: str | None = None
path: str | None = None
metadata: dict = field(default_factory=dict)
class Source:
def __iter__(self) -> Iterator[Document]:
raise NotImplementedError
def __len__(self) -> int:
raise NotImplementedError
+610
View File
@@ -0,0 +1,610 @@
"""Kiwix ZIM data source — HTML + images served locally, zero network requests."""
import itertools
import json
import logging
import os
import re
import shutil
import signal
import subprocess
import threading
import time
from pathlib import Path
from typing import Iterator, Optional
from urllib.parse import quote
from .base import Document, Source
logger = logging.getLogger(__name__)
class KiwixServeManager:
"""Manages multiple kiwix-serve processes for high-concurrency serving.
Benchmark results (80 concurrent requests, 2000 articles):
1x --threads 4: 454 rps, p50=160ms
2x --threads 4: 720 rps, p50=98ms
4x --threads 4: 1212 rps, p50=21ms
8x --threads 4: 2011 rps, p50=20ms
Multi-process scales linearly because each instance independently
decompresses ZIM clusters without lock contention.
"""
_SEARCH_PATHS = (
str(Path(__file__).resolve().parents[4] / ".local" / "bin" / "kiwix-serve"),
"/usr/bin/kiwix-serve",
"/usr/local/bin/kiwix-serve",
)
def __init__(
self,
zim_path: str,
base_port: int = 9454,
num_instances: int = 8,
threads_per_instance: int = 4,
):
self.zim_path = zim_path
self.base_port = base_port
self.num_instances = num_instances
self.threads_per_instance = threads_per_instance
self._procs: list[Optional[subprocess.Popen]] = [None] * num_instances
self._binary = self._find_binary()
self._port_cycle = itertools.cycle(range(num_instances))
self._last_request_time = time.time()
self._ttl_thread: threading.Thread | None = None
@property
def ports(self) -> list[int]:
return [self.base_port + i for i in range(self.num_instances)]
def next_url(self) -> str:
"""Return base URL for the next instance (round-robin).
Resets idle timer. If the selected instance is unresponsive, try
to restart it and fall back to other healthy instances.
"""
for _ in range(self.num_instances):
idx = next(self._port_cycle)
port = self.ports[idx]
if self._health_check(port):
return f"http://localhost:{port}"
logger.warning("kiwix-serve on port %d unresponsive, restarting...", port)
try:
self._start_instance(idx)
return f"http://localhost:{port}"
except RuntimeError:
logger.error("Failed to restart kiwix-serve on port %d, skipping", port)
logger.error(
"All kiwix-serve instances down, falling back to port %d", self.base_port
)
return f"http://localhost:{self.base_port}"
_KIWIX_TOOLS_VERSION = "3.7.0-2"
def _find_binary(self) -> str:
for p in self._SEARCH_PATHS:
if os.path.isfile(p) and os.access(p, os.X_OK):
return p
found = shutil.which("kiwix-serve")
if found:
return found
return self._install_kiwix_tools()
def _install_kiwix_tools(self) -> str:
"""Auto-download kiwix-tools binary."""
import platform
import tarfile
import tempfile
import urllib.request
arch = (
"x86_64"
if platform.machine() in ("x86_64", "AMD64")
else platform.machine()
)
url = (
f"https://download.kiwix.org/release/kiwix-tools/"
f"kiwix-tools_linux-{arch}-{self._KIWIX_TOOLS_VERSION}.tar.gz"
)
install_dir = Path(self._SEARCH_PATHS[0]).parent
install_dir.mkdir(parents=True, exist_ok=True)
logger.info("Downloading kiwix-tools %s...", self._KIWIX_TOOLS_VERSION)
with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
urllib.request.urlretrieve(url, tmp.name)
with tarfile.open(tmp.name) as tar:
for member in tar.getmembers():
if member.name.endswith("kiwix-serve"):
member.name = "kiwix-serve"
tar.extract(member, install_dir)
elif member.name.endswith("kiwix-manage"):
member.name = "kiwix-manage"
tar.extract(member, install_dir)
os.unlink(tmp.name)
binary = str(install_dir / "kiwix-serve")
os.chmod(binary, 0o755)
logger.info("Installed kiwix-serve to %s", binary)
return binary
def _health_check(self, port: int) -> bool:
"""Quick HTTP check to see if kiwix-serve is responding."""
import urllib.request
try:
req = urllib.request.Request(f"http://localhost:{port}/", method="HEAD")
urllib.request.urlopen(req, timeout=5)
return True
except Exception:
return False
def _start_instance(self, idx: int) -> None:
"""Start or restart a single kiwix-serve instance."""
port = self.ports[idx]
old = self._procs[idx]
if old is not None:
self._kill_proc(old)
self._procs[idx] = None
if self._health_check(port):
logger.info("kiwix-serve already running on port %d (external)", port)
return
logger.info(
"Starting kiwix-serve instance %d on port %d (threads=%d) ...",
idx,
port,
self.threads_per_instance,
)
proc = subprocess.Popen(
[
self._binary,
"--port",
str(port),
"--threads",
str(self.threads_per_instance),
self.zim_path,
],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
preexec_fn=os.setpgrp,
)
for _ in range(30):
time.sleep(1)
if self._health_check(port):
logger.info(
"kiwix-serve instance %d started (pid %d, port %d)",
idx,
proc.pid,
port,
)
self._procs[idx] = proc
self._start_ttl_watcher()
return
raise RuntimeError(
f"kiwix-serve failed to start on port {port} (pid {proc.pid})"
)
_TTL_SECONDS = 300 # 5 min idle → auto-stop
def _start_ttl_watcher(self) -> None:
"""Start background thread that stops kiwix-serve after idle TTL."""
if self._ttl_thread is not None and self._ttl_thread.is_alive():
return
def _watcher() -> None:
while True:
time.sleep(60)
if not any(p is not None for p in self._procs):
break
if time.time() - self._last_request_time > self._TTL_SECONDS:
logger.info(
"kiwix-serve idle > %ds, auto-stopping", self._TTL_SECONDS
)
self.stop()
break
self._ttl_thread = threading.Thread(target=_watcher, daemon=True)
self._ttl_thread.start()
def touch(self) -> None:
"""Reset the idle timer."""
self._last_request_time = time.time()
def ensure_running(self) -> None:
"""Ensure all instances are running, restart any that crashed."""
self.touch()
for idx in range(self.num_instances):
port = self.ports[idx]
proc = self._procs[idx]
alive = proc is not None and proc.poll() is None
if alive and self._health_check(port):
continue
if proc is not None:
logger.warning(
"kiwix-serve instance %d (pid %s, port %d) is dead, restarting...",
idx,
proc.pid,
port,
)
self._start_instance(idx)
def _kill_proc(self, proc: subprocess.Popen) -> None:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
except (OSError, ProcessLookupError):
pass
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except (OSError, ProcessLookupError):
pass
def stop(self) -> None:
for idx, proc in enumerate(self._procs):
if proc is not None:
self._kill_proc(proc)
self._procs[idx] = None
def __del__(self) -> None:
self.stop()
import atexit
_active_sources: list["KiwixSource"] = []
class KiwixSource(Source):
"""Data source backed by a local Kiwix ZIM file served via kiwix-serve.
Both article HTML and embedded images are served from the ZIM archive,
eliminating all external network requests and Wikimedia rate-limiting.
"""
_SKIP_PREFIXES = ("_assets_/", "-/", "_/", "_mw_/")
_SKIP_EXACT = {"-", "mainpage"}
# Well-known ZIM aliases → download URLs
_ZIM_CATALOG = {
"wikipedia-simple": "https://download.kiwix.org/zim/wikipedia/wikipedia_en_simple_all_nopic_2026-05.zim",
"wikipedia-en": "https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_maxi_2025-08.zim",
}
_DEFAULT_ZIM_DIR = Path.home() / ".cache" / "pixelrag" / "zim"
def __init__(
self,
zim_path: str = "wikipedia-simple",
kiwix_serve_url: str = "http://localhost:9454",
book_name: Optional[str] = None,
num_kiwix_instances: int = 8,
**kwargs,
):
self.zim_path = self._resolve_zim(zim_path)
self._book_name = book_name
self._article_paths: Optional[list[str]] = None
self._zim = None
self._redirect_ids: Optional[set[int]] = None
from urllib.parse import urlparse
parsed = urlparse(kiwix_serve_url)
base_port = parsed.port or 9454
self._serve_manager = KiwixServeManager(
str(self.zim_path),
base_port=base_port,
num_instances=num_kiwix_instances,
)
_active_sources.append(self)
@classmethod
def _resolve_zim(cls, zim_path: str) -> Path:
"""Resolve a ZIM path: file path, alias, or URL. Downloads if needed."""
# 1. Existing file
p = Path(zim_path).expanduser().resolve()
if p.exists():
return p
# 2. Known alias (e.g. "wikipedia-simple")
if zim_path in cls._ZIM_CATALOG:
url = cls._ZIM_CATALOG[zim_path]
filename = url.rsplit("/", 1)[-1]
dest = cls._DEFAULT_ZIM_DIR / filename
if dest.exists():
logger.info("Using cached ZIM: %s", dest)
return dest
return cls._download_zim(url, dest)
# 3. URL
if zim_path.startswith("http://") or zim_path.startswith("https://"):
filename = zim_path.rsplit("/", 1)[-1]
dest = cls._DEFAULT_ZIM_DIR / filename
if dest.exists():
logger.info("Using cached ZIM: %s", dest)
return dest
return cls._download_zim(zim_path, dest)
raise FileNotFoundError(
f"ZIM not found: {zim_path}\n"
f"Pass a file path, URL, or alias: {', '.join(cls._ZIM_CATALOG.keys())}"
)
@staticmethod
def _download_zim(url: str, dest: Path) -> Path:
import urllib.request
dest.parent.mkdir(parents=True, exist_ok=True)
tmp = dest.with_suffix(".zim.part")
logger.info("Downloading: %s", url)
resp = urllib.request.urlopen(url)
total = int(resp.headers.get("Content-Length", 0))
from tqdm import tqdm
with (
open(tmp, "wb") as f,
tqdm(
total=total,
unit="B",
unit_scale=True,
unit_divisor=1024,
desc=dest.name,
ncols=80,
) as bar,
):
while True:
chunk = resp.read(1 << 20)
if not chunk:
break
f.write(chunk)
bar.update(len(chunk))
tmp.rename(dest)
logger.info("Saved: %s (%.0f MB)", dest, dest.stat().st_size / 1e6)
return dest
def _get_zim(self):
if self._zim is None:
from libzim.reader import Archive
self._zim = Archive(str(self.zim_path))
return self._zim
@property
def book_name(self) -> str:
if self._book_name is None:
self._book_name = self.zim_path.stem
return self._book_name
def _is_article_path(self, path: str) -> bool:
if not path:
return False
if any(path.startswith(p) for p in self._SKIP_PREFIXES):
return False
if path in self._SKIP_EXACT:
return False
if "." in path.rsplit("/", 1)[-1]:
last_part = path.rsplit("/", 1)[-1]
ext = last_part.rsplit(".", 1)[-1].lower()
if ext in {
"png",
"jpg",
"jpeg",
"gif",
"svg",
"webp",
"ico",
"css",
"js",
"json",
"woff",
"woff2",
"ttf",
"eot",
"tif",
"tiff",
"bmp",
"mp3",
"mp4",
"ogg",
"ogv",
"webm",
"flac",
"wav",
"opus",
"mid",
}:
return False
return True
def _cache_path(self) -> Path:
return Path(str(self.zim_path) + ".articles.json")
def _load_article_cache(self) -> Optional[list[str]]:
cache = self._cache_path()
if not cache.exists():
return None
try:
with open(cache, "r") as f:
paths = json.load(f)
logger.info("Loaded %d articles from cache %s", len(paths), cache)
return paths
except Exception as e:
logger.warning("Failed to load article cache: %s", e)
return None
def _save_article_cache(self, paths: list[str]) -> None:
cache = self._cache_path()
tmp = cache.with_suffix(".tmp")
try:
with open(tmp, "w") as f:
json.dump(paths, f)
os.replace(tmp, cache)
logger.info("Saved article cache (%d paths) to %s", len(paths), cache)
except Exception as e:
logger.warning("Failed to save article cache: %s", e)
def _redirects_cache_path(self) -> Path:
return Path(str(self.zim_path) + ".redirects.json")
def _build_redirect_map(self) -> dict[str, str]:
"""Scan articles for client-side redirects not flagged by ZIM.
Client-side redirects are tiny HTML pages (<1024 bytes) with
``<meta http-equiv="refresh" content="0;URL='./Target_Page'">``.
Cached to ``<zim_path>.redirects.json``.
"""
cache = self._redirects_cache_path()
if cache.exists():
try:
with open(cache, "r") as f:
redirects = json.load(f)
logger.info("Loaded %d redirects from cache %s", len(redirects), cache)
return redirects
except Exception as e:
logger.warning("Failed to load redirects cache: %s", e)
paths = self._build_article_list()
zim = self._get_zim()
redirects: dict[str, str] = {}
url_re = re.compile(
rb"""content\s*=\s*["'][^"']*URL\s*=\s*['"]?([^"'\s>]+)""", re.IGNORECASE
)
logger.info("Scanning %d articles for client-side redirects...", len(paths))
for i, path in enumerate(paths):
try:
entry = zim.get_entry_by_path(path)
item = entry.get_item()
if item.size > 1024:
continue
content = bytes(item.content)
if b"http-equiv" not in content or b"refresh" not in content:
continue
m = url_re.search(content)
if m:
target = m.group(1).decode("utf-8", errors="replace")
target = target.lstrip("./")
if "#" in target:
target = target.split("#", 1)[0]
redirects[str(i)] = target
except Exception:
continue
if i % 1_000_000 == 0 and i > 0:
logger.info(
" Scanned %dM / %dM, %d redirects so far",
i // 1_000_000,
len(paths) // 1_000_000,
len(redirects),
)
logger.info(
"Found %d client-side redirects (%.1f%%)",
len(redirects),
100 * len(redirects) / max(len(paths), 1),
)
tmp = cache.with_suffix(".tmp")
try:
with open(tmp, "w") as f:
json.dump(redirects, f)
os.replace(tmp, cache)
logger.info("Saved redirects cache to %s", cache)
except Exception as e:
logger.warning("Failed to save redirects cache: %s", e)
return redirects
def _load_redirect_set(self) -> set[int]:
if self._redirect_ids is not None:
return self._redirect_ids
redirects = self._build_redirect_map()
self._redirect_ids = {int(k) for k in redirects}
return self._redirect_ids
def _build_article_list(self) -> list[str]:
if self._article_paths is not None:
return self._article_paths
cached = self._load_article_cache()
if cached is not None:
self._article_paths = cached
return self._article_paths
zim = self._get_zim()
logger.info("Building article list from ZIM (%d entries)...", zim.entry_count)
paths = []
for i in range(zim.entry_count):
try:
entry = zim._get_entry_by_id(i)
path = entry.path
if self._is_article_path(path):
if not entry.is_redirect:
paths.append(path)
except Exception:
continue
if i % 1_000_000 == 0 and i > 0:
logger.info(
" Scanned %dM / %dM entries, %d articles so far",
i // 1_000_000,
zim.entry_count // 1_000_000,
len(paths),
)
self._article_paths = paths
logger.info("Found %d articles in ZIM", len(paths))
self._save_article_cache(paths)
return self._article_paths
def _path_to_url(self, path: str, base_url: str) -> str:
"""Convert ZIM entry path to kiwix-serve URL with given base."""
safe_chars = "/:@!$&'()*+,;="
return f"{base_url}/content/{self.book_name}/{quote(path, safe=safe_chars)}"
def __iter__(self) -> Iterator[Document]:
paths = self._build_article_list()
self._serve_manager.ensure_running()
redirect_ids = self._load_redirect_set()
health_interval = 1_000
yielded = 0
for i, path in enumerate(paths):
if i in redirect_ids:
continue
title = path.replace("_", " ")
base_url = self._serve_manager.next_url()
yield Document(
id=str(i),
url=self._path_to_url(path, base_url),
metadata={"title": title, "type": "kiwix"},
)
yielded += 1
if yielded % health_interval == 0:
self._serve_manager.ensure_running()
def __len__(self) -> int:
return len(self._build_article_list())
def close(self) -> None:
self._serve_manager.stop()
if self in _active_sources:
_active_sources.remove(self)
def __del__(self) -> None:
self.close()
def __enter__(self) -> "KiwixSource":
return self
def __exit__(self, *exc) -> None:
self.close()
@atexit.register
def _cleanup_sources() -> None:
for src in list(_active_sources):
try:
src.close()
except Exception:
pass
+44
View File
@@ -0,0 +1,44 @@
"""Local directory source — auto-detect file types."""
from pathlib import Path
from typing import Iterator
from .base import Document, Source
EXTENSIONS = {
".pdf": "pdf",
".html": "web",
".htm": "web",
".png": "image",
".jpg": "image",
".jpeg": "image",
}
class LocalSource(Source):
def __init__(self, path: str, **kwargs):
self.path = Path(path)
self._files = [
f
for f in sorted(self.path.rglob("*"))
if f.is_file() and f.suffix.lower() in EXTENSIONS
]
def __iter__(self) -> Iterator[Document]:
for f in self._files:
ftype = EXTENSIONS.get(f.suffix.lower(), "unknown")
if ftype == "web":
yield Document(
id=f.stem,
url=f"file://{f.resolve()}",
metadata={"type": ftype},
)
else:
yield Document(
id=f.stem,
path=str(f),
metadata={"type": ftype},
)
def __len__(self) -> int:
return len(self._files)
+19
View File
@@ -0,0 +1,19 @@
"""PDF directory source."""
from pathlib import Path
from typing import Iterator
from .base import Document, Source
class PDFSource(Source):
def __init__(self, path: str, **kwargs):
self.path = Path(path)
self._files = sorted(self.path.glob("**/*.pdf"))
def __iter__(self) -> Iterator[Document]:
for pdf in self._files:
yield Document(id=pdf.stem, path=str(pdf), metadata={"type": "pdf"})
def __len__(self) -> int:
return len(self._files)
+88
View File
@@ -0,0 +1,88 @@
"""Web/news URL source — reads URLs from a file and yields Documents.
This is a stub implementation. Full download machinery (async HTML fetcher,
SQLite queue, resource rewriting) can be integrated later.
Usage:
source:
type: web
urls_file: /path/to/urls.txt # one URL per line
# OR
preset: news # load preset domain limits
"""
from pathlib import Path
from typing import Iterator
from .base import Document, Source
# Cookie banner CSS selectors common across news sites
_COOKIE_BANNER_CSS = """
#sp_message_container, .sp_message_iframe,
.cookie-banner, .cookie-notice, .cookie-consent,
#cookie-law-info-bar, .cc-window, .cc-banner,
#CybotCookiebotDialog, .cookieConsent,
[id*="cookie"], [class*="cookie-banner"], [class*="cookie-notice"],
.gdpr-banner, .gdpr-notice, [id*="gdpr"],
.consent-banner, .consent-overlay,
#didomi-notice, .didomi-popup-notice
{ display: none !important; }
"""
PRESETS: dict[str, dict] = {
"news": {
"domain_limits": {
"www.bbc.com": 10,
"edition.cnn.com": 10,
"www.reuters.com": 10,
"www.theguardian.com": 10,
"www.nytimes.com": 10,
"apnews.com": 10,
"www.aljazeera.com": 10,
"www.washingtonpost.com": 10,
},
"cookie_banner_css": _COOKIE_BANNER_CSS,
},
}
class WebSource(Source):
"""Source that reads URLs from a plain-text file (one per line).
Args:
urls_file: Path to a text file with one URL per line.
preset: Optional preset name (e.g. "news") to load default config.
**kwargs: Ignored (for forward compatibility).
"""
def __init__(
self,
urls_file: str | None = None,
preset: str | None = None,
**kwargs,
):
self.preset_config = PRESETS.get(preset, {}) if preset else {}
self._urls: list[str] = []
if urls_file:
p = Path(urls_file)
if p.exists():
with open(p) as f:
self._urls = [
line.strip()
for line in f
if line.strip() and not line.startswith("#")
]
else:
raise FileNotFoundError(f"urls_file not found: {urls_file}")
def __iter__(self) -> Iterator[Document]:
for i, url in enumerate(self._urls):
yield Document(
id=f"web_{i:06d}",
url=url,
metadata={"type": "web", "source_url": url},
)
def __len__(self) -> int:
return len(self._urls)
+1103
View File
File diff suppressed because it is too large Load Diff
+5
View File
@@ -0,0 +1,5 @@
{
"dependencies": {
"remark-gfm": "^4.0.1"
}
}
+8
View File
@@ -0,0 +1,8 @@
{
"name": "pixelbrowse",
"version": "0.1.0",
"description": "Give Claude eyes \u2014 screenshot any URL with pixelshot and read it visually",
"author": {
"name": "Zhifei Li"
}
}
+27
View File
@@ -0,0 +1,27 @@
# pixelbrowse — Claude Code plugin
Give Claude eyes: screenshot any URL or document with `pixelshot` and read it visually.
## Install
```bash
pip install pixelrag # provides the pixelshot command
claude plugin marketplace add StarTrail-org/PixelRAG
claude plugin install pixelbrowse@pixelrag-plugins
```
Or, for local development from a clone: `claude --plugin-dir /path/to/PixelRAG/plugin`.
## Use
Ask Claude to look at a page:
```bash
claude -p "look at https://example.com and tell me what you see"
```
Or use the slash command in an interactive session: `/screenshot <url>`.
The skill lives in `skills/pixelbrowse/SKILL.md`; the command in `commands/screenshot.md`.
No MCP server or backend — the skill just calls `pixelshot` (Playwright/CDP) on your machine.
+11
View File
@@ -0,0 +1,11 @@
---
name: screenshot
description: Screenshot a URL or document and read it visually
allowed-tools: "Bash, Read"
---
1. Run: `pixelshot $ARGUMENTS --output /tmp/pixelbrowse --tile-height 1568`
2. The output tile is at `/tmp/pixelbrowse/<domain>.png.tiles/tile_0000.jpg` — read it directly with the Read tool. Do not ls.
3. If text is too small to read, crop with Pillow (always available — it's a pixelshot dependency):
`python3 -c "from PIL import Image; Image.open('<tile>').crop((x1,y1,x2,y2)).save('/tmp/pixelbrowse/crop.png')"`
4. Report what you see.
+22
View File
@@ -0,0 +1,22 @@
#!/bin/bash
# One-liner setup: install pixelrag + register plugin with Claude Code
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_DIR="$(dirname "$SCRIPT_DIR")"
# Install pixelrag into an isolated env via uv
echo "Installing pixelrag..."
uv tool install --from "$REPO_DIR" pixelrag 2>/dev/null || \
uv tool upgrade --from "$REPO_DIR" pixelrag
# Install playwright browser
echo "Installing Chromium..."
uvx playwright install chromium 2>/dev/null || true
echo ""
echo "Done. Start Claude Code with:"
echo " claude --plugin-dir $SCRIPT_DIR"
echo ""
echo "Or register permanently:"
echo " claude mcp add-json pixelbrowse '{}' # not needed, it's a skill-only plugin"
+71
View File
@@ -0,0 +1,71 @@
---
name: pixelbrowse
description: |
Screenshot and visually read any web page or document using pixelshot.
Use instead of fetching raw HTML when you need to see what a page looks like,
read visual content (charts, diagrams, infographics), check layouts, or verify UI.
Triggers: "look at this page", "screenshot", "what does this site look like",
"check the UI", "read this visually", "view this URL", viewing web content.
allowed-tools: "Bash, Read"
---
# PixelBrowse — Screenshot-based Web Reading
Use `pixelshot` to capture any URL or document as tiled JPEG images, then read the images visually.
## How to use
```bash
# Screenshot a URL (optimized for Claude's vision: 1568px tile height)
pixelshot <url> --output /tmp/pixelbrowse --tile-height 1568
# Screenshot multiple URLs in parallel
pixelshot <url1> <url2> --output /tmp/pixelbrowse --tile-height 1568 --workers 4
# Wider viewport for desktop layouts
pixelshot <url> --output /tmp/pixelbrowse --tile-height 1568 --viewport-width 1280
# Render a PDF
pixelshot document.pdf --output /tmp/pixelbrowse
```
IMPORTANT: Always use `--tile-height 1568` for screenshots you will read visually.
Claude's vision model downscales images with long edge > 1568px (Sonnet/Haiku) or 2576px (Opus).
The default 8192px tile height will be downscaled and text becomes unreadable.
After rendering, read the tile images from the output directory to visually understand the content.
## Workflow
1. Run `pixelshot <url> --output /tmp/pixelbrowse`
2. Read `/tmp/pixelbrowse/<domain>.png.tiles/tile_0000.jpg` directly (no need to ls — the naming is deterministic)
3. If the page is long, also read tile_0001.jpg, tile_0002.jpg, etc.
Output path pattern: `/tmp/pixelbrowse/<sanitized-url>.png.tiles/tile_NNNN.jpg`
- For `https://news.ycombinator.com``/tmp/pixelbrowse/news.ycombinator.com.png.tiles/tile_0000.jpg`
- For `https://example.com/page``/tmp/pixelbrowse/example.com_page.png.tiles/tile_0000.jpg`
Do NOT run `ls` — just read tile_0000.jpg. If it doesn't exist, the page had no content.
## Crop & Zoom
If text or details are too small to read, crop the region of interest and re-read at full resolution.
Pillow is always available (it's a pixelshot dependency):
```bash
python3 -c "from PIL import Image; Image.open('<tile_path>').crop((x1, y1, x2, y2)).save('/tmp/pixelbrowse/crop.png')"
```
- Coordinates are in pixels from the top-left corner of the tile
- Crop to roughly 800x800 or smaller for maximum clarity
- You can crop multiple times to inspect different regions
- Read the cropped image with the Read tool just like any other image
Use this whenever you see content but can't make out the details — tables, small labels, fine print, chart axes, etc.
## Tips
- Output is tiled JPEG images — tile_0000.jpg is the top, higher numbers go down the page
- Use `--viewport-width 1280` for desktop layouts, default 875 for mobile/article width
- Supports URLs (http/https), local HTML files, PDFs, and images
- Backend options: `--backend cdp` (default, fastest) or `--backend playwright`
+123
View File
@@ -0,0 +1,123 @@
[project]
name = "pixelrag"
version = "0.2.1"
description = "Visual Retrieval-Augmented Generation — render, embed, index, search"
requires-python = ">=3.12"
authors = [{ name = "Zhifei Li", email = "andylizf@outlook.com" }]
license = "Apache-2.0"
readme = "README.md"
keywords = ["rag", "retrieval", "vision-language", "screenshots", "embeddings", "faiss"]
classifiers = [
"License :: OSI Approved :: Apache Software License",
"Programming Language :: Python :: 3.12",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
# Core stays light — rendering / screenshots only (no torch). The `pixelshot` command and
# the `pixelrag` umbrella work on this alone. Heavy ML stages are opt-in extras below.
dependencies = [
"pillow>=10.0.0",
"websockets>=12.0",
"pymupdf>=1.27.2.3",
"pyturbojpeg>=2.2.0",
"cef-capi-py>=131.3.5",
"anthropic>=0.102.0",
]
[project.optional-dependencies]
embed = [
"torch>=2.9.0",
"torchvision>=0.24.0",
"transformers>=4.57.0",
"faiss-cpu>=1.9.0",
"numpy>=1.26.0",
"tqdm>=4.60.0",
]
serve = [
"fastapi>=0.115.0",
"uvicorn>=0.30.0",
"numpy>=1.26.0",
"faiss-cpu>=1.9.0",
"transformers>=4.57.0",
"torch>=2.9.0",
"qwen-vl-utils",
"pydantic>=2.0.0",
]
index = ["pixelrag[embed]", "pyyaml>=6.0"]
all = ["pixelrag[embed,serve,index]"]
gpu = ["faiss-gpu-cu12>=1.13.2"]
playwright = ["playwright>=1.40.0"]
pdf = ["pdf2image>=1.16.0"]
kiwix = ["libzim>=3.6.0"]
distributed = ["boto3>=1.42.0"]
eval = [
"pandas>=2.0",
"Pillow>=10.0",
"tqdm>=4.60",
"trafilatura>=1.6",
"openai>=1.0",
"aiohttp>=3.9",
"datasets>=2.14",
"huggingface-hub>=0.20",
]
dev = ["pytest>=8.0"]
[project.scripts]
pixelshot = "pixelrag_render.render:main"
pixelrag = "pixelrag.cli:main"
[project.urls]
Homepage = "https://pixelrag.ai"
Repository = "https://github.com/StarTrail-org/PixelRAG"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
# One distribution bundles the umbrella + every stage module.
packages = [
"src/pixelrag",
"render/src/pixelrag_render",
"embed/src/pixelrag_embed",
"index/src/pixelrag_index",
"serve/src/pixelrag_serve",
]
[tool.hatch.build.targets.sdist]
# Repo root holds multi-GB data dirs (.venv, tiles, arxiv, …); ship only the package sources.
include = [
"/src/pixelrag",
"/render/src",
"/embed/src",
"/index/src",
"/serve/src",
"/README.md",
"/LICENSE",
]
[tool.uv]
override-dependencies = ["nvidia-cudnn-cu12==9.20.0.48"]
environments = ["sys_platform == 'linux'"]
[tool.uv.sources]
torch = [{ index = "pytorch-cu129" }]
torchvision = [{ index = "pytorch-cu129" }]
[[tool.uv.index]]
name = "pytorch-cu129"
url = "https://download.pytorch.org/whl/cu129"
explicit = true
[tool.ruff]
exclude = ["tmp"]
[tool.ruff.lint]
ignore = [
"E402", # module-level import not at top — common in ML codebases
"E741", # ambiguous variable name (l, O, I)
]
[tool.pyright]
extraPaths = ["render/src", "embed/src", "index/src", "serve/src", "train/src"]
venvPath = "."
venv = ".venv"
+8
View File
@@ -0,0 +1,8 @@
"""pixelshot: Document to image tiles.
Renders web pages, PDFs, and local files as tiled screenshots.
"""
from .render import render_url, render_pdf, render_file
__all__ = ["render_url", "render_pdf", "render_file"]
@@ -0,0 +1 @@
"""pixelrag_render backends: cdp, playwright, pdf."""
@@ -0,0 +1,636 @@
"""Fast CDP backend: raw BGRA capture → async JPEG compression.
Architecture:
Chrome workers (n_workers) Compression pool (n_compressors procs)
↓ rawFilePath ↓ read /dev/shm
/dev/shm/pixelrag_render/raw/ → JPEG compress
(28MB × n_workers slots) → output/tiles/
Capture and compression are fully decoupled. Chrome writes raw BGRA to
/dev/shm via Page.captureScreenshot rawFilePath. A background asyncio task
drains the compression queue and submits work to a ProcessPoolExecutor.
Capture never waits for compression.
Requirements: pillow, websockets (no playwright needed)
Usage:
from pixelrag_render.backends.fast_cdp import render_articles
result = render_articles(articles, "./tiles")
# result: {"total_tiles": N, "wall_s": T, "tiles_per_s": tps}
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import signal
import struct
import subprocess
import time
import urllib.request
from pathlib import Path
logger = logging.getLogger("pixelrag_render.backends.fast_cdp")
VIEWPORT_WIDTH = 875
TILE_HEIGHT = 8192
CHROME_ARGS = [
"--no-sandbox",
"--disable-dev-shm-usage",
"--enable-gpu-rasterization",
"--force-gpu-rasterization",
"--disable-renderer-backgrounding",
"--disable-backgrounding-occluded-windows",
"--disable-background-networking",
"--disable-features=Translate,MediaRouter,OptimizationHints",
]
# JS: wait for fonts + eager images, then return scrollHeight
_WAIT_FONTS_IMGS = """new Promise(resolve => {
const waitEagerImgs = Promise.all(
Array.from(document.images)
.filter(i => !i.complete && i.loading !== 'lazy')
.map(i => new Promise(r => {
i.addEventListener('load', r, {once: true});
i.addEventListener('error', r, {once: true});
}))
);
const timeout = new Promise(r => setTimeout(r, 2000));
Promise.race([
Promise.all([document.fonts.ready, waitEagerImgs]),
timeout
]).then(() => {
requestAnimationFrame(() => {
requestAnimationFrame(() => {
document.documentElement.style.scrollBehavior = 'auto';
const sh = document.documentElement.scrollHeight;
const body = document.body;
resolve(body
? Math.min(sh, Math.max(Math.ceil(body.getBoundingClientRect().bottom), 1))
: sh);
});
});
});
})"""
# ---------------------------------------------------------------------------
# Subprocess: JPEG compression (runs in ProcessPoolExecutor worker)
# ---------------------------------------------------------------------------
def compress_tile(raw_path: str, out_path: str, quality: int = 85) -> None:
"""Read raw BGRA file, compress to JPEG, delete raw file.
Raw file layout (written by Chrome rawFilePath):
bytes 0-3: width (uint32 LE)
bytes 4-7: height (uint32 LE)
bytes 8-11: rowBytes (uint32 LE)
bytes 12+: BGRA pixels
"""
from PIL import Image
data = open(raw_path, "rb").read()
w, h, rb = struct.unpack_from("<III", data, 0)
img = Image.frombuffer("RGBA", (w, h), data[12:], "raw", "BGRA", rb, 1)
img = img.convert("RGB")
img.save(out_path, "JPEG", quality=quality)
os.unlink(raw_path)
# ---------------------------------------------------------------------------
# Chrome connection helpers (inlined to avoid circular deps)
# ---------------------------------------------------------------------------
_port_counter = 0
def _next_base_port() -> int:
global _port_counter
_port_counter += 1
return 12000 + (_port_counter - 1) * 500
async def _launch_chrome(chrome_path: str, port: int) -> tuple:
"""Launch a headless Chrome and return (websocket, proc)."""
import websockets
args = (
[chrome_path, f"--remote-debugging-port={port}", "--headless"]
+ CHROME_ARGS
+ ["about:blank"]
)
proc = subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
for attempt in range(10):
await asyncio.sleep(1)
try:
data = urllib.request.urlopen(
f"http://localhost:{port}/json", timeout=3
).read()
targets = json.loads(data)
ws = await websockets.connect(
targets[0]["webSocketDebuggerUrl"],
open_timeout=10,
max_size=50 * 1024 * 1024,
)
return ws, proc
except Exception:
if attempt == 9:
proc.kill()
raise ConnectionError(f"Failed to connect to Chrome on port {port}")
class _Conn:
"""Minimal CDP connection with a receive loop."""
def __init__(self, ws, proc):
self._ws = ws
self._proc = proc
self._msg_id = 0
self._pending: dict[int, asyncio.Future] = {}
self._event_listeners: dict[str, list] = {}
self._recv_task: asyncio.Task | None = None
def _ensure_recv(self):
if self._recv_task is None or self._recv_task.done():
self._recv_task = asyncio.get_event_loop().create_task(self._recv_loop())
async def _recv_loop(self):
try:
async for raw in self._ws:
msg = json.loads(raw)
mid = msg.get("id")
if mid is not None:
fut = self._pending.pop(mid, None)
if fut and not fut.done():
fut.set_result(msg)
else:
method = msg.get("method", "")
listeners = self._event_listeners.get(method, [])
remaining = []
for fut, filter_fn in listeners:
if fut.done():
continue
params = msg.get("params", {})
matched = filter_fn(params) if filter_fn else True
if matched:
fut.set_result(params)
else:
remaining.append((fut, filter_fn))
self._event_listeners[method] = remaining
except Exception:
exc = ConnectionError("WebSocket receive loop ended")
for fut in self._pending.values():
if not fut.done():
fut.set_exception(exc)
for listeners in self._event_listeners.values():
for fut, _ in listeners:
if not fut.done():
fut.set_exception(exc)
async def cdp(self, method: str, params: dict | None = None) -> dict:
self._ensure_recv()
self._msg_id += 1
mid = self._msg_id
msg = {"id": mid, "method": method}
if params:
msg["params"] = params
loop = asyncio.get_event_loop()
fut: asyncio.Future = loop.create_future()
self._pending[mid] = fut
await self._ws.send(json.dumps(msg))
return await asyncio.wait_for(fut, timeout=180)
async def wait_for_event(
self, method: str, timeout: float = 30.0, filter_fn=None
) -> dict:
self._ensure_recv()
loop = asyncio.get_event_loop()
fut: asyncio.Future = loop.create_future()
self._event_listeners.setdefault(method, []).append((fut, filter_fn))
try:
return await asyncio.wait_for(asyncio.shield(fut), timeout=timeout)
except asyncio.TimeoutError:
listeners = self._event_listeners.get(method, [])
self._event_listeners[method] = [
(f, fn) for f, fn in listeners if f is not fut
]
if not fut.done():
fut.cancel()
raise
async def close(self):
try:
await self._ws.close()
except Exception:
pass
self._proc.send_signal(signal.SIGTERM)
try:
self._proc.wait(timeout=5)
except subprocess.TimeoutExpired:
self._proc.kill()
# ---------------------------------------------------------------------------
# Core render logic
# ---------------------------------------------------------------------------
async def _run_render(
articles: list[dict],
output_dir: Path,
chrome_path: str,
n_workers: int,
tile_height: int,
jpeg_quality: int,
n_compressors: int,
) -> dict:
raw_dir = Path("/dev/shm/pixelrag_render/raw")
raw_dir.mkdir(parents=True, exist_ok=True)
output_dir.mkdir(parents=True, exist_ok=True)
# Semaphore: limit concurrent captures (CPU-bound) to n_workers // 2
capture_limit = max(1, n_workers // 2)
capture_sem = asyncio.Semaphore(capture_limit)
# Compression: dedicated thread with its own multiprocessing.Pool.
# Tiles are pushed to a thread-safe queue from capture workers.
# The thread runs pool.starmap in batches, fully independent of asyncio.
from multiprocessing import Pool as MPPool
import queue as _queue
import threading
metrics = {
"total_tiles": 0,
"total_capture_ms": 0.0,
"errors": 0,
}
compress_inbox: _queue.Queue = _queue.Queue()
compress_done = threading.Event()
n_cpus = os.cpu_count() or 128
compress_cores = set(range(max(0, n_cpus - n_compressors), n_cpus))
def _pool_init():
try:
os.sched_setaffinity(0, compress_cores)
except OSError:
pass
def _compressor_thread():
pool = MPPool(processes=n_compressors, initializer=_pool_init)
# Warm up: ensure all workers are forked and idle before capture starts
pool.map(int, range(n_compressors))
async_results = []
while True:
item = compress_inbox.get() # block until item available
if item is None:
break
async_results.append(pool.apply_async(compress_tile, item))
# Wait for all remaining
for ar in async_results:
try:
ar.get(timeout=60)
except Exception:
pass
pool.close()
pool.join()
compress_done.set()
compress_thread = threading.Thread(target=_compressor_thread, daemon=True)
compress_thread.start()
base_port = _next_base_port()
# Work-stealing queue
work_q: asyncio.Queue = asyncio.Queue()
for art in articles:
work_q.put_nowait(art)
# Launch Chrome workers
connections: list[_Conn] = []
frame_ids: list[str] = []
logger.info(
"Launching %d Chrome workers on ports %d-%d",
n_workers,
base_port,
base_port + n_workers - 1,
)
for i in range(n_workers):
ws, proc = await _launch_chrome(chrome_path, base_port + i)
conn = _Conn(ws, proc)
connections.append(conn)
for i, conn in enumerate(connections):
await conn.cdp("Page.enable")
await conn.cdp(
"Emulation.setDeviceMetricsOverride",
{
"width": VIEWPORT_WIDTH,
"height": tile_height,
"deviceScaleFactor": 1,
"mobile": False,
},
)
ft = await conn.cdp("Page.getFrameTree")
frame_ids.append(ft["result"]["frameTree"]["frame"]["id"])
logger.info("All workers ready. Processing %d articles.", len(articles))
t_start = time.monotonic()
async def worker_task(wi: int):
conn = connections[wi]
main_fid = frame_ids[wi]
while True:
try:
article = work_q.get_nowait()
except asyncio.QueueEmpty:
break
art_path = article["path"]
raw_file_val = article.get("file", "")
target_url = (
raw_file_val
if raw_file_val.startswith("http")
else f"file://{raw_file_val}"
)
# Make output tile dir: use path slug
slug = art_path.replace("/", "_").replace(" ", "_")[:200] or "article"
tile_dir = output_dir / f"{slug}.tiles"
tile_dir.mkdir(parents=True, exist_ok=True)
try:
# --- NAV (outside semaphore — I/O bound) ---
nav_event_fut = asyncio.ensure_future(
conn.wait_for_event(
"Page.frameStoppedLoading",
timeout=30.0,
filter_fn=lambda p: p.get("frameId") == main_fid,
)
)
try:
await conn.cdp("Page.navigate", {"url": target_url})
except Exception as e:
nav_event_fut.cancel()
logger.warning("[w%d] nav failed for %s: %s", wi, art_path, e)
metrics["errors"] += 1
continue
try:
await nav_event_fut
except asyncio.TimeoutError:
logger.warning(
"[w%d] frameStoppedLoading timeout for %s", wi, art_path
)
metrics["errors"] += 1
continue
# Wait for fonts + images, get page height
try:
r = await conn.cdp(
"Runtime.evaluate",
{
"expression": _WAIT_FONTS_IMGS,
"awaitPromise": True,
"returnByValue": True,
},
)
page_h = r["result"]["result"]["value"]
if not page_h or page_h <= 0:
page_h = tile_height
except Exception:
page_h = tile_height
n_tiles = max(1, (page_h + tile_height - 1) // tile_height)
n_written = 0
tile_names = []
for t in range(n_tiles):
clip_h = min(tile_height, page_h - t * tile_height)
if clip_h <= 28:
break
# Scroll + wait in-viewport images (outside semaphore)
if t > 0:
y = t * tile_height
try:
await conn.cdp(
"Runtime.evaluate",
{
"expression": f"""new Promise(resolve => {{
window.scrollTo(0, {y});
requestAnimationFrame(() => requestAnimationFrame(() => {{
const imgs = Array.from(document.images).filter(i => {{
if (i.complete) return false;
const r = i.getBoundingClientRect();
return r.bottom > 0 && r.top < window.innerHeight;
}});
if (imgs.length === 0) return resolve();
const timeout = new Promise(r => setTimeout(r, 500));
const loaded = Promise.all(imgs.map(i => new Promise(r => {{
i.addEventListener('load', r, {{once: true}});
i.addEventListener('error', r, {{once: true}});
}})));
Promise.race([loaded, timeout]).then(resolve);
}}));
}})""",
"awaitPromise": True,
},
)
except Exception:
pass
# Acquire semaphore → capture → release (fine-grained)
raw_path = str(raw_dir / f"w{wi}_{slug}_{t}.raw")
out_path = str(tile_dir / f"tile_{t:04d}.jpg")
await capture_sem.acquire()
try:
t0 = time.monotonic()
r = await conn.cdp(
"Page.captureScreenshot",
{
"fromSurface": True,
"optimizeForSpeed": True,
"rawFilePath": raw_path,
"clip": {
"x": 0,
"y": t * tile_height,
"width": VIEWPORT_WIDTH,
"height": clip_h,
"scale": 1,
},
},
)
shot_ms = (time.monotonic() - t0) * 1000
except Exception as e:
logger.warning(
"[w%d] capture failed tile %d of %s: %s", wi, t, art_path, e
)
metrics["errors"] += 1
continue
finally:
capture_sem.release()
if "error" in r.get("result", {}):
logger.warning(
"[w%d] CDP error tile %d of %s: %s",
wi,
t,
art_path,
r["result"]["error"],
)
metrics["errors"] += 1
continue
metrics["total_capture_ms"] += shot_ms
# Enqueue compression (non-blocking — capture continues)
compress_inbox.put((raw_path, out_path, jpeg_quality))
n_written += 1
tile_names.append(f"tile_{t:04d}.jpg")
# Write manifest
manifest = {
"path": art_path,
"url": target_url,
"page_height": page_h,
"tiles": tile_names,
"complete": True,
}
with open(tile_dir / "tiles.json", "w") as f:
json.dump(manifest, f)
metrics["total_tiles"] += n_written
logger.info(
"[w%d] %s%d tiles (%.0f ms capture)",
wi,
art_path,
n_written,
shot_ms if n_tiles == 1 else 0,
)
except Exception as e:
logger.warning("[w%d] unexpected error for %s: %s", wi, art_path, e)
metrics["errors"] += 1
# Run all workers concurrently
await asyncio.gather(*[worker_task(i) for i in range(n_workers)])
capture_wall_s = time.monotonic() - t_start
total = metrics["total_tiles"]
capture_tps = total / capture_wall_s if capture_wall_s > 0 else 0.0
logger.info(
"Capture done: %d tiles in %.1fs (%.1f tiles/s)",
total,
capture_wall_s,
capture_tps,
)
# Wait for compression — run in thread to avoid blocking asyncio event loop
import threading
# Signal compression thread to finish, teardown Chrome in parallel
compress_inbox.put(None)
for conn in connections:
try:
await conn.close()
except Exception:
pass
# Wait for compression to finish (runs in its own thread, no asyncio)
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, compress_done.wait, 120)
wall_s = time.monotonic() - t_start
tps = total / wall_s if wall_s > 0 else 0.0
logger.info(
"Done: %d tiles in %.1fs (%.1f tiles/s, capture=%.1f tiles/s)",
total,
wall_s,
tps,
capture_tps,
)
return {
"total_tiles": total,
"wall_s": wall_s,
"capture_wall_s": capture_wall_s,
"capture_tiles_per_s": capture_tps,
"tiles_per_s": tps,
"errors": metrics["errors"],
"avg_capture_ms": (metrics["total_capture_ms"] / total if total > 0 else 0.0),
}
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
async def render_articles(
articles: list[dict],
output_dir: str,
chrome_path: str = None,
n_workers: int = 48,
tile_height: int = TILE_HEIGHT,
jpeg_quality: int = 85,
n_compressors: int = 4,
) -> dict:
"""Render articles to JPEG tiles with async compression.
Capture (Chrome → /dev/shm raw BGRA) and compression (raw → JPEG on disk)
are fully decoupled. Chrome workers never wait for compression.
Args:
articles: List of dicts with keys ``path`` (article ID) and
``file`` (URL, http:// or absolute filesystem path).
output_dir: Directory for output tile subdirectories.
chrome_path: Path to Chrome binary. Auto-detected if None.
n_workers: Number of parallel Chrome processes (default 48).
tile_height: Max tile height in pixels (default 8192).
jpeg_quality: JPEG quality 1100 (default 85).
n_compressors: ProcessPoolExecutor workers for compression (default 4).
Returns:
dict with keys:
``total_tiles`` number of tiles written
``wall_s`` total wall-clock time in seconds
``tiles_per_s`` throughput
``errors`` count of capture/nav errors
``avg_capture_ms`` average per-tile capture time (ms)
"""
if not articles:
return {
"total_tiles": 0,
"wall_s": 0.0,
"tiles_per_s": 0.0,
"errors": 0,
"avg_capture_ms": 0.0,
}
if chrome_path is None:
from ..chrome import find_chrome
chrome_path = find_chrome()
actual_workers = min(n_workers, len(articles))
return await _run_render(
articles=articles,
output_dir=Path(output_dir),
chrome_path=chrome_path,
n_workers=actual_workers,
tile_height=tile_height,
jpeg_quality=jpeg_quality,
n_compressors=n_compressors,
)
+103
View File
@@ -0,0 +1,103 @@
"""PDF backend for pixelshot.
Renders PDF pages to JPEG tiles using pdf2image (poppler).
Requires: pdf2image>=1.16.0 (install pixelrag-render[pdf])
"""
import json
import logging
from pathlib import Path
from typing import Optional
logger = logging.getLogger("pixelrag_render.backends.pdf")
def render_pdf(
path: str | Path,
output_dir: str | Path,
*,
dpi: int = 200,
pages: Optional[list[int]] = None,
quality: int = 85,
) -> list[Path]:
"""Render a PDF to JPEG tiles.
Each page is written as ``{stem}.png.tiles/tile_NNNN.jpg`` with a
``tiles.json`` manifest alongside.
Args:
path: Path to the source PDF file.
output_dir: Directory to write the tile subdirectory into.
dpi: Resolution for rendering (default 200 gives ~1650×2200px for A4).
pages: 1-based list of page numbers to render. ``None`` renders all pages.
quality: JPEG quality 1-100 (default 85).
Returns:
List containing the single tile directory Path on success.
Raises:
ImportError: If pdf2image is not installed.
FileNotFoundError: If the PDF file does not exist.
"""
try:
from pdf2image import convert_from_path
except ImportError as e:
raise ImportError(
"pdf2image is required for PDF rendering. "
"Install with: pip install 'pixelrag-render[pdf]'"
) from e
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"PDF not found: {path}")
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
stem = path.stem
tile_dir = output_dir / f"{stem}.png.tiles"
tile_dir.mkdir(parents=True, exist_ok=True)
logger.info("Rendering PDF: %s (dpi=%d)", path, dpi)
convert_kwargs: dict = {
"pdf_path": str(path),
"dpi": dpi,
"fmt": "jpeg",
"jpegopt": {"quality": quality, "progressive": True},
"thread_count": 4,
}
if pages is not None:
# pdf2image uses 1-based page numbers
convert_kwargs["first_page"] = min(pages)
convert_kwargs["last_page"] = max(pages)
images = convert_from_path(**convert_kwargs)
saved_tiles: list[str] = []
for idx, img in enumerate(images):
# If caller provided a sparse page list, skip pages not in the list
if pages is not None:
page_num = min(pages) + idx
if page_num not in pages:
continue
tile_name = f"tile_{idx:04d}.jpg"
tile_path = tile_dir / tile_name
img.save(str(tile_path), "JPEG", quality=quality)
saved_tiles.append(tile_name)
logger.debug(" Page %d%s (%dx%d)", idx, tile_name, *img.size)
manifest = {
"source": str(path),
"dpi": dpi,
"total_pages": len(saved_tiles),
"tiles": saved_tiles,
"complete": True,
}
with open(tile_dir / "tiles.json", "w") as f:
json.dump(manifest, f)
logger.info("PDF rendered: %d pages → %s", len(saved_tiles), tile_dir)
return [tile_dir]
@@ -0,0 +1,399 @@
"""Direct websocket CDP backend for pixelshot.
No Playwright dependency — uses subprocess to launch Chrome and websockets
to communicate via CDP directly. ~35% faster than the Playwright-based cdp.py
backend due to eliminating the Node.js IPC layer.
Requirements: websockets, pillow (no playwright needed)
Usage:
from pixelrag_render.backends.websocket import render_urls
tile_dirs = render_urls(["https://example.com"], "./tiles", workers=4)
"""
import asyncio
import base64
import io
import json
import logging
import signal
import subprocess
import time
import urllib.request
from pathlib import Path
from PIL import Image
logger = logging.getLogger("pixelrag_render.backends.websocket")
VIEWPORT_W = 875
VIEWPORT_H = 1080
BROWSER_ARGS = [
"--disable-dev-shm-usage",
"--no-sandbox",
"--disable-renderer-backgrounding",
"--disable-backgrounding-occluded-windows",
"--disable-background-networking",
"--disable-features=Translate,MediaRouter,OptimizationHints",
"--enable-gpu-rasterization",
"--force-gpu-rasterization",
]
def _find_chrome() -> str:
from ..chrome import find_chrome
return find_chrome()
async def _connect_cdp(port: int, retries: int = 5, delay: float = 1.0):
"""Connect to Chrome's CDP websocket endpoint."""
import websockets
for attempt in range(retries):
try:
data = urllib.request.urlopen(
f"http://localhost:{port}/json", timeout=3
).read()
targets = json.loads(data)
ws = await websockets.connect(
targets[0]["webSocketDebuggerUrl"],
open_timeout=10,
max_size=50 * 1024 * 1024,
)
return ws
except Exception:
if attempt < retries - 1:
await asyncio.sleep(delay)
raise ConnectionError(f"Failed to connect to Chrome on port {port}")
async def _cdp_send(ws, msg_id_ref: list, method: str, params: dict | None = None):
"""Send a CDP command and wait for its response."""
msg_id_ref[0] += 1
mid = msg_id_ref[0]
msg = {"id": mid, "method": method}
if params:
msg["params"] = params
await ws.send(json.dumps(msg))
while True:
r = json.loads(await asyncio.wait_for(ws.recv(), timeout=180))
if r.get("id") == mid:
if "error" in r:
raise RuntimeError(f"CDP error: {r['error']}")
return r.get("result", {})
async def capture_url(
ws,
msg_id_ref: list,
url: str,
tile_dir: Path,
*,
tile_h: int = 8192,
quality: int = 85,
viewport_w: int = VIEWPORT_W,
image_format: str = "jpeg",
from_surface: bool = True,
) -> int:
"""Capture a URL as tiled images via direct CDP websocket.
Returns the number of tiles written.
"""
tile_dir.mkdir(parents=True, exist_ok=True)
await _cdp_send(ws, msg_id_ref, "Page.navigate", {"url": url})
# Wait for fonts + layout to stabilize, return scrollHeight in one call
result = await _cdp_send(
ws,
msg_id_ref,
"Runtime.evaluate",
{
"expression": """new Promise(resolve => {
document.fonts.ready.then(() => {
requestAnimationFrame(() => {
requestAnimationFrame(() => {
document.documentElement.style.scrollBehavior = 'auto';
const sh = document.documentElement.scrollHeight;
const body = document.body;
if (body) {
const bottom = Math.ceil(body.getBoundingClientRect().bottom);
resolve(Math.min(sh, Math.max(bottom, 1)));
} else {
resolve(sh);
}
});
});
});
})""",
"awaitPromise": True,
"returnByValue": True,
},
)
try:
page_height = result["result"]["value"]
except (KeyError, TypeError):
page_height = tile_h
tiles = []
y = 0
idx = 0
while y < page_height:
clip_h = min(tile_h, page_height - y)
if clip_h <= 0:
break
params = {
"format": image_format,
"fromSurface": from_surface,
"optimizeForSpeed": True,
"clip": {
"x": 0,
"y": y,
"width": viewport_w,
"height": clip_h,
"scale": 1,
},
}
if image_format == "jpeg":
params["quality"] = quality
result = await _cdp_send(ws, msg_id_ref, "Page.captureScreenshot", params)
img_bytes = base64.b64decode(result["data"])
tile_path = (
tile_dir / f"tile_{idx:04d}.{'jpg' if image_format == 'jpeg' else 'png'}"
)
if clip_h < tile_h:
img = Image.open(io.BytesIO(img_bytes))
w, h = img.size
if h > clip_h:
img = img.crop((0, 0, w, clip_h))
img.save(
tile_path, "JPEG" if image_format == "jpeg" else "PNG", quality=quality
)
else:
tile_path.write_bytes(img_bytes)
tiles.append(tile_path.name)
idx += 1
y += tile_h
manifest = {
"url": url,
"page_height": page_height,
"tiles": tiles,
"complete": True,
}
with open(tile_dir / "tiles.json", "w") as f:
json.dump(manifest, f)
return len(tiles)
async def _worker(
chrome_path: str,
port: int,
work_queue: asyncio.Queue,
output_dir: Path,
tile_height: int,
quality: int,
viewport_w: int,
image_format: str,
from_surface: bool,
worker_id: int,
stats: dict,
results: list,
):
"""Async worker: owns a Chrome process, pulls URLs from queue."""
proc = subprocess.Popen(
[chrome_path, f"--remote-debugging-port={port}", "--headless"]
+ BROWSER_ARGS
+ ["about:blank"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
try:
await asyncio.sleep(3)
ws = await _connect_cdp(port)
msg_id_ref = [0]
await _cdp_send(ws, msg_id_ref, "Page.enable")
await _cdp_send(
ws,
msg_id_ref,
"Emulation.setDeviceMetricsOverride",
{
"width": viewport_w,
"height": tile_height,
"deviceScaleFactor": 1,
"mobile": False,
},
)
while True:
try:
item = work_queue.get_nowait()
except asyncio.QueueEmpty:
break
url = item["url"]
stem = item["stem"]
tile_dir = output_dir / f"{stem}.png.tiles"
t0 = time.monotonic()
try:
n_tiles = await capture_url(
ws,
msg_id_ref,
url,
tile_dir,
tile_h=tile_height,
quality=quality,
viewport_w=viewport_w,
image_format=image_format,
from_surface=from_surface,
)
stats["done"] += 1
elapsed = time.monotonic() - t0
logger.info(
"[w%d] %s%d tiles (%.1fs)", worker_id, url, n_tiles, elapsed
)
results.append(tile_dir)
except Exception as e:
stats["failed"] += 1
logger.warning("[w%d] FAIL %s: %s", worker_id, url, str(e)[:200])
await ws.close()
finally:
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
proc.kill()
async def _run_batch(
urls: list[str],
output_dir: Path,
num_workers: int,
tile_height: int,
quality: int,
viewport_w: int,
image_format: str,
from_surface: bool,
stems: list[str] | None,
chrome_path: str,
) -> list[Path]:
work_queue: asyncio.Queue = asyncio.Queue()
seen_stems: dict[str, int] = {}
for i, url in enumerate(urls):
if stems and i < len(stems):
stem = str(stems[i])
else:
from urllib.parse import urlparse
parsed = urlparse(url)
raw = (parsed.netloc + parsed.path).rstrip("/")
stem = (
raw.replace("/", "_")
.replace(":", "_")
.replace("?", "_")
.replace("&", "_")
)
stem = stem[:200] or "page"
count = seen_stems.get(stem, 0)
seen_stems[stem] = count + 1
if count > 0:
stem = f"{stem}_{count}"
work_queue.put_nowait({"url": url, "stem": stem})
stats = {"done": 0, "failed": 0}
results: list[Path] = []
base_port = 9400
actual_workers = min(num_workers, len(urls))
workers = [
_worker(
chrome_path,
base_port + wid,
work_queue,
output_dir,
tile_height,
quality,
viewport_w,
image_format,
from_surface,
wid,
stats,
results,
)
for wid in range(actual_workers)
]
await asyncio.gather(*workers, return_exceptions=True)
logger.info("Batch complete: done=%d failed=%d", stats["done"], stats["failed"])
return results
def render_urls(
urls: list[str],
output_dir: str | Path,
*,
stems: list[str] | None = None,
tile_height: int = 8192,
quality: int = 85,
viewport_width: int = VIEWPORT_W,
workers: int = 4,
image_format: str = "jpeg",
from_surface: bool = True,
chrome_path: str | None = None,
) -> list[Path]:
"""Render URLs to tiled images using direct CDP websocket.
No Playwright dependency. Each worker launches its own Chrome process
and communicates via CDP over websocket.
Args:
urls: URLs to capture.
output_dir: Output directory for tile subdirectories.
stems: Optional output directory name per URL.
tile_height: Max tile height in pixels (default 8192).
quality: JPEG quality 1-100 (default 85).
viewport_width: Browser viewport width (default 875).
workers: Number of parallel Chrome processes (default 4).
image_format: 'jpeg' or 'png' (default 'jpeg').
from_surface: CDP fromSurface param. True for batch (throughput),
False for serve (low latency). Default True.
chrome_path: Path to Chrome binary. Auto-detected if None.
Returns:
List of Path objects for created tile directories.
"""
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
if not urls:
return []
chrome = chrome_path or _find_chrome()
return asyncio.run(
_run_batch(
urls,
output_dir,
workers,
tile_height,
quality,
viewport_width,
image_format,
from_surface,
stems,
chrome,
)
)
@@ -0,0 +1,16 @@
"""pixelshot benchmark harness.
Usage:
from pixelrag_render.strategies import CDPSequentialStrategy
from pixelrag_render.bench import Bench
bench = Bench(zim_path="...", chrome_path="...", output_dir="./results")
result = await bench.run(CDPSequentialStrategy(chrome_path=..., n_workers=32, fmt="raw"))
"""
from .bench_throughput import (
Bench as Bench,
prepare_articles as prepare_articles,
generate_ground_truth as generate_ground_truth,
run_and_verify as run_and_verify,
)
@@ -0,0 +1,509 @@
#!/usr/bin/env python3
"""
Screenshot benchmark with correctness verification.
Bench is a clean measurement harness. It takes a strategy object, runs it,
verifies results against cached GT, and dumps config + results.
Strategies live in pixelrag_render.strategies — bench does NOT know about
specific strategy implementations or naming conventions.
Usage (programmatic):
from pixelrag_render.strategies import CDPSequentialStrategy
from pixelrag_render.bench.bench_throughput import Bench
bench = Bench(zim_path="...", chrome_path="...", output_dir="./results")
strategy = CDPSequentialStrategy(chrome_path=..., n_workers=32, fmt="raw")
result = await bench.run(strategy)
"""
from __future__ import annotations
import hashlib
import io
import json
import os
import random
import struct
import tempfile
import time
from pathlib import Path
import numpy as np
from PIL import Image
from pixelrag_render.strategies.base import TileCapture
from pixelrag_render.strategies.cdp_sequential import (
CDPSequentialStrategy,
VIEWPORT_WIDTH,
)
CORRECT_THRESHOLD = 99.0
JPEG_MAX_MEAN_DIFF = 5.0
LOSSLESS_MAX_MEAN_DIFF = 3.0
# ---------------------------------------------------------------------------
# Article preparation
# ---------------------------------------------------------------------------
def prepare_articles(
zim_path: str, n: int, seed: int = 42, kiwix_url: str | None = None
) -> list[dict]:
"""Sample articles from ZIM.
kiwix_url can be:
- None: write HTML to temp files (file:// mode)
- "http://host:port": single kiwix-serve instance
- "http://host:9461,http://host:9462,...": multiple instances (round-robin)
"""
from libzim.reader import Archive
from urllib.parse import quote
archive = Archive(zim_path)
# Support multiple kiwix URLs (comma-separated)
kiwix_urls = kiwix_url.split(",") if kiwix_url else []
# Detect book_name from first URL or ZIM filename
if kiwix_urls:
# Extract book_name from URL: http://host:port/content/{book_name}/...
# For symlinks like wiki_1.zim, book_name = wiki_1
# We need to figure out the right book_name for each URL
pass
book_name = Path(zim_path).stem
rng = random.Random(seed)
articles = []
tried = 0
while len(articles) < n and tried < n * 20:
idx = rng.randint(0, archive.all_entry_count - 1)
tried += 1
try:
e = archive._get_entry_by_id(idx)
if e.is_redirect or e.path.startswith("-/") or len(e.path) <= 2:
continue
entry = archive.get_entry_by_path(e.path)
item = entry.get_item()
if "html" not in item.mimetype:
continue
html = bytes(item.content).decode("utf-8")
if 'http-equiv="refresh"' in html.lower() or len(html) < 300:
continue
if kiwix_urls:
safe = "/:@!$&'()*+,;="
# Round-robin across kiwix instances
base = kiwix_urls[len(articles) % len(kiwix_urls)]
# Detect book_name from the symlink/ZIM each instance serves
parts = base.rstrip("/").rsplit(":", 1)
port = int(parts[1]) if len(parts) > 1 else 9454
# Each instance may have different book_name (wiki_1, wiki_2, etc.)
bname = f"wiki_{port - 9460}" if port > 9460 else book_name
url = f"{base}/content/{bname}/{quote(e.path, safe=safe)}"
articles.append({"path": e.path, "file": url})
else:
tmp = tempfile.NamedTemporaryFile(
suffix=".html", delete=False, dir="/tmp", prefix="bench_"
)
tmp.write(html.encode())
tmp.close()
articles.append({"path": e.path, "file": tmp.name})
except Exception:
continue
return articles
def cleanup_articles(articles: list[dict]):
for a in articles:
if a["file"].startswith("http"):
continue
try:
os.unlink(a["file"])
except OSError:
pass
# ---------------------------------------------------------------------------
# Ground truth (cached)
# ---------------------------------------------------------------------------
def gt_cache_key(articles: list[dict], seed: int) -> str:
paths = sorted(a["path"] for a in articles)
content = f"seed={seed}\n" + "\n".join(paths)
return hashlib.sha256(content.encode()).hexdigest()[:16]
async def generate_ground_truth(
articles: list[dict],
chrome_path: str,
cache_dir: Path,
seed: int,
timeout_ms: int = 5000,
) -> dict[str, list[Path]]:
cache_key = gt_cache_key(articles, seed)
manifest_path = cache_dir / f"gt_{cache_key}.json"
if manifest_path.exists():
manifest = json.loads(manifest_path.read_text())
all_exist = all(Path(p).exists() for paths in manifest.values() for p in paths)
if all_exist:
result = {k: [Path(p) for p in v] for k, v in manifest.items()}
total = sum(len(v) for v in result.values())
print(
f"Ground truth cache hit: {len(result)} articles, {total} tiles",
flush=True,
)
return result
cache_dir.mkdir(parents=True, exist_ok=True)
strategy = _make_gt_strategy(chrome_path, timeout_ms)
await strategy.setup()
try:
results = await strategy.capture_articles(articles)
finally:
await strategy.teardown()
ground_truth = {}
for ac in results:
tile_paths = []
for tc in ac.tiles:
tile_path = (
cache_dir
/ f"gt_{cache_key}_{ac.article_path.replace('/', '_')}_{tc.tile_index:02d}.png"
)
if tc.image_bytes:
tile_path.write_bytes(tc.image_bytes)
tile_paths.append(tile_path)
ground_truth[ac.article_path] = tile_paths
manifest = {k: [str(p) for p in v] for k, v in ground_truth.items()}
manifest_path.write_text(json.dumps(manifest))
total = sum(len(v) for v in ground_truth.values())
print(
f"Ground truth generated: {len(ground_truth)} articles, {total} tiles",
flush=True,
)
return ground_truth
def _make_gt_strategy(chrome_path: str, timeout_ms: int):
"""GT uses the most conservative strategy: 1 worker, PNG, long timeout.
Uses port 9222 to avoid TIME_WAIT conflicts with test strategies (9300+).
"""
s = CDPSequentialStrategy(
chrome_path=chrome_path, n_workers=1, fmt="png", from_surface=True
)
s._base_port = 9222
return s
def validate_gt(ground_truth: dict[str, list[Path]]) -> tuple[int, int, list[str]]:
"""Validate GT tiles are non-degenerate (not blank, not tiny, readable).
Returns (ok, bad, bad_examples).
"""
ok = 0
bad = 0
examples = []
for article_path, tile_paths in ground_truth.items():
for tp in tile_paths:
try:
img = Image.open(tp)
arr = np.array(img)
if arr.std() < 1.0:
bad += 1
if len(examples) < 10:
examples.append(
f"{article_path} {tp.name}: blank (std={arr.std():.1f})"
)
continue
ok += 1
except Exception as e:
bad += 1
if len(examples) < 10:
examples.append(f"{article_path} {tp.name}: {e}")
return ok, bad, examples
# ---------------------------------------------------------------------------
# Decode + verify (NOT timed)
# ---------------------------------------------------------------------------
def decode_tile(tc: TileCapture) -> Image.Image | None:
try:
if tc.raw_file_path and os.path.exists(tc.raw_file_path):
data = open(tc.raw_file_path, "rb").read()
w, h, rb = struct.unpack_from("<III", data, 0)
img = Image.frombuffer(
"RGBA", (w, h), data[12:], "raw", "BGRA", rb, 1
).convert("RGB")
return img
elif tc.image_bytes:
return Image.open(io.BytesIO(tc.image_bytes)).convert("RGB")
except Exception:
return None
return None
def verify_tile(
captured: Image.Image, gt_path: Path, is_lossy: bool
) -> tuple[bool, float]:
gt = Image.open(gt_path).convert("RGB")
cap_arr = np.array(captured, dtype=np.float32)
gt_arr = np.array(gt, dtype=np.float32)
if cap_arr.shape != gt_arr.shape:
return False, 999.0
diff = np.abs(cap_arr - gt_arr)
mean_diff = float(diff.mean())
threshold = JPEG_MAX_MEAN_DIFF if is_lossy else LOSSLESS_MAX_MEAN_DIFF
return mean_diff <= threshold, mean_diff
# ---------------------------------------------------------------------------
# Run one strategy: time capture, then verify separately
# ---------------------------------------------------------------------------
async def run_and_verify(strategy, articles, ground_truth) -> dict:
await strategy.setup()
t0 = time.monotonic()
try:
article_captures = await strategy.capture_articles(articles)
finally:
wall_s = time.monotonic() - t0
await strategy.teardown()
# --- UNTIMED: decode + verify ---
tiles_ok = 0
tiles_bad = 0
tiles_total = 0
total_shot_ms = 0.0
total_nav_ms = 0.0
total_pixels = 0
total_height_px = 0
per_tile_shot_ms = []
per_tile_nav_ms = []
bad_examples = []
is_lossy = strategy.fmt in ("jpeg",)
for ac in article_captures:
gt_tiles = ground_truth.get(ac.article_path, [])
total_height_px += ac.page_height
total_shot_ms += ac.total_shot_ms
total_nav_ms += ac.total_nav_ms
for tc in ac.tiles:
tiles_total += 1
total_pixels += VIEWPORT_WIDTH * tc.clip_h
per_tile_shot_ms.append(tc.shot_ms)
if tc.nav_ms > 0:
per_tile_nav_ms.append(tc.nav_ms)
if tc.tile_index >= len(gt_tiles):
tiles_bad += 1
bad_examples.append(f"{ac.article_path} tile {tc.tile_index}: no GT")
continue
img = decode_tile(tc)
if img is None:
tiles_bad += 1
bad_examples.append(
f"{ac.article_path} tile {tc.tile_index}: decode failed"
)
continue
ok, mean_diff = verify_tile(img, gt_tiles[tc.tile_index], is_lossy)
if ok:
tiles_ok += 1
else:
tiles_bad += 1
if len(bad_examples) < 10:
bad_examples.append(
f"{ac.article_path} tile {tc.tile_index}: mean_diff={mean_diff:.2f}"
)
for ac in article_captures:
for tc in ac.tiles:
if tc.raw_file_path:
try:
os.unlink(tc.raw_file_path)
except OSError:
pass
correct_pct = tiles_ok / tiles_total * 100 if tiles_total > 0 else 0
tps = tiles_total / wall_s if wall_s > 0 else 0
ms_per_tile = total_shot_ms / tiles_total if tiles_total > 0 else 0
articles_per_s = len(article_captures) / wall_s if wall_s > 0 else 0
mpix_per_s = (total_pixels / 1_000_000) / wall_s if wall_s > 0 else 0
shot_share = (
total_shot_ms / (total_shot_ms + total_nav_ms)
if (total_shot_ms + total_nav_ms) > 0
else 0
)
# Latency percentiles
sorted_shots = sorted(per_tile_shot_ms) if per_tile_shot_ms else [0]
sorted_navs = sorted(per_tile_nav_ms) if per_tile_nav_ms else [0]
def percentile(arr, p):
idx = int(len(arr) * p / 100)
return arr[min(idx, len(arr) - 1)]
return {
"name": strategy.name,
"tiles_total": tiles_total,
"tiles_ok": tiles_ok,
"tiles_bad": tiles_bad,
"correct_pct": correct_pct,
"wall_s": wall_s,
"tiles_per_s": tps,
"ms_per_tile": ms_per_tile,
"articles_per_s": articles_per_s,
"mpix_per_s": mpix_per_s,
"height_kpx_per_s": (total_height_px / 1000) / wall_s if wall_s > 0 else 0,
"shot_pct": shot_share * 100,
"bad_examples": bad_examples,
# Latency distribution
"shot_min": sorted_shots[0],
"shot_p50": percentile(sorted_shots, 50),
"shot_p95": percentile(sorted_shots, 95),
"shot_p99": percentile(sorted_shots, 99),
"shot_max": sorted_shots[-1],
"nav_avg": sum(sorted_navs) / len(sorted_navs),
"nav_p95": percentile(sorted_navs, 95),
}
# ---------------------------------------------------------------------------
# Bench: clean harness that takes any CaptureStrategy
# ---------------------------------------------------------------------------
class Bench:
"""Benchmark harness. Measures throughput/latency and verifies correctness.
Usage:
bench = Bench(zim_path="...", chrome_path="...", output_dir="./results")
result = await bench.run(strategy, articles=200, seed=42)
"""
def __init__(
self,
zim_path: str,
chrome_path: str,
output_dir: str = "./bench_results",
kiwix_url: str | None = None,
gt_timeout_ms: int = 5000,
):
self.zim_path = zim_path
self.chrome_path = chrome_path
self.output_dir = Path(output_dir)
self.kiwix_url = kiwix_url
self.gt_timeout_ms = gt_timeout_ms
self._articles: list[dict] | None = None
self._gt: dict[str, list[Path]] | None = None
def prepare(self, n_articles: int = 200, seed: int = 42) -> list[dict]:
if self._articles is None:
self._articles = prepare_articles(
self.zim_path, n_articles, seed, kiwix_url=self.kiwix_url
)
return self._articles
async def ensure_gt(
self, n_articles: int = 200, seed: int = 42
) -> dict[str, list[Path]]:
if self._gt is not None:
return self._gt
articles = self.prepare(n_articles, seed)
gt_dir = self.output_dir / "ground_truth"
self._gt = await generate_ground_truth(
articles, self.chrome_path, gt_dir, seed, timeout_ms=self.gt_timeout_ms
)
ok, bad, examples = validate_gt(self._gt)
total = ok + bad
print(f"GT validation: {ok}/{total} OK, {bad} bad", flush=True)
if examples:
for ex in examples:
print(f" GT BAD: {ex}", flush=True)
if bad > 0:
pct = ok / total * 100 if total else 0
if pct < CORRECT_THRESHOLD:
raise RuntimeError(
f"GT itself is only {pct:.1f}% valid ({bad} bad tiles). "
f"Fix image loading or increase gt_timeout_ms."
)
return self._gt
async def run(self, strategy, n_articles: int = 200, seed: int = 42) -> dict:
articles = self.prepare(n_articles, seed)
gt = await self.ensure_gt(n_articles, seed)
result = await run_and_verify(strategy, articles, gt)
exp = self._build_experiment(strategy, result, n_articles, seed)
try:
self._dump_experiment(exp)
except Exception as e:
print(f"Warning: failed to save experiment: {e}", flush=True)
return result
def _build_experiment(
self, strategy, result: dict, n_articles: int, seed: int
) -> dict:
config = {
"strategy_class": type(strategy).__name__,
"strategy_name": strategy.name,
"n_workers": getattr(strategy, "n_workers", None),
"fmt": strategy.fmt,
"launcher": getattr(strategy, "launcher", None),
"from_surface": getattr(strategy, "from_surface", None),
"chrome_path": getattr(strategy, "chrome_path", None),
"n_articles": n_articles,
"seed": seed,
"zim_path": self.zim_path,
"kiwix_url": self.kiwix_url,
"gt_timeout_ms": self.gt_timeout_ms,
}
return {
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
"config": config,
"results": {k: v for k, v in result.items() if k != "bad_examples"},
"bad_examples": result.get("bad_examples", []),
}
def _dump_experiment(self, exp: dict):
exp_dir = self.output_dir / "experiments"
exp_dir.mkdir(parents=True, exist_ok=True)
ts = exp["timestamp"].replace(":", "")
name = exp["config"]["strategy_name"].replace(" ", "_").replace("/", "-")
path = exp_dir / f"{ts}_{name}.json"
path.write_text(json.dumps(exp, indent=2, default=str))
print(f"Experiment saved: {path}", flush=True)
# ---------------------------------------------------------------------------
# Display
# ---------------------------------------------------------------------------
def format_result_line(r: dict) -> str:
status = "PASS" if r["correct_pct"] >= CORRECT_THRESHOLD else "FAIL"
ok = f"{r['tiles_ok']}/{r['tiles_total']}"
return (
f" {r['name']:<25} {ok:>7} {r['correct_pct']:>5.1f}% "
f"{r['tiles_per_s']:>6.1f} {r['ms_per_tile']:>5.0f} "
f"{r['shot_pct']:>4.0f}% {status}"
)
def print_results(results: list[dict]):
for r in results:
print(format_result_line(r), flush=True)
if r["bad_examples"]:
for ex in r["bad_examples"][:3]:
print(f" BAD: {ex}", flush=True)
+216
View File
@@ -0,0 +1,216 @@
"""Chrome binary management for pixelshot.
Downloads and manages a patched headless Chrome binary with rawFilePath
support. Similar to `playwright install chromium`.
Usage:
pixelshot install-chrome # download patched headless_shell
pixelshot which-chrome # print path to active binary
Programmatic:
from pixelrag_render.chrome import find_chrome, install_chrome
path = find_chrome() # auto-detect best available
path = install_chrome() # download if needed
"""
import json
import os
import platform
import subprocess
import sys
import tarfile
import tempfile
import urllib.request
from pathlib import Path
INSTALL_DIR = Path.home() / ".cache" / "pixelrag" / "chrome"
VERSION_FILE = "version.json"
# Update these when releasing a new build
CHROME_VERSION = "150.0.7844.0"
RELEASE_URL_TEMPLATE = (
"https://github.com/StarTrail-org/PixelRAG/releases/download/"
"chrome-{version}/headless_shell-linux-x64.tar.zst"
)
# Search order for find_chrome()
_SEARCH_PATHS = [
lambda: os.environ.get("CHROME_PATH", ""),
lambda: str(INSTALL_DIR / "headless_shell"),
lambda: os.path.expanduser(
"~/.cache/ms-playwright/chromium-1217/chrome-linux64/chrome"
),
lambda: "/usr/bin/google-chrome",
lambda: "/usr/bin/google-chrome-stable",
lambda: "/usr/bin/chromium-browser",
lambda: "/usr/bin/chromium",
]
def find_chrome(auto_install: bool = True) -> str:
"""Find the best available Chrome binary. Auto-installs if none found.
Search order:
1. CHROME_PATH env var
2. pixelrag-installed headless_shell (~/.cache/pixelrag/chrome/)
3. Playwright's Chrome
4. System Chrome/Chromium
5. Auto-install patched headless_shell (if auto_install=True)
Returns:
Path to Chrome binary.
Raises:
FileNotFoundError: No Chrome binary found and auto_install=False.
"""
for path_fn in _SEARCH_PATHS:
path = path_fn()
if path and os.path.isfile(path) and os.access(path, os.X_OK):
return path
if auto_install:
print("No Chrome found. Installing headless_shell...", flush=True)
return str(install_chrome())
raise FileNotFoundError(
"No Chrome binary found. Run 'pixelshot install-chrome' or set CHROME_PATH."
)
def get_installed_version() -> str | None:
"""Return version string of installed headless_shell, or None."""
version_path = INSTALL_DIR / VERSION_FILE
if version_path.exists():
try:
data = json.loads(version_path.read_text())
return data.get("version")
except Exception:
pass
return None
def install_chrome(version: str | None = None, force: bool = False) -> Path:
"""Download and install the patched headless_shell binary.
Args:
version: Chrome version to install. Defaults to CHROME_VERSION.
force: Re-download even if already installed.
Returns:
Path to the installed headless_shell binary.
"""
version = version or CHROME_VERSION
binary_path = INSTALL_DIR / "headless_shell"
if binary_path.exists() and not force:
installed = get_installed_version()
if installed == version:
print(f"Already installed: headless_shell {version}")
return binary_path
if platform.system() != "Linux" or platform.machine() != "x86_64":
raise RuntimeError(
f"Pre-built headless_shell only available for linux-x64, "
f"got {platform.system()}-{platform.machine()}"
)
url = RELEASE_URL_TEMPLATE.format(version=version)
print(f"Downloading headless_shell {version}...")
print(f" URL: {url}")
INSTALL_DIR.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(suffix=".tar.zst", delete=False) as tmp:
tmp_path = tmp.name
try:
urllib.request.urlretrieve(url, tmp_path, _progress_hook)
print()
# Decompress: zstd → tar → extract
print("Extracting...")
# Try zstd decompression
decomp_path = tmp_path + ".tar"
try:
subprocess.run(
["zstd", "-d", tmp_path, "-o", decomp_path],
check=True,
capture_output=True,
)
except (FileNotFoundError, subprocess.CalledProcessError):
# Fallback: try python zstandard
try:
import zstandard
with open(tmp_path, "rb") as f_in, open(decomp_path, "wb") as f_out:
dctx = zstandard.ZstdDecompressor()
dctx.copy_stream(f_in, f_out)
except ImportError:
raise RuntimeError(
"zstd not found. Install with: apt install zstd (or pip install zstandard)"
)
with tarfile.open(decomp_path) as tar:
tar.extractall(INSTALL_DIR)
os.unlink(decomp_path)
# Set executable permission
binary_path.chmod(0o755)
# Write version file
version_data = {"version": version, "binary": str(binary_path)}
(INSTALL_DIR / VERSION_FILE).write_text(json.dumps(version_data))
print(
f"Installed: {binary_path} ({binary_path.stat().st_size / 1024 / 1024:.0f}MB)"
)
return binary_path
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
def _progress_hook(block_num, block_size, total_size):
downloaded = block_num * block_size
if total_size > 0:
pct = min(100, downloaded * 100 // total_size)
mb = downloaded / 1024 / 1024
total_mb = total_size / 1024 / 1024
print(f"\r {mb:.0f}/{total_mb:.0f} MB ({pct}%)", end="", flush=True)
def main():
"""CLI entry point for chrome management."""
import argparse
parser = argparse.ArgumentParser(description="Manage Chrome for pixelshot")
sub = parser.add_subparsers(dest="command")
sub.add_parser("install", help="Download patched headless_shell")
sub.add_parser("which", help="Print path to active Chrome binary")
sub.add_parser("version", help="Print installed version")
args = parser.parse_args()
if args.command == "install":
install_chrome()
elif args.command == "which":
try:
print(find_chrome())
except FileNotFoundError as e:
print(str(e), file=sys.stderr)
sys.exit(1)
elif args.command == "version":
v = get_installed_version()
if v:
print(v)
else:
print("Not installed", file=sys.stderr)
sys.exit(1)
else:
parser.print_help()
if __name__ == "__main__":
main()

Some files were not shown because too many files have changed in this diff Show More