使用 Milvus 和 BentoML 进行检索增强生成 (RAG)

介绍

本指南演示如何在 BentoCloud 上使用开源嵌入模型和大语言模型，结合 Milvus 向量数据库构建 RAG (检索增强生成) 应用程序。

BentoCloud 是一个为快速发展的 AI 团队提供的 AI 推理平台，提供专门为模型推理定制的完全托管基础设施。它与 BentoML（一个开源模型服务框架）配合使用，以促进高性能模型服务的轻松创建和部署。在本演示中，我们使用 Milvus Lite 作为向量数据库，它是可以嵌入到您的 Python 应用程序中的 Milvus 轻量级版本。

开始之前

Milvus Lite 在 PyPI 上可用。您可以通过 pip 为 Python 3.8+ 安装它：

$ pip install -U pymilvus bentoml

如果您使用 Google Colab，为了启用刚安装的依赖项，您可能需要重启运行时（点击屏幕顶部的"Runtime"菜单，并从下拉菜单中选择"Restart session"）。

使用 BentoML/BentoCloud 服务嵌入

要使用此端点，导入 bentoml 并通过指定端点和可选的 token（如果您在 BentoCloud 上开启了 Endpoint Authorization）使用 SyncHTTPClient 设置 HTTP 客户端。或者，您可以使用通过 BentoML 服务的相同模型，使用其 Sentence Transformers Embeddings 代码仓库。

import bentoml

BENTO_EMBEDDING_MODEL_END_POINT = "BENTO_EMBEDDING_MODEL_END_POINT"
BENTO_API_TOKEN = "BENTO_API_TOKEN"

embedding_client = bentoml.SyncHTTPClient(
    BENTO_EMBEDDING_MODEL_END_POINT, token=BENTO_API_TOKEN
)

一旦我们连接到 embedding_client，我们需要处理我们的数据。我们提供了几个函数来执行数据分割和嵌入。

读取文件并将文本预处理为字符串列表。

# 简单地按换行符分块
def chunk_text(filename: str) -> list:
    with open(filename, "r") as f:
        text = f.read()
    sentences = text.split("\n")
    return sentences

首先我们需要下载城市数据。

import os
import requests
import urllib.request

# 设置数据源
repo = "ytang07/bento_octo_milvus_RAG"
directory = "data"
save_dir = "./city_data"
api_url = f"https://api.github.com/repos/{repo}/contents/{directory}"


response = requests.get(api_url)
data = response.json()

if not os.path.exists(save_dir):
    os.makedirs(save_dir)

for item in data:
    if item["type"] == "file":
        file_url = item["download_url"]
        file_path = os.path.join(save_dir, item["name"])
        urllib.request.urlretrieve(file_url, file_path)

接下来，我们处理我们拥有的每个文件。

# 请在此文件的文件夹下上传您的数据目录
cities = os.listdir("city_data")
# 将每个城市的分块文本存储在字典列表中
city_chunks = []
for city in cities:
    chunked = chunk_text(f"city_data/{city}")
    cleaned = []
    for chunk in chunked:
        if len(chunk) > 7:
            cleaned.append(chunk)
    mapped = {"city_name": city.split(".")[0], "chunks": cleaned}
    city_chunks.append(mapped)

将字符串列表分割为嵌入列表，每组 25 个文本字符串。

def get_embeddings(texts: list) -> list:
    if len(texts) > 25:
        splits = [texts[x : x + 25] for x in range(0, len(texts), 25)]
        embeddings = []
        for split in splits:
            embedding_split = embedding_client.encode(sentences=split)
            embeddings += embedding_split
        return embeddings
    return embedding_client.encode(
        sentences=texts,
    )

现在，我们需要匹配嵌入和文本块。由于嵌入列表和句子列表应该按索引匹配，我们可以通过 enumerate 任一列表来匹配它们。

entries = []
for city_dict in city_chunks:
    # 如果 get_embeddings 已经返回列表的列表，则不需要嵌入列表
    embedding_list = get_embeddings(city_dict["chunks"])  # 返回列表的列表
    # 现在匹配文本与嵌入和城市名称
    for i, embedding in enumerate(embedding_list):
        entry = {
            "embedding": embedding,
            "sentence": city_dict["chunks"][
                i
            ],  # 假设 "chunks" 具有嵌入的相应文本
            "city": city_dict["city_name"],
        }
        entries.append(entry)
    print(entries)

将数据插入向量数据库以进行检索

准备好嵌入和数据后，我们可以将向量与元数据一起插入 Milvus Lite 中，以便稍后进行向量搜索。本节的第一步是通过连接到 Milvus Lite 来启动客户端。

我们只需导入 MilvusClient 模块并初始化连接到您的 Milvus Lite 向量数据库的 Milvus Lite 客户端。维度大小来自嵌入模型的大小，例如 Sentence Transformer 模型 all-MiniLM-L6-v2 产生 384 维的向量。

from pymilvus import MilvusClient

