使用Qdrant扩展PDF检索

时间: 30 分钟级别: 中级输出: GitHubOpen In Colab

高效的PDF文档检索是(代理)检索增强生成(RAG)等任务以及许多其他基于搜索的应用程序中的常见需求。同时，设置PDF文档检索通常也会面临一些额外的挑战。

许多传统的PDF检索解决方案依赖于光学字符识别（OCR）以及特定用例的启发式方法来处理视觉复杂的元素，如表格、图像和图表。这些算法通常不可转移——即使在同一领域内——因为它们具有任务定制的解析和分块策略，劳动密集型，容易出错，并且难以扩展。

最近在视觉大语言模型（VLLMs）方面的进展，如ColPali及其继任者ColQwen，开启了PDF检索的变革。这些多模态模型直接以PDF页面作为输入，无需预处理。任何可以转换为图像的内容（将PDF视为文档页面的截图）都可以被这些模型有效处理。VLLMs在使用上更为简单，在PDF检索基准测试中达到了最先进的性能，如视觉文档检索（ViDoRe）基准测试。

VLLMs 如何用于 PDF 检索

像ColPali和ColQwen这样的VLLMs为每个PDF页面生成多向量表示；这些表示被存储并索引在向量数据库中。在检索过程中，模型动态地为（文本）用户查询创建多向量表示，并通过延迟交互机制实现PDF页面和查询之间的精确匹配。

扩展VLLMs的挑战

由VLLMs生成的重型多向量表示使得大规模PDF检索在计算上非常密集。如果不进行优化，这些模型在大规模PDF检索任务中效率低下。

缩放背后的数学原理

ColPali 每页PDF生成超过 1,000个向量，而其继任者 ColQwen 生成的向量略少，最多为 768个向量，根据图像大小动态调整。通常，ColQwen 每页生成 约700个向量。

为了理解其影响，考虑构建一个HNSW索引，这是向量数据库的一种常见索引算法。让我们粗略估计将一个新的PDF页面插入索引所需的比较次数。

  • 每页向量数： ~700 (ColQwen) 或 ~1,000 (ColPali)
  • ef_construct: 100（默认）

向量比较次数的下界估计将是：

$$ 700 \times 700 \times 100 = 49 \ \text{百万} $$

现在想象一下在20,000页上建立索引需要多少时间！

对于ColPali，这个数字会翻倍。结果是极其缓慢的索引构建时间。

我们的解决方案

我们建议减少PDF页面表示中的向量数量以进行第一阶段检索。在使用减少的向量数量进行第一阶段检索后，我们建议使用原始的未压缩表示对检索到的子集进行重新排序。

向量的减少可以通过对多向量VLLM生成的输出应用均值池化操作来实现。均值池化平均了选定子组内所有向量的值，将多个向量压缩成一个代表性的向量。如果操作得当，它可以在显著减少向量数量的同时，保留原始页面的重要信息。

VLLMs生成对应于表示PDF页面不同部分的补丁的向量。这些补丁可以分组为PDF页面的列和行。

例如：

  • ColPali 将 PDF 页面分成 1,024 个补丁。
  • 通过对该补丁矩阵的行（或列）应用均值池化，将页面表示减少到仅32个向量。

ColPali patching of a PDF page

我们使用ColPali模型测试了这种方法，通过PDF页面行对其多向量进行平均池化。结果显示：

  • 索引时间快一个数量级
  • 检索质量与原始模型相当

有关此实验的详细信息，请参阅我们的gitHub 仓库、ColPali 优化博客文章或网络研讨会“大规模 PDF 检索”

本教程的目标

在本教程中，我们将展示一种使用Qdrant和ColPali & ColQwen2 VLLMs进行PDF检索的可扩展方法。 强烈推荐使用所展示的方法，以避免长时间索引和检索速度慢的常见问题。

在以下部分中，我们将展示一个通过我们成功实验得出的优化检索算法：

第一阶段检索使用均值池化向量：

  • 使用仅均值池化向量构建HNSW索引。
  • 将它们用于第一阶段检索。

使用原始模型多向量重新排序：

  • 使用来自ColPali或ColQwen2的原始多向量重新排序第一阶段检索到的结果。

设置

安装并导入所需的库

