使用嵌入列表進行檢索

本頁說明如何使用 Milvus 中的結構體陣列建立 ColBERT 文字檢索系統和 ColPali 文字檢索系統，這可讓您在嵌入清單中儲存文件及其向量化區塊。

概述

在建立文字擷取系統時，您可能需要將文件分割成區塊，並將每個區塊連同其嵌入作為一個實體儲存在向量資料庫中，以確保精確度和準確性，尤其是對於長文件而言，全文嵌入可能會削弱語義的特定性或超出模型輸入的限制。

然而，將資料儲存在區塊中會導致以區塊為單位的搜尋結果，這表示檢索最初會識別相關的區段，而非連貫的文件。為了解決這個問題，您應該執行額外的搜尋後處理。

ColBERT (arXiv:2004.12832) 是一個文本-文本檢索系統，透過 BERT 上的上下文化後期互動，提供有效率且有效的段落搜尋。它可以對查詢和文件進行獨立的代號化編碼，並計算它們的相似性。

令牌式編碼

在 ColBERT 的資料擷取過程中，每個文件都會被分割成標記，然後將其向量化並儲存為嵌入清單，如$\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d$\to$ E d = $[e$ $,$ ed2,$...$, e dn]$\in$ R ${}^n$×d。當一個查詢到達時，它也會被標記化、向量化，並儲存為一個嵌入列表，如$\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q$\to$ E q = $[e$ q1, e q2,$...$, e qm]$\in$ R ${}^m$×d。

在上述公式中

• $dd$ d：文件

• $qq$ q：查詢

• $EdE_d$ Ed ：表示該文件的嵌入列表。

• $EqE_q$ Eq: 代表查詢的嵌入清單。

• $∈Rn×d[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$.$.., edn]$\in$ R${}^n$×d：表示文件的嵌入列表中向量嵌入的數量在$Rn×d\R^{n×d}$ R${}^{n\times d}$ 的範圍內。

• $∈Rm×d[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1, eq2,$.$.$.$, eqm]$\in$ R${}^m$×d：嵌入列表中代表查詢的向量嵌入數目在$Rm×d\R^{m×d}$ R${}^{m\times d}$ 的範圍內。

後期互動

一旦向量化完成，查詢的嵌入清單就會與每個文件的嵌入清單逐一比較，以決定最終的相似性分數。

後期互動

如上圖所示，查詢包含兩個標記，即machine 和learning ，而視窗中的文件有四個標記：neural,network,python, 和tutorial 。當這些標記被向量化之後，每個查詢標記的向量嵌入與文件中的向量嵌入進行比較，以得到相似性分數清單。然後將每個分數列表中最高的分數相加，得出最終分數。決定文件最終得分的過程稱為最大相似性(MAX_SIM)。有關最大相似性的詳細資訊，請參閱最大相似性。

ColPali 延伸

在 ColBERT 的基礎上，ColPali (arXiv:2407.01449) 提出了一種利用視覺語言模型 (Vision-Language Models, VLMs) 進行視覺豐富的文件檢索的新方法。在資料擷取過程中，每個文件頁面會被渲染成高解析度影像，然後再分割成不同的片段，而不是標記化。例如，一張 448 x 448 像素的文件頁面影像可產生 1,024 個修補區，每個修補區的大小為 14 x 14 像素。

此方法可保留非文字資訊，例如文件排版、影像和表格結構，這些資訊在使用純文字檢索系統時都會遺失。

Copali 延伸

ColPali 使用的 VLM 稱為 PaliGemma (arXiv：2407.07726)，它包含一個影像編碼器(SigLIP-400M)、一個僅解碼的語言模型(Gemma2-2B)，以及一個將圖像編碼器的輸出投射到語言模型向量空間的線性層，如上圖所示。

在資料擷取的過程中，以原始影像表示的文件頁面會被分割成多個視覺斑塊，每個斑塊都會被嵌入以產生向量嵌入清單。然後將它們投射到語言模型的向量空間，以獲得最終的嵌入列表，如$\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d$\to$ E d = $[e$ $,$ ed2,$...$, e dn]$\in$ R ${}^n$×d。當一個查詢到達時，它會被標記化，每個標記會被嵌入以產生一個向量嵌入列表，如$\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q$\to$ E q = $[e$ q1, e q2,$...$, e qm]$\in$ R ${}^m$×d。然後運用MAX_SIM來比較兩個嵌入清單，並得到查詢頁面與文件頁面之間的最終得分。

ColBERT 文本檢索系統

本節將以 Milvus 的 Array of Structs 建立 ColBERT 文字檢索系統。在此之前，建立一個與 Milvus v2.6.x 相容的 Milvus v2.6.x instanceZilliz Cloud 叢集，取得 Cohere access token。

步驟 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	A 記錄 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 ，它是一個 Structs 陣列。

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 中的 tokens，並嵌入到浮點向量清單中。接著就可以使用 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]",
    search_params={
        "metric_type": "MAX_SIM_COSINE"
    },
    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 ，上述輸出中的相似性分數清楚地展示了多個標記級相似性分數的總和。

ColPali 文字檢索系統

在本節中，我們將使用 Milvus 的 Array of Structs 建立一個以 ColPali 為基礎的文字檢索系統。在此之前，請先設定一個與 Milvus v2.6.x 相容的 Milvus v2.6.x instanceZilliz Cloud 叢集。

步驟 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 Benchmark 的 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 iterrow() generator to get the first row
row = next(df.iterrows())[1]

# Include the image in the above row in a list
images = [ Image.open(row['image']['bytes'] ]
patches = processor.process_images(images).to(model.device)
patches_embeddings = model(**patches_in_pixels)[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：

data = []

for index, row in df.iterrows():
  row = next(df.iterrows())[1]
  corpus_id = row['corpus_id']
  
  images = [Image.open(BytesIO(row['image']['bytes']))]
  batch_images = processor.process_images(images).to(model.device)
  patches = model(**batch_images)[0]

  doc_id = row['doc_id']
  markdown = row['markdown']
  page_number_in_doc = row['page_number_in_doc']

  data.append({
      "corpus_id": corpus_id,
      "patches": [ {"emb": emb} for emb in patches ],
      "doc_id": markdown,
      "page_number_in_doc": row['page_number_in_doc']
  })

步驟 4：為財務報告資料集建立集合

資料準備就緒後，我們將建立一個集合。在集合中，一個名為patches 的欄位是一個 Structs 陣列。

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].cpu())

results = client.search(
    collection_name="financial_reports",
    data=[query_emb_list],
    anns_field="patches[emb]",
    search_params={
        "metric_type": "MAX_SIM_COSINE"
    },
    limit=10,
    output_fields=["doc_id", "page_number_in_doc"]
)