Blog
Inferencia de embeddings a escala para aplicaciones RAG con Ray Data y Milvus

Inferencia de embeddings a escala para aplicaciones RAG con Ray Data y Milvus

Apr 12, 20248 min read

La Generación Aumentada por Recuperación (RAG) es uno de los casos de uso más populares para la IA generativa empresarial. La mayoría de los tutoriales de RAG muestran cómo usar la API de OpenAI tanto para el modelo de embeddings como para la inferencia del modelo de lenguaje grande (LLM). ¿Por qué deberías pagar para acceder a tus propios datos, especialmente durante el proceso de desarrollo? Puedes acceder a tus propios datos e iterar tan rápido como quieras usando código abierto.

Uno de los descubrimientos más intrigantes fue el notable aumento de rendimiento logrado con Ray Data durante el paso de embeddings, donde los datos se transforman en vectores.

Ejecutar embeddings de código abierto usando solicitudes de inferencia por lotes agrupadas con herramientas como Ray Data ahorró recursos y tiempo en comparación con Pandas. Usando solo cuatro workers en una laptop Mac M2 con 16GB de RAM, Ray Data fue 60 veces más rápido; más detalles más adelante en este blog.

Nuestra pila RAG de código abierto:

Nuevo modelo de embeddings BGM-M3 (genera 3 tipos de vectores en una sola ronda: dispersos, densos y multi-vector)
Ray Data para una inferencia de embeddings rápida y distribuida
AWS S3 para almacenar temporalmente el resultado de la inferencia
Base de datos vectorial Milvus o Zilliz Cloud

Datos de ejemplo descargados de Kaggle IMDB poster

Nuestra pila RAG de código abierto

Modelo de embeddings BGM-M3

Embeddings dispersos, densos y multi-vector potentes. El modelo de embeddings BGE-M3 recibe su apodo por sus capacidades “multi”: multilingüismo, multifuncionalidad y multigranularidad. Puede trabajar con más de 100 idiomas, calcular simultáneamente embeddings para los tres métodos comunes de recuperación: embeddings densos, dispersos y multi-vector. También funciona con varias longitudes de texto, desde oraciones cortas hasta documentos largos (hasta 8.192 tokens). Puedes obtener más información en este Paper, o en esta página de HuggingFace sobre el model.

Desde la versión 2.4, Milvus tiene soporte integrado para BGE M3.

Ray Data

¿Tareas de transformación de datos de larga duración?

El procesamiento de datos escalable de Ray Data hace que sea más fácil y rápido procesar enormes cantidades de datos en paralelo en múltiples máquinas (CPU, GPU, etc.). Ray Data es especialmente útil cuando los datos pueden dividirse en procesos paralelos, como muchas transformaciones simultáneas de fragmentación y embeddings. Internamente, Ray Data cuenta con un potente motor de ejecución en streaming para maximizar la utilización de la GPU en el clúster. En comparación con ejecutar embeddings con un servicio en línea (como la API de embeddings de OpenAI), ejecutar un trabajo de embeddings sin conexión con Ray Data puede ahorrar la mayor parte del coste.

Anyscale es una plataforma gestionada para Ray. Puedes escalar fácilmente los trabajos de embeddings en Anyscale para aprovechar cientos de máquinas con GPU.

Milvus y Zilliz

¡La salsa secreta detrás de una app RAG ultrarrápida es una potente base de datos vectorial!_ Milvus está diseñado para manejar cantidades masivas de datos para uso a gran escala por parte de las empresas. A diferencia de algunas bases de datos vectoriales, Milvus puede crecer de forma flexible a medida que aumentan tus necesidades de datos: sus capas arquitectónicas para almacenamiento, indexación y consulta están diseñadas para escalar de manera independiente hacia arriba y/o hacia afuera. Esto hace que tu app RAG sea rápida, ya que Milvus realiza inteligentemente cálculos sin conexión antes y mientras las consultas llegan en tiempo real. Además, Milvus incluye otras prestaciones importantes para las empresas, como mantener tus datos seguros y organizados (multi-tenancy y control de acceso basado en roles) y garantizar que siempre estén disponibles (alta disponibilidad).

Zilliz es un producto cloud gestionado y utiliza Milvus de código abierto.

Configura tus herramientas RAG

Usaremos el SDK de Python para Milvus, Ray Data, Amazon S3 y Zilliz.

Para Amazon S3, tendrás que registrarte para obtener una cuenta de AWS.

En tu navegador, ve a console.aws.amazon.com > IAM > My security credentials > Create access key. Copia y guarda de forma segura localmente tu clave y clave secreta.

Instala las bibliotecas y ejecuta aws config. Esto colocará las variables de AWS en un archivo de credenciales.

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

Instala Ray Data:

pip install -U "ray[data]"

Instala Pymilvus:

pip install -U pymilvus "pymilvus[model]" langchain

El modelo de embeddings BGE-M3 ya viene empaquetado con Pymilvus desde la versión 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

Para usar el nivel gratuito de Zilliz (hasta 2 colecciones, 1 millón de vectores cada una), regístrate para obtener una cuenta y crea un clúster Starter.

Preparar los datos

El código de este blog usa los conocidos datos públicos de pósteres de IMDB de Kaggle. Contiene unas 48.000 películas, reseñas, enlaces a pósteres y más metadatos.

Copié todos los campos de texto (nombre de la película, descripción, texto de la reseña) en una nueva columna llamada ‘text’ y la guardé en formato Parquet, ya que es más eficiente que CSV.

