MilvusとVoyageAIによるセマンティック検索
このガイドでは、MilvusベクトルデータベースとVoyageAIのEmbedding APIを使用して、テキストをセマンティック検索する方法を紹介します。
はじめに
始める前に、Voyage APIキーが用意されているか、VoyageAIのウェブサイトから取得できることを確認してください。
この例で使用するデータは本のタイトルである。データセットはこちらからダウンロードでき、以下のコードを実行するのと同じディレクトリに置くことができる。
まず、MilvusとVoyage AI用のパッケージをインストールする：
$ pip install --upgrade voyageai pymilvus
Google Colabを使用している場合、インストールした依存関係を有効にするために、ランタイムを再起動する必要があるかもしれません。(画面上部の "Runtime "メニューをクリックし、ドロップダウンメニューから "Restart session "を選択する）。
これで、埋め込みデータを生成し、ベクターデータベースを使ってセマンティック検索を行う準備が整いました。
VoyageAI & Milvusによる書籍タイトルの検索
以下の例では、ダウンロードしたCSVファイルから書籍のタイトルデータを読み込み、VoyageAIの埋め込みモデルを使ってベクトル表現を生成し、Milvusのベクトルデータベースに格納して意味検索を行います。
import voyageai
from pymilvus import MilvusClient
MODEL_NAME = "voyage-law-2" # Which model to use, please check https://docs.voyageai.com/docs/embeddings for available models
DIMENSION = 1024 # Dimension of vector embedding
# Connect to VoyageAI with API Key.
voyage_client = voyageai.Client(api_key="<YOUR_VOYAGEAI_API_KEY>")
docs = [
"Artificial intelligence was founded as an academic discipline in 1956.",
"Alan Turing was the first person to conduct substantial research in AI.",
"Born in Maida Vale, London, Turing was raised in southern England.",
]
vectors = voyage_client.embed(texts=docs, model=MODEL_NAME, truncation=False).embeddings
# Prepare data to be stored in Milvus vector database.
# We can store the id, vector representation, raw text and labels such as "subject" in this case in Milvus.
data = [
{"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
for i in range(len(docs))
]
# Connect to Milvus, all data is stored in a local file named "milvus_voyage_demo.db"
# in current directory. You can also connect to a remote Milvus server following this
# instruction: https://milvus.io/docs/install_standalone-docker.md.
milvus_client = MilvusClient(uri="milvus_voyage_demo.db")
COLLECTION_NAME = "demo_collection" # Milvus collection name
# Create a collection to store the vectors and text.
if milvus_client.has_collection(collection_name=COLLECTION_NAME):
milvus_client.drop_collection(collection_name=COLLECTION_NAME)
milvus_client.create_collection(collection_name=COLLECTION_NAME, dimension=DIMENSION)
# Insert all data into Milvus vector database.
res = milvus_client.insert(collection_name="demo_collection", data=data)
print(res["insert_count"])
引数としてMilvusClient を指定する：
./milvus.dbのようにuriをローカルファイルとして設定する方法が最も便利である。- データ規模が大きい場合は、dockerやkubernetes上に、よりパフォーマンスの高いMilvusサーバを構築することができます。このセットアップでは、
http://localhost:19530などのサーバ uri をuriとして使用してください。 - MilvusのフルマネージドクラウドサービスであるZilliz Cloudを使用する場合は、Zilliz CloudのPublic EndpointとApi keyに対応する
uriとtokenを調整してください。
Milvusのベクトルデータベースに全てのデータが登録されたので、クエリに対するベクトル埋め込みを生成してベクトル検索を行うことで、セマンティック検索を行うことができる。
queries = ["When was artificial intelligence founded?"]
query_vectors = voyage_client.embed(
texts=queries, model=MODEL_NAME, truncation=False
).embeddings
res = milvus_client.search(
collection_name=COLLECTION_NAME, # target collection
data=query_vectors, # query vectors
limit=2, # number of returned entities
output_fields=["text", "subject"], # specifies fields to be returned
)
for q in queries:
print("Query:", q)
for result in res:
print(result)
print("\n")
Query: When was artificial intelligence founded?
[{'id': 0, 'distance': 0.7196218371391296, 'entity': {'text': 'Artificial intelligence was founded as an academic discipline in 1956.', 'subject': 'history'}}, {'id': 1, 'distance': 0.6297335028648376, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.', 'subject': 'history'}}]
