Inferenza di embedding su larga scala per applicazioni RAG con Ray Data e Milvus
Retrieval Augmented Generation (RAG) è uno dei casi d'uso più popolari per l'IA generativa aziendale. La maggior parte dei tutorial RAG mostra come usare l'API OpenAI sia per il modello di embedding sia per l'inferenza del modello linguistico di grandi dimensioni (LLM). Perché dovresti pagare per accedere ai tuoi dati, soprattutto durante il processo di sviluppo? Puoi accedere ai tuoi dati e iterare rapidamente quanto vuoi usando l'open source.
Una delle scoperte più intriganti è stato il notevole aumento delle prestazioni ottenuto con Ray Data durante la fase di embedding, in cui i dati vengono trasformati in vettori.
L'esecuzione di embedding open-source usando richieste di inferenza batch raggruppate con strumenti come Ray Data ha risparmiato risorse e tempo rispetto a Pandas. Usando solo quattro worker su un laptop Mac M2 con 16 GB di RAM, Ray Data è stato 60 volte più veloce; maggiori dettagli più avanti in questo blog.
Il nostro stack RAG open source:
Nuovo modello di embedding BGM-M3 (genera 3 tipi di vettori in un'unica esecuzione: sparsi, densi e multi-vettore)
Ray Data per un'inferenza di embedding veloce e distribuita
AWS S3 per archiviare temporaneamente il risultato dell'inferenza
Database vettoriale Milvus o Zilliz Cloud
Dati di esempio scaricati da Kaggle IMDB poster
Il nostro stack RAG open source
Modello di embedding BGM-M3
Potenti embedding sparsi, densi e multi-vettore. Il modello di embedding BGE-M3 deve il suo soprannome alle sue capacità “multi”: multilinguismo, multifunzionalità e multi-granularità. Può funzionare con oltre 100 lingue, calcolare simultaneamente embedding per i tre metodi di retrieval comuni: embedding densi, sparsi e multi-vettore. Funziona inoltre con testi di varie lunghezze, da frasi brevi a documenti lunghi (fino a 8.192 token). Puoi saperne di più da questo paper, o da questa pagina HuggingFace sul modello.
Dalla versione 2.4, Milvus ha supporto integrato per BGE M3.
Ray Data
Attività di trasformazione dei dati di lunga durata?
L'elaborazione scalabile dei dati di Ray Data rende più semplice e veloce elaborare enormi quantità di dati in parallelo su più macchine (CPU, GPU, ecc.). Ray Data è particolarmente utile quando i dati possono essere suddivisi in processi paralleli, come molte trasformazioni simultanee di chunking ed embedding! Dietro le quinte, Ray Data dispone di un potente motore di esecuzione in streaming, per massimizzare l'utilizzo della GPU nel cluster. Rispetto all'esecuzione di embedding con un servizio online (come l'API di embedding di OpenAI), eseguire un job di embedding offline con Ray Data può far risparmiare la maggior parte dei costi.
Anyscale è una piattaforma gestita per Ray. Puoi scalare facilmente i job di embedding su Anyscale per sfruttare centinaia di macchine GPU.
Milvus e Zilliz
La salsa segreta dietro un'app RAG velocissima è un potente database vettoriale!_ Milvus è progettato per gestire enormi quantità di dati per l'uso su larga scala da parte delle aziende. A differenza di alcuni database vettoriali, Milvus può crescere in modo flessibile man mano che le tue esigenze di dati aumentano: i suoi livelli architetturali per archiviazione, indicizzazione e query sono progettati per scalare verso l'alto indipendentemente e/o verso l'esterno. Questo rende la tua app RAG veloce, poiché Milvus esegue in modo intelligente calcoli offline prima e mentre le query arrivano in tempo reale. Inoltre, Milvus include altre funzionalità importanti per le aziende, come mantenere i dati sicuri e organizzati (multi-tenancy e controllo degli accessi basato sui ruoli) e garantire che siano sempre disponibili (alta disponibilità).
Zilliz è un prodotto cloud gestito e utilizza Milvus open-source.
Configura i tuoi strumenti RAG
Useremo l’SDK Python per Milvus, Ray Data, Amazon S3 e Zilliz.
Per Amazon S3, dovrai registrarti per un account AWS.
Nel tuo browser, vai su console.aws.amazon.com > IAM > My security credentials > Create access key. Copia e salva localmente in modo sicuro la tua chiave e la chiave segreta.
Installa le librerie ed esegui aws config. Questo inserirà le variabili AWS in un file di credenziali.
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
Installa Ray Data:
pip install -U "ray[data]"
Installa Pymilvus:
pip install -U pymilvus "pymilvus[model]" langchain
Il modello di embedding BGE-M3 è già incluso in Pymilvus dalla v2.4.
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
Per usare il piano gratuito Zilliz (fino a 2 raccolte, 1 milione di vettori ciascuna), registrati per un account e crea un cluster Starter.
Preparare i dati
Il codice in questo blog utilizza il noto dataset pubblico Kaggle dei poster IMDB. Contiene circa 48.000 film, recensioni, link ai poster e altri metadati.
Ho copiato tutti i campi di testo (nome del film, descrizione, testo della recensione) in una nuova colonna chiamata ‘text’ e l’ho salvata in formato Parquet, poiché è più efficiente del CSV.
Generare embedding
I passaggi per creare embedding sono:
- Suddividere i dati: dividere il testo di input in blocchi, per mantenere insieme le parti di testo semanticamente correlate.
- Chiamare un modello di embedding in modalità inferenza per generare rappresentazioni vettoriali dei blocchi.
Ray Data è in grado di parallelizzare queste operazioni sui dati usando:
- flat_map() per suddividere i dati, poiché l’output avrà più righe dell’input.
- map_batches() per chiamare il modello di embedding dall’interno di un metodo Class callable.
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')
Per eseguirlo, lo invierai come job Ray:
Salva il codice in un file script Python. Io l’ho chiamato ray_data_demo.py
Per eseguirlo localmente dal tuo laptop, crea una directory pulita, con solo il file script .py e il file dati .parquet. Metti solo il minimo indispensabile in questa directory pulita. Io ho chiamato la mia ‘ray_cluster’.
Esegui lo script Python. Questo avvierà un cluster Ray e invierà automaticamente un job.
**Vai a
http://127.0.0.1:8265. Visualizza i tempi di Cluster e Jobs.
Latenza degli embedding - 60 volte più veloce su un laptop
| Approccio | Dimensione dei dati di input | Tempo totale | Screenshot |
| Pandas | 100 righe | 23 sec |
|
| Ray Data | 100 righe | 50 sec |
|
| Pandas | 45K righe | >4 ore |
|
| Ray Data | 45K righe | 4 min |
|
Tabella: Tempi per l’embedding dei dati su un laptop M2 con 16 GB. L’elaborazione batch di Ray Data è stata eseguita su un cluster Ray a nodo singolo, concurrency= 4 worker. Pandas era lento perché aveva un solo processore, mentre Ray Data aveva 4 processori. Entrambi funzionerebbero più velocemente su un cluster più grande.
Inserisci in blocco i dati embedded da S3 direttamente in Milvus o Zilliz
Sia Milvus sia Zilliz offrono l’inserimento in blocco per importare dati già embedded direttamente da AWS, GCP o Azure. Oltre alla console web (mostrata sotto), Zilliz offre anche una API RESTful e SDK.
Per un ampio corpus di embedding generati in batch, l’uso dell’importazione in blocco può far risparmiare significativamente risorse macchina e ridurre il tempo di inserimento rispetto all’inserimento incrementale. Ancora più importante, l’indice di ricerca vettoriale creato dall’importazione in blocco è molto più efficiente rispetto a quello creato dall’inserimento incrementale (pensa all’ottimizzazione globale vs. l’ottimizzazione locale).
Vediamo come eseguire comodamente l’importazione in blocco con pochi semplici clic sulla console web di Zilliz Cloud. Partendo dal Cluster in cui vuoi creare la nuova collection, crea una nuova collection con AutoID, solo la colonna “vector” con EMBEDDING_DIMENSION corretto, usa la comoda opzione “Dynamic Field” e fai clic su “Create Collection”.
Quindi, fai clic su “Import Data” e segui le istruzioni sullo schermo per copiare il percorso dei file parquet scritti dal job Ray Data. (Nota che devi anche specificare la Access Key e la Secret Key se il tuo bucket S3 è privato, in modo che Zilliz Cloud possa leggere i dati al suo interno). Sono supportate sorgenti cloud come Amazon S3, Google Cloud Storage o Azure Blob Storage. Fai clic su “Import” per iniziare a importare tutti i dati nella collection del database vettoriale.
Una volta importati, puoi facoltativamente fare clic su build index sulla collection per rendere la ricerca vettoriale più efficiente nel passaggio Query your data.
Immagine: Screenshot delle schermate di bulk insert di Zilliz.
Interroga i tuoi dati
Per testare la collection appena importata, poniamo una domanda e recuperiamo le risposte dai nostri dati sui film.
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)
Scorrendo i primi 2 risultati unici, possiamo vedere il seguente contenuto restituito in modo molto pertinente dalla query di ricerca sopra:
Lo script completo di Ray Data è disponibile su GitHub.
Conclusione
Questo blog ha mostrato come utilizzare le funzionalità Ray Data e Milvus Bulk Import per accelerare significativamente la generazione di vettori e caricarli in batch in modo efficiente in un database vettoriale. Ad esempio, l'embedding di 102K righe di dati usando Ray Data ha richiesto 4 minuti rispetto a 4 ore usando un approccio Pandas ingenuo! ****Inoltre, l'uso di Bulk Import in Milvus può creare un indice vettoriale altamente efficiente e far risparmiare risorse e tempo rispetto al normale inserimento incrementale. Consulta Ray Data e le funzionalità di Bulk Import in Milvus e Zilliz Cloud per maggiori dettagli!
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.
Continua a leggere
Announcing the General Availability of Single Sign-On (SSO) on Zilliz Cloud
SSO is GA on Zilliz Cloud, delivering the enterprise-grade identity management capabilities your teams need to deploy vectorDB with confidence.
Why I’m Against Claude Code’s Grep-Only Retrieval? It Just Burns Too Many Tokens
Learn how vector-based code retrieval cuts Claude Code token consumption by 40%. Open-source solution with easy MCP integration. Try claude-context today.
What is the K-Nearest Neighbors (KNN) Algorithm in Machine Learning?
KNN is a supervised machine learning technique and algorithm for classification and regression. This post is the ultimate guide to KNN.