Open In Colab GitHub Repository

RAG mit Milvus und EmbedAnything aufbauen

EmbedAnything ist eine blitzschnelle, leichtgewichtige Einbettungspipeline in Rust, die Text, PDFs, Bilder, Audio und mehr unterstützt.

In diesem Tutorial zeigen wir, wie man eine Retrieval-Augmented Generation (RAG) Pipeline mit EmbedAnything und Milvus erstellt. EmbedAnything ist nicht eng mit einer bestimmten Datenbank gekoppelt, sondern verwendet ein steckbares Adaptersystem - Adapter dienen als Wrapper, die definieren, wie Einbettungen formatiert, indiziert und im Zielvektorspeicher gespeichert werden.

Indem Sie EmbedAnything mit einem Milvus-Adapter verbinden, können Sie Einbettungen aus verschiedenen Dateitypen generieren und sie effizient in Milvus speichern - und das mit nur wenigen Zeilen Code.

⚠️ Hinweis: Während der Adapter in EmbedAnything das Einfügen in Milvus handhabt, unterstützt er die Suche nicht von Haus aus. Um eine vollständige RAG-Pipeline zu erstellen, müssen Sie einen separaten MilvusClient instanziieren und die Abfragelogik (z. B. Ähnlichkeitssuche über Vektoren) als Teil Ihrer Anwendung implementieren.

Vorbereitung

Abhängigkeiten und Umgebung

$ pip install -qU pymilvus milvus-lite openai embed_anything

Wenn Sie Google Colab verwenden, müssen Sie eventuell die Runtime neu starten, um die soeben installierten Abhängigkeiten zu aktivieren (klicken Sie auf das Menü "Runtime" am oberen Rand des Bildschirms und wählen Sie "Restart session" aus dem Dropdown-Menü).

Klonen des Repositorys und Laden des Adapters

Als nächstes klonen wir das EmbedAnything-Repository und fügen das Verzeichnis examples/adapters zum Python-Pfad hinzu. Dort speichern wir die benutzerdefinierte Milvus-Adapter-Implementierung, die es EmbedAnything ermöglicht, mit Milvus für die Vektoreinfügung zu kommunizieren.

import sys

# Clone the EmbedAnything repository if not already cloned
![ -d "EmbedAnything" ] || git clone https://github.com/StarlightSearch/EmbedAnything.git

# Add the `examples/adapters` directory to the Python path
sys.path.append("EmbedAnything/examples/adapters")
print("✅ EmbedAnything cloned and adapter path added.")
✅ EmbedAnything cloned and adapter path added.

Wir werden OpenAI als LLM in dieser RAG-Pipeline verwenden. Sie sollten den api-Schlüssel OPENAI_API_KEY als Umgebungsvariable vorbereiten.

import os
from openai import OpenAI

os.environ["OPENAI_API_KEY"] = "sk-***********"

openai_client = OpenAI()

RAG erstellen

Milvus initialisieren

Bevor wir irgendwelche Dateien einbetten, müssen wir zwei Komponenten vorbereiten, die mit Milvus interagieren:

  1. MilvusVectorAdapter - Dies ist der Milvus-Adapter für EmbedAnything und wird nur für die Aufnahme von Vektoren verwendet (d.h. das Einfügen von Einbettungen und das Erstellen von Indizes). Er unterstützt derzeit keine Suchoperationen.
  2. MilvusClient - Dies ist der offizielle Client von pymilvus, der den vollen Zugriff auf die Milvus-Funktionen einschließlich Vektorsuche, Filterung und Sammlungsverwaltung ermöglicht.

Um Verwirrung zu vermeiden:

  • Betrachten Sie MilvusVectorAdapter als Ihr "schreibgeschütztes" Werkzeug zum Speichern von Vektoren.
  • Betrachten Sie MilvusClient als Ihre "Lese- und Suchmaschine", um Abfragen durchzuführen und Dokumente für RAG abzurufen.
import embed_anything
from embed_anything import (
    WhichModel,
    EmbeddingModel,
)
from milvus_db import MilvusVectorAdapter
from pymilvus import MilvusClient

# Official Milvus client for full operations
milvus_client = MilvusClient(uri="./milvus.db", token="")

# EmbedAnything adapter for pushing embeddings into Milvus
index_name = "embed_anything_milvus_collection"
milvus_adapter = MilvusVectorAdapter(
    uri="./milvus.db", token="", collection_name=index_name
)

# Delete existing collection if it exists
if milvus_client.has_collection(index_name):
    milvus_client.drop_collection(index_name)

# Create a new collection with dimension matching the embedding model later used
milvus_adapter.create_index(dimension=384)
Ok - Milvus DB connection established.
Collection 'embed_anything_milvus_collection' created with index.

