Масштабируемый инференс эмбеддингов для RAG-приложений с Ray Data и Milvus
Retrieval Augmented Generation (RAG) — один из самых популярных сценариев использования корпоративного генеративного ИИ. Большинство руководств по RAG показывают, как использовать OpenAI API как для модели эмбеддингов, так и для инференса большой языковой модели (LLM). Почему вы должны платить за доступ к собственным данным, особенно в процессе разработки? Вы можете получать доступ к собственным данным и выполнять итерации так быстро, как захотите, используя open source.
Одним из самых интригующих открытий стал значительный прирост производительности, достигнутый с Ray Data на этапе создания эмбеддингов, когда данные преобразуются в векторы.
Запуск open-source эмбеддингов с использованием объединённых пакетных запросов инференса с такими инструментами, как Ray Data, сэкономил ресурсы и время по сравнению с Pandas. Используя всего четыре воркера на ноутбуке Mac M2 с 16 ГБ RAM, Ray Data оказался в 60 раз быстрее; подробнее об этом далее в этом блоге.
Наш open source RAG-стек:
Новая модель эмбеддингов BGM-M3 (генерирует 3 типа векторов за один проход: разреженные, плотные и multi-vector)
Ray Data для быстрого распределённого инференса эмбеддингов
AWS S3 для временного хранения результата инференса
Векторная база данных Milvus или Zilliz Cloud
Пример данных загружен из Kaggle IMDB poster
Наш open source RAG-стек
Модель эмбеддингов BGM-M3
Мощные разреженные, плотные и multi-vector эмбеддинги. Модель эмбеддингов BGE-M3 получила своё прозвище благодаря своим «multi» возможностям: многоязычности, многофункциональности и многоуровневой гранулярности. Она может работать более чем со 100 языками, одновременно вычислять эмбеддинги для трёх распространённых методов поиска: плотные, разреженные и multi-vector эмбеддинги. Она также работает с текстами различной длины — от коротких предложений до длинных документов (до 8 192 токенов). Подробнее можно узнать из этой статьи или на этой странице HuggingFace о модели.
Начиная с версии 2.4 Milvus имеет встроенную поддержку BGE M3.
Ray Data
Долгие задачи преобразования данных?
Масштабируемая обработка данных Ray Data упрощает и ускоряет параллельную обработку огромных объёмов данных на нескольких машинах (CPU, GPU и т. д.). Ray Data особенно полезен, когда данные можно разделить на параллельные процессы, например множество одновременных преобразований чанкинга и эмбеддингов! Под капотом Ray Data имеет мощный потоковый движок выполнения, чтобы максимально использовать GPU в кластере. По сравнению с запуском эмбеддингов через онлайн-сервис (например, OpenAI embedding API), запуск офлайн-задачи эмбеддингов с Ray Data может сэкономить большую часть затрат.
Anyscale — это управляемая платформа для Ray. Вы можете легко масштабировать задачи эмбеддингов на Anyscale, чтобы задействовать сотни GPU-машин.
Milvus and Zilliz
Секретный ингредиент молниеносного RAG-приложения — мощная векторная база данных!_ Milvus создан для обработки огромных объёмов данных при крупномасштабном использовании бизнесом. В отличие от некоторых векторных баз данных, Milvus может гибко расти по мере увеличения ваших потребностей в данных — его архитектурные уровни хранения, индексирования и запросов спроектированы так, чтобы масштабироваться независимо вертикально и/или горизонтально. Это делает ваше RAG-приложение быстрым, поскольку Milvus разумно выполняет офлайн-вычисления до поступления запросов и во время их обработки в реальном времени. Кроме того, Milvus включает и другие важные для бизнеса возможности, например обеспечение безопасности и упорядоченности ваших данных (мультитенантность и управление доступом на основе ролей), а также гарантию их постоянной доступности (высокая доступность).
Zilliz — это управляемый облачный продукт, использующий open-source Milvus.
Настройте свои RAG-инструменты
Мы будем использовать Python SDK для Milvus, Ray Data, Amazon S3 и Zilliz.
Для Amazon S3 вам нужно будет зарегистрировать учетную запись AWS.
В браузере перейдите по пути console.aws.amazon.com > IAM > My security credentials > Create access key. Скопируйте и надежно сохраните локально ваш ключ и секретный ключ.
Установите библиотеки и запустите aws config. Это поместит переменные AWS в файл учетных данных.
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
Установите Ray Data:
pip install -U "ray[data]"
Установите Pymilvus:
pip install -U pymilvus "pymilvus[model]" langchain
Модель эмбеддингов BGE-M3 уже поставляется вместе с Pymilvus начиная с 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
Чтобы использовать бесплатный тариф Zilliz (до 2 коллекций, по 1 миллиону векторов в каждой), зарегистрируйтесь и создайте Starter cluster.
Подготовка данных
Код в этом блоге использует хорошо известные публичные данные Kaggle с постерами IMDB. Они содержат около 48 000 фильмов, обзоры, ссылки на постеры и другие метаданные.
Я скопировал все текстовые поля (название фильма, описание, текст обзора) в новый столбец под названием ‘text’ и сохранил его в формате Parquet, поскольку он эффективнее CSV.
Генерация эмбеддингов
Шаги для создания эмбеддингов:
- Разбиение данных на фрагменты: разделите входной текст на фрагменты, чтобы семантически связанные части текста оставались вместе.
- Вызов модели эмбеддингов в режиме инференса для генерации векторных представлений фрагментов.
Ray Data способен распараллеливать эти операции с данными, используя:
- flat_map() для разбиения данных на фрагменты, поскольку на выходе будет больше строк, чем на входе.
- map_batches() для вызова модели эмбеддингов из метода вызываемого Class.
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')
Чтобы запустить это, вы отправите его как задание Ray:
Сохраните код в файл Python-скрипта. Я назвал его ray_data_demo.py
Чтобы запустить локально со своего ноутбука, создайте чистую директорию, в которой будут только файл скрипта .py и файл данных .parquet. Поместите в эту чистую директорию только самый необходимый минимум. Я назвал свою ‘ray_cluster’.
Запустите Python-скрипт. Это запустит кластер Ray и автоматически отправит задание.
**Перейдите по адресу
http://127.0.0.1:8265. Просмотрите тайминги Cluster и Jobs.
Задержка эмбеддинга — в 60 раз быстрее на ноутбуке
| Подход | Размер входных данных | Общее время | Скриншот |
| Pandas | 100 строк | 23 сек |
|
| Ray Data | 100 строк | 50 сек |
|
| Pandas | 45 тыс. строк | >4 часов |
|
| Ray Data | 45 тыс. строк | 4 мин |
|
Таблица: Тайминги для эмбеддинга данных на ноутбуке M2 16GB. Пакетная обработка Ray Data выполнялась на одноузловом кластере Ray, concurrency= 4 workers. Pandas был медленным, потому что у него был только один процессор, тогда как у Ray Data было 4 процессора. Оба варианта работали бы быстрее на более крупном кластере.
Массово вставьте эмбеддированные данные из S3 напрямую в Milvus или Zilliz
И Milvus, и Zilliz предлагают bulk-insert для импорта уже эмбеддированных данных напрямую из AWS, GCP или Azure. Помимо web console (показанной ниже), Zilliz также предлагает restful API and SDK.
Для большого корпуса эмбеддингов, сгенерированных пакетно, использование массового импорта может значительно сэкономить машинные ресурсы и сократить время вставки по сравнению с инкрементальной вставкой. Что еще важнее, индекс векторного поиска, построенный с помощью массового импорта, намного эффективнее, чем индекс, созданный при инкрементальной вставке (подумайте о глобальной оптимизации по сравнению с локальной оптимизацией).
Давайте посмотрим, как удобно выполнить массовый импорт несколькими простыми кликами в web console Zilliz Cloud. Начиная с кластера, в котором вы хотите создать новую коллекцию, создайте новую коллекцию с AutoID, только столбцом “vector” с корректным EMBEDDING_DIMENSION, используйте удобную опцию “Dynamic Field” и нажмите “Create Collection”.
Затем нажмите «Import Data» и следуйте инструкциям на экране, чтобы скопировать путь к parquet-файлам, записанным заданием Ray Data. (Обратите внимание, что вам также нужно указать Access Key и Secret Key, если ваш бакет S3 является приватным, чтобы Zilliz Cloud мог читать данные в нем). Поддерживаются облачные источники Amazon S3, Google Cloud Storage и Azure Blob Storage. Нажмите «Import», чтобы начать импорт всех данных в коллекцию векторной базы данных.
После импорта при желании вы можете нажать на build index в коллекции, чтобы сделать векторный поиск более эффективным на этапе Query your data.
Изображение: снимок экрана интерфейсов Zilliz bulk insert.
Запрос к вашим данным
Чтобы протестировать недавно импортированную коллекцию, давайте зададим вопрос и получим ответы из наших данных о фильмах.
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)
Проходя циклом по 2 верхним уникальным результатам, мы можем увидеть следующий контент, близко соответствующий приведенному выше поисковому запросу:
Полный скрипт Ray Data доступен на GitHub.
Заключение
В этом блоге показано, как использовать функции Ray Data и Milvus Bulk Import, чтобы значительно ускорить генерацию векторов и эффективно пакетно загружать их в векторную базу данных. Например, создание эмбеддингов для 102 тыс. строк данных с помощью Ray Data заняло 4 минуты по сравнению с 4 часами при использовании наивного подхода на Pandas! ****Кроме того, использование Bulk Import в Milvus позволяет построить высокоэффективный векторный индекс и сэкономить ресурсы и время по сравнению с обычной инкрементальной вставкой. Ознакомьтесь с Ray Data и функциями Bulk Import в Milvus и 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.
Читать далее
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.
Zilliz Cloud BYOC Now Available Across AWS, GCP, and Azure
Zilliz Cloud BYOC is now generally available on all three major clouds. Deploy fully managed vector search in your own AWS, GCP, or Azure account — your data never leaves your VPC.
How to Use Anthropic MCP Server with Milvus
MCP + Milvus: Streamline AI agent development with standardized data access, eliminating integration hassles while enhancing context and flexibility.