# pip install colpali_engine>=0.3.1
from colpali_engine.models import ColPali, ColPaliProcessor
# pip install qdrant-client>=1.12.0
from qdrant_client import QdrantClient, models

要运行这些实验，我们使用的是Qdrant集群。如果你刚开始，可以设置一个免费层级的集群进行测试和探索。请按照文档中的说明操作“如何创建一个免费层级的Qdrant集群”

client = QdrantClient(
    url=<YOUR CLUSTER URL>,
    api_key=<YOUR API KEY>
)

下载ColPali模型及其输入处理器。请确保选择适合您设置的后端。

colpali_model = ColPali.from_pretrained(
        "vidore/colpali-v1.3",
        torch_dtype=torch.bfloat16,
        device_map="mps",  # Use "cuda:0" for GPU, "cpu" for CPU, or "mps" for Apple Silicon
    ).eval()

colpali_processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.3")
For ColQwen model
from colpali_engine.models import ColQwen2, ColQwen2Processor

colqwen_model = ColQwen2.from_pretrained(
        "vidore/colqwen2-v0.1",
        torch_dtype=torch.bfloat16,
        device_map="mps", # Use "cuda:0" for GPU, "cpu" for CPU, or "mps" for Apple Silicon
    ).eval()

colqwen_processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v0.1")

创建Qdrant集合

我们现在可以在Qdrant中创建一个集合，用于存储由ColPali或ColQwen生成的PDF页面的多向量表示。

集合将包括PDF页面的行和列表示的均值池化，以及原始多向量表示。

client.create_collection(
    collection_name=collection_name,
    vectors_config={
        "original": 
            models.VectorParams( #switch off HNSW
                    size=128,
                    distance=models.Distance.COSINE,
                    multivector_config=models.MultiVectorConfig(
                        comparator=models.MultiVectorComparator.MAX_SIM
                    ),
                    hnsw_config=models.HnswConfigDiff(
                        m=0 #switching off HNSW
                    )
            ),
        "mean_pooling_columns": models.VectorParams(
                size=128,
                distance=models.Distance.COSINE,
                multivector_config=models.MultiVectorConfig(
                    comparator=models.MultiVectorComparator.MAX_SIM
                )
            ),
        "mean_pooling_rows": models.VectorParams(
                size=128,
                distance=models.Distance.COSINE,
                multivector_config=models.MultiVectorConfig(
                    comparator=models.MultiVectorComparator.MAX_SIM
                )
            )
    }
)

选择一个数据集

我们将使用Daniel van Strien提供的UFO数据集进行本教程。该数据集可在Hugging Face上获取；您可以直接从那里下载。

from datasets import load_dataset
ufo_dataset = "davanstrien/ufo-ColPali"
dataset = load_dataset(ufo_dataset, split="train")

嵌入和平均池化

我们将使用一个函数，该函数批量生成每个PDF页面（即图像）的多向量表示及其平均池化版本。 为了完全理解，重要的是要考虑ColPali和ColQwen的以下具体细节：

ColPali: 理论上，ColPali 被设计为每页 PDF 生成 1,024 个向量，但实际上，它生成了 1,030 个向量。这种差异是由于 ColPali 的预处理器，它在每个输入前附加了文本 Describe the image.。这段额外的文本生成了额外的 6 个多向量。

ColQwen: ColQwen 根据PDF页面的大小动态确定其“行和列”中的补丁数量。因此，多向量的数量可能会因输入而异。ColQwen 预处理器会在前面添加 <|im_start|>user<|vision_start|> 并在后面追加 <|vision_end|>描述图像。<|im_end|><|endoftext|>。

例如，这就是ColQwen多向量输出的形成方式。

that’s how ColQwen multivector output is formed

get_patches 函数用于获取 ColPali/ColQwen2 模型将 PDF 页面划分为的 x_patches（行）和 y_patches（列）的数量。 对于 ColPali，数量将始终为 32 乘 32；ColQwen 将根据 PDF 页面大小动态定义它们。

x_patches, y_patches = model_processor.get_n_patches(
    image_size, 
    patch_size=model.patch_size
)
For ColQwen model
model_processor.get_n_patches(
    image_size, 
    patch_size=model.patch_size,
    spatial_merge_size=model.spatial_merge_size
)

