快速掌握大语言模型+向量数据库

本文链接：https://blog.csdn.net/qq_36254517/article/details/147641969

一、前言

结合前面掌握的vLLM部署Qwen7B模型、通过Embedding模型（bdg-large-zh模型）提取高质量作文内容并预先存储到Milvus向量数据库中，我们很容易实现RAG方案进一步提高写作内容的生成质量。

本篇要实现的目标是：通过FlaskAPI对上游业务系统提供一个开放接口，调用该接口将返回写作内容（非流式）。该接口优先通过dbg-large-zh模型生成ebedding向量，然后从Milvus向量数据库中查询和向量相似度最高的1条数据，若查到数据且相似度达到指定阈值（如0.75或0.9,具体数值由业务决定），则直接返回Milvus中存储的内容，否则调用Qwen2-7B-Instruct模型生成写作内容。

要实现的架构为：

往期文章：

快速掌握大语言模型-Qwen2-7B-Instruct落地1-CSDN博客

快速掌握大语言模型-Qwen2-7B-Instruct落地2-CSDN博客

快速掌握向量数据库-Milvus探索1-CSDN博客

快速掌握向量数据库-Milvus探索2_集成Embedding模型-CSDN博客

二、术语

2.1 vLLM

vLLM采用连续批处理,吞吐量相比传统框架（如Hugging Face Transformers）可提升5-10倍;

分布式推理支持允许在多GPU环境下并行处理;

通过PagedAttention技术，vLLM将注意力计算的键值（KV Cache）分页存储在显存中，动态分配空间以减少碎片；

内置了对 OpenAI API 的兼容支持，可直接启动符合 OpenAI 标准的 API 服务；

2.2 FastApi

基于Python的Web框架，主要用于构建高性能、易维护的API和Web应用程序，通过FastApi可以对上游业务系统暴露API，上游业务系统通过http请求即调用相关接口。

2.3 Milvus

Milvus是一个专为处理高维向量数据设计的开源向量数据库，支持数百亿级数据规模，在多数开源向量数据库中综合表现突出（一般是其他的2~5倍）。

提供三种部署方式：本地调试Milvus Lite、企业级小规模数据的Milvus Standalone（一亿以内向量）、企业级大规模数据的Milvus Distributed （数百亿向量）。

2.4 bge-large-zh

bge-large-zh模型是专为中文文本设计的，同尺寸模型中性能优异的开源Embedding模型，凭借其中文优化、高效检索、长文本支持和低资源消耗等特性，成为中文Embedding领域的标杆模型。

2.5 Qwen2-7B-Instruct

Qwen2-7B-Instruct 是阿里云推出的开源指令微调大语言模型，属于 Qwen2 系列中的 70 亿参数版本。该模型基于 Transformer 架构，通过大规模预训练和指令优化，展现出强大的语言理解、生成及多任务处理能力

三、代码

3.1 代码目录

vllm.py ：提供调用Qwen7B模型的方法；
app.py ：提供通过FastAPI暴露给上游业务系统调用的API；
embedding.py ：提供调用bge-large-zh模型生成内容的embedding向量的方法；
milvus_server.py ：提供查询Milvus向量数据库的方法；
generate_text_request.py : 模拟上游业务系统通过http调用app.py提供的生成写作内容的API

3.2 完整代码

3.2.1 安装依赖

# vllm需要的依赖
pip install vllm

# milvus需要的依赖
python -m pip install -U pymilvus

3.2.2 app.py

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import logging
from vllm import Model
from embedding import search_writing_embeddings

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
app = FastAPI()
model = Model()

# 流式响应
@app.get("/chat")
async def chat(prompt: str):
    return StreamingResponse(model.generate_stream(prompt), media_type="text/plain")


