使用 EmbeddingLists 进行搜索：ColBERT 和 ColPali

本教程演示如何在 Milvus 中利用 EmbeddingList 搜索功能，针对 StructArray 向量子字段构建 ColBERT 风格和 ColPali 风格的检索系统。当您的查询和存储数据均以向量列表形式表示，且您希望实现基于实体级晚期交互的检索并采用MAX_SIM* 指标时，可使用本教程。

有关本教程所涉及的 StructArray 基础知识，请参阅《创建 StructArray 字段》、《为 StructArray 字段建立索引》以及《使用 StructArray 进行基础向量搜索》。本教程重点介绍 ColBERT 和 ColPali 工作流，而非通用的 StructArray 语法。

概述

要构建文本检索系统，您可能需要将文档拆分为多个片段，并将每个片段与其Embeddings一起作为实体存储在向量数据库中，以确保精确度和准确度——对于长文档而言尤其如此，因为全文本Embeddings可能会稀释语义特异性或超出模型输入限制。

然而，按片段存储数据会导致搜索结果以片段为单位，这意味着检索过程最初识别的是相关片段，而非连贯的文档。为解决此问题，您应执行额外的检索后处理。

ColBERT（arXiv:2004.12832）是一种文本-文本检索系统，通过基于BERT的上下文化后期交互，提供高效且有效的段落检索。它支持对查询和文档进行独立的词元级编码，并计算它们之间的相似度。

令牌级编码

在 ColBERT 的数据摄入过程中，每篇文档会被拆分为令牌，随后进行向量化并存储为向量列表，例如$∈Rn×dd \rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ E d = [$e$ d1 $,{}^{}e$ d2 $,\dots ,e$ dn $]$ $\in$ R${}^{n\times d}$. 当查询到达时，它也会被分词、向量化，并作为嵌入列表存储，例如$∈R^{m×d} q \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ $\to$ E q $=$ [e q1 $,e$ q2 $,\dots ,e$ qm $]\in$ R${}^{m\times d}$.

在上述公式中，

• $dd$：一份文档

• $qq$：查询

• $EdE_d$ ：表示该文档的嵌入向量列表。

• $EqE_q$ ：表示查询的嵌入列表。

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$$,e_{}$,…,以及dn$]\in$ R${}^{n\times d}$：表示该文档的嵌入列表中向量嵌入的数量在$\R^{n×d}$ 的范围内。

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1,eq2,…,eqm$]\in$ R${}^{m\times d}$：表示查询的嵌入列表中向量嵌入的数量在$\R^{m×d}$ 的范围内。

后期交互

向量化完成后，将查询嵌入列表与每个文档嵌入列表进行逐令牌比对，以确定最终的相似度得分。

如上图所示，查询包含两个令牌，即machine 和learning ，而窗口中的文档包含四个令牌：neural 、network 、python 和tutorial 。将这些令牌向量化后，会将每个查询令牌的向量嵌入与文档中的向量嵌入进行比较，从而获得一组相似度得分。 随后，将每个相似度列表中得分最高的项求和，从而得出最终得分。确定文档最终得分的过程称为最大相似度（MAX_SIM）。有关最大相似度的详细信息，请参阅“最大相似度”。

ColPali 扩展

基于 ColBERT，ColPali（arXiv:2407.01449）提出了一种利用视觉语言模型（VLMs）进行视觉丰富文档检索的新方法。 在数据摄入过程中，每个文档页面会被渲染成高分辨率图像，然后分割成图像块，而不是进行分词。例如，一张 448 x 448 像素的文档页面图像可以生成 1,024 个图像块，每个图像块尺寸为 14 x 14 像素。

该方法保留了非文本信息，例如文档布局、图片和表格结构，而这些信息在使用纯文本检索系统时会丢失。

