使用 Milvus、vLLM 和 Llama 3.1 搭建 RAG 应用

vLLM 是一个简单易用的 LLM 推理服务库。加州大学伯克利分校于 2024 年 7 月将 vLLM 作为孵化项目正式捐赠给 LF AI & Data Foundation 基金会。欢迎 vLLM 加入 LF AI & Data 大家庭！🎉

在主流的 AI 应用架构中，大语言模型（LLM）通常与向量数据库配套使用，用于构建检索增强生成（RAG）应用，从而解决 AI 幻觉问题。本文将介绍如何使用 Milvus、vLLM 和 Llama 3.1 构建并运行RAG 应用。我们将详细展示如何将文本信息转换为 Embedding 向量并存储到 Milvus 向量数据库中、如何将 Milvus 作为知识库有效检索与用户问题相关的文本块。最后，我们将通过 vLLM 使用 Meta的Llama 3.1-8B 模型生成答案。

Milvus、vLLM 和 Llama 3.1 简介

Milvus 向量数据库

Milvus 是一款开源的分布式向量数据库，可用于存储、索引和搜索向量数据，适用于生成式 AI（GenAI）应用。Milvus 支持 hybrid search、元数据过滤、重排（Reranking），能够高效处理万亿规模的向量，助力开发者搭建 AI 和 ML 应用。您可以在本地运行 Milvus standalone 或 cluster 版本，或者使用全托管的 Milvus 服务——Zilliz Cloud。

vLLM

vLLM 是加州大学伯克利分校 SkyLab 推出的一个开源项目，专注于优化 LLM 服务性能。通过高效的内存管理技术，如 PagedAttention、持续批处理和优化 CUDA 内核，vLLm 与传统方法相比将服务性能提高了多达 24 倍，同时将 GPU 内存用量减少了一半。

根据论文《Efficient Memory Management for Large Language Model Serving with PagedAttention》，KV 缓存使用约 30% 的 GPU 内存，可能会导致内存问题。KV 缓存存储在连续内存（contiguous memory）中，但内存变化可能导致内存碎片化，不利于计算效率。

通过使用虚拟内存缓存 KV，vLLM 只需要在必要时分配物理 GPU 内存，有效避免了内存碎片化和内存预分配。在测试中，vLLM 的吞吐量比 HuggingFace Transformers (HF) 高出多达 24 倍，比基于NVIDIA A10G 和 A100 GPU 的 Text Generation Inference (TGI) 高出 3.5 倍。

Meta Llama 3.1

Meta 于 2024 年 7 月 23 日宣布推出 Llama 3.1，允许用于多种商业用途。其 405B 模型（4050 亿参数）在多个公开的性能测试中均展示出了最出色的性能，并支持 128,000 个输入 Token 的上下文窗口。除了 405B 模型外，Meta 还发布了 Llama3.1 70B（700 亿参数）和 8B（80 亿参数）模型。您可以通过 Meta 官网下载模型权重（model weight）。

需要注意微调生成的数据可以提高模型性能，但低质量的数据可能会降低模型性能。Llama 团队已不断识别和去除这些低质量的数据，使用模型本身及其他辅助模型和工具，进一步优化模型。

使用 Milvus 搭建 RAG-Retrieval部分

准备数据

本教程将使用 Milvus 文档作为数据集。我们需要先下载并本地保存 Milvus 文档。

from langchain.document_loaders import DirectoryLoader# Load HTML files already saved in a local directorypath = "../../RAG/rtdocs_new/"global_pattern = '*.html'loader = DirectoryLoader(path=path, glob=global_pattern)docs = loader.load()# Print num documents and a preview.print(f"loaded {len(docs)} documents")print(docs[0].page_content)pprint.pprint(docs[0].metadata)

下载 Embedding 模型

接着，从 HuggingFace 上下载一个免费的开源 Embedding 模型。

import torchfrom sentence_transformers import SentenceTransformer# Initialize torch settings for device-agnostic code.N_GPU = torch.cuda.device_count()DEVICE = torch.device('cuda:N_GPU' if torch.cuda.is_available() else 'cpu')# Download the model from huggingface model hub.model_name = "BAAI/bge-large-en-v1.5"encoder = SentenceTransformer(model_name, device=DEVICE)# Get the model parameters and save for later.EMBEDDING_DIM = encoder.get_sentence_embedding_dimension()MAX_SEQ_LENGTH_IN_TOKENS = encoder.get_max_seq_length()# Inspect model parameters.print(f"model_name: {model_name}")print(f"EMBEDDING_DIM: {EMBEDDING_DIM}")print(f"MAX_SEQ_LENGTH: {MAX_SEQ_LENGTH}")

切分数据并编码为向量

将文档数据切分成固定长度（512 个字符）的文本块，并将切分 overlap 设置为 10%。

from langchain.text_splitter import RecursiveCharacterTextSplitterCHUNK_SIZE = 512chunk_overlap = np.round(CHUNK_SIZE * 0.10, 0)print(f"chunk_size: {CHUNK_SIZE}, chunk_overlap: {chunk_overlap}")# Define the splitter.child_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE,chunk_overlap=chunk_overlap)# Chunk the docs.chunks = child_splitter.split_documents(docs)print(f"{len(docs)} docs split into {len(chunks)} child documents.")# Encoder input is doc.page_content as strings.list_of_strings = [doc.page_content for doc in chunks if hasattr(doc, 'page_content')]# Embedding inference using HuggingFace encoder.embeddings = torch.tensor(encoder.encode(list_of_strings))# Normalize the embeddings.embeddings = np.array(embeddings / np.linalg.norm(embeddings))# Milvus expects a list of `numpy.ndarray` of `numpy.float32` numbers.converted_values = list(map(np.float32, embeddings))# Create dict_list for Milvus insertion.dict_list = []for chunk, vector in zip(chunks, converted_values):# Assemble embedding vector, original text chunk, metadata.chunk_dict = {'chunk': chunk.page_content,'source': chunk.metadata.get('source', ""),'vector': vector,}dict_list.append(chunk_dict)

