埋め込みリストによる検索

このページでは、Milvusの構造体配列を使って、ColBERTテキスト検索システムとColPaliテキスト検索システムを構築する方法を説明します。Milvusの構造体配列は、文書をベクトル化されたチャンクとともに埋め込みリストに格納することができます。

概要

テキスト検索システムを構築するには、ドキュメントをチャンクに分割し、各チャンクをエンベッディングと共にエンティティとしてベクトルデータベースに格納し、精度と正確性を確保する必要があります。

しかし、データをチャンク単位で保存すると、検索結果がチャンク単位になり、まとまりのあるドキュメントではなく、関連するセグメントを検索結果として特定することになります。これに対処するためには、検索後の処理を追加で行う必要がある。

ColBERT（arXiv:2004.12832）は、BERT上の文脈に基づく後期インタラクションを通じて、効率的かつ効果的な通過検索を提供するテキスト-テキスト検索システムである。クエリと文書を独立にトークン単位で符号化し、それらの類似度を計算することができる。

トークン単位のエンコーディング

ColBERT におけるデータ取り込みの際、各文書はトークンに分割され、それらはベクトル化され、$d\to \mathrm{Ed=}[{\mathrm{ed}1}_{},{\mathrm{ed}2}_{},...]$のように埋め込みリストとして格納される。$$e$d1$,$\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d$\to$ E d = $[e$ d1, e d2,$...$, e dn]$\in$ R ${}^{n\times d}$.クエリが到着すると、それはまた、$q\to \mathrm{Eq=}[{\mathrm{eq}}_1,{\mathrm{eq}}_2,...]$のように、トークン化、ベクトル化され、埋め込みリストとして格納される。$e_{q1}, e_{q2}, e_{qm}] ∈Rm×dq \rightarrow E_q = [e_{q1}, e_{q2}, e_{dots, e_{qm}] ∈R^{m×d}$ q$\to$ E q = $[e$ q1, e q2,$...$, e ${}_{}q18$ ]∈ R ${}^{m\times d}$.

上記の式では

• $dd$ d: 文書

• $qq$ q: クエリー

• $EdE_d$ Ed: 文書を表す埋め込みリスト。

• $EqE_q$ Eq: クエリを表す埋め込みリスト。

• [$[e_{d1}, e_{d2}, ୧dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$...$, edn]$\in$ R${}^{n\times d}$：文書を表現する埋め込みリストのベクトル埋め込み数が$Rn×dR^{n×d}$ R${}^{n\times d}$ の範囲内にある。

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1, eq2,$...$, eqm]$\in$ R${}^{m\times d}$：クエリを表現する埋め込みリストのベクトル埋め込み数が$Rm×dR$ R${}^{m\times dの}$範囲内にある。

遅い相互作用

ベクトル化が完了すると、クエリ埋め込みリストと各文書埋め込みリストがトークンごとに比較され、最終的な類似度スコアが決定される。

遅い相互作用

上の図に示すように、クエリには2つのトークン、すなわちmachine とlearning が含まれ、ウィンドウ内の文書には4つのトークンがある：neural network 、python 、tutorial の4つのトークンがある。これらのトークンがベクトル化されると、各クエリトークンのベクトル埋め込みがドキュメント内のものと比較され、類似度スコアのリストが得られる。そして、各スコアリストから最も高いスコアが合計され、最終的なスコアが生成される。文書の最終スコアを決定するプロセスは最大類似度（MAX_SIM）として知られている。最大類似度の詳細については、最大類似度を参照のこと。

ColPali 拡張機能

ColBERTをベースにしたColPali（arXiv:2407.01449）は、視覚言語モデル（VLM）を活用した、ビジュアルリッチな文書検索への新しいアプローチを提案している。データを取り込む際、各文書ページは高解像度の画像にレンダリングされ、トークン化されるのではなく、パッチに分割される。例えば、448 x 448ピクセルの文書ページ画像は、それぞれ14 x 14ピクセルの1,024パッチを生成することができる。

この方法では、文書のレイアウトや画像、表構造など、テキストのみの検索システムでは失われてしまう非テキスト情報が保存されます。

コパリ拡張

ColPaliで使用されるVLMはPaliGemma（arXiv:2407.07726）と呼ばれ、上図に示すように、画像エンコーダ（SigLIP-400M）、デコーダのみの言語モデル（Gemma2-2B）、画像エンコーダの出力を言語モデルのベクトル空間に投影する線形レイヤから構成される。

データ取り込みの際、生画像として表現された文書ページは複数の視覚的パッチに分割され、それぞれが埋め込まれてベクトル埋め込みリストが生成される。そして、$d\to \mathrm{Ed=}[{\mathrm{ed}1}_{},{\mathrm{ed}2}_{},...,edn$]のように、最終的な埋め込みリストを得るために、言語モデルのベクトル空間に投影される。$$e$d1$,$\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d$\to$ E d = $[e$ d1, e d2,$...$, e dn]$\in$ R ${}^{n\times d}$.クエリが到着すると、それはトークン化され、各トークンは$q\to \mathrm{Eq=}[{\mathrm{eq}}_1,{\mathrm{eq}}_2,...]$のようにベクトル埋め込みリストを生成するために埋め込まれる。$e_{q1}, e_{q2}, e_{qm}] ∈Rm×dq \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈R^{m×d}$ q$\to$ E q = $[e$ q1, e q2,$...$, e ${}_{}q18$ ]∈ R ${}^{m\times d}$.次に、MAX_SIMを適用して、2つの埋め込みリストを比較し、クエリと文書ページの間の最終スコアを求める。

ColBERT テキスト検索システム

本節では、MilvusのArray of Structsを用いたColBERTテキスト検索システムを構築する。その前に、Milvus v2.6.xインスタンスのセットアップMilvus v2.6.xと互換性のあるZilliz Cloudクラスタ、Cohereアクセストークンの取得を行う。

ステップ1: 依存関係のインストール

以下のコマンドを実行して依存関係をインストールします。

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

ステップ 2: Cohere データセットのロード

この例では、Cohere の Wikipedia データセットを使用し、最初の 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 というフィールドがあり、これは構造体の配列である。

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がWikipediaデータセット内の段落の埋め込みを生成するために使用したのと同じモデルを使用する。

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]",
    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と互換性のあるZilliz Cloudクラスタをセットアップする。

ステップ1: 依存関係のインストール

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

ステップ2：Vidoreデータセットのロード

ここでは、vidore_v2_finance_jaというVidoreデータセットを使用する。このデータセットは銀行セクターの年次報告書のコーパスであり、長文文書理解タスクを対象としている。ViDoRe v3 Benchmarkを構成する10のコーパスのうちの一つである。このデータセットの詳細はこちらのページを参照されたい。

from datasets import load_dataset

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

上記のスクリプトを実行すると、データセットがローカルにない場合はダウンロードされる。データセットの各レコードは、財務報告書の1ページである。次の表は、このデータセットの構造を示している。

カラム名	説明
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ピクセルのパッチに分割する。最後に、これらのパッチはそれぞれ128次元の1,031個の埋め込みに埋め込まれます。

以下のようなループを使って、すべての画像の埋め込みを生成することができます：

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 というフィールドは、構造体の配列である。

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