使用 Hugging Face 和 Milvus 构建 RAG 系统

前言

Milvus 是一个广受欢迎的开源向量数据库，为人工智能应用提供高性能和可扩展的向量相似性搜索。在本教程中，我们将向您展示如何使用 Hugging Face 和 Milvus 构建 RAG（检索增强生成）流程。

RAG 系统将检索系统与 LLM 相结合。该系统首先使用向量数据库 Milvus 从语料库中检索相关文档，然后使用 Hugging Face 的 LLM 根据检索到的文档生成回答。

01.

准备

依赖环境关系

! pip install --upgrade pymilvus sentence-transformers huggingface-hub langchain_community langchain-text-splitters pypdf tqdm

如果您使用 Google Colab，要启用刚刚安装的依赖项，您可能需要 重新启动运行时（单击屏幕顶部的“运行时”菜单，然后从下拉菜单中选择“重新启动会话”） 。

另外，我们建议您配置您的Hugging Face User Access Token，并将其设置在您的环境变量中，因为我们将使用 Hugging Face Hub 的 LLM 。如果不设置 token 环境变量，您可能会受到较低的请求速率限制。

import os

os.environ["HF_TOKEN"] = "hf_..."

准备数据

我们使用AI Act PDF作为我们 RAG 中的私有知识。这是一个针对人工智能的监管框架，具有不同风险级别，对应或多或少的监管。

%%bash

if [ ! -f "The-AI-Act.pdf" ]; then
    wget -q https://artificialintelligenceact.eu/wp-content/uploads/2021/08/The-AI-Act.pdf
fi

我们使用 LangChain 的 PyPDFLoader 从 PDF 中提取文本，然后将文本拆分为更小的块。默认情况下，我们将块大小设置为 1000，重叠设置为 200，这意味着每个块将有近 1000 个字符，两个块之间的重叠将是 200 个字符。

from langchain_community.document_loaders import PyPDFLoader

loader = PyPDFLoader("The-AI-Act.pdf")
docs = loader.load()
print(len(docs))

from langchain_text_splitters import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
chunks = text_splitter.split_documents(docs)

text_lines = [chunk.page_content for chunk in chunks]

准备 embedding 模型

定义一个函数来生成文本 embedding。我们使用BGE embedding模型 作为示例，但您也可以将其更改为任何其他 embedding 模型，例如在MTEB排行榜上选择。

from sentence_transformers import SentenceTransformer

embedding_model = SentenceTransformer("BAAI/bge-small-en-v1.5")

def emb_text(text):
    return embedding_model.encode([text], normalize_embeddings=True).tolist()[0]

生成测试 embedding 并打印其大小和前几个元素。

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

02.

将数据加载到Milvus中

创建 collection

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./hf_milvus_demo.db")

collection_name = "rag_collection"

对于 MilvusClient 的参数：

• 将 uri 设置为本地文件，例如 ./hf_milvus_demo.db ，是最方便的方法，因为它会自动使用 Milvus Lite 将所有数据存储在此文件中。
• 如果您有大量数据，例如超过一百万个向量，您可以在 Docker 或 Kubernetes 上设置性能更高的 Milvus 服务器。在此设置中，请使用服务器 uri，例如 http://localhost:19530 作为您的 uri 。
• 如果您想使用 Milvus 的全托管云服务 Zilliz Cloud ，请调整 uri 和 token，分别对应 Zilliz Cloud 中 Public Endpoint 和 Api key 。

检查 collection 是否已存在，如果存在则将其删除。

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

使用指定参数创建一个新 collection。

如果我们不指定任何字段信息，Milvus 会自动创建一个默认的 id 字段作为主键，并创建一个 vector 字段来存储向量数据。保留的 JSON 字段用于存储未定义 schema 的字段及其值。

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Strong consistency level
)

插入数据

迭代文本行，创建 embedding，然后将数据插入到 Milvus 中。

这是一个新字段 text，它是 collection schema 中的未定义字段。它将自动添加到保留的 JSON 动态字段中，从更高的层次看来，它可被视为普通字段。

from tqdm import tqdm

data = []

for i, line in enumerate(tqdm(text_lines, desc="Creating embeddings")):
    data.append({"id": i, "vector": emb_text(line), "text": line})

insert_res = milvus_client.insert(collection_name=collection_name, data=data)
insert_res["insert_count"]

03.

构建RAG

检索查询数据

让我们具体指定一个有关该语料的问题。

question = "What is the legal basis for the proposal?"

在 collection 中搜索问题并检索语义前 3 个匹配项。

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[
        emb_text(question)
    ],  # Use the `emb_text` function to convert the question to an embedding vector
    limit=3,  # Return top 3 results
    search_params={"metric_type": "IP", "params": {}},  # Inner product distance
    output_fields=["text"],  # Return the text field
)

我们来看看查询的搜索结果

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

使用 LLM 获取 RAG 回答

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

定义语言模型的提示词。该提示与从 Milvus 检索到的文档组合在一起。

PROMPT = """
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>
"""

我们使用部署在 Hugging Face 推理服务上的 Mixtral-8x7B-Instruct-v0.1 根据提示词生成响应。

from huggingface_hub import InferenceClient

repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"

llm_client = InferenceClient(model=repo_id, timeout=120)

我们设置提示词格式并生成最后的回答。

prompt = PROMPT.format(context=context, question=question)

answer = llm_client.text_generation(
    prompt,
    max_new_tokens=1000,
).strip()
print(answer)

恭喜！您已经使用 Hugging Face 和 Milvus 构建了 RAG 流程。