ColPali中使用的VLM名为PaliGemma（arXiv:2407.07726），它由图像编码器（SigLIP-400M）、仅解码器的语言模型（Gemma2-2B） 以及一个将图像编码器输出投影到语言模型向量空间的线性层，如上图所示。

在数据摄入过程中，以原始图像形式呈现的文档页面会被划分为多个视觉片段，每个片段均经过嵌入处理以生成一组向量Embeddings。 随后，这些嵌入向量被投影到语言模型的向量空间中，以获得最终的嵌入向量列表，例如$∈Rn×dd \rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ E d = [$e$ d1 $,e$ d2 $,\dots ,e$ dn $]\in$ R${}^{n\times d}$. 当查询到达时，将其分词，并将每个词向量进行嵌入以生成向量Embeddings列表，例如$∈R^{m×d} q \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ $\to$ E q $=$ [e q1 $,e$ q2 $,\dots ,e$ qm $]\in$ R${}^{m\times d}$. 随后，应用了MAX_SIM算法来比较这两组Embeddings列表，并计算出查询与文档页面之间的最终得分。

ColBERT 文本检索系统

在本节中，我们将使用 StructArray 搭建一个 ColBERT 文本检索系统。在此之前，请先部署一个 Milvus v2.6.x 实例，并获取一个 Cohere 访问令牌。

步骤 1：安装依赖项

运行以下命令安装依赖项。

pip install --upgrade huggingface-hub transformers datasets pymilvus cohere

步骤 2：加载 Cohere 数据集

在本示例中，我们将使用 Cohere 的维基百科数据集，并检索前 10,000 条记录。您可以在此页面上找到有关此数据集的信息。

from datasets import load_dataset

lang = "simple"
docs = load_dataset(
    "Cohere/wikipedia-2023-11-embed-multilingual-v3",
    lang,
    split="train[:10000]"
)

如果本地没有该数据集，运行上述脚本将自动下载数据集。数据集中的每条记录都来自维基百科页面中的一个段落。下表展示了该数据集的结构。

列名	描述
_id	记录 ID
url	当前记录的 URL。
title	源文档的标题。
text	源文档中的一段文字。
emb	源文档中文本的Embeddings。

步骤 3：按标题对段落进行分组

若要搜索文档而非段落，应按标题对段落进行分组。

df = docs.to_pandas()
groups = df.groupby('title')

data = []

for title, group in groups:
  data.append({
      "title": title,
      "paragraphs": [{
          "text": row['text'],
          'emb': row['emb']
      } for _, row in group.iterrows()]
  })

在此代码中，我们将分组后的段落作为文档存储，并将其加入data 列表。每个文档都包含一个paragraphs 键，该键对应一个段落列表；每个段落对象包含text 和emb 两个键。

步骤 4：为 Cohere 数据集创建 Collection

数据准备就绪后，我们将创建一个 Collection。在该 Collection 中，paragraphs 是一个 StructArray 字段。有关 StructArray Schema 的一般说明，请参阅《创建 StructArray 字段》。

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri="http://localhost:19530",
    token="root:Milvus"
)

# Create collection schema
schema = client.create_schema()

schema.add_field('id', DataType.INT64, is_primary=True, auto_id=True)
schema.add_field('title', DataType.VARCHAR, max_length=512)

# Create struct schema
struct_schema = client.create_struct_field_schema()
struct_schema.add_field('text', DataType.VARCHAR, max_length=65535)
struct_schema.add_field('emb', DataType.FLOAT_VECTOR, dim=512)

schema.add_field('paragraphs', DataType.ARRAY,
                 element_type=DataType.STRUCT,
                 struct_schema=struct_schema, max_capacity=200)

# Create index parameters
index_params = client.prepare_index_params()
index_params.add_index(
    field_name="paragraphs[emb]",
    index_type="AUTOINDEX",
    metric_type="MAX_SIM_COSINE"
)

# Create a collection
client.create_collection(
    collection_name='wiki_documents',
    schema=schema,
    index_params=index_params
)