# 非流式响应（RAG增强检索）
@app.get("/generate")
async def generate(prompt: str):
    # 优先通过embedding向量从milvus中搜索相似度最高的1条数据，如果相似度达到90%以上，则直接返回该数据的content_full字段作为结果，否则调用llm模型生成
    # 这里的90%是一个阈值，可以根据实际情况调整
    res = search_writing_embeddings(prompt)
    print(f"search_writing_embeddings:{res}")

    threshold = 0.75
    if res and res[0] and res[0][0]['distance'] >= threshold:
        print("使用milvus中的数据")
        result = res[0][0]['entity']['content_full']
    else:
        print("使用llm模型生成")
        result = await model.generate_text(prompt)

    return result


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

3.2.3 vllm.py

# 大语言模型LLM文件

import logging
from openai import OpenAI
from pydantic_core import Url

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
model_path = "/mnt/workspace/models/Qwen2-7B-Instruct"  
base_url = "http://localhost:9000/v1"

class Model:
    def __init__(self):
        # 创建 OpenAI 客户端实例,这里的密钥可以随意填写，vLLM 不需要验证
        self.client = OpenAI(api_key="EMPTY", base_url= base_url)   

    # 流式处理
    async def generate_stream(self, prompt: str):
        try:
            logging.info(f"开始处理提示: {prompt}")
            # 使用新客户端调用 API
            # 确保使用类的实例属性 self.client
            chat_response = self.client.chat.completions.create(
                model= model_path,
                messages=[{"role": "user", "content": prompt}],
                stream=True,
                max_tokens=512,
                temperature=0.7,
                top_p=0.9
            )
            
            for chunk in chat_response:
                msg = chunk.choices[0].delta.content
                if msg is not None:
                    logging.info(f"生成文本块: {msg}")
                    yield msg
            logging.info("流式输出结束")
        except Exception as e:
            logging.error(f"处理提示时出错: {e}", exc_info=True)

    # 非流式处理
    async def generate_text(self, prompt: str):
        try:
            logging.info(f"开始处理提示: {prompt}")
            # 使用新客户端调用 API
            response = self.client.chat.completions.create(
                model= model_path,
                messages=[{"role": "user", "content": prompt}],
                stream=False,
                max_tokens=512,
                temperature=0.7,
                top_p=0.9
            )
            # 获取完整响应
            content = response.choices[0].message.content
            logging.info(f"文本生成完成:{content}")
            return content
        except Exception as e:
            logging.error(f"处理提示时出错: {e}", exc_info=True)
            return ""

3.2.4 embedding.py

from transformers import AutoTokenizer, AutoModel
import torch
import random
from milvus_server import WritingDTO, insert_data_to_milvus,search_data_from_milvus

# 加载模型和分词器
model_path = "/mnt/workspace/models/bge-large-zh"  # 根据实际情况修改
model = AutoModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
# 推理模式
model.eval()  

# 写作生成合集名称
writing_collection_name = "writing"


def get_embedding_list(text_list):
    """
    输入文本列表,返回文本的embedding向量
    :param text_list: 待获取embedding的文本集合
    :return: ebedding向量集合
    """
    
     # 编码输入（自动截断和填充）
    encoded_input = tokenizer(text_list, padding=True, truncation=True, return_tensors='pt')

    # 调用大模型得到文本的embedding
    with torch.no_grad():
        model_output = model(**encoded_input)
        sentence_embeddings = model_output[0][:, 0] # bdge-large-zh 模型的输出是一个元组，第一个元素是句子的嵌入向量,1024维
        print(f"sentence_embeddings:{ len(sentence_embeddings[0])}")
        # 归一化处理:可以提高模型的稳定性和收敛速度，尤其在处理向量相似度计算时非常有用
        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)

        # 转换嵌入向量为列表
        embeddings_list = sentence_embeddings.tolist()
        return embeddings_list


