使用 EmbeddingLists 進行搜尋：ColBERT 與 ColPali

本教學將展示如何在 Milvus 中，透過對 StructArray 向量子欄位進行 EmbeddingList 搜尋，建置 ColBERT 風格與 ColPali 風格的檢索系統。當您的查詢與儲存資料皆以向量清單形式表示，且您希望採用MAX_SIM* 指標進行實體層級的延遲交互檢索時，可參考本教學。

關於本教學背後 StructArray 的基礎知識，請參閱《建立 StructArray 欄位》、《為 StructArray 欄位建立索引》以及《使用 StructArray 進行基本向量搜尋》。本教學重點在於 ColBERT 和 ColPali 的工作流程，而非一般的 StructArray 語法。

概述

要建置文字檢索系統，您可能需要將文件分割成多個區塊，並將每個區塊連同其嵌入向量一併儲存為向量資料庫中的實體，以確保精確度與準確度；對於長篇文件而言，這點尤為重要，因為全文字嵌入向量可能會稀釋語義特異性，或超出模型輸入限制。

然而，以片段形式儲存資料會導致搜尋結果以片段為單位，這意味著檢索過程最初僅能識別相關片段，而非完整的文檔。為解決此問題，您應執行額外的搜尋後處理。

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$ $,\dots ,{}^e$ qm $]\in$ ${}^{Rm\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）， 以及一個將圖像編碼器輸出投影至語言模型向量空間的線性層，如上圖所示。

在資料導入過程中，以原始影像形式呈現的文件頁面會被劃分為多個視覺區塊，每個區塊皆會進行嵌入以產生一組向量嵌入值。 接著，這些嵌入向量會被投影至語言模型的向量空間中，以獲得最終的嵌入向量清單，例如$∈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}$。 接著，應用MAX_SIM演算法來比較這兩組嵌入向量清單，並計算查詢與文檔頁面之間的最終分數。

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	來源文件中文字的嵌入向量。

步驟 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 資料集建立集合

資料準備就緒後，我們將建立一個集合。在該集合中，paragraphs 是一個 StructArray 欄位。關於 StructArray 架構的一般說明，請參閱《建立 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 資料集插入集合

現在，我們可以將已準備好的資料插入到上述建立的集合中。

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

步驟 6：在 Cohere 資料集中進行搜尋

根據 ColBERT 的設計，查詢文字應先進行分詞，然後嵌入至 EmbeddingList 中。在此步驟中，我們將使用與 Cohere 相同的模型，為維基百科資料集中的段落生成嵌入向量。

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：為頁面圖像生成嵌入向量

如「概述」章節所示，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 維。

您可以透過以下迴圈為所有圖片產生嵌入向量：

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：為財務報表資料集建立集合

資料準備就緒後，我們將建立一個集合。在此集合中，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：將財務報告插入集合

現在，我們可以將已準備好的財務報告插入到集合中。

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

從輸出結果中，您可以發現 Vidore 資料集的所有頁面均已成功插入。

步驟 6：在財務報告中進行搜尋

資料準備就緒後，我們可以按照以下方式對集合中的資料進行搜尋：

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"]
)