ColPali: 基于PaliGemma-3B的视觉检索器与ColBERT策略

ColPali: 基于PaliGemma-3B的视觉检索器与ColBERT策略

colpali 项目地址: https://gitcode.com/mirrors/vidore/colpali

ColPali是一个基于视觉语言模型（VLMs）的新型模型架构和训练策略，能够高效地索引文档的视觉特征。它是PaliGemma-3B的扩展，生成ColBERT风格的文本和图像的多向量表示。

ColPali在论文ColPali: 基于视觉语言模型的文档检索中首次提出，并在此代码库中首次发布。

模型介绍

ColPali模型从现成的SigLIP模型开始构建，我们对其进行了微调以创建BiSigLIP，并将SigLIP输出的patch-embeddings输入到大型语言模型PaliGemma-3B中，以创建BiPali。

将图像patch embeddings输入语言模型的一个好处是它们被自然地映射到与文本输入（查询）类似的潜在空间。这使得可以利用ColBERT策略来计算文本token和图像patch之间的交互，与BiPali相比，这带来了性能上的显著提升。

模型训练

数据集

我们的训练数据集由127,460个查询-页面对组成，包括63%的公开可用的学术数据集的训练集和一个由网络爬取的PDF文档页面组成的合成数据集，以及由VLM生成的(Claude-3 Sonnet)伪问题增强(37%)。我们的训练集完全是英文的，旨在研究对非英语语言的零样本泛化。我们明确验证没有多页PDF文档同时用于ViDoRe和训练集中，以防止评估污染。使用样本的2%创建一个验证集，以调整超参数。

注意：多语言数据存在于语言模型（Gemma-2B）的预训练语料库中，并在PaliGemma-3B的多模态训练中可能发生。

参数

所有模型都在训练集上训练了1个epoch。除非另有说明，我们以bfloat16格式训练模型，使用低秩适配器(LoRA)，在语言模型的transformer层以及最终的随机初始化投影层上使用alpha=32和r=32，并使用paged_adamw_8bit优化器。我们在8个GPU上使用数据并行性进行训练，学习率为5e-5，线性衰减，2.5%的预热步骤，批处理大小为32。

使用方法

为了获得最佳性能，有更新的模型可用(vidore/colpali-v1.2)

# 此模型检查点兼容版本0.1.1，但不兼容推理库的较新版本
pip install colpali_engine==0.1.1

import torch
import typer
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import AutoProcessor
from PIL import Image

from colpali_engine.models.paligemma_colbert_architecture import ColPali
from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
from colpali_engine.utils.colpali_processing_utils import process_images, process_queries
from colpali_engine.utils.image_from_page_utils import load_from_dataset


def main() -> None:
    """使用ColPali进行推理的示例脚本"""

    # 加载模型
    model_name = "vidore/colpali"
    model = ColPali.from_pretrained("vidore/colpaligemma-3b-mix-448-base", torch_dtype=torch.bfloat16, device_map="cuda").eval()
    model.load_adapter(model_name)
    processor = AutoProcessor.from_pretrained(model_name)

    # 选择图像 -> load_from_pdf(<pdf_path>),  load_from_image_urls(["<url_1>"]), load_from_dataset(<path>)
    images = load_from_dataset("vidore/docvqa_test_subsampled")
    queries = ["James V. Fiorca来自哪所大学？", "日本的首相是谁？"]

    # 运行推理 - 文档
    dataloader = DataLoader(
        images,
        batch_size=4,
        shuffle=False,
        collate_fn=lambda x: process_images(processor, x),
    )
    ds = []
    for batch_doc in tqdm(dataloader):
        with torch.no_grad():
            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
            embeddings_doc = model(**batch_doc)
        ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))

    # 运行推理 - 查询
    dataloader = DataLoader(
        queries,
        batch_size=4,
        shuffle=False,
        collate_fn=lambda x: process_queries(processor, x, Image.new("RGB", (448, 448), (255, 255, 255))),
    )

    qs = []
    for batch_query in dataloader:
        with torch.no_grad():
            batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
            embeddings_query = model(**batch_query)
        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))

    # 运行评估
    retriever_evaluator = CustomEvaluator(is_multi_vector=True)
    scores = retriever_evaluator.evaluate(qs, ds)
    print(scores.argmax(axis=1))


if __name__ == "__main__":
    typer.run(main)

局限性

• 焦点：该模型主要关注PDF类型文档和高资源语言，可能限制其在其他文档类型或不太代表语言的泛化。
• 支持：该模型依赖于来自ColBERT晚期交互机制的多元向量检索，这可能需要工程工作才能适应缺乏本地多元向量支持的广泛使用的向量检索框架。

许可证

ColPali的视觉语言骨干模型（PaliGemma）根据其模型卡中指定的gemma许可证进行授权。附加到模型上的适配器根据MIT许可证进行授权。

联系方式

• Manuel Faysse: manuel.faysse@illuin.tech
• Hugues Sibille: hugues.sibille@illuin.tech
• Tony Wu: tony.wu@illuin.tech

引用

如果您在研究中使用此组织提供的任何数据集或模型，请按以下方式引用原始数据集：

@misc{faysse2024colpaliefficientdocumentretrieval,
  title={ColPali: Efficient Document Retrieval with Vision Language Models}, 
  author={Manuel Faysse and Hugues Sibille and Tony Wu and Bilel Omrani and Gautier Viaud and Céline Hudelot and Pierre Colombo},
  year={2024},
  eprint={2407.01449},
  archivePrefix={arXiv},
  primaryClass={cs.IR},
  url={https://arxiv.org/abs/2407.01449}, 
}

colpali 项目地址: https://gitcode.com/mirrors/vidore/colpali