COLLECTION_NAME = "Bento_Milvus_RAG"  # Collection 的随机名称
DIMENSION = 384

# 初始化 Milvus Lite 客户端
milvus_client = MilvusClient("milvus_demo.db")

关于 MilvusClient 的参数：

将 uri 设置为本地文件，例如 ./milvus.db，是最方便的方法，因为它会自动利用 Milvus Lite 将所有数据存储在此文件中。
如果您有大规模数据，可以在 docker 或 kubernetes 上设置更高性能的 Milvus 服务器。在此设置中，请使用服务器 uri，例如 http://localhost:19530，作为您的 uri。
如果您想使用 Zilliz Cloud，Milvus 的完全托管云服务，请调整 uri 和 token，它们对应于 Zilliz Cloud 中的公共端点和 API 密钥。

或者使用旧的 connections.connect API（不推荐）：

from pymilvus import connections

connections.connect(uri="milvus_demo.db")

创建您的 Milvus Lite Collection

使用 Milvus Lite 创建 Collection 涉及两个步骤：首先，定义模式，其次，定义索引。对于本节，我们需要一个模块：DataType 告诉我们 Field 中将包含什么类型的数据。我们还需要使用两个函数来创建模式和添加 Field。create_schema()：创建 Collection 模式，add_field()：向 Collection 的模式添加 Field。

from pymilvus import MilvusClient, DataType, Collection

# 创建模式
schema = MilvusClient.create_schema(
    auto_id=True,
    enable_dynamic_field=True,
)

# 3.2. 向模式添加 Field
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=DIMENSION)

现在我们已经创建了模式并成功定义了数据 Field，我们需要定义索引。在搜索方面，"索引"定义了我们将如何映射数据以进行检索。我们使用默认选择 AUTOINDEX 为此项目索引我们的数据。

接下来，我们使用之前给定的名称、模式和索引创建 Collection。最后，我们插入之前处理的数据。

# 准备索引参数
index_params = milvus_client.prepare_index_params()

# 添加索引
index_params.add_index(
    field_name="embedding",
    index_type="AUTOINDEX",  # 使用 autoindex 而不是其他复杂的索引方法
    metric_type="COSINE",  # L2, COSINE, 或 IP
)

# 创建 Collection
if milvus_client.has_collection(collection_name=COLLECTION_NAME):
    milvus_client.drop_collection(collection_name=COLLECTION_NAME)
milvus_client.create_collection(
    collection_name=COLLECTION_NAME, schema=schema, index_params=index_params
)

# 在循环外，现在您一次性 upsert 所有条目
milvus_client.insert(collection_name=COLLECTION_NAME, data=entries)

为 RAG 设置您的 LLM

要构建 RAG 应用程序，我们需要在 BentoCloud 上部署 LLM。让我们使用最新的 Llama3 LLM。一旦它启动并运行，只需复制此模型服务的端点和 token 并为其设置客户端。

BENTO_LLM_END_POINT = "BENTO_LLM_END_POINT"

llm_client = bentoml.SyncHTTPClient(BENTO_LLM_END_POINT, token=BENTO_API_TOKEN)

LLM 指令

现在，我们使用提示、上下文和问题设置 LLM 指令。这是一个作为 LLM 的函数，然后以字符串格式从客户端返回输出。

def dorag(question: str, context: str):

    prompt = (
        f"You are a helpful assistant. The user has a question. Answer the user question based only on the context: {context}. \n"
        f"The user question is {question}"
    )

    results = llm_client.generate(
        max_tokens=1024,
        prompt=prompt,
    )

    res = ""
    for result in results:
        res += result

    return res

RAG 示例

现在我们准备好提问了。此函数只需接受一个问题，然后执行 RAG 以从背景信息生成相关上下文。然后，我们将上下文和问题传递给 dorag() 并获得结果。

question = "What state is Cambridge in?"


def ask_a_question(question):
    embeddings = get_embeddings([question])
    res = milvus_client.search(
        collection_name=COLLECTION_NAME,
        data=embeddings,  # 搜索作为列表的列表返回的一个 (1) 嵌入
        anns_field="embedding",  # 跨嵌入搜索
        limit=5,  # 给我前 5 个结果
        output_fields=["sentence"],  # 获取句子/块和城市
    )

    sentences = []
    for hits in res:
        for hit in hits:
            print(hit)
            sentences.append(hit["entity"]["sentence"])
    context = ". ".join(sentences)
    return context


context = ask_a_question(question=question)
print(context)

实现 RAG

print(dorag(question=question, context=context))

对于询问剑桥在哪个州的示例问题，我们可以打印来自 BentoML 的整个响应。但是，如果我们花时间解析它，它看起来更好，它应该告诉我们剑桥位于马萨诸塞州。

介绍​

开始之前​

使用 BentoML/BentoCloud 服务嵌入​

将数据插入向量数据库以进行检索​

创建您的 Milvus Lite Collection​

为 RAG 设置您的 LLM​

LLM 指令​

RAG 示例​

介绍