为了账号安全,请及时绑定邮箱和手机立即绑定

使用Milvus、vLLM和Llama 3.1搭建检索增强生成系统

Blog Cover

加州大学伯克利分校于2024年7月将vLLM,一个快速且易于使用的大型语言模型推理和提供服务的库(vLLM),捐赠给LFAI & 数据基金会作为孵化项目的成员之一。作为基金会成员项目之一,我们非常高兴欢迎vLLM加入LF AI & Data大家庭!🎉

大型语言模型(LLMs)和向量数据库通常结合起来构建检索增强生成(RAG),它是一种流行的AI应用架构,用于解决AI幻觉的问题。这篇博客将向你展示如何使用Milvus、vLLM和Llama 3.1构建并运行一个RAG。具体来说,我将向你展示如何在Milvus中存储并嵌入文本信息作为向量嵌入。这个向量存储将作为知识库,高效检索与用户问题相关的文本片段。最后,我们将利用vLLM来提供服务Meta的Llama 3.1–8B模型,生成由检索到的文本增强的回答。让我们开始吧!

Milvus、VLLM 和 Meta 的 Llama 3.1 简介 Milvus 向量数据库

Milvus 是一个开源的、专门为处理向量而设计的分布式向量数据库,用于存储、索引和搜索 生成式人工智能 (GenAI) 工作负载中的向量。它能够执行 混合搜索元数据过滤、重排序以及高效处理数万亿的向量,因此对于 AI 和机器学习的任务来说,Milvus 是一个理想的选择。Milvus 可以在本地运行,也可以在集群中运行,或者托管于 Zilliz Cloud。

vLLM

vLLM 是一个由加州伯克利大学 SkyLab 发起的开源项目,专注于优化大型语言模型 (LLM) 的服务性能。它采用高效的内存管理技术,包括分页注意力 (PagedAttention)、连续批处理和优化的 CUDA 内核。与传统方法相比,vLLM 可将服务性能提升高达 24 倍,同时将 GPU 内存使用量减半。

根据论文“Efficient Memory Management for Large Language Model Serving with PagedAttention”,KV缓存大约占用了GPU内存的30%,可能会引发内存问题。KV缓存存储在连续的内存中,但其大小的调整会导致内存碎片,这对计算效率不利。

图1.现有系统的KV缓存的内存管理(2023 Paged Attention这篇论文)

通过使用虚拟内存来管理KV缓存,vLLM仅在需要时分配物理GPU内存,从而避免了内存碎片和预分配。在实际测试中,vLLM的性能超过了HuggingFace Transformers(HF)和Text Generation Inference(TGI),在NVIDIA A10G和A100 GPU上,其吞吐量比HF高出24倍,比TGI高出3.5倍。

