MilvusとOpenAIによるセマンティック検索

このガイドでは、OpenAIのEmbedding APIをMilvusベクトルデータベースとどのように組み合わせて、テキストのセマンティック検索を行うかを紹介します。

はじめに

始める前に、OpenAIのAPIキーを用意するか、OpenAIのウェブサイトから取得してください。

この例で使用するデータは本のタイトルです。データセットはここからダウンロードでき、以下のコードを実行するのと同じディレクトリに置くことができる。

まず、MilvusとOpenAIのパッケージをインストールする：

pip install --upgrade openai pymilvus

これで、埋め込みデータを生成し、ベクターデータベースを使ってセマンティック検索を行う準備ができました。

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

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",
        },
    },
]