我们选择保留前缀和后缀多向量。我们的池化操作根据模型确定的行数和列数（ColPali为静态32x32，ColQwen为动态XxY）压缩表示图像标记的多向量。该函数保留并将模型产生的额外多向量整合回池化表示中。

ColPali模型的简化版本池化：

（请参阅完整版本——同样适用于ColQwen——在教程笔记本中）


processed_images = model_processor.process_images(image_batch) 
# Image embeddings of shape (batch_size, 1030, 128)
image_embeddings = model(**processed_images)

# (1030, 128)
image_embedding = image_embeddings[0] # take the first element of the batch

# Now we need to identify vectors that correspond to the image tokens
# It can be done by selecting tokens corresponding to special `image_token_id`

# (1030, ) - boolean mask (for the first element in the batch), True for image tokens 
mask = processed_images.input_ids[0] == model_processor.image_token_id

# For convenience, we now select only image tokens 
#   and reshape them to (x_patches, y_patches, dim)

# (x_patches, y_patches, 128)
image_tokens = image_embedding[mask].view(x_patches, y_patches, model.dim)

# Now we can apply mean pooling by rows and columns

# (x_patches, 128)
pooled_by_rows = image_tokens.mean(dim=0)

# (y_patches, 128)
pooled_by_columns = image_tokens.mean(dim=1)

# [Optionally] we can also concatenate special tokens to the pooled representations, 
# For ColPali, it's only postfix

# (x_patches + 6, 128)
pooled_by_rows = torch.cat([pooled_by_rows, image_embedding[~mask]])

# (y_patches + 6, 128)
pooled_by_columns = torch.cat([pooled_by_columns, image_embedding[~mask]])

上传到Qdrant

上传过程非常简单；唯一需要注意的是ColPali和ColQwen2模型的计算成本。 在资源有限的环境中，建议使用较小的批量大小进行嵌入和平均池化。

上传代码的完整版本可在教程笔记本中找到

查询PDF

在索引PDF文档之后，我们可以继续使用我们的两阶段检索方法来查询它们。

query = "Lee Harvey Oswald's involvement in the JFK assassination"
processed_queries = model_processor.process_queries([query]).to(model.device)

# Resulting query embedding is a tensor of shape (22, 128)
query_embedding = model(**processed_queries)[0]

现在让我们设计一个用于两阶段检索的函数，该函数使用由VLLMs生成的多向量：

  • 步骤 1: 使用压缩的多向量表示和HNSW索引预取结果。
  • 步骤 2: 使用原始的多向量表示对预取的结果进行重新排序。

让我们使用组合平均池化表示来查询我们的集合，以进行检索的第一阶段。

# Final amount of results to return
search_limit = 10
# Amount of results to prefetch for reranking
prefetch_limit = 100

response = client.query_points(
    collection_name=collection_name,
    query=query_embedding,
    prefetch=[
        models.Prefetch(
            query=query_embedding,
            limit=prefetch_limit,
            using="mean_pooling_columns"
        ),
        models.Prefetch(
            query=query_embedding,
            limit=prefetch_limit,
            using="mean_pooling_rows"
        ),
    ],
    limit=search_limit,
    with_payload=True,
    with_vector=False,
    using="original"
)

并检查我们查询的顶部检索结果“李·哈维·奥斯瓦尔德在肯尼迪遇刺事件中的参与”。

dataset[response.points[0].payload['index']]['image']

Results, ColPali

结论

在本教程中，我们展示了一种使用Qdrant进行大规模PDF检索的优化方法，该方法利用VLLMs生成复杂的多向量表示，如ColPali和ColQwen2。

如果没有这样的优化，检索系统的性能可能会严重下降，无论是在索引时间还是查询延迟方面，尤其是在数据集大小增长时。

我们强烈推荐在你的工作流程中实施这种方法，以确保高效且可扩展的PDF检索。忽视优化检索过程可能会导致性能慢到无法接受，从而阻碍系统的可用性。

立即开始扩展您的PDF检索能力！

这个页面有用吗？

感谢您的反馈！🙏

我们很抱歉听到这个消息。😔 你可以在GitHub上编辑这个页面，或者创建一个GitHub问题。

本页内容：