Milvus와 OpenAI를 사용한 시맨틱 검색

이 가이드에서는 OpenAI의 임베딩 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"])

이제 밀버스 벡터 데이터베이스에 있는 모든 데이터를 가지고 쿼리에 대한 벡터 임베딩을 생성하여 시맨틱 검색을 수행하고 벡터 검색을 수행할 수 있습니다.

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