Создайте RAG с помощью Milvus и Docling

Docling упрощает разбор и понимание документов различных форматов для приложений искусственного интеллекта. Благодаря расширенному пониманию PDF и унифицированному представлению документов Docling делает неструктурированные данные документов готовыми для последующих рабочих процессов.

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

Подготовка

Зависимости и среда

Для начала установите необходимые зависимости, выполнив следующую команду:

$ pip install --upgrade pymilvus milvus-lite docling openai

Настройка ключей API

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

import os

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

Подготовка LLM и модели встраивания

Мы инициализируем клиент OpenAI, чтобы подготовить модель встраивания.

from openai import OpenAI

openai_client = OpenAI()

Определите функцию для генерации текстовых вкраплений с помощью клиента OpenAI. В качестве примера мы используем модель text-embedding-3-small.

def emb_text(text):
    return (
        openai_client.embeddings.create(input=text, model="text-embedding-3-small")
        .data[0]
        .embedding
    )

Сгенерируйте тестовое вкрапление и выведите его размерность и первые несколько элементов.

test_embedding = emb_text("This is a test")
embedding_dim = len(test_embedding)
print(embedding_dim)
print(test_embedding[:10])

1536
[0.00988506618887186, -0.005540902726352215, 0.0068014683201909065, -0.03810417652130127, -0.018254263326525688, -0.041231658309698105, -0.007651153020560741, 0.03220026567578316, 0.01892443746328354, 0.00010708322952268645]

Обработка данных с помощью Docling

Docling может анализировать различные форматы документов в единое представление (Docling Document), которое затем можно экспортировать в различные выходные форматы. Полный список поддерживаемых входных и выходных форматов можно найти в официальной документации.

В этом руководстве мы будем использовать файл Markdown(исходный) в качестве входного. Мы обработаем документ с помощью HierarchicalChunker, предоставленного Docling, чтобы создать структурированные, иерархические фрагменты, подходящие для последующих задач RAG.

from docling.document_converter import DocumentConverter
from docling_core.transforms.chunker import HierarchicalChunker

converter = DocumentConverter()
chunker = HierarchicalChunker()

# Convert the input file to Docling Document
source = "https://milvus.io/docs/overview.md"
doc = converter.convert(source).document

# Perform hierarchical chunking
texts = [chunk.text for chunk in chunker.chunk(doc)]

for i, text in enumerate(texts[:5]):
    print(f"Chunk {i+1}:\n{text}\n{'-'*50}")

Chunk 1:
Milvus is a high-performance, highly scalable vector database that runs efficiently across a wide range of environments, from a laptop to large-scale distributed systems. It is available as both open-source software and a cloud service.
--------------------------------------------------
Chunk 2:
Milvus is an open-source project under LF AI & Data Foundation distributed under the Apache 2.0 license. Most contributors are experts from the high-performance computing (HPC) community, specializing in building large-scale systems and optimizing hardware-aware code. Core contributors include professionals from Zilliz, ARM, NVIDIA, AMD, Intel, Meta, IBM, Salesforce, Alibaba, and Microsoft.
--------------------------------------------------
Chunk 3:
Unstructured data, such as text, images, and audio, varies in format and carries rich underlying semantics, making it challenging to analyze. To manage this complexity, embeddings are used to convert unstructured data into numerical vectors that capture its essential characteristics. These vectors are then stored in a vector database, enabling fast and scalable searches and analytics.
--------------------------------------------------
Chunk 4:
Milvus offers robust data modeling capabilities, enabling you to organize your unstructured or multi-modal data into structured collections. It supports a wide range of data types for different attribute modeling, including common numerical and character types, various vector types, arrays, sets, and JSON, saving you from the effort of maintaining multiple database systems.
--------------------------------------------------
Chunk 5:
Untructured data, embeddings, and Milvus
--------------------------------------------------

Загрузка данных в Milvus

Создайте коллекцию

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./milvus_demo.db")
collection_name = "my_rag_collection"

Проверьте, не существует ли уже коллекция, и удалите ее, если она существует.

if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

Создайте новую коллекцию с указанными параметрами.

Если мы не укажем информацию о полях, Milvus автоматически создаст поле по умолчанию id для первичного ключа и поле vector для хранения векторных данных. Зарезервированное поле JSON используется для хранения не определенных схемой полей и их значений.

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
)

Вставка данных

from tqdm import tqdm

data = []

for i, chunk in enumerate(tqdm(texts, desc="Processing chunks")):
    embedding = emb_text(chunk)
    data.append({"id": i, "vector": embedding, "text": chunk})

milvus_client.insert(collection_name=collection_name, data=data)

Processing chunks: 100%|██████████| 36/36 [00:18<00:00,  1.96it/s]





{'insert_count': 36, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35], 'cost': 0}

Построить RAG

Получение данных для запроса

Давайте зададим вопрос о сайте, который мы только что соскоблили.

question = (
    "What are the three deployment modes of Milvus, and what are their differences?"
)

Найдем этот вопрос в коллекции и получим семантический топ-3 совпадений.

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[emb_text(question)],
    limit=3,
    search_params={"metric_type": "IP", "params": {}},
    output_fields=["text"],
)

Давайте посмотрим на результаты поиска по этому запросу.

import json

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]
print(json.dumps(retrieved_lines_with_distances, indent=4))

[
    [
        "Milvus offers three deployment modes, covering a wide range of data scales\u2014from local prototyping in Jupyter Notebooks to massive Kubernetes clusters managing tens of billions of vectors:",
        0.6503741145133972
    ],
    [
        "Milvus Lite is a Python library that can be easily integrated into your applications. As a lightweight version of Milvus, it\u2019s ideal for quick prototyping in Jupyter Notebooks or running on edge devices with limited resources. Learn more.\nMilvus Standalone is a single-machine server deployment, with all components bundled into a single Docker image for convenient deployment. Learn more.\nMilvus Distributed can be deployed on Kubernetes clusters, featuring a cloud-native architecture designed for billion-scale or even larger scenarios. This architecture ensures redundancy in critical components. Learn more.",
        0.6281254291534424
    ],
    [
        "What is Milvus?\nUnstructured Data, Embeddings, and Milvus\nWhat Makes Milvus so Fast\uff1f\nWhat Makes Milvus so Scalable\nTypes of Searches Supported by Milvus\nComprehensive Feature Set",
        0.6117545962333679
    ]
]

Использование LLM для получения ответа RAG

Преобразуйте полученные документы в строковый формат.

context = "\n".join(
    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]
)

Определите системные и пользовательские подсказки для модели Lanage. Эта подсказка собрана из документов, полученных из Milvus.

SYSTEM_PROMPT = """
Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.
"""
USER_PROMPT = f"""
Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.
<context>
{context}
</context>
<question>
{question}
</question>
"""

Используйте OpenAI ChatGPT для генерации ответа на основе подсказок.

response = openai_client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ],
)
print(response.choices[0].message.content)

The three deployment modes of Milvus are Milvus Lite, Milvus Standalone, and Milvus Distributed. 

1. **Milvus Lite**: This is a Python library designed for easy integration into applications. It is lightweight and ideal for quick prototyping in Jupyter Notebooks or for use on edge devices with limited resources.

2. **Milvus Standalone**: This deployment mode involves a single-machine server with all components bundled into a single Docker image for convenient deployment.

3. **Milvus Distributed**: This mode can be deployed on Kubernetes clusters and is built for larger-scale scenarios, including managing billions of vectors. It features a cloud-native architecture that ensures redundancy in critical components, making it suited for extensive scalability.