将向量数据存储在 Milvus 中

将向量存储到 Milvus 向量数据库中。

# Connect a client to the Milvus Lite server.from pymilvus import MilvusClientmc = MilvusClient("milvus_demo.db")# Create a collection with flexible schema and AUTOINDEX.COLLECTION_NAME = "MilvusDocs"mc.create_collection(COLLECTION_NAME,EMBEDDING_DIM,consistency_level="Eventually",auto_id=True, overwrite=True)# Insert data into the Milvus collection.print("Start inserting entities")start_time = time.time()mc.insert(COLLECTION_NAME,data=dict_list,progress_bar=True)end_time = time.time()print(f"Milvus insert time for {len(dict_list)} vectors: ", end="")print(f"{round(end_time - start_time, 2)} seconds")

进行向量搜索

输入问题，并在 Milvus 知识库中搜索与问题最相似的文本块。

SAMPLE_QUESTION = "What do the parameters for HNSW mean?"# Embed the question using the same encoder.query_embeddings = torch.tensor(encoder.encode(SAMPLE_QUESTION))# Normalize embeddings to unit length.query_embeddings = F.normalize(query_embeddings, p=2, dim=1)# Convert the embeddings to list of list of np.float32.query_embeddings = list(map(np.float32, query_embeddings))# Define metadata fields you can filter on.OUTPUT_FIELDS = list(dict_list[0].keys())OUTPUT_FIELDS.remove('vector')# Define how many top-k results you want to retrieve.TOP_K = 2# Run semantic vector search using your query and the vector database.results = mc.search(COLLECTION_NAME,data=query_embeddings,output_fields=OUTPUT_FIELDS,limit=TOP_K,consistency_level="Eventually")

搜索结果如下所示：

使用 vLLM 和 Llama 3.1-8B 搭建 RAG-Generation 部分

安装 vLLM 与 HuggingFace 模型

vLLM 默认从 HuggingFace 下载大语言模型。通常情况下，如果您想使用 HuggingFace 上的新模型，需要执行 pip install --update 或 -U。此外，我们还需要 GPU 通过 vLLM 来运行 Meta 的 Llama 3.1 推理模型。

# (Recommended) Create a new conda environment.conda create -n myenv python=3.11 -yconda activate myenv# Install vLLM with CUDA 12.1.pip install -U vllm transformers torch

import vllm, torchfrom vllm import LLM, SamplingParams# Clear the GPU memory cache.torch.cuda.empty_cache()# Check the GPU.!nvidia-smi

获取 HuggingFace token

HuggingFace 上的部分模型（如 Meta Llama 3.1）要求用户在下载前接受其许可证。因此，您必须先创建一个 HuggingFace 帐户，接受模型的许可证，并生成一个 Token。

在 HuggingFace 的 Llama3.1 页上，您会收到一条消息要求您同意条款。单击"Accept License"以接受 Meta 条款，然后再下载模型权重。审批流程通常可以在一天内完成。

审批通过后，需要生成一个新的 HuggingFace token。旧 Token 无法使用。

在安装 vLLM 之前，请使用您的新 Token 登录 HuggingFace。以下示例代码中使用 Colab Secrets 来存储 Token。

# Login to HuggingFace using your new token.from huggingface_hub import loginfrom google.colab import userdatahf_token = userdata.get('HF_TOKEN')login(token = hf_token, add_to_git_credential=True)

运行 RAG-Generation 部分

我们需要 GPU 和较大的内存来运行 Llama-3.1-8B 模型。以下示例是在 Google Colab Pro上使用 A100 GPU 运行的。

# 1. Choose a modelMODELTORUN = "meta-llama/Meta-Llama-3.1-8B-Instruct"# 2. Clear the GPU memory cache, you're going to need it all!torch.cuda.empty_cache()# 3. Instantiate a vLLM model instance.llm = LLM(model=MODELTORUN,enforce_eager=True,dtype=torch.bfloat16,gpu_memory_utilization=0.5,max_model_len=1000,seed=415,max_num_batched_tokens=3000)

# Separate all the context together by space.contexts_combined = ' '.join(contexts)# Lance Martin, LangChain, says put the best contexts at the end.contexts_combined = ' '.join(reversed(contexts))# Separate all the unique sources together by comma.source_combined = ' '.join(reversed(list(dict.fromkeys(sources))))SYSTEM_PROMPT = f"""First, check if the provided Context is relevant tothe user's question.  Second, only if the provided Context is strongly relevant, answer the question using the Context.  Otherwise, if the Context is not strongly relevant, answer the question without using the Context. Be clear, concise, relevant.  Answer clearly, in fewer than 2 sentences.Grounding sources: {source_combined}Context: {contexts_combined}User's question: {SAMPLE_QUESTION}"""prompts = [SYSTEM_PROMPT]

使用从 Milvus 中检索获得的上下文和原始提问来编写提示，并生成回答。

# Sampling parameterssampling_params = SamplingParams(temperature=0.2, top_p=0.95)# Invoke the vLLM model.outputs = llm.generate(prompts, sampling_params)# Print the outputs.for output in outputs:prompt = output.promptgenerated_text = output.outputs[0].text# !r calls repr(), which prints a string inside quotes.print()print(f"Question: {SAMPLE_QUESTION!r}")pprint.pprint(f"Generated text: {generated_text!r}")