Semantische Suche mit Milvus und OpenAI

Dieser Leitfaden zeigt, wie die Einbettungs-API von OpenAI mit der Vektordatenbank Milvus verwendet werden kann, um eine semantische Suche im Text durchzuführen.

Erste Schritte

Bevor Sie beginnen, stellen Sie sicher, dass Sie den OpenAI-API-Schlüssel bereithalten, oder Sie erhalten ihn von der OpenAI-Website.

Die in diesem Beispiel verwendeten Daten sind Buchtitel. Sie können den Datensatz hier herunterladen und ihn in das gleiche Verzeichnis legen, in dem Sie den folgenden Code ausführen.

Installieren Sie zunächst das Paket für Milvus und OpenAI:

pip install --upgrade openai pymilvus

Damit sind wir bereit, Einbettungen zu erzeugen und die Vektordatenbank für die semantische Suche zu verwenden.

Durchsuchen von Buchtiteln mit OpenAI & Milvus

Im folgenden Beispiel laden wir Buchtiteldaten aus der heruntergeladenen CSV-Datei, verwenden das OpenAI-Einbettungsmodell, um Vektordarstellungen zu erzeugen, und speichern sie in der Milvus-Vektordatenbank für die semantische Suche.

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

Mit allen Daten in der Milvus-Vektordatenbank können wir nun eine semantische Suche durchführen, indem wir eine Vektoreinbettung für die Abfrage erzeugen und eine Vektorsuche durchführen.

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

Sie sollten die folgende Ausgabe sehen:

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