Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/StarTrail-org/PixelRAG/llms.txt

Use this file to discover all available pages before exploring further.

The train/ directory is a separate uv project (wiki-screenshot-training) that LoRA fine-tunes Qwen/Qwen3-VL-Embedding-2B on Wikipedia screenshot data. Pre-trained adapters are already published — you only need this if you want to retrain from scratch or adapt the model to a new domain or dataset.

Pre-trained adapters

The trained LoRA adapters are published on Hugging Face and ready to use without any retraining:

LoRA Adapters

Chrisyichuan/wiki-screenshot-embedding-lora — Best checkpoint: lora_vit/ckpt200 (ViT-LoRA at step 200, ~102 MB)

Training Dataset

Chrisyichuan/screenshot-training-natural-filtered-v2 — 104K train pairs with hard negatives, LLM-augmented queries
Load the best checkpoint with PEFT:
from peft import PeftModel
from transformers import AutoModel

base = AutoModel.from_pretrained("Qwen/Qwen3-VL-Embedding-2B")
model = PeftModel.from_pretrained(
    base,
    "Chrisyichuan/wiki-screenshot-embedding-lora",
    subfolder="lora_vit/ckpt200",
)
To use the adapter with a running PixelRAG server, pass the --peft-adapter flag to pixelrag serve:
pixelrag serve --index-dir ./index --peft-adapter ./lora_vit/ckpt200 --port 30001

Training dataset

The full training set is published at Chrisyichuan/screenshot-training-natural-filtered-v2. It contains 104K training query–image pairs with 2 hard negatives each, built using a data curation pipeline of LLM-augmented query generation, filtering, and hard-negative mining. Because the dataset is released independently of the PixelRAG model, you can use it to adapt other backbones — a larger Qwen variant, or any other visual embedding model — without needing to reproduce the full PixelRAG training pipeline.

Requirements for retraining

Before starting a training run, make sure you have the following resources available:
ResourcePurpose
1× GPU ≥ 40 GB (H100 or A100)Fine-tuning
1× GPU for vLLM inferenceQA evaluation at each eval step
~95 GB fast local storageDataset images
OpenAI API key (gpt-4.1 grader)Eval QA scoring
W&B API keyLogging training curves online
The Python environment is a separate uv project in train/ with strictly pinned dependencies:
PackagePinned version
PyTorch2.9.1+cu129
transformers4.57.1
cuDNN9.20.0.48
Without a valid OpenAI API key, the QA score metric will silently read 0 for every evaluation step — the grader swallows errors rather than raising them. This makes a correctly running training job appear broken. Verify your key works before launching a multi-hour run.

Setting up the training environment

1

Change into the train/ directory and sync the environment

The train/ project has its own lockfile and must be installed from inside the directory:
cd train && uv sync
The train/ environment is intentionally separate from the root pyproject.toml. Do not install the training dependencies from the repo root — the pinned versions (torch==2.9.1+cu129, transformers==4.57.1, cuDNN 9.20.0.48) are only locked inside train/uv.lock.
2

Set your API keys

Both keys must be exported before you start training so they are available when the first evaluation step runs:
export OPENAI_API_KEY=sk-...
export WANDB_API_KEY=...
If you don’t need online W&B logging, you can substitute WANDB_MODE=offline in place of a W&B key — metrics will still be written to local eval_step*.jsonl files in the output directory.
3

Follow the complete training recipe in train/README.md

The full recipe — dataset downloads, image extraction, vLLM reader setup, and the exact training command — is documented in train/README.md. See that file for the step-by-step instructions.

Results

The fine-tuned model is evaluated on the miniv8 test set: 400 SimpleQA questions against 7,426 candidate screenshot tiles.
ModelQA Score
Base Qwen3-VL-Embedding-2B~0.715 – 0.730
LoRA lora_vit/ckpt200≈ 0.785
The QA score peaks around step 150–250 and may decay slightly thereafter due to overfitting — checkpoint at the peak. The original training run is logged at W&B: wandb.ai/yichuan_wang-uc-berkeley-electrical-engineering-computer/wiki-screenshot-training/runs/2y39owix

Build docs developers (and LLMs) love