Ricerca semantica con Milvus e VoyageAI

Questa guida illustra come l 'API di embedding di VoyageAI 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 Voyage o di ottenerne una dal sito web di VoyageAI.

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

Innanzitutto, installare il pacchetto per Milvus e Voyage AI:

$ pip install --upgrade voyageai pymilvus

A questo punto siamo pronti a generare embeddings e a utilizzare il database vettoriale per effettuare la ricerca semantica.

Ricerca di titoli di libri con VoyageAI e Milvus

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

import voyageai
from pymilvus import MilvusClient

MODEL_NAME = "voyage-law-2"  # Which model to use, please check https://docs.voyageai.com/docs/embeddings for available models
DIMENSION = 1024  # Dimension of vector embedding

# Connect to VoyageAI with API Key.
voyage_client = voyageai.Client(api_key="<YOUR_VOYAGEAI_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 = voyage_client.embed(texts=docs, model=MODEL_NAME, truncation=False).embeddings

# 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_voyage_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_voyage_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 = voyage_client.embed(
    texts=queries, model=MODEL_NAME, truncation=False
).embeddings

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

Query: When was artificial intelligence founded?
[{'id': 0, 'distance': 0.7196218371391296, 'entity': {'text': 'Artificial intelligence was founded as an academic discipline in 1956.', 'subject': 'history'}}, {'id': 1, 'distance': 0.6297335028648376, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.', 'subject': 'history'}}]