步骤 5：将 Cohere 数据集插入 Collection

现在，我们可以将准备好的数据插入到上面创建的Collection中。

client.insert(
    collection_name='wiki_documents',
    data=data
)

步骤 6：在 Cohere 数据集中进行搜索

根据 ColBERT 的设计，查询文本应先进行分词，然后嵌入到 EmbeddingList 中。在此步骤中，我们将使用 Cohere 用于为维基百科数据集中的段落生成 Embeddings 的同一模型。

import cohere

co = cohere.ClientV2("COHERE_API_KEY")

query_inputs = [
    {
        'content': [
            {'type': 'text', 'text': 'Adobe'},
        ]
    },
    {
        'content': [
            {'type': 'text', 'text': 'software'}
        ]
    }
]

embeddings = co.embed(
    inputs=query_inputs,
    model='embed-multilingual-v3.0',
    input_type="classification",
    embedding_types=["float"],
)

在代码中，查询文本会被拆分为query_inputs 中的令牌，并嵌入到一个浮点向量列表中。然后，您可以使用Milvus的EmbeddingList进行相似度搜索，具体如下。

from pymilvus.client.embedding_list import EmbeddingList

query_emb_list = EmbeddingList()

if (embeddings.embeddings.float):
  query_emb_list.add_batch(embeddings.embeddings.float)

results = client.search(
    collection_name="wiki_documents",
    data=[query_emb_list],
    anns_field="paragraphs[emb]",
    limit=10,
    output_fields=["title"]
)

for hit in results[0]:
  print(f"Document {hit['entity']['title']}: {hit['distance']:.4f}")

上述代码的输出结果类似于以下内容：

# Document Software: 2.3035
# Document Application: 2.1875
# Document Adobe Illustrator: 2.1167
# Document Open source: 2.0542
# Document Computer: 1.9811
# Document Microsoft: 1.9784
# Document Web browser: 1.9655
# Document Program: 1.9627
# Document Website: 1.9594
# Document Computer science: 1.9460

每对的余弦相似度得分范围在-1 到1 之间。最终的MAX_SIM_COSINE 得分可能大于1 ，因为它聚合了多个令牌级别的最大相似度得分。

ColPali 文档检索系统

在本节中，我们将使用 StructArray 构建一个基于 ColPali 的文档检索系统。在此之前，请先部署一个 Milvus v2.6.x 实例。

步骤 1：安装依赖项

pip install --upgrade huggingface-hub transformers datasets pymilvus 'colpali-engine>=0.3.0,<0.4.0'

步骤 2：加载 Vidore 数据集

在本节中，我们将使用名为vidore_v2_finance_en 的 Vidore 数据集。该数据集是由银行业年度报告组成的语料库，旨在用于长文档理解任务。它是构成 ViDoRe v3 基准测试的 10 个语料库之一。您可以在此页面上查看有关此数据集的详细信息。

from datasets import load_dataset

ds = load_dataset("vidore/vidore_v3_finance_en", "corpus")
df = ds['test'].to_pandas()

若本地未存储该数据集，运行上述脚本将自动下载数据集。数据集中的每条记录均来自一份财务报告的单页内容。下表展示了该数据集的结构。

列名	描述
corpus_id	语料库中的一个记录
image	以字节为单位的页面图像。
doc_id	描述性文档 ID。
page_number_in_doc	文档中当前页面的页码。

步骤 3：为页面图像生成Embeddings

如“概述”部分所示，ColPali 模型是一种将图像投影到文本模型向量空间中的视觉语言模型（VLM）。在此步骤中，我们将使用最新的 ColPali 模型vidore/colpali-v1.3。您可以在此页面上查看该模型的详细信息。

import torch
from typing import cast
from colpali_engine.models import ColPali, ColPaliProcessor

model_name = "vidore/colpali-v1.3"

model = ColPali.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda:0",  # or "mps" if on Apple Silicon
).eval()

