Pixel-Native RAG: คู่มือสร้างระบบ Indexing เอกสารเชิงภาพแบบครบวงจร

import os
import sys
import io
import re
import json
import time
import math
import shutil
import hashlib
import asyncio
import logging
import argparse
import threading
import subprocess
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Any, Optional, Tuple
@dataclass
class Config:
urls: List[str] = field(default_factory=lambda: [
"https://en.wikipedia.org/wiki/Retrieval-augmented_generation",
"https://en.wikipedia.org/wiki/Vector_database",
"https://en.wikipedia.org/wiki/Transformer_(deep_learning_architecture)",
"https://en.wikipedia.org/wiki/Photosynthesis",
"https://en.wikipedia.org/wiki/Delhi",
])
include_synthetic_pdf: bool = True
tile_width: int = 1024
tile_height: int = 1024
tile_overlap: int = 128
device_scale: float = 1.0
max_page_height: int = 24000
max_tiles_per_doc: int = 12
min_tile_height: int = 200
blank_std_threshold: float = 6.0
dedup_hamming: int = 4
nav_timeout_ms: int = 60000
headless_args: List[str] = field(default_factory=lambda: [
"--no-sandbox", "--disable-dev-shm-usage", "--hide-scrollbars",
"--disable-gpu", "--force-color-profile=srgb", "--font-render-hinting=none",
])
backend: str = "siglip"
model_id: str = "google/siglip-base-patch16-224"
qwen_model_id: str = "Qwen/Qwen3-VL-Embedding-2B"
embed_batch_size: int = 8
embed_image_size: Optional[int] = None
index_dir: str = "./pixel_index"
ivf_threshold: int = 2000
ivf_nprobe: int = 16
top_k_tiles: int = 20
n_docs: int = 5
use_ocr_hybrid: bool = True
rrf_k: int = 60
dense_weight: float = 1.0
sparse_weight: float = 1.0
enable_server: bool = True
server_port: int = 8000
enable_eval: bool = True
enable_adapter_train: bool = True
enable_vlm_answer: bool = False
vlm_model_id: str = "Qwen/Qwen2.5-VL-3B-Instruct"
show_plots: bool = True
work_dir: str = "./pixelrag_work"
seed: int = 0
CFG = Config()
EVAL_QUERIES: List[Tuple[str, str]] = [
("how do plants convert sunlight into chemical energy", "Photosynthesis"),
("chlorophyll light dependent reactions", "Photosynthesis"),
("converting scanned images of text into machine readable characters", "Optical_character"),
("approximate nearest neighbour search over embeddings", "Vector_database"),
("self-attention multi-head architecture", "Transformer"),
("grounding a language model with retrieved documents", "Retrieval-augmented"),
("capital territory of india red fort", "Delhi"),
]
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)-7s | %(message)s",
datefmt="%H:%M:%S")
log = logging.getLogger("pixelrag")
for noisy in ("urllib3", "PIL", "matplotlib", "httpx", "asyncio", "uvicorn.error"):
logging.getLogger(noisy).setLevel(logging.WARNING)
IN_COLAB = "google.colab" in sys.modules
def _pip(*pkgs: str) -> None:
"""Install quietly; never explode the notebook on a single bad wheel."""
cmd = [sys.executable, "-m", "pip", "install", "-q", "--disable-pip-version-check", *pkgs]
subprocess.run(cmd, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
def _have(mod: str) -> bool:
import importlib.util
return importlib.util.find_spec(mod) is not None
def ensure_deps(cfg: Config) -> None:
log.info("Installing dependencies (first run only, ~2-4 min)...")
wanted = []
for mod, pkg in [
("PIL", "pillow"), ("numpy", "numpy"), ("faiss", "faiss-cpu"),
("fitz", "pymupdf"), ("transformers", "transformers"),
("fastapi", "fastapi"), ("uvicorn", "uvicorn"), ("requests", "requests"),
("matplotlib", "matplotlib"), ("tqdm", "tqdm"), ("rank_bm25", "rank-bm25"),
("playwright", "playwright"), ("sentencepiece", "sentencepiece"),
]:
if not _have(mod):
wanted.append(pkg)
if cfg.use_ocr_hybrid and not _have("pytesseract"):
wanted.append("pytesseract")
if wanted:
_pip(*wanted)
if not _have("torch"):
log.warning("torch not found — installing CPU wheel (Colab normally ships torch).")
_pip("torch", "torchvision")
if cfg.use_ocr_hybrid and shutil.which("tesseract") is None:
log.info("Installing tesseract-ocr system package...")
subprocess.run("apt-get -qq update && apt-get -qq install -y tesseract-ocr",
shell=True, check=False,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
if shutil.which("tesseract") is None:
log.warning("tesseract unavailable -> hybrid retrieval will run dense-only.")
cfg.use_ocr_hybrid = False
marker = Path(cfg.work_dir) / ".chromium_ok"
if not marker.exists():
log.info("Downloading Playwright Chromium...")
r = subprocess.run([sys.executable, "-m", "playwright", "install", "--with-deps", "chromium"],
capture_output=True, text=True)
if r.returncode != 0:
r = subprocess.run([sys.executable, "-m", "playwright", "install", "chromium"],
capture_output=True, text=True)
if r.returncode == 0:
marker.parent.mkdir(parents=True, exist_ok=True)
marker.write_text("ok")
else:
log.warning("Chromium install failed -> falling back to the text renderer.\n%s",
(r.stderr or "")[-600:])
log.info("Dependencies ready.")
def run_async(coro):
"""
Run a coroutine from a Jupyter/Colab cell.
Colab already owns a running event loop, which makes Playwright's *sync*
API raise. Rather than monkey-patching with nest_asyncio, we hand the
coroutine to a private loop on a private thread — the most robust option.
"""
box: Dict[str, Any] = {}
def _runner():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
box["value"] = loop.run_until_complete(coro)
except BaseException as exc:
box["error"] = exc
finally:
try:
loop.run_until_complete(loop.shutdown_asyncgens())
finally:
loop.close()
t = threading.Thread(target=_runner, daemon=True)
t.start()
t.join()
if "error" in box:
raise box["error"]
return box.get("value")

@dataclass
class Tile:
tile_id: str
doc_id: str
source: str
kind: str
page: int
seq: int
y0: int
y1: int
path: str
ocr_text: str = ""
title: str = ""
def _doc_id_from_source(src: str) -> str:
tail = src.rstrip("/").split("/")[-1] or src
tail = re.sub(r"\.(html?|pdf|png|jpg)$", "", tail, flags=re.I)
return re.sub(r"[^A-Za-z0-9_.\-()]+", "_", tail)[:80] or hashlib.md5(src.encode()).hexdigest()[:10]
def _ahash(img, size: int = 8) -> int:
"""64-bit average hash — cheap near-duplicate detection for repeated headers."""
import numpy as np
g = img.convert("L").resize((size, size))
a = np.asarray(g, dtype="float32")
bits = (a > a.mean()).flatten()
out = 0
for b in bits:
out = (out << 1) | int(b)
return out
def _hamming(a: int, b: int) -> int:
return bin(a ^ b).count("1")
def _is_informative(img, cfg: Config) -> bool:
"""Reject blank / solid-colour tiles before they ever reach the GPU."""
import numpy as np
a = np.asarray(img.convert("L"), dtype="float32")
return float(a.std()) >= cfg.blank_std_threshold
def _save_tile(img, out_dir: Path, name: str) -> str:
out_dir.mkdir(parents=True, exist_ok=True)
p = out_dir / f"{name}.png"
img.convert("RGB").save(p, format="PNG", optimize=True)
return str(p)
def slice_image_to_tiles(img, cfg: Config, *, doc_id: str, source: str, kind: str,
page: int, out_dir: Path, start_seq: int = 0,
seen_hashes: Optional[List[int]] = None,
title: str = "") -> List[Tile]:
"""Vertical sliding window with overlap. Used for PDFs and text fallback."""
from PIL import Image
seen_hashes = seen_hashes if seen_hashes is not None else []
W, H = img.size
if W != cfg.tile_width:
new_h = max(1, int(H * cfg.tile_width / W))
img = img.resize((cfg.tile_width, new_h))
W, H = img.size
step = max(1, cfg.tile_height - cfg.tile_overlap)
tiles: List[Tile] = []
y, seq = 0, start_seq
while y < H and (seq - start_seq) < cfg.max_tiles_per_doc:
h = min(cfg.tile_height, H - y)
if h < cfg.min_tile_height and seq > start_seq:
break
crop = img.crop((0, y, W, y + h))
if _is_informative(crop, cfg):
hsh = _ahash(crop)
if all(_hamming(hsh, s) > cfg.dedup_hamming for s in seen_hashes):
seen_hashes.append(hsh)
tid = f"{doc_id}__p{page}__t{seq}"
tiles.append(Tile(
tile_id=tid, doc_id=doc_id, source=source, kind=kind, page=page,
seq=seq, y0=y, y1=y + h, title=title,
path=_save_tile(crop, out_dir, tid),
))
seq += 1
y += step
return tiles
_JS_AUTOSCROLL = """
async () => {
await new Promise((resolve) => {
let y = 0;
const timer = setInterval(() => {
window.scrollBy(0, 800);
y += 800;
if (y >= document.body.scrollHeight || y > 40000) {
clearInterval(timer);
window.scrollTo(0, 0);
setTimeout(resolve, 250);
}
}, 40);
});
}
"""
_JS_FLATTEN = """
() => {
document.querySelectorAll('*').forEach((el) => {
const s = getComputedStyle(el);
if (s.position === 'fixed' || s.position === 'sticky') el.style.position = 'absolute';
});
document.querySelectorAll('[role="dialog"], .cookie, #cookie-banner, .cc-banner')
.forEach((el) => el.remove());
}
"""
_CSS_CLEANUP = """
* { animation: none !important; transition: none !important;
scroll-behavior: auto !important; }
html { -webkit-font-smoothing: antialiased; }
video, iframe[src*="youtube"] { visibility: hidden !important; }
"""
_UA = ("Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/124.0 Safari/537.36 PixelRAG-Tutorial/1.0")
async def _render_urls_async(urls: List[str], cfg: Config, out_dir: Path) -> List[Tile]:
from playwright.async_api import async_playwright
from PIL import Image
all_tiles: List[Tile] = []
async with async_playwright() as pw:
browser = await pw.chromium.launch(headless=True, args=cfg.headless_args)
ctx = await browser.new_context(
viewport={"width": cfg.tile_width, "height": cfg.tile_height},
device_scale_factor=cfg.device_scale,
user_agent=_UA,
java_script_enabled=True,
)
for url in urls:
doc_id = _doc_id_from_source(url)
page = await ctx.new_page()
try:
await page.goto(url, wait_until="domcontentloaded", timeout=cfg.nav_timeout_ms)
try:
await page.wait_for_load_state("networkidle", timeout=12000)
except Exception:
pass
await page.evaluate(_JS_AUTOSCROLL)
await page.add_style_tag(content=_CSS_CLEANUP)
await page.evaluate(_JS_FLATTEN)
title = (await page.title()) or doc_id
height = await page.evaluate(
"() => Math.max(document.body.scrollHeight, "
"document.documentElement.scrollHeight)")
height = int(min(height, cfg.max_page_height))
step = max(1, cfg.tile_height - cfg.tile_overlap)
seen: List[int] = []
y, seq = 0, 0
while y < height and seq < cfg.max_tiles_per_doc:
h = min(cfg.tile_height, height - y)
if h < cfg.min_tile_height and seq > 0:
break
buf = await page.screenshot(
full_page=True, type="png",
clip={"x": 0, "y": y, "width": cfg.tile_width, "height": h})
img = Image.open(io.BytesIO(buf)).convert("RGB")
if img.size[0] != cfg.tile_width:
img = img.resize((cfg.tile_width,
max(1, int(img.size[1] * cfg.tile_width / img.size[0]))))
if _is_informative(img, cfg):
hsh = _ahash(img)
if all(_hamming(hsh, s) > cfg.dedup_hamming for s in seen):
seen.append(hsh)
tid = f"{doc_id}__p0__t{seq}"
all_tiles.append(Tile(
tile_id=tid, doc_id=doc_id, source=url, kind="web",
page=0, seq=seq, y0=y, y1=y + h, title=title,
path=_save_tile(img, out_dir, tid)))
seq += 1
y += step
log.info(" rendered %-34s -> %2d tiles (page %dpx)", doc_id, seq, height)
except Exception as exc:
log.warning(" FAILED %s (%s)", url, type(exc).__name__)
finally:
await page.close()
await ctx.close()
await browser.close()
return all_tiles
def render_urls(urls: List[str], cfg: Config, out_dir: Path) -> List[Tile]:
"""Screenshot every URL into tiles; degrade to the text renderer on failure."""
try:
tiles = run_async(_render_urls_async(urls, cfg, out_dir))
if tiles:
return tiles
log.warning("Browser produced no tiles — using text-render fallback.")
except Exception as exc:
log.warning("Playwright unavailable (%s: %s) — using text-render fallback.",
type(exc).__name__, str(exc)[:160])
return [t for u in urls for t in render_url_as_text(u, cfg, out_dir)]
def _strip_html(html: str) -> str:
html = re.sub(r"(?is)<(script|style|nav|footer|header|noscript).*?</\1>", " ", html)
html = re.sub(r"(?s)<!--.*?-->", " ", html)
html = re.sub(r"(?i)</(p|div|h[1-6]|li|tr|br)>", "\n", html)
text = re.sub(r"(?s)<[^>]+>", " ", html)
for a, b in [(" ", " "), ("&", "&"), ("<", "<"), (">", ">"), ("\"", '"')]:
text = text.replace(a, b)
text = re.sub(r"\[\d+\]", "", text)
text = re.sub(r"[ \t]+", " ", text)
return re.sub(r"\n{2,}", "\n", text).strip()
def _mono_font(size: int = 20):
from PIL import ImageFont
for cand in ("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf"):
if os.path.exists(cand):
return ImageFont.truetype(cand, size)
try:
import matplotlib.font_manager as fm
return ImageFont.truetype(fm.findfont("DejaVu Sans"), size)
except Exception:
return ImageFont.load_default()
def text_to_image(text: str, cfg: Config, title: str = "") -> Any:
"""Render plain text onto a tall white canvas — a browser-free stand-in."""
from PIL import Image, ImageDraw
font, tfont = _mono_font(20), _mono_font(30)
pad, lh, wrap = 40, 30, max(20, (cfg.tile_width - 80) // 11)
lines: List[str] = []
for para in text.split("\n"):
para = para.strip()
if not para:
continue
while len(para) > wrap:
cut = para.rfind(" ", 0, wrap)
cut = cut if cut > 0 else wrap
lines.append(para[:cut])
para = para[cut:].lstrip()
lines.append(para)
lines = lines[:900]
height = pad * 2 + 60 + lh * len(lines)
img = Image.new("RGB", (cfg.tile_width, max(cfg.tile_height, height)), "white")
d = ImageDraw.Draw(img)
d.text((pad, pad), title[:60], font=tfont, fill=(15, 15, 15))
for i, ln in enumerate(lines):
d.text((pad, pad + 60 + i * lh), ln, font=font, fill=(35, 35, 35))
return img
def render_url_as_text(url: str, cfg: Config, out_dir: Path) -> List[Tile]:
import requests
doc_id = _doc_id_from_source(url)
try:
r = requests.get(url, timeout=30, headers={"User-Agent": _UA})
r.raise_for_status()
body = _strip_html(r.text)
m = re.search(r"(?is)<title>(.*?)</title>", r.text)
title = m.group(1).strip() if m else doc_id
except Exception as exc:
log.warning(" fetch failed for %s (%s)", url, type(exc).__name__)
return []
img = text_to_image(body, cfg, title=title)
log.info(" text-rendered %-30s -> canvas %dpx", doc_id, img.size[1])
return slice_image_to_tiles(img, cfg, doc_id=doc_id, source=url, kind="text",
page=0, out_dir=out_dir, title=title)
def render_pdf(pdf_path: str, cfg: Config, out_dir: Path, dpi: int = 150) -> List[Tile]:
import fitz
from PIL import Image
doc_id = _doc_id_from_source(pdf_path)
tiles: List[Tile] = []
with fitz.open(pdf_path) as doc:
title = (doc.metadata or {}).get("title") or doc_id
n_pages = doc.page_count
for pno in range(n_pages):
pix = doc[pno].get_pixmap(dpi=dpi)
img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
tiles += slice_image_to_tiles(img, cfg, doc_id=doc_id, source=pdf_path,
kind="pdf", page=pno, out_dir=out_dir,
title=title)
log.info(" rendered %-34s -> %2d tiles (%d pages)", doc_id, len(tiles), n_pages)
return tiles
def make_synthetic_pdf(path: Path) -> str:
"""A tiny PDF so the tutorial always exercises the PDF path, offline or not."""
import fitz
body = [
("PixelRAG Internal Note", 22),
("", 12),
("Why pixel-native retrieval?", 16),
("Parsers are per-site glue code. A renderer is one code path for every", 11),
("document type: HTML, PDF, scanned fax, spreadsheet export, dashboard.", 11),
("", 11),
("Tiling policy", 16),
("Tiles are 1024x1024 with 128px of vertical overlap. Overlap keeps a", 11),
("sentence or table row from being split across two embeddings, which is", 11),
("the single biggest source of recall loss in naive screenshot pipelines.", 11),
("", 11),
("Serving", 16),
("FAISS inner-product over L2-normalised vectors equals cosine similarity.", 11),
("Tile scores are max-pooled per document so one strong tile can surface", 11),
("a long page, mirroring late-interaction retrieval behaviour.", 11),
("", 11),
("The mitochondria reference is a joke; the overlap advice is not.", 11),
]
doc = fitz.open()
page = doc.new_page()
y = 72
for line, size in body:
page.insert_text((72, y), line, fontsize=size, fontname="helv")
y += size + 8
doc.save(str(path))
doc.close()
return str(path)

def ocr_tiles(tiles: List[Tile], cfg: Config) -> None:
if not cfg.use_ocr_hybrid:
return
try:
import pytesseract
from PIL import Image
except Exception:
log.warning("pytesseract missing -> dense-only retrieval.")
cfg.use_ocr_hybrid = False
return
from tqdm.auto import tqdm
t0 = time.time()
for t in tqdm(tiles, desc="OCR", unit="tile"):
try:
raw = pytesseract.image_to_string(Image.open(t.path), config="--psm 6")
t.ocr_text = re.sub(r"\s+", " ", raw).strip()[:4000]
except Exception:
t.ocr_text = ""
log.info("OCR over %d tiles in %.1fs", len(tiles), time.time() - t0)
def torch_device() -> str:
import torch
if torch.cuda.is_available():
return "cuda"
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
return "mps"
return "cpu"
class DualEncoderBackend:
"""
SigLIP / CLIP image-text dual encoder.
Honest caveat: these encoders were trained on natural images with short
captions (64-77 token text towers). They understand a screenshot's *gist*
— layout, topic, figures — not its fine print. That is exactly why upstream
PixelRAG uses Qwen3-VL-Embedding-2B plus a LoRA trained on screenshots.
Sections §9 (OCR hybrid) and §10 (adapter) exist to close part of the gap
on hardware that can't host a 2B VLM.
"""
def __init__(self, cfg: Config):
import torch
from transformers import AutoModel, AutoProcessor
self.cfg = cfg
self.device = torch_device()
self.dtype = torch.float16 if self.device == "cuda" else torch.float32
self.model_id = cfg.model_id if cfg.backend != "clip" else "openai/clip-vit-base-patch32"
log.info("Loading embedding model %s on %s (%s)", self.model_id, self.device,
str(self.dtype).replace("torch.", ""))
self.processor = AutoProcessor.from_pretrained(self.model_id)
self.model = AutoModel.from_pretrained(self.model_id, torch_dtype=self.dtype)
self.model.to(self.device).eval()
self.is_siglip = "siglip" in self.model_id.lower()
self.dim = int(getattr(self.model.config, "projection_dim", 0) or
getattr(self.model.config.text_config, "hidden_size", 512))
self.name = f"{'siglip' if self.is_siglip else 'clip'}:{self.model_id}"
@staticmethod
def _l2(x):
import numpy as np
n = np.linalg.norm(x, axis=-1, keepdims=True)
return (x / np.clip(n, 1e-12, None)).astype("float32")
def embed_images(self, images: List[Any], bs: Optional[int] = None):
import torch, numpy as np
from tqdm.auto import tqdm
bs = bs or self.cfg.embed_batch_size
out = []
for i in tqdm(range(0, len(images), bs), desc="embed:image", unit="batch"):
batch = images[i:i + bs]
inputs = self.processor(images=batch, return_tensors="pt")
inputs = {k: v.to(self.device, self.dtype if v.is_floating_point() else v.dtype)
for k, v in inputs.items()}
with torch.no_grad():
feats = self.model.get_image_features(**inputs)
out.append(feats.float().cpu().numpy())
return self._l2(np.concatenate(out, 0)) if out else np.zeros((0, self.dim), "float32")
def embed_texts(self, texts: List[str], bs: Optional[int] = None):
import torch, numpy as np
bs = bs or max(16, self.cfg.embed_batch_size)
out = []
for i in range(0, len(texts), bs):
batch = [t if t.strip() else " " for t in texts[i:i + bs]]
kw = dict(text=batch, return_tensors="pt", truncation=True)
kw.update(padding="max_length", max_length=64) if self.is_siglip else kw.update(padding=True)
inputs = self.processor(**kw)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
feats = self.model.get_text_features(**inputs)
out.append(feats.float().cpu().numpy())
return self._l2(np.concatenate(out, 0)) if out else np.zeros((0, self.dim), "float32")
class Qwen3VLEmbeddingBackend:
"""
Opt-in backend matching upstream (Qwen/Qwen3-VL-Embedding-2B).
Needs a recent transformers (>= 4.57) and ~8 GB of VRAM in fp16. It embeds
text and images into one space by mean-pooling the last hidden state of a
VLM prompt, which is why it handles dense screenshot text far better than
a CLIP-style tower.
"""
def __init__(self, cfg: Config):
import torch
from transformers import AutoModel, AutoProcessor
self.cfg = cfg
self.device = torch_device()
self.dtype = torch.float16 if self.device == "cuda" else torch.float32
mid = cfg.qwen_model_id
log.info("Loading %s (this is a large download)...", mid)
self.processor = AutoProcessor.from_pretrained(mid, trust_remote_code=True)
self.model = AutoModel.from_pretrained(mid, torch_dtype=self.dtype,
trust_remote_code=True).to(self.device).eval()
self.dim = int(self.model.config.hidden_size)
self.name = f"qwen3vl:{mid}"
def _pool(self, hidden, mask):
import torch
m = mask.unsqueeze(-1).to(hidden.dtype)
return (hidden * m).sum(1) / m.sum(1).clamp(min=1e-6)
def _encode(self, **proc_kwargs):
import torch, numpy as np
inputs = self.processor(return_tensors="pt", padding=True, **proc_kwargs)
inputs = {k: (v.to(self.device) if hasattr(v, "to") else v) for k, v in inputs.items()}
with torch.no_grad():
out = self.model(**inputs, output_hidden_states=True)
hidden = out.hidden_states[-1] if getattr(out, "hidden_states", None) is not None \
else out.last_hidden_state
vec = self._pool(hidden, inputs["attention_mask"]).float().cpu().numpy()
return DualEncoderBackend._l2(vec)
def embed_images(self, images: List[Any], bs: Optional[int] = None):
import numpy as np
from tqdm.auto import tqdm
bs = bs or max(1, self.cfg.embed_batch_size // 4)
chunks = []
for i in tqdm(range(0, len(images), bs), desc="embed:image", unit="batch"):
batch = images[i:i + bs]
prompt = ["Describe this document screenshot for retrieval."] * len(batch)
chunks.append(self._encode(text=prompt, images=batch))
return np.concatenate(chunks, 0)
def embed_texts(self, texts: List[str], bs: Optional[int] = None):
import numpy as np
bs = bs or 8
chunks = [self._encode(text=[f"Query: {t}" for t in texts[i:i + bs]])
for i in range(0, len(texts), bs)]
return np.concatenate(chunks, 0) if chunks else np.zeros((0, self.dim), "float32")
def build_backend(cfg: Config):
if cfg.backend == "qwen3vl":
try:
return Qwen3VLEmbeddingBackend(cfg)
except Exception as exc:
log.warning("Qwen3-VL backend failed (%s: %s) -> falling back to SigLIP.",
type(exc).__name__, str(exc)[:200])
cfg.backend = "siglip"
return DualEncoderBackend(cfg)
def embed_tiles(tiles: List[Tile], backend, cfg: Config):
from PIL import Image
import numpy as np
vecs = []
bs = cfg.embed_batch_size
for i in range(0, len(tiles), bs):
imgs = [Image.open(t.path).convert("RGB") for t in tiles[i:i + bs]]
vecs.append(backend.embed_images(imgs, bs=bs))
for im in imgs:
im.close()
return np.concatenate(vecs, 0) if vecs else np.zeros((0, backend.dim), "float32")

class PixelIndex:
"""
Inner-product FAISS index over L2-normalised vectors (== cosine similarity).
Flat below `ivf_threshold` vectors (exact, no training); IVF above it
(sub-linear, needs training + nprobe tuning). Raw vectors are also kept in
memory so §10 can re-project them after adapter training without re-running
the encoder.
"""
def __init__(self, dim: int, cfg: Config):
self.dim, self.cfg = dim, cfg
self.index = None
self.metas: List[Dict[str, Any]] = []
self.vectors = None
self._bm25 = None
self._bm25_corpus: List[List[str]] = []
def build(self, vectors, tiles: List[Tile]) -> "PixelIndex":
import faiss, numpy as np
vectors = np.ascontiguousarray(vectors.astype("float32"))
n = vectors.shape[0]
if n == 0:
raise RuntimeError("No vectors to index — did rendering produce any tiles?")
if n >= self.cfg.ivf_threshold:
nlist = max(4, min(4096, int(4 * math.sqrt(n))))
quant = faiss.IndexFlatIP(self.dim)
base = faiss.IndexIVFFlat(quant, self.dim, nlist, faiss.METRIC_INNER_PRODUCT)
base.train(vectors)
base.nprobe = self.cfg.ivf_nprobe
log.info("FAISS IndexIVFFlat n=%d nlist=%d nprobe=%d", n, nlist, base.nprobe)
else:
base = faiss.IndexFlatIP(self.dim)
log.info("FAISS IndexFlatIP n=%d dim=%d (exact search)", n, self.dim)
self.index = faiss.IndexIDMap2(base)
self.index.add_with_ids(vectors, np.arange(n).astype("int64"))
self.vectors = vectors
self.metas = [asdict(t) for t in tiles]
self._fit_bm25()
return self
def _fit_bm25(self) -> None:
if not self.cfg.use_ocr_hybrid:
return
try:
from rank_bm25 import BM25Okapi
except Exception:
return
self._bm25_corpus = [re.findall(r"[a-z0-9]+", (m.get("ocr_text", "") + " " +
m.get("title", "")).lower())
for m in self.metas]
if any(self._bm25_corpus):
self._bm25 = BM25Okapi([c or ["_"] for c in self._bm25_corpus])
log.info("BM25 fitted over OCR sidecar (%d docs)", len(self._bm25_corpus))
def search_dense(self, qvecs, k: int):
import numpy as np
scores, ids = self.index.search(np.ascontiguousarray(qvecs.astype("float32")), k)
return scores, ids
def search_sparse(self, query: str, k: int) -> List[Tuple[int, float]]:
if self._bm25 is None:
return []
import numpy as np
toks = re.findall(r"[a-z0-9]+", query.lower())
if not toks:
return []
s = np.asarray(self._bm25.get_scores(toks))
top = np.argsort(-s)[:k]
return [(int(i), float(s[i])) for i in top if s[i] > 0]
def save(self, out_dir: str) -> None:
import faiss, numpy as np
p = Path(out_dir)
p.mkdir(parents=True, exist_ok=True)
faiss.write_index(self.index, str(p / "tiles.faiss"))
np.save(p / "vectors.npy", self.vectors)
(p / "metas.jsonl").write_text("\n".join(json.dumps(m) for m in self.metas))
(p / "manifest.json").write_text(json.dumps(
{"dim": self.dim, "n": len(self.metas), "created": time.time(),
"config": asdict(self.cfg)}, indent=2))
log.info("Index saved to %s (%d tiles)", p.resolve(), len(self.metas))
@classmethod
def load(cls, out_dir: str, cfg: Config) -> "PixelIndex":
import faiss, numpy as np
p = Path(out_dir)
man = json.loads((p / "manifest.json").read_text())
obj = cls(man["dim"], cfg)
obj.index = faiss.read_index(str(p / "tiles.faiss"))
obj.vectors = np.load(p / "vectors.npy")
obj.metas = [json.loads(l) for l in (p / "metas.jsonl").read_text().splitlines() if l]
obj._fit_bm25()
return obj
def reproject(self, new_vectors) -> None:
"""Swap in re-embedded vectors (used after adapter training in §10)."""
tiles = [Tile(**m) for m in self.metas]
self.build(new_vectors, tiles)
def build_index(cfg: Config) -> Tuple[PixelIndex, Any, List[Tile]]:
work = Path(cfg.work_dir)
tiles_dir = work / "tiles"
tiles_dir.mkdir(parents=True, exist_ok=True)
log.info("=" * 74)
log.info("STAGE 1/4 RENDER (documents -> image tiles)")
log.info("=" * 74)
tiles: List[Tile] = render_urls(cfg.urls, cfg, tiles_dir)
if cfg.include_synthetic_pdf:
pdf_path = make_synthetic_pdf(work / "pixelrag_note.pdf")
tiles += render_pdf(pdf_path, cfg, tiles_dir)
if not tiles:
raise RuntimeError("Rendering produced zero tiles. Check network access.")
log.info("Total tiles: %d across %d documents",
len(tiles), len({t.doc_id for t in tiles}))
log.info("=" * 74)
log.info("STAGE 2/4 OCR SIDECAR (for hybrid retrieval + pair mining)")
log.info("=" * 74)
ocr_tiles(tiles, cfg)
log.info("=" * 74)
log.info("STAGE 3/4 EMBED (tiles -> vectors)")
log.info("=" * 74)
backend = build_backend(cfg)
t0 = time.time()
vecs = embed_tiles(tiles, backend, cfg)
log.info("Embedded %d tiles -> %s in %.1fs (%.2f tiles/s)",
vecs.shape[0], vecs.shape, time.time() - t0,
vecs.shape[0] / max(1e-6, time.time() - t0))
log.info("=" * 74)
log.info("STAGE 4/4 INDEX (vectors -> FAISS)")
log.info("=" * 74)
index = PixelIndex(vecs.shape[1], cfg).build(vecs, tiles)
index.save(cfg.index_dir)
return index, backend, tiles

def search(query: str, index: PixelIndex, backend, cfg: Config,
n_docs: Optional[int] = None) -> List[Dict[str, Any]]:
import numpy as np
n_docs = n_docs or cfg.n_docs
k = min(cfg.top_k_tiles, len(index.metas))
qv = backend.embed_texts([query])
dscores, dids = index.search_dense(qv, k)
dense = [(int(i), float(s)) for i, s in zip(dids[0], dscores[0]) if i >= 0]
fused: Dict[int, float] = {}
for rank, (tid, _) in enumerate(dense):
fused[tid] = fused.get(tid, 0.0) + cfg.dense_weight / (cfg.rrf_k + rank + 1)
sparse = index.search_sparse(query, k) if cfg.use_ocr_hybrid else []
for rank, (tid, _) in enumerate(sparse):
fused[tid] = fused.get(tid, 0.0) + cfg.sparse_weight / (cfg.rrf_k + rank + 1)
dense_lookup = dict(dense)
tile_hits = sorted(fused.items(), key=lambda kv: -kv[1])
per_doc: Dict[str, Dict[str, Any]] = {}
for tid, fscore in tile_hits:
m = index.metas[tid]
d = per_doc.setdefault(m["doc_id"], {
"doc_id": m["doc_id"], "title": m.get("title") or m["doc_id"],
"source": m["source"], "kind": m["kind"], "score": 0.0,
"dense_score": 0.0, "tiles": [],
})
d["score"] = max(d["score"], fscore)
d["dense_score"] = max(d["dense_score"], dense_lookup.get(tid, 0.0))
if len(d["tiles"]) < 3:
d["tiles"].append({
"tile_id": m["tile_id"], "path": m["path"], "seq": m["seq"],
"page": m["page"], "y0": m["y0"], "y1": m["y1"],
"rrf": round(fscore, 6),
"cosine": round(dense_lookup.get(tid, 0.0), 4),
"snippet": (m.get("ocr_text", "") or "")[:220],
})
return sorted(per_doc.values(), key=lambda d: -d["score"])[:n_docs]
def pretty_print(query: str, results: List[Dict[str, Any]]) -> None:
print(f"\n\033[1mQ: {query}\033[0m")
if not results:
print(" (no hits)")
return
for i, r in enumerate(results, 1):
print(f" {i}. [{r['score']:.4f} rrf | {r['dense_score']:.3f} cos] "
f"{r['title'][:64]} ({r['kind']})")
top = r["tiles"][0]
print(f" tile {top['tile_id']} y={top['y0']}-{top['y1']}")
if top["snippet"]:
print(f" \033[2m{top['snippet'][:150]}...\033[0m")
class SearchServer:
"""FastAPI + uvicorn on a background thread, mirroring upstream's POST /search."""
def __init__(self, index: PixelIndex, backend, cfg: Config):
from fastapi import FastAPI
from pydantic import BaseModel
class Query(BaseModel):
text: str
class SearchRequest(BaseModel):
queries: List[Query]
n_docs: int = cfg.n_docs
app = FastAPI(title="PixelRAG (tutorial)", version="1.0")
@app.get("/health")
def health():
return {"status": "ok", "tiles": len(index.metas),
"docs": len({m["doc_id"] for m in index.metas}),
"backend": getattr(backend, "name", "unknown")}
@app.post("/search")
def do_search(req: SearchRequest):
return {"results": [
{"query": q.text, "docs": search(q.text, index, backend, cfg, req.n_docs)}
for q in req.queries]}
self.app, self.cfg = app, cfg
self.thread: Optional[threading.Thread] = None
self.server = None
def start(self) -> bool:
import uvicorn, requests
config = uvicorn.Config(self.app, host="127.0.0.1", port=self.cfg.server_port,
log_level="error")
self.server = uvicorn.Server(config)
self.thread = threading.Thread(target=self.server.run, daemon=True)
self.thread.start()
for _ in range(40):
time.sleep(0.25)
try:
if requests.get(f"http://127.0.0.1:{self.cfg.server_port}/health",
timeout=2).ok:
log.info("Search API live on http://127.0.0.1:%d", self.cfg.server_port)
return True
except Exception:
continue
log.warning("Server did not come up in time.")
return False
def stop(self) -> None:
if self.server:
self.server.should_exit = True
if self.thread:
self.thread.join(timeout=5)

def evaluate(index: PixelIndex, backend, cfg: Config,
queries: List[Tuple[str, str]] = EVAL_QUERIES,
label: str = "eval", quiet: bool = False) -> Dict[str, float]:
ranks: List[Optional[int]] = []
for q, want in queries:
docs = search(q, index, backend, cfg, n_docs=10)
hit = next((i for i, d in enumerate(docs) if want.lower() in d["doc_id"].lower()), None)
ranks.append(hit)
if not quiet:
got = docs[0]["doc_id"] if docs else "-"
mark = "OK " if hit == 0 else (f"@{hit + 1}" if hit is not None else "MISS")
print(f" [{mark:>4}] {q[:56]:<58} -> {got[:32]}")
n = len(ranks)
m = {
"recall@1": sum(r == 0 for r in ranks) / n,
"recall@3": sum(r is not None and r < 3 for r in ranks) / n,
"recall@5": sum(r is not None and r < 5 for r in ranks) / n,
"mrr": sum(1.0 / (r + 1) for r in ranks if r is not None) / n,
}
print(f" \033[1m{label}\033[0m R@1={m['recall@1']:.2f} R@3={m['recall@3']:.2f} "
f"R@5={m['recall@5']:.2f} MRR={m['mrr']:.3f}")
return m
def mine_training_pairs(tiles: List[Tile], max_per_tile: int = 2) -> List[Tuple[str, int]]:
"""Weak supervision: pseudo-queries from a tile's own OCR text / title."""
import random
rng = random.Random(0)
pairs: List[Tuple[str, int]] = []
for idx, t in enumerate(tiles):
text = (t.ocr_text or "").strip()
cands: List[str] = []
if len(text) > 80:
words = text.split()
for _ in range(max_per_tile):
if len(words) <= 14:
break
s = rng.randint(0, len(words) - 14)
span = " ".join(words[s:s + rng.randint(8, 14)])
if len(span) > 30:
cands.append(span)
if t.title:
cands.append(t.title)
for c in cands[:max_per_tile]:
pairs.append((c, idx))
return pairs
class ResidualAdapter:
"""Shared two-layer residual MLP applied to both query and tile vectors."""
def __init__(self, dim: int, hidden: int = 512, device: str = "cpu"):
import torch
import torch.nn as nn
self.device = device
self.net = nn.Sequential(
nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim)
).to(device)
for p in self.net[-1].parameters():
torch.nn.init.zeros_(p)
self.logit_scale = torch.nn.Parameter(torch.tensor(2.996, device=device))
self.dim = dim
def forward_t(self, x):
import torch
y = x + self.net(x)
return torch.nn.functional.normalize(y, dim=-1)
def apply_np(self, arr):
import torch, numpy as np
with torch.no_grad():
t = torch.from_numpy(np.ascontiguousarray(arr.astype("float32"))).to(self.device)
return self.forward_t(t).cpu().numpy().astype("float32")
def train_adapter(index: PixelIndex, backend, tiles: List[Tile], cfg: Config,
epochs: int = 12, batch: int = 24, lr: float = 1e-4):
import torch, numpy as np
pairs = mine_training_pairs(tiles)
if len(pairs) < 32:
log.warning("Only %d mined pairs — skipping adapter training "
"(enable OCR or add documents).", len(pairs))
return None
log.info("Mined %d (pseudo-query, tile) pairs from %d tiles", len(pairs), len(tiles))
q_texts = [p[0] for p in pairs]
t_idx = np.array([p[1] for p in pairs], dtype="int64")
log.info("Pre-embedding pseudo-queries (frozen encoder, done once)...")
Q = torch.from_numpy(backend.embed_texts(q_texts))
V = torch.from_numpy(index.vectors)
device = "cuda" if torch.cuda.is_available() else "cpu"
ad = ResidualAdapter(index.dim, device=device)
Q, V = Q.to(device), V.to(device)
opt = torch.optim.AdamW(list(ad.net.parameters()) + [ad.logit_scale], lr=lr, weight_decay=1e-2)
n = len(pairs)
doc_ids = torch.from_numpy(t_idx).to(device)
for ep in range(epochs):
perm = torch.randperm(n, device=device)
total, steps = 0.0, 0
for i in range(0, n, batch):
sel = perm[i:i + batch]
if sel.numel() < 4:
continue
qb = ad.forward_t(Q[sel])
docs = doc_ids[sel]
vb = ad.forward_t(V[docs])
logits = ad.logit_scale.exp().clamp(max=100) * qb @ vb.T
same = docs[:, None] == docs[None, :]
eye = torch.eye(len(sel), dtype=torch.bool, device=device)
logits = logits.masked_fill(same & ~eye, float("-inf"))
labels = torch.arange(len(sel), device=device)
loss = 0.5 * (torch.nn.functional.cross_entropy(logits, labels) +
torch.nn.functional.cross_entropy(logits.T, labels))
opt.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(ad.net.parameters(), 1.0)
opt.step()
total += loss.detach().item()
steps += 1
if ep % 3 == 0 or ep == epochs - 1:
log.info(" epoch %2d/%d InfoNCE loss %.4f", ep + 1, epochs, total / max(steps, 1))
return ad
class AdaptedBackend:
"""Wraps a frozen backend so queries pass through the trained adapter."""
def __init__(self, backend, adapter: ResidualAdapter):
self.backend, self.adapter = backend, adapter
self.dim = backend.dim
self.name = f"{getattr(backend, 'name', 'backend')}+adapter"
def embed_texts(self, texts, bs=None):
return self.adapter.apply_np(self.backend.embed_texts(texts, bs=bs))
def embed_images(self, images, bs=None):
return self.adapter.apply_np(self.backend.embed_images(images, bs=bs))
def answer_with_vlm(query: str, results: List[Dict[str, Any]], cfg: Config,
max_tiles: int = 3) -> str:
"""
Retrieval returns pixels, so generation must accept pixels. Any VLM works;
Qwen2.5-VL-3B is a reasonable Colab-sized default (~7 GB download).
"""
try:
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText
except Exception as exc:
return f"[VLM unavailable: {exc}]"
paths = [t["path"] for r in results for t in r["tiles"]][:max_tiles]
if not paths:
return "[no retrieved tiles]"
log.info("Loading VLM %s ...", cfg.vlm_model_id)
proc = AutoProcessor.from_pretrained(cfg.vlm_model_id)
model = AutoModelForImageTextToText.from_pretrained(
cfg.vlm_model_id,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto")
images = [Image.open(p).convert("RGB") for p in paths]
content = [{"type": "image"} for _ in images] + [{"type": "text", "text":
f"These are screenshots retrieved for the question. Answer using only what "
f"is visible, and say so if the answer is not shown.\n\nQuestion: {query}"}]
prompt = proc.apply_chat_template([{"role": "user", "content": content}],
add_generation_prompt=True, tokenize=False)
inputs = proc(text=[prompt], images=images, return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
text = proc.batch_decode(out[:, inputs["input_ids"].shape[1]:],
skip_special_tokens=True)[0]
return text.strip()
def show_results(query: str, results: List[Dict[str, Any]], max_tiles: int = 3) -> None:
try:
import matplotlib.pyplot as plt
from PIL import Image
except Exception:
return
tiles = [(r, t) for r in results for t in r["tiles"][:1]][:max_tiles]
if not tiles:
return
fig, axes = plt.subplots(1, len(tiles), figsize=(5 * len(tiles), 6))
axes = [axes] if len(tiles) == 1 else list(axes)
for ax, (r, t) in zip(axes, tiles):
ax.imshow(Image.open(t["path"]))
ax.set_title(f"{r['title'][:34]}\nrrf={r['score']:.4f} cos={t['cosine']:.3f}",
fontsize=9)
ax.axis("off")
fig.suptitle(f"Q: {query}", fontsize=12)
plt.tight_layout()
plt.show()

def main(cfg: Config = CFG) -> Dict[str, Any]:
Path(cfg.work_dir).mkdir(parents=True, exist_ok=True)
ensure_deps(cfg)
import numpy as np
np.random.seed(cfg.seed)
banner = """
██████╗ ██╗██╗ ██╗███████╗██╗ ██████╗ █████╗ ██████╗
██╔══██╗██║╚██╗██╔╝██╔════╝██║ ██╔══██╗██╔══██╗██╔════╝
██████╔╝██║ ╚███╔╝ █████╗ ██║ ██████╔╝███████║██║ ███╗
██╔═══╝ ██║ ██╔██╗ ██╔══╝ ██║ ██╔══██╗██╔══██║██║ ██║
██║ ██║██╔╝ ██╗███████╗███████╗██║ ██║██║ ██║╚██████╔╝
╚═╝ ╚═╝╚═╝ ╚═╝╚══════╝╚══════╝╚═╝ ╚═╝╚═╝ ╚═╝ ╚═════╝
pixel-native retrieval: render -> tile -> embed -> FAISS -> serve
"""
print(banner)
index, backend, tiles = build_index(cfg)
print("\n" + "=" * 74)
print("SEARCH DEMO — text query against a pixel index")
print("=" * 74)
demo_queries = [
"how do plants turn light into sugar",
"what does a vector database store",
"why use overlapping tiles when screenshotting a page",
]
for q in demo_queries:
res = search(q, index, backend, cfg)
pretty_print(q, res)
if cfg.show_plots:
show_results(q, res)
metrics_before = None
if cfg.enable_eval:
print("\n" + "=" * 74)
print("EVALUATION — baseline")
print("=" * 74)
metrics_before = evaluate(index, backend, cfg, label="baseline")
if cfg.use_ocr_hybrid:
cfg.use_ocr_hybrid = False
print("\n -- ablation: dense only (OCR/BM25 disabled) --")
evaluate(index, backend, cfg, label="dense-only", quiet=True)
cfg.use_ocr_hybrid = True
active_backend = backend
if cfg.enable_adapter_train:
print("\n" + "=" * 74)
print("ADAPTER TRAINING — contrastive head over frozen embeddings")
print("=" * 74)
adapter = train_adapter(index, backend, tiles, cfg)
if adapter is not None:
index.reproject(adapter.apply_np(index.vectors))
active_backend = AdaptedBackend(backend, adapter)
if cfg.enable_eval:
print("\n -- after adapter --")
after = evaluate(index, active_backend, cfg, label="adapted")
if metrics_before:
d = after["mrr"] - metrics_before["mrr"]
print(f" MRR delta: {d:+.3f} "
f"({'improved' if d > 0 else 'no gain — expected on a corpus this small'})")
server = None
if cfg.enable_server:
print("\n" + "=" * 74)
print("SERVE — FastAPI, upstream-compatible POST /search")
print("=" * 74)
server = SearchServer(index, active_backend, cfg)
if server.start():
import requests
r = requests.post(f"http://127.0.0.1:{cfg.server_port}/search",
json={"queries": [{"text": "what is retrieval augmented generation"}],
"n_docs": 3}, timeout=120)
payload = r.json()
for res in payload["results"]:
print(f"\n POST /search query={res['query']!r}")
for d in res['docs']:
print(f" - {d['score']:.4f} {d['title'][:56]} <{d['source'][:48]}>")
print("\n Equivalent curl:")
print(f" curl -X POST http://127.0.0.1:{cfg.server_port}/search \\")
print(" -H 'Content-Type: application/json' \\")
print(" -d '{\"queries\":[{\"text\":\"capital of india\"}],\"n_docs\":3}'")
if cfg.enable_vlm_answer:
print("\n" + "=" * 74)
print("GENERATION — answering from retrieved pixels")
print("=" * 74)
q = "According to the retrieved screenshots, what is photosynthesis?"
res = search(q, index, active_backend, cfg, n_docs=2)
print(answer_with_vlm(q, res, cfg))
else:
print("\n[i] Set CFG.enable_vlm_answer = True (GPU) to generate answers "
"directly from the retrieved tiles.")
n_docs = len({m['doc_id'] for m in index.metas})
print("\n" + "=" * 74)
print("DONE")
print("=" * 74)
print(f" tiles indexed : {len(index.metas)} across {n_docs} documents")
print(f" embedding dim : {index.dim} backend: {getattr(active_backend, 'name', '?')}")
print(f" index on disk : {Path(cfg.index_dir).resolve()}")
print(f" tiles on disk : {Path(cfg.work_dir).resolve() / 'tiles'}")
print("""
Try next:
* CFG.urls -> point at your own pages, then re-run main()
* CFG.backend = "qwen3vl" -> upstream's Qwen3-VL-Embedding-2B (needs a big GPU)
* CFG.device_scale = 2.0 -> sharper tiles, better small-text retrieval
* CFG.tile_overlap = 256 -> higher recall on prose, more vectors to store
* render_pdf("/content/your.pdf", CFG, Path(CFG.work_dir)/"tiles")
* The real deal: git clone https://github.com/StarTrail-org/PixelRAG
uv sync --package pixelrag-index && pixelrag-index build
""")
return {"index": index, "backend": active_backend, "tiles": tiles, "server": server,
"search": lambda q, k=5: pretty_print(q, search(q, index, active_backend, cfg, k))}
if __name__ == "__main__":
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument("--no-server", action="store_true")
parser.add_argument("--no-train", action="store_true")
parser.add_argument("--backend", default=None)
args, _ = parser.parse_known_args()
if args.no_server:
CFG.enable_server = False
if args.no_train:
CFG.enable_adapter_train = False
if args.backend:
CFG.backend = args.backend
STATE = main(CFG)