MilvusとOpenAIによるセマンティック検索
このガイドでは、OpenAIのEmbedding APIをMilvusベクトルデータベースとどのように組み合わせて、テキストのセマンティック検索を行うかを紹介します。
はじめに
始める前に、OpenAIのAPIキーを用意するか、OpenAIのウェブサイトから取得してください。
この例で使用するデータは本のタイトルです。データセットはここからダウンロードでき、以下のコードを実行するのと同じディレクトリに置くことができる。
まず、MilvusとOpenAIのパッケージをインストールする:
pip install --upgrade openai pymilvus
Google Colabを使用している場合、インストールした依存関係を有効にするために、ランタイムを再起動する必要があるかもしれません。(画面上部の "Runtime "メニューをクリックし、ドロップダウンメニューから "Restart session "を選択する)。
これで、埋め込みデータを生成し、ベクターデータベースを使ってセマンティック検索を行う準備ができました。
OpenAI & Milvusで本のタイトルを検索する
以下の例では、ダウンロードしたCSVファイルから書名データを読み込み、OpenAIの埋め込みモデルを使ってベクトル表現を生成し、Milvusのベクトルデータベースに格納して意味検索を行います。
from openai import OpenAI
from pymilvus import MilvusClient
MODEL_NAME = "text-embedding-3-small" # Which model to use, please check https://platform.openai.com/docs/guides/embeddings for available models
DIMENSION = 1536 # Dimension of vector embedding
# Connect to OpenAI with API Key.
openai_client = OpenAI(api_key="<YOUR_OPENAI_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 = [
vec.embedding
for vec in openai_client.embeddings.create(input=docs, model=MODEL_NAME).data
]
# 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_openai_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_openai_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 = [
vec.embedding
for vec in openai_client.embeddings.create(input=queries, model=MODEL_NAME).data
]
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")
出力は以下のようになる:
[
{
"id": 0,
"distance": -0.772376537322998,
"entity": {
"text": "Artificial intelligence was founded as an academic discipline in 1956.",
"subject": "history",
},
},
{
"id": 1,
"distance": -0.58596271276474,
"entity": {
"text": "Alan Turing was the first person to conduct substantial research in AI.",
"subject": "history",
},
},
]