processor = ColPaliProcessor.from_pretrained(model_name)

模型准备就绪后，您可以按照以下方式尝试为特定图像生成图像块。

from PIL import Image
from io import BytesIO

# Use the iterrows() generator to get the first row.
row = next(df.iterrows())[1]

# Decode the image bytes and generate patch embeddings.
images = [Image.open(BytesIO(row["image"]["bytes"]))]
batch_images = processor.process_images(images).to(model.device)

with torch.no_grad():
    patches_embeddings = model(**batch_images)[0]

# Check the shape of the embeddings generated for the patches.
print(patches_embeddings.shape)

# [1031, 128]

在上面的代码中，ColPali 模型将图像调整为 448 x 448 像素，然后将其划分为每个尺寸为 14 x 14 像素的片段。最后，这些片段被嵌入为 1,031 个维度均为 128 的 Embeddings。

您可以通过如下循环为所有图像生成Embeddings：

data = []

for _, row in df.iterrows():
    corpus_id = row["corpus_id"]
    images = [Image.open(BytesIO(row["image"]["bytes"]))]
    batch_images = processor.process_images(images).to(model.device)

    with torch.no_grad():
        patches = model(**batch_images)[0]

    doc_id = row["doc_id"]
    page_number_in_doc = row["page_number_in_doc"]

    data.append({
        "corpus_id": corpus_id,
        "patches": [
            {"emb": emb.float().cpu().tolist()}
            for emb in patches
        ],
        "doc_id": doc_id,
        "page_number_in_doc": page_number_in_doc,
    })

步骤 4：为财务报告数据集创建Collection

数据准备就绪后，我们将创建一个 Collection。在该 Collection 中，patches 是一个 StructArray 字段。每个 Struct 元素存储一个片段嵌入。有关 StructArray 向量字段的索引要求，请参阅《索引 StructArray 字段》。

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri=YOUR_CLUSTER_ENDPOINT,
    token=YOUR_API_KEY
)

schema = client.create_schema()

schema.add_field(
    field_name="corpus_id",
    datatype=DataType.INT64,
    is_primary=True
)

patch_schema = client.create_struct_field_schema()

patch_schema.add_field(
    field_name="emb",
    datatype=DataType.FLOAT_VECTOR,
    dim=128
)

schema.add_field(
    field_name="patches",
    datatype=DataType.ARRAY,
    element_type=DataType.STRUCT,
    struct_schema=patch_schema,
    max_capacity=1031
)

schema.add_field(
    field_name="doc_id",
    datatype=DataType.VARCHAR,
    max_length=512
)

schema.add_field(
    field_name="page_number_in_doc",
    datatype=DataType.INT64
)

index_params = client.prepare_index_params()

index_params.add_index(
    field_name="patches[emb]",
    index_type="AUTOINDEX",
    metric_type="MAX_SIM_COSINE"
)

client.create_collection(
    collection_name="financial_reports",
    schema=schema,
    index_params=index_params
)

步骤 5：将财务报告插入 Collection

现在，我们可以将准备好的财务报告插入到Collection中。

client.insert(
    collection_name="financial_reports",
    data=data
)

从输出结果中，您可以发现 Vidore 数据集的所有页面均已插入。

步骤 6：在财务报告中进行搜索

数据准备就绪后，我们可以按照以下方式对Collection中的数据进行搜索：

from pymilvus.client.embedding_list import EmbeddingList

queries = [
    "quarterly revenue growth chart"
]

batch_queries = processor.process_queries(queries).to(model.device)

with torch.no_grad():
    query_embeddings = model(**batch_queries)

query_emb_list = EmbeddingList()
query_emb_list.add_batch(query_embeddings[0].float().cpu().tolist())

results = client.search(
    collection_name="financial_reports",
    data=[query_emb_list],
    anns_field="patches[emb]",
    limit=10,
    output_fields=["doc_id", "page_number_in_doc"]
)