Создание RAG с помощью Milvus и EmbedAnything

EmbedAnything - это молниеносный, легкий конвейер встраивания, построенный на Rust, который поддерживает текст, PDF, изображения, аудио и многое другое.

В этом руководстве мы покажем, как построить конвейер Retrieval-Augmented Generation (RAG), используя EmbedAnything вместе с Milvus. Вместо жесткой связи с конкретной базой данных EmbedAnything использует подключаемую систему адаптеров - адаптеры служат обертками, которые определяют, как эмбеддинги форматируются, индексируются и хранятся в целевом векторном хранилище.

Соединив EmbedAnything с адаптером Milvus, вы сможете генерировать эмбеддинги из различных типов файлов и эффективно хранить их в Milvus всего за несколько строк кода.

⚠️ Примечание: Хотя адаптер в EmbedAnything обрабатывает вставку в Milvus, он не поддерживает поиск из коробки. Чтобы построить полноценный конвейер RAG, вам нужно будет отдельно создать MilvusClient и реализовать логику поиска (например, поиск сходства по векторам) в своем приложении.

Подготовка

Зависимости и окружение

$ pip install -qU pymilvus milvus-lite openai embed_anything

Клонирование репозитория и загрузка адаптера

Далее мы клонируем репозиторий EmbedAnything и добавим каталог examples/adapters в путь Python. Здесь хранится пользовательская реализация адаптера Milvus, которая позволяет EmbedAnything взаимодействовать с Milvus для вставки векторов.

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.

Мы будем использовать OpenAI в качестве LLM в этом RAG-конвейере. Вы должны подготовить api ключ OPENAI_API_KEY в качестве переменной окружения.

import os
from openai import OpenAI

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

openai_client = OpenAI()

Сборка RAG

Инициализация Milvus

Прежде чем встраивать какие-либо файлы, нам нужно подготовить два компонента, которые взаимодействуют с Milvus:

1. MilvusVectorAdapter - Это адаптер Milvus для EmbedAnything, который используется только для встраивания векторов (т.е. вставки вложений и создания индексов). В настоящее время он не поддерживает операции поиска.
2. MilvusClient - Это официальный клиент с сайта pymilvus, который обеспечивает полный доступ к возможностям Milvus, включая векторный поиск, фильтрацию и управление коллекциями.

Чтобы избежать путаницы:

• Считайте, что MilvusVectorAdapter - это инструмент для хранения векторов "только для записи".
• Думайте о MilvusClient как о механизме "чтения и поиска" для выполнения запросов и получения документов для RAG.

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.

Инициализация модели встраивания и встраивание PDF-документа

Теперь мы инициализируем модель встраивания. Мы воспользуемся моделью all-MiniLM-L12-v2 model из библиотеки sentence-transformers, которая представляет собой легкую, но мощную модель для генерации вкраплений текста. Она создает 384-мерные вкрапления, так что это соответствует размерности нашей коллекции Milvus, установленной на 384. Это выравнивание очень важно и обеспечивает совместимость между размерами векторов, хранящихся в Milvus, и теми, которые генерирует модель.

EmbedAnything поддерживает гораздо больше моделей встраивания. За более подробной информацией обращайтесь к официальной документации.

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

Теперь давайте встроим PDF-файл. EmbedAnything позволяет легко обрабатывать PDF (и многие другие) документы и хранить их вложения непосредственно в Milvus.

# 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.

Получение и генерация ответа

Опять же, MilvusVectorAdapter от EmbedAnything в настоящее время представляет собой легковесную абстракцию, предназначенную только для векторного ввода и индексирования. Она не поддерживает запросы на поиск или извлечение информации. Поэтому для поиска релевантных документов для построения нашего конвейера RAG мы должны напрямую использовать экземпляр MilvusClient (milvus_client) для запроса к нашему векторному хранилищу Milvus.

Определите функцию для получения соответствующих документов из 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

Определите функцию для генерации ответа с использованием полученных документов в конвейере RAG.

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

Давайте протестируем конвейер RAG на примере вопроса.

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.