多模态检索增强生成 (RAG) 与文档检索 (ColPali) 和视觉语言模型 (VLM)

🚨 警告：此笔记本资源密集型，需要强大的计算能力。如果你在 Colab 中运行它，它将使用 A100 GPU。

在此笔记本中，我们演示了如何通过将 ColPali 文档检索器与 Qwen2-VL 视觉语言模型 (VLM) 相结合来构建一个 **多模态检索增强生成 (RAG)** 系统。这些模型共同构成一个功能强大的 RAG 系统，能够通过文本文档和视觉数据来增强查询响应。

我们不会依赖于复杂的文档处理器管道，该管道通过 OCR 提取数据，而是将利用文档检索模型来根据特定的用户查询有效地检索相关文档。

我还建议查看并收藏 smol-vision 存储库，它启发了此笔记本 - 尤其是此笔记本。有关 RAG 的介绍，你可以查看另一个食谱！

此图表灵感来自 Aymeric Roucher 在高级 RAG 或 RAG 评估食谱中的工作。

1. 安装依赖项

让我们从安装项目所需的必要库开始！🚀

!pip install -U -q byaldi pdf2image qwen-vl-utils transformers
# Tested with byaldi==0.0.4, pdf2image==1.17.0, qwen-vl-utils==0.0.8, transformers==4.45.0

我们还将安装 **poppler-utils** 以便于 PDF 操作。此实用程序提供了用于处理 PDF 文件的基本工具，确保我们可以有效地处理项目中的任何与文档相关的任务。

!sudo apt-get install -y poppler-utils

2. 加载数据集 📁

在本节中，我们将使用宜家组装说明作为我们的数据集。这些 PDF 包含有关组装各种家具的逐步指南。想象一下，你可以在组装新的宜家家具时向我们的助手寻求帮助！🛋

要下载组装说明，你可以按照这些步骤。

对于此笔记本，我选择了一些示例，但在现实情况下，我们可以使用大量的 PDF 来增强模型的功能。

import requests
import os

pdfs = {
    "MALM": "https://www.ikea.com/us/en/assembly_instructions/malm-4-drawer-chest-white__AA-2398381-2-100.pdf",
    "BILLY": "https://www.ikea.com/us/en/assembly_instructions/billy-bookcase-white__AA-1844854-6-2.pdf",
    "BOAXEL": "https://www.ikea.com/us/en/assembly_instructions/boaxel-wall-upright-white__AA-2341341-2-100.pdf",
    "ADILS": "https://www.ikea.com/us/en/assembly_instructions/adils-leg-white__AA-844478-6-2.pdf",
    "MICKE": "https://www.ikea.com/us/en/assembly_instructions/micke-desk-white__AA-476626-10-100.pdf",
}

output_dir = "data"
os.makedirs(output_dir, exist_ok=True)

for name, url in pdfs.items():
    response = requests.get(url)
    pdf_path = os.path.join(output_dir, f"{name}.pdf")

    with open(pdf_path, "wb") as f:
        f.write(response.content)

    print(f"Downloaded {name} to {pdf_path}")

print("Downloaded files:", os.listdir(output_dir))

下载组装说明后，我们将把 PDF 转换为图像。此步骤至关重要，因为它允许文档检索模型 (ColPali) 有效地处理和操作视觉内容。

import os
from pdf2image import convert_from_path


def convert_pdfs_to_images(pdf_folder):
    pdf_files = [f for f in os.listdir(pdf_folder) if f.endswith(".pdf")]
    all_images = {}

    for doc_id, pdf_file in enumerate(pdf_files):
        pdf_path = os.path.join(pdf_folder, pdf_file)
        images = convert_from_path(pdf_path)
        all_images[doc_id] = images

    return all_images


all_images = convert_pdfs_to_images("/content/data/")

让我们可视化一个组装指南示例，以了解这些说明是如何呈现的！这将帮助我们了解我们将要处理的内容的格式和布局。👀

>>> import matplotlib.pyplot as plt

>>> fig, axes = plt.subplots(1, 8, figsize=(15, 10))