def insert_writing_embeddings(text_list):
    """
    获取写作生成内容的文本的embedding向量,并写入到milvus中
    :param text_list: 待获取embedding的写作生成内容集合
    :return: 写入到到milvus结果,ebedding向量集合
    """
    # 获取embedding 
    embeddings_list = get_embedding_list(text_list)
    print(f"embeddings_list:{len(embeddings_list)}")
    # 写作记录合集
    writing_list = []
    for i, embedding in enumerate(embeddings_list):
        writingDTO = WritingDTO(embedding, text_list[i], random.randint(500,9999))  #组装对象数据,其中biz_id是业务ID,这里这里方便说明暂设为随机数字
        writing_list.append(writingDTO.to_dict()) # 将对象转换为字典,并添加到集合中
    # 插入数据到 Milvus
    res = insert_data_to_milvus(writing_collection_name,writing_list)
    return res,embeddings_list

def search_writing_embeddings(text):
    """
    获取写作内容的embedding向量,并从milvus中搜索
    """
    # 获取embedding
    embeddings_list = get_embedding_list([text])
    # 搜索数据
    res = search_data_from_milvus(writing_collection_name,embeddings_list,["id","content_full","biz_id"],1)
    return res


if __name__ == "__main__":
    # 输入文本
    content = "推荐广州有哪些好玩的地方"
    # 将文本转eembedding向量并写入到milvus中
    res = search_writing_embeddings(content)
    print("Search result:", res)

3.2.5 milvus_server.py

from pymilvus import MilvusClient, db
import numpy as np
from pymilvus.orm import collection
from typing import Iterable

# 定义 Milvus 服务的主机地址
host = "阿里云公网IP"

# 创建一个 Milvus 客户端实例，连接到指定的 Milvus 服务
client = MilvusClient(uri=f"http://{host}:19530",db_name="db001") # 连接到 Milvus 服务并选择数据库 "db001"
# collection_name = "writing" # 指定要连接的集合名称 "writing"

# 写作生成对象
class WritingDTO:
    def __init__(self, content_vector, content_full, biz_id):
        self.content_vector = content_vector
        self.content_full = content_full
        self.biz_id = biz_id

    def to_dict(self):
            """
            将 WritingDTO 对象转换为字典。

            :return: 包含 WritingDTO 对象属性的字典
            """
            return {
                "content_vector": self.content_vector,
                "content_full": self.content_full,
                "biz_id": self.biz_id
            }


def insert_data_to_milvus(collection_name,data):
    """
    将对象集合中的数据插入到 Milvus 集合中。

    :param data: 字典对象集合
    :return: 插入操作的结果
    """

    print(f"insert_data_to_milvus:{collection_name}")
    print(f"insert_data_to_milvus:{len(data)}")
    res = client.insert(
        collection_name=collection_name,
        data=data
    )
    return res

def search_data_from_milvus(collection_name,query_vector,output_fields, top_k=10):
    """
    从 Milvus 集合中搜索与查询向量最相似的向量。

    :param query_vector: 查询向量
    :param top_k: 返回的最相似向量的数量
    :return: 搜索结果
    """
    res = client.search(
        collection_name= collection_name, # 合集名称
        data=query_vector, # 查询向量
        search_params={
            "metric_type": "COSINE", # 向量相似性度量方式，COSINE 表示余弦相似度(适用于文本/语义相似性场景); 可选 IP/COSINE/L2 
            "params": {"level":1}, 
        }, # 搜索参数
        limit=top_k, # 查询结果数量
        output_fields= output_fields, # 查询结果需要返回的字段
        consistency_level="Bounded" # 数据一致性级别,Bounded允许在有限时间窗口内读取旧数据,相比强一致性（STRONG）提升 20 倍查询性能，适合高吞吐场景; 
    )
    return res

3.2.6 generate_text_request.py

import requests

url = "http://0.0.0.0:8000/generate?prompt=广州有什么好玩的"
response = requests.get(url, stream=True)
print(response.text)

四、调用结果

4.1 启动vllm

python -m vllm.entrypoints.openai.api_server  --model  /mnt/workspace/models/Qwen2-7B-Instruct  --swap-space 10 --disable-log-requests --max-num-seqs 256 --host 0.0.0.0 --port 9000  --dtype float16 --max-parallel-loading-workers 1  --max-model-len 10240 --enforce-eager