RAG 数据检索

数据检索我们前面已经学习完了知识库的构建那我们接下来来看一下数据检索是如何实现的。在大模型回答问题之前先从私有知识库 / 文档库里快速找出和用户问题最相关、最匹配的一段 / 多段真实原文内容喂给大模型再让模型基于这些真实资料 自身知识来回答。简单一句话先找资料再答题不让模型凭空瞎编。数据检索这一部分我会从以下 4 个部分给大家展开讲解。 这部分整体是比较简单的。大家回想一下我们前期完成了知识库的构建也体验过向量数据库的创建如何往里面插入数据以及如何根据文本、根据用户查询检索得到相似度较高的文本内容。那么在本章节我们要实现的场景就是已经建好知识库的前提下如何处理用户输入的查询内容经过向量化转换把它作为查询条件在向量数据库中完成检索获取我们需要的内容主要分为以下四个部分。用户查询第一部分就是用户查询也就是我们的query。 在使用 RAG 的任意场景下都必须要有用户输入也就是用户的查询条件。用户查询User Query是用户在检索系统或对话系统中输入的原始问题、指令或关键词用于表达其信息需求。在 RAG 系统中它是整个检索与生成流程的起点。简单来说就是用户向应用发起提问。比如我们问一个 AI 应用的内容就是用户查询例如我们问公司的 AI 助理“公司新发布的休假政策是什么”。查询向量化第二部分是查询向量化。 简单来说就是把用户输入的query通过embedding model嵌入模型转换成向量数组也就是浮点数组。⚠️重点注意这里使用的模型必须和构建知识库时使用的模型保持完全一致。 比如我们使用 BGE‑M3 模型在构建知识库、把分块文本向量化入库的时候用的是 BGE‑M3做查询的时候也必须使用同一个 BGE‑M3 模型。 如果使用不同的模型生成的向量会处在不同的语义空间检索出来的结果就会出现错误所以模型一致性非常关键。用户输入问题后系统使用与索引阶段相同的嵌入模型将查询文本转换为向量表示使其与知识库中的文档向量处于同一语义空间。这个步骤和我们将语句块转为向量化的过程一样我们不再重复赘述。向量检索第三部分是向量检索。 把用户query经过嵌入模型转换成向量之后拿到向量数据库中执行查询。 将查询向量和向量库内存储的全部向量做相似度比对同时可以通过limit参数控制返回结果数量比如设置返回 10 条、20 条相似度最高的数据这就是我们之前讲过的limit参数的作用。将用户查询转换为向量后系统需要在向量数据库中快速找到语义最接近的文本块。这个过程本质上是向量相似度检索用搜索算法检索查询向量与库中所有向量的“距离”距离越近语义越相似。距离之前我们已经介绍过了这里就不重复介绍了。代码实现第四部分就是以上三个流程的代码实现。大家要注意数据检索整体逻辑比较简单。 我们的重点工作其实都在数据加载、数据清洗、各类文本分块方案以及分块之后的文本向量化、入库。 数据检索本质只是查询环节所以本章的实现完全是基于前面已经完成的工作。讲完本篇整体的目录我们就先来学习用户查询。用户查询也很好理解。我们使用 RAG 的时候首先接收用户查询在知识库中检索出相关内容再把检索到的内容连同用户的原始问题一起送入大模型。 在提示词环节做内容限定划定模型回答的边界避免模型凭空生成内容也就是解决大模型的幻觉问题。当知识库中没有对应资料时模型容易胡乱编造答案RAG 就是通过检索真实知识库内容来规避这个问题这就是用户查询在 RAG 里的作用。接下来是查询向量化。 复用构建知识库时使用的那套嵌入模型对用户输入的查询做向量表示。把用户输入的query传入嵌入模型得到对应的vector向量。然后就是向量检索。 把用户查询转换得到向量之后直接在向量数据库执行相似度比对取出语义相近的文本块即可。概念部分都比较简单接下来我们着手编写代码实现。我们打开编辑器新建一个 Python 文件命名为vectorSearch.py用来实现向量检索功能。回顾向量检索的完整流程首先对用户查询做向量化同时需要构建向量数据库提升检索效率。 所以第一步我们先构建向量数据库。首先声明数据库名称DB_NAME我们命名为vectorSearch。 数据库定义完成之后接着创建集合collection。定义好数据库名、集合名之后我们编写init()初始化函数完成向量数据库的初始化这里可以复用我们之前写好的milvus模块内的工具方法。初始化第一步先判断数据库是否已经存在。Milvus 本地数据库本质是磁盘上的文件夹。我们构造db_path路径config.base_path拼接上{DB_NAME}.db这个就是数据库存放的目录。判断这个路径是否存在如果已经存在就调用shutil.rmtree()把旧数据库文件夹删除并且加上time.sleep(0.5)等待保证删除操作完成再执行后续逻辑。删除旧库完成之后第二步创建数据库。调用milvus.create_db()传入我们定义好的DB_NAME。 数据库创建完毕再基于该数据库创建集合调用milvus.create_collection()第一个参数传入集合名称COLLECTION因为我们使用 BGE‑M3 模型向量维度是 1024我们把向量维度定义为常量DIMENSION 1024作为第二个参数传入。到这里向量数据库的初始化就完成了。数据库初始化完成之后下一步文本分块将分块内容向量化之后入库。 拿到原始文档使用我们之前学过的分块方案句子分块、语义分块、结构分块把文档切分成文本块再将每一个文本块向量化插入向量数据库。这里我们复用之前写好的分块逻辑选择句子分块方案。 我们新建函数index_to_milvus()返回值为int类型。调用sentence_chunk_documents()做句子分块传入config.structure_path文档路径设置最大句子数参数为1也就是一个句子作为一个文本块。分块完成之后做校验如果chunks分块结果为空直接抛出ValueError异常提示分块结果为空请检测数据集内容。接下来准备入库数据定义空列表rows。 循环遍历每一个分块chunk 调用milvus.embed_model.get_text_embedding(chunk.text)把分块文本转为向量vec。 把每一条数据组装成字典追加到rows列表 字典内包含三个字段idi1作为主键 idvector刚才生成的向量vectext原始分块文本chunk.text所有分块全部处理完毕调用milvus.insert()传入集合名COLLECTION和rows完成批量入库。 打印日志提示入库完成输出入库记录条数和向量维度函数返回入库记录数量。# 导入需要用到的标准库 import shutil import time # 导入项目自定义模块 import config import milvus from sentenceChunk import sentence_chunk_documents # 向量数据库配置常量 # 数据库名称 DB_NAME vectorSearch # 集合(表)名称 COLLECTION vectorSearch # 向量维度和embedding模型输出维度保持一致 DIMENSION 1024 def init(): 初始化Milvus向量数据库 功能 1. 如果旧的数据库文件存在先删除旧库清理残留数据 2. 创建新的数据库 3. 创建指定维度的集合(数据表) # 拼接本地数据库文件路径 dp_path config.base_path / f{DB_NAME}.db # 判断数据库文件是否已经存在 if dp_path.exists(): # 递归删除整个旧数据库目录清理历史数据 shutil.rmtree(dp_path) # 短暂休眠确保文件系统删除操作完成避免文件占用报错 time.sleep(0.5) # 创建Milvus数据库 milvus.create_db(DB_NAME) # 在数据库内创建集合指定向量维度 milvus.create_collection(COLLECTION, DIMENSION) def index_to_milvus() - int: 文档分块 - 生成embedding向量 - 批量插入Milvus向量库 :return: 成功插入的文档块数量 :raises ValueError: 如果分块结果为空抛出异常 # 读取文档路径执行句子级分块第二个参数为分块参数 chunks sentence_chunk_documents(config.structure_path, 1) # 校验分块结果如果没有得到任何文本块抛出异常终止流程 if not chunks: raise ValueError(分块结果为空请检查数据集内容) # 用于存放待插入Milvus的数据行 rows [] # 遍历所有分块生成向量组装插入数据 for i, chunk in enumerate(chunks): # 使用嵌入模型把文本转换成向量数组 vec milvus.embed_model.get_text_embedding(chunk.text) # 组装单条记录id、向量、原始文本 rows.append( { id: i 1, # 主键id从1开始计数 vector: vec, # embedding向量 text: chunk.text, # 原始分块文本 } ) # 将组装好的全部数据批量插入Milvus集合 milvus.insert(COLLECTION, rows) # 打印入库统计信息 print(f已经完成入库, {len(rows)} 个记录, 向量维度{DIMENSION}) # 返回插入的记录总数 return len(rows)以上就是知识库构建、分块、向量化、入库的完整代码流程接下来就可以实现用户查询、查询向量化、向量检索的业务逻辑。完成向量入库之后接下来我们来实现查询功能也就是文本检索功能。我们定义一个函数search_similar这个函数的第一个入参就是query也就是用户输入的查询内容。第二个参数我们声明为limit用来控制返回结果的条数这里给它设置默认值为2。接下来我们来实现函数内部逻辑。 首先需要加载集合调用milvus.get_client().load_collection()传入我们定义好的集合名称COLLECTION。 集合加载完成之后执行查询操作调用milvus.search_by_text()这就是我们之前向量数据库里封装好的向量查询方法。传入参数第一个参数集合名COLLECTIONtextquery传入用户的查询文本limitlimit设置返回结果数量fields[id, text]指定需要返回的字段拿到查询返回结果之后我们要对返回的数据做简单的整理处理。 我们观察原始返回结果里面包含id、text还有distance相似度距离原始的distance小数位数过长可读性不好所以我们要对输出做格式化处理。我们用变量hits接收结果hits results[0] if results else []做判断如果hits为空说明没有检索到匹配内容打印未检索到相关结果然后直接return退出函数。如果检索到结果我们先打印用户的原始查询print(f\n用户查询:{query})接着打印检索到结果的总条数print(f共{len(hits)}条相似的结果)然后遍历hits列表使用enumerate循环设置start1拿到每条结果的序号rank和结果对象hit。 从hit里面取出entity实体对象再从entity中拿到text原始文本。注意数据结构里原始文本存放在entity内部。接下来格式化打印输出内容打印排名rank打印相似度距离distance。原始distance小数位数很多我们做格式化保留小数点后 4 位。打印id这里要注意外层结果对象hit有identity内部也有id两个都打印出来方便对照。打印检索出来的原始文本text。def search_similar(query: str, limit: int 2): 根据用户查询文本执行向量相似度检索 :param query: 用户输入查询问题 :param limit: 返回相似度最高的结果条数默认2条 # 加载集合到内存保证可以执行检索 milvus.get_client().load_collection(collection_nameCOLLECTION) # 执行文本向量检索返回匹配结果 results milvus.search_by_text( COLLECTION, textquery, limitlimit, fields[id, text], ) # 取出检索命中结果列表 hits results[0] if results else [] # 无命中结果处理 if not hits: print(未检索到相关结果) return # 打印查询与结果统计 print(f\n用户查询:{query}) print(f共{len(hits)}条相似的结果) # 遍历打印每条命中结果输出排名、距离、id、原文 for rank, hit in enumerate(hits, start1): text hit.get(text) entity hit.get(entity) print(f\n【rank】{rank} 距离 distance{hit.get(distance):.4f}) print(f【id】{entity.get(id)}) print(f【text】{text})到此search_similar查询函数就编写完成。接下来我们要实现交互逻辑程序运行之后可以循环接收用户输入查询持续检索知识库内容。 我们新建一个search()函数实现循环交互。函数内部先打印提示信息向量检索系统请输入用户查询问题。 开启while True无限循环使用input()获取用户输入调用.strip()去除首尾空格赋值给query。判断如果用户输入q代表退出程序打印已经退出执行break跳出循环。如果用户输入为空执行continue继续等待下一次输入。输入合法查询文本调用search_similar(query)使用默认limit2执行检索。最后编写程序入口if __name__ __main__:程序执行的顺序分为三步调用init()完成向量数据库初始化创建数据库和集合调用index_to_milvus()完成文档分块、向量化、数据入库调用search()启动交互循环接收用户查询执行检索。def search(): 交互式检索循环持续接收用户输入查询调用相似度检索 输入 q / quit 即可退出程序 print(向量检索系统请输入用户查询问题) while True: # 获取用户输入去除首尾空白字符 query input(请输入查询问题).strip() # 退出条件输入q或者quit if query in [q, quit]: print(已经退出) break # 空输入直接跳过 if not query: continue # 执行相似度检索 search_similar(query) if __name__ __main__: # 程序入口初始化库 - 文档入库 - 开启交互式检索 init() index_to_milvus() search()写到这里我们就完整实现了用户查询、向量检索的整套流程。那我们把这个代码写完了之后我们来测一下它这个效果。这块我们先来执行代码看看它能否按照我们的预期运行。我们执行程序观察两件事第一能否把这个数据创建出来第二能否正常进行查询。 首先程序会先把模型先加载到内存。好可以看到已经完成入库一共16个记录向量维度是1024维。因为我们原始数据集采用按句子分块的策略每一个句子单独定为一个块最终得到 16 条记录。Loading weights: 100%|██████████| 391/391 [00:0000:00, 48135.39it/s] 已经完成入库,16个记录向量维度1024 向量检索系统请输入用户查询问题 请输入查询问题接下来我们输入问题测试检索效果我们选取数据集里面存在的业务场景来提问PDF 文件乱码的原因是什么执行程序观察向量数据库返回结果一共查出来了两条。这里要注意返回的数值并不是相似度而是距离。距离数值越小代表向量越接近。 第一条距离为0.2距离比较近返回内容就是 RAG 数据集里面原始文本完整原样返回。 第二条对应内容是内存溢出相关距离达到0.49向量距离就相对更远。我们再换一个问题测试文本分块有哪几种策略。 执行检索可以成功命中对应的原始句子文档 ID 也正常返回。 第二条返回结果是检索增强策略因为关键词 “策略” 发生匹配向量距离相对更近。到这里本章数据检索的全部内容就实现完毕。整体难度不高大部分底层能力在前面章节知识库构建阶段就已经铺垫完成。