>>> for i, ax in enumerate(axes.flat):
...     img = all_images[0][i]
...     ax.imshow(img)
...     ax.axis("off")

>>> plt.tight_layout()
>>> plt.show()

3. 初始化 ColPali 多模态文档检索模型 🤖

现在我们的数据集已准备好，我们将初始化文档检索模型，该模型将负责从原始图像中提取相关信息，并根据我们的查询为我们提供合适的文档。

通过使用此模型，我们可以显着增强我们的对话能力。

对于此任务，我们将使用 **Byaldi**。开发人员将该库描述如下：“Byaldi 是 RAGatouille 的迷你姐妹项目。它是一个围绕 ColPali 存储库的简单包装器，可以轻松地使用熟悉 API 的 ColPALI 等后期交互多模态模型。”

在此项目中，我们将特别关注 **ColPali**。

ColPali architecture

此外，你可以探索 **ViDore（视觉文档检索基准）** 以了解顶级检索器的实际应用。

首先，我们将从检查点加载模型。

from byaldi import RAGMultiModalModel

docs_retrieval_model = RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2")

接下来，我们可以通过指定 PDF 存储的文件夹，直接使用文档检索模型索引我们的文档。这将允许模型处理和组织文档，以便根据我们的查询有效地检索它们。

docs_retrieval_model.index(
    input_path="data/", index_name="image_index", store_collection_with_index=False, overwrite=True
)

4. 使用文档检索模型检索文档 🤔

在初始化文档检索模型后，我们可以通过提交用户查询并检查检索到的相关文档来测试其功能。

该模型将返回结果，根据其与查询的相关性直接进行排名。

让我们试一试！

text_query = "How many people are needed to assemble the Malm?"

results = docs_retrieval_model.search(text_query, k=3)
results

现在，让我们检查模型检索到的特定文档（图像）。这将使我们能够看到与我们的查询相对应的视觉内容，并了解模型如何选择相关信息。

def get_grouped_images(results, all_images):
    grouped_images = []

    for result in results:
        doc_id = result["doc_id"]
        page_num = result["page_num"]
        grouped_images.append(
            all_images[doc_id][page_num - 1]
        )  # page_num are 1-indexed, while doc_ids are 0-indexed. Source https://github.com/AnswerDotAI/byaldi?tab=readme-ov-file#searching

    return grouped_images


grouped_images = get_grouped_images(results, all_images)

让我们仔细看看检索到的文档，以了解它们包含的信息。这项检查将帮助我们评估检索到的内容与我们的查询的相关性和质量。

>>> import matplotlib.pyplot as plt

>>> fig, axes = plt.subplots(1, 3, figsize=(15, 10))

>>> for i, ax in enumerate(axes.flat):
...     img = grouped_images[i]
...     ax.imshow(img)
...     ax.axis("off")

>>> plt.tight_layout()
>>> plt.show()

5. 初始化用于问答的视觉语言模型 🙋

接下来，我们将初始化用于问答的视觉语言模型 (VLM)。在本例中，我们将使用 Qwen2_VL。

Qwen2_VL architecture

您可以查看 Open VLM 的排行榜，以了解最新进展此处。

首先，我们将从预训练的检查点加载模型，并将其移动到 GPU 上以获得最佳性能。请查看此处的模型。

from transformers import Qwen2VLForConditionalGeneration, Qwen2VLProcessor
from qwen_vl_utils import process_vision_info
import torch

vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-7B-Instruct",
    torch_dtype=torch.bfloat16,
)
vl_model.cuda().eval()

接下来，我们将初始化 VLM 处理器。在此步骤中，我们将指定最小和最大像素大小，以优化更多图像拟合到 GPU 内存的过程。

有关优化图像分辨率以提高性能的更多详细信息，您可以参考此处的文档。

min_pixels = 224 * 224
max_pixels = 1024 * 1024
vl_model_processor = Qwen2VLProcessor.from_pretrained(
    "Qwen/Qwen2-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels
)

6. 组装 VLM 模型并测试系统 🔧

加载完所有组件后，我们现在可以组装系统进行测试。首先，我们将通过向系统提供三个检索到的图像以及用户查询来创建聊天结构。此步骤可以根据您的特定需求进行自定义，从而允许您更灵活地与模型交互！

