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:
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
# Markdown
|
||||
*.md linguist-detectable=true
|
||||
*.md linguist-documentation=false
|
||||
|
||||
# JSON
|
||||
*.json linguist-detectable=true
|
||||
|
||||
# YAML
|
||||
*.yml linguist-detectable=true
|
||||
@@ -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
|
||||
@@ -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/
|
||||
@@ -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) }}
|
||||
@@ -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 }}
|
||||
@@ -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
@@ -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/
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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 † 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> ·
|
||||
<a href="#give-claude-eyes">Give Claude eyes</a> ·
|
||||
<a href="#how-it-works">How it works</a> ·
|
||||
<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
|
||||
@@ -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 |
@@ -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");
|
||||
@@ -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.
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
@@ -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
|
||||
@@ -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()
|
||||
@@ -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>
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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: Retrieval–Reader 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 84–88 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`.
|
||||
@@ -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 30–90 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 2–5 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 60–120 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 ~200–300MB, 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 | 2–5 min |
|
||||
| Task 2: `fetch --no-history chromium` | 30–90 min (network-dependent) |
|
||||
| Task 3: patch transfer + apply | 2 min |
|
||||
| Task 4: args.gn setup | 1 min |
|
||||
| Task 5: `gn gen` | 3–7 min |
|
||||
| Task 6: `autoninja` build | 60–120 min (224 cores, no debug) |
|
||||
| Task 7: verification | 2 min |
|
||||
| **Total** | **~2–4 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)
|
||||
@@ -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
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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`
|
||||
@@ -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.6–1.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 ~2–6pp 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>
|
||||
```
|
||||
@@ -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%
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
@@ -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
|
||||
@@ -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
@@ -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
|
||||
@@ -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
|
||||
@@ -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.
|
||||
Executable
+118
@@ -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
File diff suppressed because it is too large
Load Diff
+1233
File diff suppressed because it is too large
Load Diff
+1350
File diff suppressed because it is too large
Load Diff
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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>
|
||||
@@ -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>
|
||||
@@ -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>
|
||||
@@ -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": {}
|
||||
}
|
||||
]
|
||||
@@ -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>
|
||||
@@ -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>
|
||||
@@ -0,0 +1 @@
|
||||
"""pixelrag-index: orchestration layer for building searchable FAISS indexes."""
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
}
|
||||
Executable
+1027
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Generated
+1103
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"dependencies": {
|
||||
"remark-gfm": "^4.0.1"
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
Executable
+22
@@ -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"
|
||||
@@ -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
@@ -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"
|
||||
@@ -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 1–100 (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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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
Reference in New Issue
Block a user