Was das Argument von MilvusVectorAdapter und MilvusClient angeht:

  • Die Einstellung von uri als lokale Datei, z.B../milvus.db, ist die bequemste Methode, da sie automatisch Milvus Lite verwendet, um alle Daten in dieser Datei zu speichern.
  • Wenn Sie große Datenmengen haben, z. B. mehr als eine Million Vektoren, können Sie einen leistungsfähigeren Milvus-Server auf Docker oder Kubernetes einrichten. Bei dieser Einrichtung verwenden Sie bitte die Serveradresse und den Port als Uri, z. B.http://localhost:19530. Wenn Sie die Authentifizierungsfunktion auf Milvus aktivieren, verwenden Sie ":" als Token, andernfalls setzen Sie das Token nicht.
  • Wenn Sie Zilliz Cloud, den vollständig verwalteten Cloud-Service für Milvus, verwenden möchten, passen Sie uri und token an, die dem öffentlichen Endpunkt und dem Api-Schlüssel in Zilliz Cloud entsprechen.

Einbettungsmodell initialisieren und PDF-Dokument einbetten

Nun werden wir das Einbettungsmodell initialisieren. Wir verwenden all-MiniLM-L12-v2 model aus der sentence-transformers-Bibliothek, ein leichtgewichtiges und dennoch leistungsstarkes Modell zur Erzeugung von Texteinbettungen. Es erzeugt 384-dimensionale Einbettungen, so dass dies mit der Dimension unserer Milvus-Sammlung übereinstimmt, die auf 384 eingestellt ist. Dieser Abgleich ist entscheidend und gewährleistet die Kompatibilität zwischen den in Milvus gespeicherten und den vom Modell erzeugten Vektordimensionen.

EmbedAnything unterstützt eine Vielzahl weiterer Einbettungsmodelle. Für weitere Details, lesen Sie bitte die offizielle Dokumentation.

# Initialize the embedding model
model = EmbeddingModel.from_pretrained_hf(
    WhichModel.Bert, model_id="sentence-transformers/all-MiniLM-L12-v2"
)

Lassen Sie uns nun eine PDF-Datei einbetten. EmbedAnything macht es einfach, PDF-Dokumente (und viele andere) zu verarbeiten und ihre Einbettungen direkt in Milvus zu speichern.

# Embed a PDF file
data = embed_anything.embed_file(
    "./pdf_files/WhatisMilvus.pdf",
    embedder=model,
    adapter=milvus_adapter,
)
Converted 12 embeddings for insertion.
Successfully inserted 12 embeddings.

Antwort abrufen und generieren

Auch hier ist die MilvusVectorAdapter von EmbedAnything derzeit nur eine leichtgewichtige Abstraktion für die Aufnahme und Indizierung von Vektoren. Sie unterstützt keine Such- oder Retrieval-Abfragen. Daher müssen wir für die Suche nach relevanten Dokumenten zum Aufbau unserer RAG-Pipeline direkt die Instanz MilvusClient (milvus_client) verwenden, um unseren Milvus-Vektorspeicher abzufragen.

Definieren Sie eine Funktion zum Abrufen relevanter Dokumente aus Milvus.

def retrieve_documents(question, top_k=3):
    query_vector = list(
        embed_anything.embed_query([question], embedder=model)[0].embedding
    )
    search_res = milvus_client.search(
        collection_name=index_name,
        data=[query_vector],
        limit=top_k,
        output_fields=["text"],
    )
    docs = [(res["entity"]["text"], res["distance"]) for res in search_res[0]]
    return docs

Definieren Sie eine Funktion, um eine Antwort unter Verwendung der abgerufenen Dokumente in der RAG-Pipeline zu erzeugen.

def generate_rag_response(question):
    retrieved_docs = retrieve_documents(question)
    context = "\n".join([f"Text: {doc[0]}\n" for doc in retrieved_docs])
    system_prompt = (
        "You are an AI assistant. Provide answers based on the given context."
    )
    user_prompt = f"""
    Use the following pieces of information to answer the question. If the information is not in the context, say you don't know.
    
    Context:
    {context}
    
    Question: {question}
    """
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    )
    return response.choices[0].message.content

Lassen Sie uns die RAG-Pipeline mit einer Beispielfrage testen.

question = "How does Milvus search for similar documents?"
answer = generate_rag_response(question)
print(f"Question: {question}")
print(f"Answer: {answer}")
Question: How does Milvus search for similar documents?
Answer: Milvus searches for similar documents primarily through Approximate Nearest Neighbor (ANN) search, which finds the top K vectors closest to a given query vector. It also supports various other types of searches, such as filtering search under specified conditions, range search within a specified radius, hybrid search based on multiple vector fields, and keyword search based on BM25. Additionally, it can perform reranking to adjust the order of search results based on additional criteria, refining the initial ANN search results.