_当每个请求需要三个并行输出完成时的服务吞吐量。vLLM 达到了比 HF 高 8.5–15 倍,比 TGI 高 3.3–3.5 倍(2023 年 【vLLM 博客】(https://blog.vllm.ai/2023/06/20/vllm.html) )。

Meta的LLAMA 3.1

Meta的Llama 3.1于2024年7月23日宣布。该4050亿参数的模型在多个公共基准测试中表现出色的性能,并具有128,000个输入令牌的上下文窗口,允许多种商业用途。除了4050亿参数的模型外,Meta还发布了Llama3 70B(70亿参数)和8B(8亿参数)的更新版本。模型权重可以在Meta的网站上下载这里

一个关键的洞察是,微调生成的数据可以提升性能,但质量差的例子会拖累性能。Llama团队努力识别并移除这些差的例子,利用模型本身、辅助模型以及其他工具。

使用Milvus构建和执行RAG检索 先来准备一下数据集吧

我采用了官方的Milvus文档作为本次演示的资料,并将其下载保存到本地。

    从langchain.document_loaders导入DirectoryLoader模块  
        # 加载本地目录中已保存的HTML文件,路径为 "../../RAG/rtdocs_new/"  
        path = "../../RAG/rtdocs_new/"  
        定义全局模式为 '*.html'  
        global_pattern = '*.html'  
        loader = DirectoryLoader(path=path, glob=global_pattern)  
        docs = loader.load()  

        # 打印文档数量和预览内容。  
        print(f"加载了 {len(docs)} 个文档")  
        print(docs[0].page_content)  
        pprint.pprint(docs[0].metadata)

loaded 22 docs

下载嵌入模型

下一步,从HuggingFace下载一个免费的开源embedding模型

    import torch  
    from sentence_transformers import SentenceTransformer  

    # 初始化适用于设备无关代码的 torch 设置。  
    N_GPU = torch.cuda.device_count()  
    DEVICE = torch.device('cuda:N_GPU' if torch.cuda.is_available() else 'cpu')  

    # 从 Huggingface 模型中心下载这个模型。  
    model_name = "BAAI/bge-large-en-v1.5"  
    encoder = SentenceTransformer(model_name, device=DEVICE)  

    # 获取模型参数并保存以便将来使用。  
    EMBEDDING_DIM = encoder.get_sentence_embedding_dimension()  
    MAX_SEQ_LENGTH_IN_TOKENS = encoder.get_max_seq_length()  

    # 检查模型参数。  
    print(f"model_name: {model_name}")  
    print(f"EMBEDDING_DIM: {EMBEDDING_DIM}")  
    print(f"MAX_SEQ_LENGTH_IN_TOKENS: {MAX_SEQ_LENGTH_IN_TOKENS}")

将自己的自定义数据切分成小块并编码成向量。

我将使用固定长度为512个字符,并让每部分有10%的重叠。

    从langchain.text_splitter导入RecursiveCharacterTextSplitter模块  

        CHUNK_SIZE = 512  
        chunk_overlap = np.round(CHUNK_SIZE * 0.10, 0)  
        print(f"分块大小: {CHUNK_SIZE}, 分块重叠: {chunk_overlap}")  

        # 定义用于拆分文档的拆分器。  
        child_splitter = RecursiveCharacterTextSplitter(  
           chunk_size=CHUNK_SIZE,  
           chunk_overlap=chunk_overlap)  

        # 拆分文档。  
        chunks = child_splitter.split_documents(docs)  
        print(f"{len(docs)} 个文档拆分成了 {len(chunks)} 个子文档。")  

        # 编码器的输入是文档的页面内容作为字符串。  
        list_of_strings = [doc.page_content for doc in chunks if hasattr(doc, 'page_content')]  

        # 使用 HuggingFace 编码器进行嵌入向量计算。  
        embeddings = torch.tensor(encoder.encode(list_of_strings))  

        # 归一化嵌入向量。  
        embeddings = np.array(embeddings / np.linalg.norm(embeddings))  

        # Milvus 需要一个由 `numpy.float32` 类型的 `numpy.ndarray` 构成的列表。  
        converted_values = list(map(np.float32, embeddings))  

        # 创建用于 Milvus 插入的 dict_list。  
        dict_list = []  
        for chunk, vector in zip(chunks, converted_values):  
           # 组装嵌入向量、原始文本片段及其元数据。  
           chunk_dict = {  
               'chunk': chunk.page_content,  
               'source': chunk.metadata.get('source', ""),  
               'vector': vector,  
           }  
           dict_list.append(chunk_dict)

将向量保存到Milvus

将编码的向量嵌入数据存入Milvus向量数据库。

    # 连接到Milvus Lite服务器  
    from pymilvus import MilvusClient  
    mc = MilvusClient("milvus_demo.db")  

    # 使用具有灵活模式和自动索引的集合进行创建  
    COLLECTION_NAME = "MilvusDocs"  
    mc.create_collection(COLLECTION_NAME,  
               EMBEDDING_DIM,  
               consistency_level="最终",  
               auto_id=True,   
               overwrite=True)  

    # 开始插入数据  
    print("开始插入数据")  
    start_time = time.time()  
    mc.insert(  
       COLLECTION_NAME,  
       data=dict_list,  
       progress_bar=True)  

    end_time = time.time()  
    print(f"Milvus插入{len(dict_list)}向量的时间:{round(end_time - start_time, 2)}秒")

进行向量搜索

问一个问题,并在Milvus中搜索最近的相似片段。

    SAMPLE_QUESTION = "HNSW 的参数是什么意思?"  

        # 使用相同的编码器嵌入问题。  
        query_embeddings = torch.tensor(encoder.encode(SAMPLE_QUESTION))  
        # 将嵌入归一化为单位长度。  
        query_embeddings = F.normalize(query_embeddings, p=2, dim=1)  
        # 将嵌入转换为 np.float32 类型的列表。  
        query_embeddings = list(map(np.float32, query_embeddings))  

        # 定义可以用于过滤的元数据字段。  
        OUTPUT_FIELDS = list(dict_list[0].keys())  
        OUTPUT_FIELDS.remove('vector')  

        # 定义要检索的 top-k 结果数。  
        TOP_K = 2  

        # 使用查询和向量数据库执行语义向量搜索操作。  
        results = mc.search(  
            COLLECTION_NAME,  
            data=query_embeddings,  
            output_fields=OUTPUT_FIELDS,  
            limit=TOP_K,  
            consistency_level="Eventually")

检索到的结果如下。

使用vLLM和Llama 3.1-8B模型构建并运行RAG生成 安装vLLM并从HuggingFace获取模型

vLLM 默认从 HuggingFace 下载大型语言模型,默认行为是如此。如果你想使用 HuggingFace 上的新模型,你应该执行命令 pip install --upgrade-U。另外,运行 Meta 的 Llama 3.1 模型的推理需要一个 GPU。

要查看所有vLLM支持的模型的完整列表,请参看文档页面

    # (建议)创建一个新的conda环境。  
        conda create -n myenv python=3.11 -y  
        conda activate myenv  

        # 用CUDA 12.1来安装vLLM。  
        pip install -U vllm transformers torch  

        import vllm, torch  
        from vllm import LLM, SamplingParams  

        # 清空GPU缓存区。  
        torch.cuda.empty_cache()  

        # 检查一下GPU。  
        !nvidia-smi  # 运行 `!nvidia-smi` 查看GPU信息

要了解如何安装vLLM,请参阅其“如何安装”页面。

获取 HuggingFace 的 token

在HuggingFace上,例如Meta Llama 3.1这样的模型,要求用户在下载模型权重之前接受其许可。因此,您需要创建一个HuggingFace账户,接受模型的许可,并生成一个访问令牌。

访问这个Llama3.1页面时,你会看到一个消息,要求你同意条款。点击“接受许可”同意Meta的条款,然后你就可以下载模型权重了。批准过程通常需要不到一天的时间。

获得批准后,您需要生成一个新的HuggingFace令牌。您的旧令牌将不再适用于新权限。

安装 vLLM 之前,用你的新 token 登录 HuggingFace 账户。如下所示,我使用了 Colab secrets 来保存 token。

    # 使用你的新令牌登录HuggingFace。
        from huggingface_hub import 登录函数  
        from google.colab import userdata  
        hf_token = userdata.get('HF_TOKEN')  
        登录函数(token=hf_token, add_to_git_credentials=True)
运行RAG生成任务

在演示中,我们运行了 Llama-3.1-8B 模型,该模型需要 GPU 和大量内存才能启动。以下示例是在配备了 A100 GPU 的 Google Colab Pro(每月10美元)上运行的。要了解如何运行 vLLM,您可以查阅 快速入门文档

    # 1. 选择要运行的模型  
        MODELTORUN = "meta-llama/Meta-Llama-3.1-8B-Instruct"  

        # 2. 你需要清空GPU缓存,因为你将需要所有可用的内存  
        torch.cuda.empty_cache()  

        # 3. 创建一个vLLM模型实例。  
        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)

请使用从Milvus检索到的相关资料和上下文写一个提示。

    # 将所有上下文通过空格分开。  
        contexts_combined = ' '.join(contexts)  
        # Lance Martin, LangChain建议将最相关的上下文放在最后。  
        contexts_combined = ' '.join(reversed(contexts))  

        # 将所有独特的来源用逗号隔开。  
        source_combined = ' '.join(reversed(list(dict.fromkeys(sources))))  

        SYSTEM_PROMPT = f"""首先,检查提供的上下文是否与用户的提问相关。如果提供的上下文与问题高度相关,则使用上下文回答问题。否则,如果上下文不高度相关,则不使用上下文回答问题。  
        回答必须清晰、简洁且相关,不超过两句话。  
        引用来源:{source_combined}  
        参考上下文:{contexts_combined}  
        用户问题:{SAMPLE_QUESTION}  
        """  

        prompts = [SYSTEM_PROMPT]

现在,根据找到的信息和问题来生成答案。

    # 采样设置  
        sampling_params = SamplingParams(temperature=0.2, top_p=0.95)  

        # 调用vLLM模型  
        outputs = llm.generate(prompts, sampling_params)  

        # 打印输出  
        for output in outputs:  
           prompt = output.prompt  
           generated_text = output.outputs[0].text  
           # !r 调用 repr(),用于打印带引号的字符串  
           print()  
           print(f"问题: {SAMPLE_QUESTION!r}")  
           pprint.pprint(f"生成文本: {generated_text!r}")

那上面的答案看起来真是太好了!

如果你对这个演示感兴趣,不妨试试看,并告诉我们你的想法。你也可以加入我们的Milvus Discord社区,直接在Discord上交流一下,和所有GenAI开发者聊天。

参考文献
点击查看更多内容
TA 点赞

若觉得本文不错,就分享一下吧!

评论

作者其他优质文章

正在加载中
  • 推荐
  • 评论
  • 收藏
  • 共同学习,写下你的评论
感谢您的支持,我会继续努力的~
扫码打赏,你说多少就多少
赞赏金额会直接到老师账户
支付方式
打开微信扫一扫,即可进行扫码打赏哦
今天注册有机会得

100积分直接送

付费专栏免费学

大额优惠券免费领

立即参与 放弃机会
意见反馈 帮助中心 APP下载
官方微信

举报

0/150
提交
取消