十分钟跑通:byaldi + ColPali 检索一本 PDF

环境准备

# Python 3.11+,CUDA 12+
pip install byaldi pdf2image pillow

# macOS / Ubuntu 系统级依赖:pdf2image 需要 poppler
brew install poppler          # macOS
sudo apt install poppler-utils  # Ubuntu

索引一个 PDF 文件夹

from byaldi import RAGMultiModalModel

# 1. 加载模型(首次下载 ~6GB,之后走本地缓存)
model = RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2")

# 2. 把文件夹下所有 PDF 都索引
model.index(
    input_path="./pdfs/",
    index_name="annual-reports",
    store_collection_with_index=True,    # 保存页面图像,方便后续看原图
    overwrite=True,
)

# 3. 查询
results = model.search("2024 年营收同比增长多少?", k=5)
for r in results:
    print(f"{r.doc_id} 第 {r.page_num} 页,分数 {r.score:.2f}")

这 10 行代码的背后,byaldi 做了:

1. 用 pdf2image 把每页 PDF 渲染成 PIL Image(默认 DPI=150)
2. 批量送入 ColPali,每页产出 1024×128 张量
3. 向量和原图都存到本地磁盘(.byaldi/annual-reports/)
4. 查询时对 query 做同样编码,对所有页面算 MaxSim,topK 返回

不用 byaldi:裸调 transformers

byaldi 封装得太多,理解原理还得看裸版本:

import torch
from transformers import ColPaliForRetrieval, ColPaliProcessor
from pdf2image import convert_from_path

device = "cuda" if torch.cuda.is_available() else "cpu"
model = ColPaliForRetrieval.from_pretrained(
    "vidore/colpali-v1.2",
    torch_dtype=torch.bfloat16,
    device_map=device,
).eval()
processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")

# 1. 编码所有页
pages = convert_from_path("report.pdf", dpi=150)
batch = processor.process_images(pages).to(device)
with torch.no_grad():
    page_vecs = model(**batch).embeddings   # (N_pages, 1024, 128)

# 2. 编码 query
queries = ["revenue growth in 2024"]
q_batch = processor.process_queries(queries).to(device)
with torch.no_grad():
    q_vecs = model(**q_batch).embeddings    # (1, n_tok, 128)

# 3. MaxSim 打分
scores = processor.score_multi_vector(q_vecs, page_vecs)  # (1, N_pages)
top_pages = scores.squeeze().topk(5).indices
for i in top_pages:
    print(f"第 {i+1} 页,得分 {scores[0, i]:.3f}")

加速:批处理

from torch.utils.data import DataLoader

def collate(batch):
    return processor.process_images(batch)

loader = DataLoader(pages, batch_size=8, collate_fn=collate)

all_vecs = []
with torch.no_grad():
    for batch in loader:
        batch = { k: v.to(device) for k, v in batch.items() }
        all_vecs.append(model(**batch).embeddings)

page_vecs = torch.cat(all_vecs, dim=0)

A100 40GB 上 batch=16 大概 40 张 / 秒,一本 500 页的年报 12 秒索引完。

DPI 与图片尺寸

第一个实战:查一本真实财报

queries = [
    "哪一页展示 2024 年各业务线营收饼图?",
    "管理层讨论里提到的三大风险是什么",
    "ESG 报告里说温室气体排放降多少",
]
for q in queries:
    r = model.search(q, k=3)
    print(f"\n【{q}】")
    for hit in r:
        print(f"  p.{hit.page_num} score={hit.score:.2f}")

对 Apple 2024 Annual Report 测试,三条查询都在 top-1 命中正确页。对照 OCR + BGE 方案的 top-3 命中率约 60%——差距明显。

常见错误 & 排查

本章小结