VectorETLによるMilvusへの効率的なデータロード
このチュートリアルでは、ベクトルデータベース用に設計された軽量なETLフレームワークであるVectorETLを使用して、Milvusに効率的にデータをロードする方法を探ります。VectorETLは、様々なソースからデータを抽出し、AIモデルを使用してベクトル埋め込みに変換し、高速かつスケーラブルな検索のためにMilvusに格納するプロセスを簡素化します。このチュートリアルが終わるころには、ベクター検索システムを簡単に統合・管理できるETLパイプラインが完成していることでしょう。さっそく始めましょう!
準備
依存性と環境
$ pip install --upgrade vector-etl pymilvus
Google Colabを使用している場合、インストールしたばかりの依存関係を有効にするには、ランタイムを再起動する必要があるかもしれません(画面上部の "Runtime "メニューをクリックし、ドロップダウンメニューから "Restart session "を選択してください)。
VectorETLはAmazon S3, Google Cloud Storage, Local Fileなど複数のデータソースをサポートしています。サポートされているソースの全リストはこちらで確認できます。このチュートリアルでは、データソースの例として Amazon S3 を取り上げます。
Amazon S3からドキュメントをロードします。そのため、S3バケットに安全にアクセスするための環境変数として、AWS_ACCESS_KEY_ID とAWS_SECRET_ACCESS_KEY を用意する必要があります。さらに、OpenAIのtext-embedding-ada-002 エンベッディング・モデルを使って、データのエンベッディングを生成します。また、apiキー OPENAI_API_KEY も環境変数として用意する必要があります。
import os
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key-id"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-access-key"
ワークフロー
データソース(Amazon S3)の定義
今回はAmazon S3のバケットからドキュメントを抽出します。VectorETLでは、バケット名、ファイルへのパス、扱うデータの種類を指定できる。
source = {
"source_data_type": "Amazon S3",
"bucket_name": "my-bucket",
"key": "path/to/files/",
"file_type": ".csv",
"aws_access_key_id": os.environ["AWS_ACCESS_KEY_ID"],
"aws_secret_access_key": os.environ["AWS_SECRET_ACCESS_KEY"],
}
埋め込みモデルの設定(OpenAI)
データソースを設定したら、テキストデータをベクトル埋め込みに変換する埋め込みモデルを定義する必要があります。ここではOpenAIのtext-embedding-ada-002 。
embedding = {
"embedding_model": "OpenAI",
"api_key": os.environ["OPENAI_API_KEY"],
"model_name": "text-embedding-ada-002",
}
ターゲットデータベースとしてMilvusをセットアップする。
生成された埋め込みデータをMilvusに格納します。ここでは、Milvus Liteを使ってMilvusの接続パラメータを定義します。
target = {
"target_database": "Milvus",
"host": "./milvus.db", # os.environ["ZILLIZ_CLOUD_PUBLIC_ENDPOINT"] if using Zilliz Cloud
"api_key": "", # os.environ["ZILLIZ_CLOUD_TOKEN"] if using Zilliz Cloud
"collection_name": "my_collection",
"vector_dim": 1536, # 1536 for text-embedding-ada-002
}
host 、api_key :
hostをローカルファイル、例えば./milvus.dbとし、api_keyを空にしておくと、自動的にMilvus Liteが利用され、すべてのデータがこのファイルに格納されます。データ規模が大きい場合は、dockerやkubernetes上に、よりパフォーマンスの高いMilvusサーバを構築することができます。このセットアップでは、サーバのuri、例えば
http://localhost:19530をhostとして使用し、api_keyは空にしておいてください。MilvusのフルマネージドクラウドサービスであるZilliz Cloudを利用する場合は、
hostとapi_keyをZilliz CloudのPublic EndpointとApi keyに対応させてください。
埋め込みカラムの指定
次に、CSVファイルから埋め込みに変換するカラムを指定する必要があります。これにより、関連するテキストフィールドのみが処理され、効率とストレージの両方が最適化されます。
embed_columns = ["col_1", "col_2", "col_3"]
VectorETLパイプラインの作成と実行
すべての設定が完了したので、次にETLパイプラインを初期化し、データフローを設定し、実行します。
from vector_etl import create_flow
flow = create_flow()
flow.set_source(source)
flow.set_embedding(embedding)
flow.set_target(target)
flow.set_embed_columns(embed_columns)
# Execute the flow
flow.execute()
このチュートリアルに従うことで、VectorETLを使用してAmazon S3からmilvusにドキュメントを移動するエンドツーエンドのETLパイプラインを構築することに成功しました。VectorETLはデータソースに柔軟性があるので、特定のアプリケーションのニーズに基づいて好きなデータソースを選択することができます。VectorETLのモジュール設計により、このパイプラインを簡単に拡張して他のデータソースをサポートしたり、モデルを組み込んだりすることができ、AIやデータエンジニアリングワークフローのための強力なツールとなります!