chat_template = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": grouped_images[0],
            },
            {
                "type": "image",
                "image": grouped_images[1],
            },
            {
                "type": "image",
                "image": grouped_images[2],
            },
            {"type": "text", "text": text_query},
        ],
    }
]

现在，让我们应用此聊天结构。

text = vl_model_processor.apply_chat_template(chat_template, tokenize=False, add_generation_prompt=True)

接下来，我们将处理输入，以确保它们格式正确并可以作为视觉语言模型 (VLM) 的输入使用。此步骤对于使模型能够根据提供的数据有效地生成响应至关重要。

image_inputs, _ = process_vision_info(chat_template)
inputs = vl_model_processor(
    text=[text],
    images=image_inputs,
    padding=True,
    return_tensors="pt",
)
inputs = inputs.to("cuda")

现在，我们准备生成答案！让我们看看系统如何利用处理后的输入，根据用户查询和检索到的图像提供响应。

generated_ids = vl_model.generate(**inputs, max_new_tokens=500)

模型生成输出后，我们将对其进行后处理以生成最终答案。

generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
output_text = vl_model_processor.batch_decode(
    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)

>>> print(output_text[0])

The Malm requires two people to assemble it.

7. 组装所有内容！ 🧑‍🏭️

现在，让我们创建一个包含整个流水线的方法，使我们能够在将来的应用程序中轻松地重用它。

def answer_with_multimodal_rag(
    vl_model, docs_retrieval_model, vl_model_processor, grouped_images, text_query, top_k, max_new_tokens
):
    results = docs_retrieval_model.search(text_query, k=top_k)
    grouped_images = get_grouped_images(results, all_images)

    chat_template = [
        {
            "role": "user",
            "content": [{"type": "image", "image": image} for image in grouped_images]
            + [{"type": "text", "text": text_query}],
        }
    ]

    # Prepare the inputs
    text = vl_model_processor.apply_chat_template(chat_template, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(chat_template)
    inputs = vl_model_processor(
        text=[text],
        images=image_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to("cuda")

    # Generate text from the vl_model
    generated_ids = vl_model.generate(**inputs, max_new_tokens=max_new_tokens)
    generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]

    # Decode the generated text
    output_text = vl_model_processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    return output_text

让我们看看完整的 RAG 系统是如何运作的！

>>> output_text = answer_with_multimodal_rag(
...     vl_model=vl_model,
...     docs_retrieval_model=docs_retrieval_model,
...     vl_model_processor=vl_model_processor,
...     grouped_images=grouped_images,
...     text_query="How do I assemble the Micke desk?",
...     top_k=3,
...     max_new_tokens=500,
... )
>>> print(output_text[0])

To assemble the Micke desk, follow these steps:

1. **Prepare the Components**: Lay out all the components of the desk on a flat surface.

2. **Attach the Legs**: Place the legs on the bottom of the desk frame. Ensure they are securely attached.

3. **Attach the Top**: Place the top of the desk on the frame, making sure it is level and stable.

4. **Secure with Screws**: Use the provided screws to secure the top to the frame. Ensure all screws are tightened securely.

5. **Final Check**: Double-check that all parts are properly attached and the desk is stable.

Refer to the detailed instructions provided in the image for specific steps and any additional information needed for assembly.

🏆 我们现在拥有一个完全可操作的 RAG 流水线，它同时利用文档检索模型和视觉语言模型！这种强大的组合使我们能够根据用户查询和相关文档生成有见地的响应。

8. 继续旅程 🧑‍🎓️

这个方法仅仅是探索多模态 RAG 系统潜力的起点。如果您渴望更深入地研究，以下是一些想法和资源，可以指导您接下来的步骤

🔍 使用 ColPali 探索更多：

📖 其他阅读资料：

💡 有用资源集：

多模态 RAG 资源集

📝 论文和原始代码：

在您继续探索多模态检索和生成系统的旅程中，请随时浏览这些资源！

< > 更新于 GitHub

开源 AI 食谱