Ricerca semantica con Milvus e OpenAI

Questa guida mostra come l 'API Embedding di OpenAI possa essere utilizzata con il database vettoriale Milvus per condurre una ricerca semantica sul testo.

Come iniziare

Prima di iniziare, assicuratevi di avere pronta la chiave API di OpenAI o di ottenerne una dal sito web di OpenAI.

I dati utilizzati in questo esempio sono i titoli dei libri. È possibile scaricare il set di dati qui e metterlo nella stessa directory in cui si esegue il codice seguente.

Per prima cosa, installare il pacchetto per Milvus e OpenAI:

pip install --upgrade openai pymilvus

A questo punto siamo pronti a generare embeddings e a utilizzare il database vettoriale per effettuare ricerche semantiche.

Ricerca di titoli di libri con OpenAI e Milvus

Nell'esempio seguente, carichiamo i dati dei titoli dei libri dal file CSV scaricato, usiamo il modello di embedding di OpenAI per generare rappresentazioni vettoriali e le memorizziamo nel database vettoriale di Milvus per la ricerca semantica.

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

Con tutti i dati nel database vettoriale di Milvus, possiamo ora eseguire una ricerca semantica generando un'incorporazione vettoriale per la query e condurre una ricerca vettoriale.

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

L'output dovrebbe essere il seguente:

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