Generar embeddings

Los pasos para crear embeddings son:

Dividir los datos en fragmentos: Dividir el texto de entrada en fragmentos, para mantener juntas las partes de texto relacionadas semánticamente.
Llamar a un modelo de embeddings en modo de inferencia para generar representaciones vectoriales de los fragmentos.

Ray Data puede paralelizar estas operaciones de datos usando:

flat_map() para dividir los datos en fragmentos, ya que la salida tendrá más filas que la entrada.
map_batches() para llamar al modelo de embeddings desde dentro de un método de Class invocable.

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

Para ejecutar esto, lo enviarás como un trabajo de Ray:

Guarda el código en un archivo de script de Python. Yo lo llamé ray_data_demo.py
Para ejecutarlo localmente desde tu laptop, crea un directorio limpio, con solo el archivo de script .py y el archivo de datos .parquet. Pon únicamente lo mínimo indispensable en este directorio limpio. Yo llamé al mío ‘ray_cluster’.
Ejecuta el script de Python. Esto iniciará un clúster de Ray y enviará un trabajo automáticamente.
**Navega a http://127.0.0.1:8265. Consulta los tiempos de Cluster y Jobs.

Latencia de embeddings: 60 veces más rápido en una laptop


Enfoque	Tamaño de datos de entrada	Tiempo total	Captura de pantalla
Pandas	100 filas	23 s
Ray Data	100 filas	50 s
Pandas	45K filas	>4 horas
Ray Data	45K filas	4 min

Tabla: Tiempos para generar embeddings de datos en una laptop M2 de 16 GB. El procesamiento por lotes de Ray Data se realizó en un clúster de Ray de un solo nodo, concurrency= 4 workers. Pandas fue lento porque solo tenía un procesador, mientras que Ray Data tenía 4 procesadores. Ambos se ejecutarían más rápido en un clúster más grande.

Inserción masiva de los datos con embeddings desde S3 directamente en Milvus o Zilliz

Tanto Milvus como Zilliz ofrecen inserción masiva para importar datos ya vectorizados directamente desde AWS, GCP o Azure. Además de la consola web (mostrada a continuación), Zilliz también ofrece una API restful y SDK.

Para un corpus grande de embeddings generados por lotes, usar la importación masiva puede ahorrar significativamente recursos de máquina y acortar el tiempo de inserción en comparación con la inserción incremental. Más importante aún, el índice de búsqueda vectorial creado mediante importación masiva es mucho más eficiente que el de la inserción incremental (piensa en optimización global frente a optimización local).

Veamos cómo realizar cómodamente la importación masiva con unos pocos clics sencillos en la consola web de Zilliz Cloud. Comenzando desde el Cluster donde quieres crear la nueva colección, crea una nueva colección con AutoID, solo la columna “vector” con el EMBEDDING_DIMENSION correcto, usa la cómoda opción “Dynamic Field” y haz clic en “Create Collection”.

A continuación, haz clic en “Import Data” y sigue las instrucciones en pantalla para copiar la ruta a los archivos parquet escritos por el job de Ray Data. (Ten en cuenta que también debes especificar la Access Key y la Secret Key si tu bucket de S3 es privado, para que Zilliz Cloud pueda leer los datos que contiene). Se admite cualquiera de las fuentes en la nube Amazon S3, Google Cloud Storage o Azure Blob Storage. Haz clic en “Import” para empezar a importar todos los datos en la colección de la base de datos vectorial.

Una vez importados, opcionalmente puedes hacer clic en crear índice en la colección para que la búsqueda vectorial sea más eficiente en el paso Consulta tus datos.

Imagen: Captura de pantalla de las pantallas de inserción masiva de Zilliz.

Consulta tus datos

Para probar la colección recién importada, hagamos una pregunta y recuperemos respuestas de nuestros datos de películas.

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)

Al recorrer los 2 principales resultados únicos, podemos ver el siguiente contenido devuelto con alta similitud a partir de la consulta de búsqueda anterior:

El script completo de Ray Data está disponible en GitHub.

Conclusión

Este blog mostró cómo usar las funciones de Ray Data e Importación masiva de Milvus para acelerar significativamente la generación de vectores y cargarlos eficientemente por lotes en una base de datos vectorial. Por ejemplo, ¡generar embeddings para 102K filas de datos usando Ray Data tardó 4 minutos en comparación con 4 horas usando un enfoque ingenuo con Pandas! ****Además, usar la Importación masiva en Milvus puede crear un índice vectorial altamente eficiente y ahorrar recursos y tiempo en comparación con la inserción incremental regular. Consulta Ray Data y las funciones de Importación masiva en Milvus y Zilliz Cloud para más detalles!

Actualizado el Jul 13, 2026

Sigue leyendo

Zilliz Skills Breakdown: How AI Agents Master Vector Databases

Zilliz's Milvus Skill (pymilvus, 7 files) and Zilliz Cloud Skill (zilliz-cli, 14 modules) bring vector-DB dev and ops into one Claude Code session.

Vector Databases vs. Time Series Databases

Use a vector database for similarity search and semantic relationships; use a time series database for tracking value changes over time.

Why Deepseek is Waking up AI Giants Like OpenAI And Why You Should Care

Discover how DeepSeek R1's open-source AI model with superior reasoning capabilities and lower costs is disrupting the AI landscape and challenging tech giants like OpenAI.