Embedding-Inferenz im großen Maßstab für RAG-Anwendungen mit Ray Data und Milvus
Retrieval Augmented Generation (RAG) ist einer der beliebtesten Anwendungsfälle für Enterprise Generative AI. Die meisten RAG-Tutorials zeigen, wie man die OpenAI API sowohl für das Embedding-Modell als auch für die Inferenz des Large Language Model (LLM) verwendet. Warum sollten Sie für den Zugriff auf Ihre eigenen Daten bezahlen, insbesondere während des Entwicklungsprozesses? Sie können mit Open Source auf Ihre eigenen Daten zugreifen und so schnell iterieren, wie Sie möchten.
Eine der faszinierendsten Entdeckungen war der bemerkenswerte Leistungsschub, der mit Ray Data während des Embedding-Schritts erzielt wurde, bei dem Daten in Vektoren umgewandelt werden.
Das Ausführen von Open-Source-Embeddings mithilfe gepoolter Batch-Inferenzanfragen mit Tools wie Ray Data sparte im Vergleich zu Pandas Ressourcen und Zeit. Mit nur vier Workern auf einem Mac M2 Laptop mit 16 GB RAM war Ray Data 60-mal schneller, weitere Details später in diesem Blog.
Unser Open Source RAG Stack:
Neues BGM-M3-Embedding-Modell (erzeugt 3 Arten von Vektoren in einem Durchlauf: sparse, dense und multi-vector)
Ray Data für schnelle, verteilte Embedding-Inferenz
AWS S3 zur temporären Speicherung des Inferenzergebnisses
Milvus oder Zilliz Cloud Vektordatenbank
Beispieldaten heruntergeladen von Kaggle IMDB poster
Unser Open Source RAG Stack
BGM-M3-Embedding-Modell
Leistungsstarke Sparse-, Dense- und Multi-vector-Embeddings. Das BGE-M3-Embedding-Modell trägt seinen Spitznamen aufgrund seiner „Multi“-Fähigkeiten: Multi-Lingualität, Multi-Funktionalität und Multi-Granularität. Es kann mit über 100 Sprachen arbeiten und gleichzeitig Embeddings für die drei gängigen Retrieval-Methoden berechnen: dense, sparse, multi-vector embeddings. Es funktioniert außerdem mit verschiedenen Textlängen, von kurzen Sätzen bis hin zu langen Dokumenten (bis zu 8.192 Tokens). Mehr erfahren Sie in diesem Paper oder auf dieser HuggingFace-Seite über das model.
Seit Version 2.4 bietet Milvus integrierte Unterstützung für BGE M3.
Ray Data
Lang laufende Datentransformationsaufgaben?
Die skalierbare Datenverarbeitung von Ray Data macht es einfacher und schneller, riesige Datenmengen parallel über mehrere Maschinen hinweg (CPUs, GPUs usw.) zu verarbeiten. Ray Data ist besonders hilfreich, wenn die Daten in parallele Prozesse aufgeteilt werden können, wie etwa viele gleichzeitige Chunking- und Embedding-Transformationen! Unter der Haube verfügt Ray Data über eine leistungsstarke Streaming-Ausführungs-Engine, um die GPU-Auslastung im Cluster zu maximieren. Im Vergleich zum Ausführen von Embeddings mit einem Online-Dienst (wie der OpenAI embedding API) kann ein Offline-Embedding-Job mit Ray Data den Großteil der Kosten einsparen.
Anyscale ist eine verwaltete Plattform für Ray. Sie können die Embedding-Jobs auf Anyscale problemlos skalieren, um Hunderte von GPU-Maschinen zu nutzen.
Milvus und Zilliz
Die Geheimzutat hinter einer blitzschnellen RAG-App ist eine leistungsstarke Vektordatenbank!_ Milvus ist darauf ausgelegt, riesige Datenmengen für den großflächigen Einsatz in Unternehmen zu verarbeiten. Anders als einige Vektordatenbanken kann Milvus flexibel wachsen, wenn Ihre Datenanforderungen steigen – seine Architekturschichten für Speicherung, Indexierung und Abfrage sind darauf ausgelegt, unabhängig voneinander nach oben und/oder nach außen zu skalieren. Das macht Ihre RAG-App schnell, da Milvus vorab und während die Abfragen in Echtzeit eingehen, intelligent Offline-Berechnungen durchführt. Außerdem bringt Milvus weitere wichtige Funktionen für Unternehmen mit, etwa um Ihre Daten sicher und organisiert zu halten (Multi-Tenancy und rollenbasierte Zugriffskontrolle) und sicherzustellen, dass sie stets verfügbar sind (High Availability).
Zilliz ist ein verwaltetes Cloud-Produkt und verwendet das Open-Source-Milvus.
Richten Sie Ihre RAG-Tools ein
Wir verwenden das Python SDK für Milvus, Ray Data, Amazon S3 und Zilliz.
Für Amazon S3 müssen Sie sich für ein AWS-Konto registrieren.
Navigieren Sie in Ihrem Browser zu console.aws.amazon.com > IAM > My security credentials > Create access key. Kopieren Sie Ihren Schlüssel und geheimen Schlüssel und speichern Sie sie lokal an einem sicheren Ort.
Installieren Sie die Bibliotheken und führen Sie aws config aus. Dadurch werden die AWS-Variablen in einer Credentials-Datei abgelegt.
pip install boto3
pip install awscli –force-reinstall –upgrade
aws config #fill in your key and secret key
more ~/.aws/credentials #make sure this looks correct
Installieren Sie Ray Data:
pip install -U "ray[data]"
Installieren Sie Pymilvus:
pip install -U pymilvus "pymilvus[model]" langchain
Das BGE-M3-Embedding-Modell ist seit v2.4 bereits in Pymilvus enthalten.
import ray, os, pprint, time, boto3
from langchain.text_splitter import RecursiveCharacterTextSplitter
import numpy as np
import pymilvus
print(pymilvus.__version__) # must be >= 2.4.0
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
Um den kostenlosen Zilliz-Tarif zu nutzen (bis zu 2 Collections mit jeweils 1 Million Vektoren), registrieren Sie sich für ein Konto und erstellen Sie einen Starter-Cluster.
Daten vorbereiten
Der Code in diesem Blog verwendet die bekannte öffentliche Kaggle IMDB-Posterdaten. Sie enthalten etwa 48.000 Filme, Rezensionen, Poster-Links und weitere Metadaten.
Ich habe alle Textfelder (Filmname, Beschreibung, Rezensionstext) in eine neue Spalte namens „text“ kopiert und sie im Parquet-Format gespeichert, da dieses effizienter ist als CSV.
Embeddings generieren
Die Schritte zum Erstellen von Embeddings sind:
- Daten in Chunks aufteilen: Den Eingabetext in Chunks aufteilen, um semantisch zusammengehörige Textteile zusammenzuhalten.
- Ein Embedding-Modell im Inferenzmodus aufrufen, um Vektorrepräsentationen der Chunks zu generieren.
Ray Data kann diese Datenoperationen parallelisieren mit:
- flat_map() zum Aufteilen der Daten in Chunks, da die Ausgabe mehr Zeilen als die Eingabe haben wird.
- map_batches() zum Aufrufen des Embedding-Modells aus einer aufrufbaren Class-Methode heraus.
chunk_size = 512
chunk_overlap = np.round(chunk_size * 0.10, 0)
# Define a LangChain text splitter.
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
length_function=len) #len is a built-in Python function
# 1. Define a regular python function for chunking.
def chunk_row(row, splitter=text_splitter):
# Copy the row columns into metadata.
metadata = row.copy()
del metadata['text'] # Remove text from metadata
# Split the text into chunks.
chunks = splitter.create_documents(
texts=[row["text"]],
metadatas=[metadata])
chunk_list = [{
"text": chunk.page_content,
**chunk.metadata} for chunk in chunks]
return chunk_list
# 2. Define a class with a callable method to compute embeddings.
class ComputeEmbeddings:
def __init__(self):
# Initialize a Milvus built-in sparse-dense-late-interaction-reranking encoder.
# https://huggingface.co/BAAI/bge-m3
self.model = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
print(f"dense_dim: {self.model.dim['dense']}")
print(f"sparse_dim: {self.model.dim['sparse']}")
def __call__(self, batch):
# Ray data batch is a dictionary where values are array values.
# BGEM3EmbeddingFunction input is docs as a list of strings.
docs = list(batch['text'])
# Encode the documents. bge-m3 dense embeddings are already normalized.
embeddings = self.model(docs)
batch['vector_dense'] = embeddings['dense']
return batch
if __name__ == "__main__":
FILE_PATH = "s3://zilliz/kaggle_imdb.parquet"
# Load and transform data.
ds = ray.data.read_parquet(FILE_PATH)
# Chunk the input text
chunked_ds = ds.flat_map(chunk_row)
# Compute embeddings with a class that calls the embeddings model.
embeddings_ds = chunked_ds.map_batches(ComputeEmbeddings, concurrency=4)
# Save the embeddings to S3 in a folder of parquet part files.
embeddings_ds.write_parquet('s3://zilliz/kaggle_imdb_embeddings')
Um dies auszuführen, reichst du es als Ray-Job ein:
Speichere den Code in einer Python-Skriptdatei. Ich habe sie ray_data_demo.py genannt
Um ihn lokal von deinem Laptop aus auszuführen, erstelle ein sauberes Verzeichnis, das nur die .py-Skriptdatei und die .parquet-Datendatei enthält. Lege nur das absolut Notwendige in dieses saubere Verzeichnis. Ich habe meines „ray_cluster“ genannt.
Führe das Python-Skript aus. Dadurch wird ein Ray-Cluster gestartet und automatisch ein Job eingereicht.
**Navigiere zu
http://127.0.0.1:8265. Sieh dir die Zeitangaben für Cluster und Jobs an.
Embedding-Latenz – 60x schneller auf einem Laptop
| Ansatz | Eingabedatengröße | Gesamtzeit | Screenshot |
| Pandas | 100 Zeilen | 23 Sek. |
|
| Ray Data | 100 Zeilen | 50 Sek. |
|
| Pandas | 45K Zeilen | >4 Stunden |
|
| Ray Data | 45K Zeilen | 4 Min. |
|
Tabelle: Zeitmessungen für das Einbetten von Daten auf einem M2-Laptop mit 16 GB. Die Batch-Verarbeitung mit Ray Data lief auf einem Single-Node-Ray-Cluster, concurrency= 4 Worker. Pandas war langsam, weil es nur einen Prozessor hatte, während Ray Data 4 Prozessoren hatte. Beide würden auf einem größeren Cluster schneller laufen.
Die eingebetteten Daten per Bulk-Insert direkt aus S3 in Milvus oder Zilliz einfügen
Sowohl Milvus als auch Zilliz bieten Bulk-Insert, um bereits eingebettete Daten direkt aus AWS, GCP oder Azure zu importieren. Zusätzlich zur Webkonsole (unten gezeigt) bietet Zilliz auch eine restful API und SDK.
Bei einem großen Korpus batch-generierter Embeddings kann die Verwendung von Bulk Import im Vergleich zur inkrementellen Einfügung Maschinenressourcen erheblich einsparen und die Einfügezeit verkürzen. Noch wichtiger ist, dass der durch Bulk Import erstellte Vektorsuchindex wesentlich effizienter ist als der aus inkrementeller Einfügung (man denke an globale Optimierung vs. lokale Optimierung).
Sehen wir uns an, wie man den Bulk Import mit wenigen einfachen Klicks in der Zilliz Cloud-Webkonsole bequem durchführt. Ausgehend von dem Cluster, in dem du die neue Collection erstellen möchtest, erstelle eine neue Collection mit AutoID, nur der Spalte „vector“ mit korrekter EMBEDDING_DIMENSION, nutze die praktische Option „Dynamic Field“ und klicke auf „Create Collection“.
Klicken Sie als Nächstes auf „Import Data“ und folgen Sie den Bildschirmanweisungen, um den Pfad zu den Parquet-Dateien zu kopieren, die vom Ray Data-Job geschrieben wurden. (Beachten Sie, dass Sie auch den Access Key und den Secret Key angeben müssen, wenn Ihr S3-Bucket privat ist, damit Zilliz Cloud die darin enthaltenen Daten lesen kann). Alle Cloud-Quellen von Amazon S3, Google Cloud Storage oder Azure Blob Storage werden unterstützt. Klicken Sie auf „Import“, um den Import aller Daten in die Vektordatenbank-Collection zu starten.
Nach dem Import können Sie optional auf den Build-Index der Collection klicken, um die Vektorsuche im Schritt „Ihre Daten abfragen“ effizienter zu gestalten.
Bild: Screenshot der Zilliz-Bulk-Insert-Bildschirme.
Ihre Daten abfragen
Um die neu importierte Collection zu testen, stellen wir eine Frage und rufen Antworten aus unseren Filmdaten ab.
def mc_run_search(question, output_fields, top_k=2, filter_expression=""):
# Embed the question using the same encoder.
embeddings = model_bgem3([question])
query_embeddings = embeddings['dense']
# Run semantic vector search using your query and the vector database.
results = mc.search(
COLLECTION_NAME,
data=query_embeddings,
search_params=SEARCH_PARAMS,
output_fields=output_fields,
# Milvus can utilize metadata in boolean expressions to filter search.
filter=filter_expression,
limit=top_k,
consistency_level="Eventually"
)
# Assemble retrieved context and context metadata.
# The search result is in the variable `results[0]`, which is type
# 'pymilvus.orm.search.SearchResult'.
METADATA_FIELDS = [f for f in output_fields if f != 'chunk']
formatted_results, context, context_metadata = _utils.client_assemble_retrieved_context(
results, metadata_fields=METADATA_FIELDS, num_shot_answers=top_k)
return formatted_results, context, context_metadata
SAMPLE_QUESTION = "muybridge horse movie"
# Return top k unique results with HNSW index.
TOP_K = 2
# Define output fields to return.
OUTPUT_FIELDS = ["movie_id", "chunk", "PosterLink"]
formatted_results, context, context_metadata = \
mc_run_search(SAMPLE_QUESTION, OUTPUT_FIELDS, TOP_K)
Wenn wir die 2 besten eindeutigen Ergebnisse durchlaufen, sehen wir, dass der obigen Suchanfrage die folgenden Inhalte sehr nahe zurückgegeben werden:
Das vollständige Ray Data-Skript ist auf GitHub verfügbar.
Fazit
Dieser Blog zeigte, wie man Ray Data und die Milvus Bulk Import-Funktionen verwendet, um die Vektorgenerierung deutlich zu beschleunigen und sie effizient stapelweise in eine Vektordatenbank zu laden. Zum Beispiel dauerte das Einbetten von 102K Datenzeilen mit Ray Data 4 Minuten, verglichen mit 4 Stunden bei einem naiven Pandas-Ansatz! ****Darüber hinaus kann die Verwendung von Bulk Import in Milvus einen hocheffizienten Vektorindex erstellen und im Vergleich zum regulären inkrementellen Einfügen Ressourcen und Zeit sparen. Weitere Details finden Sie unter Ray Data und den Bulk Import-Funktionen in Milvus und Zilliz Cloud!
Christy BergmanChristy Bergman is a passionate Developer Advocate at Zilliz. She previously worked in distributed computing at Anyscale and as a Specialist AI/ML Solutions Architect at AWS. Christy studied applied math, is a self-taught coder, and has published papers, including one with ACM Recsys. She enjoys hiking and bird watching.
Cheng SuCheng Su is the manager of the Data team at Anyscale, a committer of the Ray open source project (https://github.com/ray-project/ray), and the code owner of the Ray Data module. Previously, Cheng Su worked in the Data Infrastructure teams at Meta, Spark, and Hadoop.
Weiterlesen
3 Easiest Ways to Use Claude Code on Your Mobile Phone
Run Claude Code from your phone with Remote Control, Happy Coder, or SSH + Tailscale. Comparison table, setup steps, and tools for typing, memory, and parallel tasks.
Why AI Databases Don't Need SQL
Whether you like it or not, here's the truth: SQL is destined for decline in the era of AI.
VidTok: Rethinking Video Processing with Compact Tokenization
VidTok tokenizes videos to reduce redundancy while preserving spatial and temporal details for efficient processing.