Semantische Suche mit Milvus und VoyageAI

Dieser Leitfaden zeigt, wie die Einbettungs-API von VoyageAI 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 Voyage-API-Schlüssel bereit haben, oder Sie erhalten einen von der VoyageAI-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 Voyage AI:

$ pip install --upgrade voyageai pymilvus milvus-lite

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

Durchsuchen von Buchtiteln mit VoyageAI & Milvus

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

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

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

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'}}]

Durchsuchen von Bildern mit VoyageAI & Milvus

import base64
import voyageai
from pymilvus import MilvusClient
import urllib.request
import matplotlib.pyplot as plt
from io import BytesIO
import urllib.request
import fitz  # PyMuPDF
from PIL import Image

def pdf_url_to_screenshots(url: str, zoom: float = 1.0) -> list[Image]:

    # Ensure that the URL is valid
    if not url.startswith("http") and url.endswith(".pdf"):
        raise ValueError("Invalid URL")

    # Read the PDF from the specified URL
    with urllib.request.urlopen(url) as response:
        pdf_data = response.read()
    pdf_stream = BytesIO(pdf_data)
    pdf = fitz.open(stream=pdf_stream, filetype="pdf")

    images = []

    # Loop through each page, render as pixmap, and convert to PIL Image
    mat = fitz.Matrix(zoom, zoom)
    for n in range(pdf.page_count):
        pix = pdf[n].get_pixmap(matrix=mat)

        # Convert pixmap to PIL Image
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)

    # Close the document
    pdf.close()

    return images


def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue())
    return img_str.decode("utf-8")

DIMENSION = 1024  # Dimension of vector embedding

Dann müssen wir die Eingabedaten für Milvus vorbereiten. Dazu verwenden wir den VoyageAI-Client, den wir im vorherigen Kapitel erstellt haben. Das verfügbare multimodale Einbettungsmodell von VoyageAI finden Sie auf dieser Seite.

pages = pdf_url_to_screenshots("https://www.fdrlibrary.org/documents/356632/390886/readingcopy.pdf", zoom=3.0)
inputs = [[img] for img in pages]

vectors = client.multimodal_embed(inputs, model="voyage-multimodal-3")

inputs = [i[0] if isinstance(i[0], str) else image_to_base64(i[0]) for i in inputs]
# 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.embeddings[i], "data": inputs[i], "subject": "fruits"}
    for i in range(len(inputs))
]

Als nächstes erstellen wir eine Milvus-Datenbankverbindung und fügen die Einbettungen in die Milvus-Datenbank ein.

milvus_client = MilvusClient(uri="milvus_voyage_multi_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"])

Nun sind wir bereit, die Bilder zu durchsuchen. Hier ist die Abfrage ein String, aber wir können auch Bilder abfragen. (siehe die Dokumentation für die multimodale API hier). Wir verwenden matplotlib, um die Ergebnisbilder anzuzeigen.

queries = [["The consequences of a dictator's peace"]]

query_vectors = client.multimodal_embed(
    inputs=queries, model="voyage-multimodal-3", truncation=False
).embeddings

res = milvus_client.search(
    collection_name=COLLECTION_NAME,  # target collection
    data=query_vectors,  # query vectors
    limit=4,  # number of returned entities
    output_fields=["data", "subject"],  # specifies fields to be returned
)

for q in queries:
    print("Query:", q)
    for result in res:
        fig, axes = plt.subplots(1, len(result), figsize=(66, 6))
        for n, page in enumerate(result):
            page_num = page['id']
            axes[n].imshow(pages[page_num])
            axes[n].axis("off")

    plt.tight_layout()
    plt.show()