大模型教程：使用 Milvus、vLLM 和 Llama 3.1 搭建 RAG 应用

下载 Embedding 模型

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

import torch
from 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 RecursiveCharacterTextSplitter

CHUNK_SIZE = 512
chunk_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 MilvusClient
mc = 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")

03. 使用 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 -y
conda activate myenv

# Install vLLM with CUDA 12.1.
pip install -U vllm transformers torch

import vllm, torch
from 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。以下示例代码中使用环境变量来存储 Token。

# Login to HuggingFace using your new token.
from huggingface_hub import login
import os
hf_token = os.getenv('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 model
MODELTORUN = "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 to
the 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 parameters
sampling_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.prompt
   generated_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}")

04. 总结与最佳实践

本文详细介绍了如何利用 Milvus 向量数据库、vLLM 推理引擎以及 Llama 3.1 大模型构建一个完整的 RAG（检索增强生成）系统。通过上述步骤，我们实现了从非结构化文本数据的加载、向量化、存储，到基于语义的检索，再到利用大模型生成精准回答的全流程。

架构优势

1. 高性能检索：Milvus 提供了毫秒级的向量搜索能力，支持海量数据的高效索引，确保在大规模知识库下仍能快速定位相关信息。
2. 高效推理：vLLM 的 PagedAttention 技术显著降低了显存占用并提升了吞吐量，使得在消费级或企业级 GPU 上部署大模型成为可能。
3. 准确生成：结合 Llama 3.1 强大的指令遵循能力和上下文理解力，配合 RAG 提供的真实上下文，有效减少了模型幻觉，提升了回答的可信度。

部署建议

• 环境配置：建议在 Linux 环境下使用 Docker 容器化部署 Milvus 和 vLLM，以便于管理和扩展。
• 监控与日志：在生产环境中，务必接入 Prometheus 和 Grafana 等监控工具，实时跟踪向量查询延迟、GPU 利用率及模型响应时间。
• 数据安全：对于敏感数据，建议在向量入库前进行脱敏处理，并确保 Milvus 集群的网络访问权限受到严格控制。
• 持续优化：定期评估 RAG 系统的召回率和准确率，根据反馈调整切片大小（Chunk Size）、重叠率（Overlap）以及检索策略（Top-K）。

通过合理组合这些组件，开发者可以构建出既具备专业知识又拥有高响应速度的智能问答系统，广泛应用于客服机器人、内部知识助手及数据分析场景。