Document RAG

Document RAG (Vision-Language Retrieval-Augmented Generation)

When to Use

Use this skill when:

• Building a retrieval-augmented generation system over PDF/image documents (papers, reports, scanned forms)
• Retrieving relevant pages without OCR, using vision-language page embeddings (ColPali-style late interaction)
• Running a VLM (Qwen2-VL, LLaVA, InternVL) to answer questions grounded in retrieved page images
• Evaluating retrieval recall@k or generation faithfulness for a document QA pipeline
• The document has complex layout (tables, figures, multi-column) where text-only RAG loses structure

Version Compatibility

• Python ≥ 3.10, torch ≥ 2.2, transformers ≥ 4.45
• colpali-engine ≥ 0.3 (ColPali / ColQwen2 checkpoints)
• pdf2image ≥ 1.17 (requires system poppler-utils), or pymupdf ≥ 1.24 as a poppler-free alternative
• faiss-cpu ≥ 1.8 or chromadb ≥ 0.5 for the vector store

Prerequisites

• pip install torch transformers colpali-engine pdf2image faiss-cpu qwen-vl-utils
• System package poppler-utils (for pdf2image) or skip it and use pymupdf
• A CUDA GPU with ≥16 GB VRAM for Qwen2-VL-7B in bf16 (or use the 2B variant / 4-bit quantization on CPU)
• Familiarity with basic RAG concepts (embed → retrieve → generate) is assumed

Core Patterns

Goal: turn a PDF into page images and page-level multi-vector embeddings. Approach: render each page at 150 DPI (enough for text legibility, cheap to embed), then embed with ColPali, which outputs one vector per image patch (not a single pooled vector) so retrieval can do token-level late interaction.

from pdf2image import convert_from_path
from colpali_engine.models import ColPali, ColPaliProcessor
import torch

model = ColPali.from_pretrained(
    "vidore/colpali-v1.2", torch_dtype=torch.bfloat16, device_map="cuda"
).eval()
processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")


def load_pages(pdf_path: str, dpi: int = 150):
    """Render each PDF page to a PIL Image. 150 DPI balances OCR-quality vs embedding speed."""
    return convert_from_path(pdf_path, dpi=dpi)


def embed_pages(pages):
    """Return a list of (n_patches, dim) tensors, one multi-vector embedding per page."""
    batch = processor.process_images(pages).to(model.device)
    with torch.no_grad():
        embeddings = model(**batch)  # (n_pages, n_patches, dim)
    return list(embeddings)


def embed_query(text: str):
    """Embed a text query into (n_tokens, dim) for late-interaction scoring against pages."""
    batch = processor.process_queries([text]).to(model.device)
    with torch.no_grad():
        return model(**batch)[0]  # (n_tokens, dim)

Goal: rank pages by relevance to a query without collapsing embeddings to a single vector. Approach: MaxSim (ColBERT-style late interaction) — for each query token, take its max similarity across all page patches, then sum over query tokens. This preserves fine-grained matches (e.g. a query token matching one table cell) that mean-pooling would wash out.

def maxsim_score(query_emb: torch.Tensor, page_emb: torch.Tensor) -> float:
    """Late-interaction similarity: sum over query tokens of the max similarity to any page patch."""
    scores = torch.einsum("qd,pd->qp", query_emb, page_emb)  # (n_query_tok, n_page_patches)
    return scores.max(dim=1).values.sum().item()


def retrieve(query_emb, page_embeddings, top_k: int = 5):
    """Return (indices, scores) of the top_k most relevant pages, highest score first."""
    scores = [maxsim_score(query_emb, pe) for pe in page_embeddings]
    ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
    return ranked, [scores[i] for i in ranked]

Goal: answer a user question grounded in the retrieved page images. Approach: feed the top-k page images plus the question directly to a VLM chat template — no OCR step needed, the model reads the rendered page.

from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

vlm = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
)
vlm_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")


def answer_from_pages(context_pages, question: str, max_new_tokens: int = 256) -> str:
    """Generate a grounded answer from a list of retrieved PIL page images."""
    content = [{"type": "image", "image": p} for p in context_pages]
    content.append({"type": "text", "text": question})
    messages = [{"role": "user", "content": content}]

    text = vlm_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, _ = process_vision_info(messages)
    inputs = vlm_processor(text=[text], images=image_inputs, return_tensors="pt").to(vlm.device)

    outputs = vlm.generate(**inputs, max_new_tokens=max_new_tokens)
    generated = outputs[:, inputs["input_ids"].shape[1]:]  # strip the prompt tokens
    return vlm_processor.decode(generated[0], skip_special_tokens=True)


# Full pipeline: index -> retrieve -> generate
pages = load_pages("document.pdf")
page_embeddings = embed_pages(pages)

query = "What is the main contribution of this paper?"
top_idx, scores = retrieve(embed_query(query), page_embeddings, top_k=3)
answer = answer_from_pages([pages[i] for i in top_idx], query)

Pitfalls

• Memory: Qwen2-VL-7B needs ~18 GB GPU RAM in bf16; use the 2B variant or 4-bit (bitsandbytes) quantization for CPU/small-GPU setups.
• DPI tradeoff: higher DPI improves legibility of small text/tables but slows embedding linearly; 150 DPI is a good default, go to 200–300 only for dense tables.
• Late interaction vs single vector: ColPali's multi-vector MaxSim beats CLIP-style single-vector retrieval on documents with tables/multi-column layouts — don't mean-pool the patch embeddings, it discards the advantage.
• Store embeddings per-patch: a FAISS/ChromaDB flat index expects fixed-size vectors, so multi-vector ColPali embeddings need either a custom late-interaction index or a library like colbert-ai's PLAID — a plain single-vector FAISS index cannot do MaxSim directly.
• Hallucination: VLMs will confidently answer from low-quality or irrelevant scans; add a confidence/faithfulness check (NLI-based or LLM-as-judge) before trusting the answer.
• Token budget: each image costs hundreds of vision tokens in Qwen2-VL; keep top_k small (3-5 pages) to stay within context and latency limits.

See Also

• ai-science-vision-rag — broader vision-language RAG patterns and evaluation
• vision-language-models — VLM architecture and inference fundamentals
• ai-science-llm-finetuning — fine-tune the generation component for domain adaptation
• pdf — general PDF parsing/text-extraction utilities