掌柜智库项目(RAG)实战5. 导入数据节点实现与测试5.6 向量化 (node_bge_embedding)文件:app/import_process/agent/nodes/node_bge_embedding.py相关工具类位置:app/lm/embedding_utils.py节点作用与实现思路节点作用: 将人类可读的文本切片转化为机器可计算的向量表示。采用 BGE-M3 模型同时生成稠密向量Dense和稀疏向量Sparse为“混合检索”提供底层数据支持。实现思路:双路编码: 利用 BGE-M3 的特性一次推理同时产出语义向量Dense捕获语义相似度和词汇向量Sparse捕获关键词匹配兼顾语义理解和精确匹配。归一化处理: 对稀疏向量进行 L2 归一化消除文本长度带来的权重偏差确保检索评分的公平性。批处理优化: 针对大量切片采用 Batch 处理模式调用模型大幅提升 GPU/CPU 的计算利用率。1. 导入与配置引入必要的依赖库配置 BGE-M3 向量化流程。importsysimportosfromtypingimportAny,List,Dictfromapp.import_process.agent.stateimportImportGraphStatefromapp.lm.embedding_utilsimportget_bge_m3_ef,generate_embeddingsfromapp.utils.task_utilsimportadd_running_task,add_done_taskfromapp.core.loggerimportlogger2. 主流程定义定义node_bge_embedding函数串联各个步骤。# # BGE-M3向量化核心节点# 核心能力将文本切片转换为稠密/稀疏双向量为Milvus向量检索提供数据基础# 依赖模型BAAI/bge-m3多语言、多粒度同时支持语义/关键词检索# 向量说明# 1. 稠密向量1024维固定长度记录文本深层语义信息用于语义相似度匹配# 2. 稀疏向量变长键值对记录文本关键词/特征位置用于关键词精准匹配# 核心设计# - 单例模型避免重复加载模型节省显存/时间# - 批量处理分批生成向量防止大批次导致的显存溢出# - 文本增强拼接商品名切片内容强化核心特征提升检索准确性# defnode_bge_embedding(state:ImportGraphState)-ImportGraphState: LangGraph核心节点BGE-M3文本向量化处理 主流程串行执行全流程异常隔离 1. 输入校验验证chunks有效性核心数据缺失则终止当前节点 2. 模型初始化获取BGE-M3单例模型实例避免重复加载 3. 批量向量化分批拼接文本、生成双向量为切片绑定向量字段 4. 状态更新将带向量的chunks更新回全局状态供下游Milvus入库节点使用 参数 state: ImportGraphState - 流程全局状态对象包含上游传入的chunks、task_id等数据 返回 ImportGraphState - 更新后的状态对象chunks字段新增dense_vector/sparse_vector 异常处理 节点内所有异常均捕获不终止整体LangGraph流程仅记录错误日志 # 获取当前节点名称用于日志和任务状态记录current_nodesys._getframe().f_code.co_name logger.info(f 开始执行LangGraph节点{current_node})# 标记任务运行状态用于任务监控/前端进度展示add_running_task(state.get(task_id,),current_node)logger.info(--- BGE-M3 文本向量化处理启动 ---)try:# 步骤1输入数据校验核心chunks无效则抛出异常texts_to_embedstep_1_validate_input(state)# 步骤2初始化BGE-M3模型单例模式仅加载一次bge_m3_efstep_2_init_model()# 步骤3批量生成双向量为切片绑定向量字段output_datastep_3_generate_embeddings(texts_to_embed,bge_m3_ef)# 步骤4更新全局状态将带向量的chunks回传下游state[chunks]output_data logger.info(f--- BGE-M3 向量化处理完成共处理{len(output_data)}条文本切片 ---)add_done_task(state.get(task_id,),current_node)exceptExceptionase:# 捕获节点所有异常记录错误堆栈不中断整体流程logger.error(fBGE-M3向量化节点执行失败{str(e)},exc_infoTrue)# 返回更新后的状态对象传递至下游节点returnstate4. 步骤 1: 校验输入检查输入数据是否有效。defstep_1_validate_input(state:ImportGraphState)-List[Dict[str,Any]]: 向量化前置步骤1输入数据有效性校验 核心作用 1. 从全局状态提取待向量化的chunks切片列表 2. 严格校验chunks类型和非空性无有效数据则终止向量化 参数 state: ImportGraphState - 流程全局状态对象 返回 List[Dict[str, Any]] - 校验通过的文本切片列表 异常 若chunks非列表/为空抛出ValueError终止当前向量化流程 # 从状态中提取切片数据texts_to_embedstate.get(chunks)# 校验必须是非空列表否则无法进行向量化ifnotisinstance(texts_to_embed,list)ornottexts_to_embed:logger.error(向量化输入校验失败chunks字段为空或非有效列表)raiseValueError(错误: 无有效文本切片数据无法执行向量化处理)logger.info(f向量化输入校验通过待处理文本切片数量{len(texts_to_embed)})returntexts_to_embed5. 步骤 2: 初始化模型加载 BGE-M3 模型客户端。可以忽略defstep_2_init_model(): 向量化步骤2初始化BGE-M3模型实例单例模式 核心作用 1. 调用单例函数get_bge_m3_ef确保模型全局仅加载一次 2. 校验模型实例有效性加载失败则抛出明确异常 返回 Any - 有效BGE-M3模型实例embedding function 异常 模型加载失败路径错误/显存不足/依赖缺失时抛出ValueError并提示配置问题 try:# 获取单例模型实例避免重复加载浪费资源efget_bge_m3_ef()# 校验模型实例是否有效ifefisNone:raiseValueError(BGE-M3模型实例为Nonepymilvus.model模块未找到或模型加载失败)logger.info(BGE-M3模型实例初始化成功单例模式)returnefexceptExceptionase:# 包装异常信息明确错误原因和排查方向error_msgfBGE-M3模型初始化失败{e}请检查模型路径/环境变量配置是否正确logger.error(error_msg)raiseValueError(error_msg)6. 步骤 3: 批量生成向量分批次为文本生成稠密和稀疏向量并进行归一化处理。defstep_3_generate_embeddings(texts_to_embed:List[Dict[str,Any]],bge_m3_ef:Any)-List[Dict[str,Any]]: 向量化核心步骤3批量生成稠密/稀疏双向量 核心逻辑分批执行每批独立异常处理 1. 文本拼接item_name商品名 换行 content切片内容强化核心特征 2. 批量调用传入拼接后的文本生成批量双向量 3. 向量绑定为每个切片复制原数据新增dense_vector/sparse_vector字段 4. 异常兜底单批次失败则保留原切片数据继续处理下一批次 参数 texts_to_embed: List[Dict[str, Any]] - 校验通过的文本切片列表含item_name/content字段 bge_m3_ef: Any - 步骤2初始化的BGE-M3模型实例 返回 List[Dict[str, Any]] - 带向量字段的文本切片列表异常批次保留原数据 关键配置 batch_size: 每批处理5条可根据服务器显存大小调整显存大则调大反之调小 # 初始化结果列表存储带向量的切片数据output_data[]# 批次大小配置平衡显存占用和处理效率建议根据实际环境调整batch_size5# 按批次遍历避免一次性处理过多数据导致显存溢出OOMtotallen(texts_to_embed)foriinrange(0,total,batch_size):# 截取当前批次的切片最后一批自动适配剩余数量【每次获取5个】batch_textstexts_to_embed[i:ibatch_size]# 计算当前批次的起止索引用于日志展示方便看从1开始也不获取下标没有影响start_idx,end_idxi1,min(ilen(batch_texts),total)try:# 构造模型输入文本拼接商品名切片内容增强核心特征input_texts[]fordocinbatch_texts:item_namedoc[item_name]contentdoc[content]# 有商品名则拼接换行分隔提升模型识别效率无则直接使用内容# 几乎所有的 Embedding 模型尤其是基于 BERT 架构的对前 128 个 token 的注意力是最集中的。越往后的词对最终向量方向的拉扯力越弱。# **“核心词前置”**的原则# 方案 1用强标点代替换行最简单、最推荐# 优化前苹果手机\n性能很好...# 优化后苹果手机。性能很好...# 方案2加一点“微量”的语义胶水适合属性明确的场景textf商品{item_name}介绍{content}ifitem_nameelsecontent# Embedding 模型是个强迫症你给它喂中文就用全套中文标点伺候给它喂英文就用全套英文标点。保持 语境纯粹 生成的向量质量最高input_texts.append(text)# 调用封装函数生成批量向量返回格式{dense: [稠密向量列表], sparse: [稀疏向量列表]}docs_embeddingsgenerate_embeddings(input_texts)ifnotdocs_embeddings:logger.warning(f第{start_idx}-{end_idx}条切片向量生成返回空保留原数据)output_data.extend(batch_texts)continue# 为当前批次每个切片绑定对应向量复制原数据避免修改上游源数据forj,docinenumerate(batch_texts):itemdoc.copy()item[dense_vector]docs_embeddings[dense][j]# 绑定稠密向量item[sparse_vector]docs_embeddings[sparse][j]# 绑定稀疏向量已归一化output_data.append(item)logger.info(f第{start_idx}-{end_idx}条切片双向量生成成功)exceptExceptionase:# 捕获单批次所有异常记录错误堆栈不终止整体批量处理logger.error(f第{start_idx}-{end_idx}条切片向量生成失败保留原数据 | 错误原因{str(e)},exc_infoTrue)# 异常批次保留原切片数据保证数据完整性后续可人工排查output_data.extend(batch_texts)continuereturnoutput_data7. 单元测试您可以在node_bge_embedding.py文件底部直接运行以下测试代码# # 本地单元测试入口# 功能独立验证向量化节点全链路逻辑无需启动整个LangGraph流程# 适用场景本地开发、调试、模型有效性验证# if__name____main__:# 加载环境变量定位项目根目录下的.env读取模型路径/设备等配置current_diros.path.dirname(os.path.abspath(__file__))project_rootos.path.dirname(os.path.dirname(current_dir))load_dotenv(os.path.join(project_root,.env))# 构造模拟测试状态模拟上游节点输出的chunks数据贴合真实业务场景test_stateImportGraphState({task_id:test_task_embedding_001,# 测试任务IDchunks:[# 模拟带item_name的文本切片上游商品名称识别节点产出{content:这是一个测试文档的内容用于验证向量化是否成功。,title:测试文档标题,item_name:测试项目,file_title:测试文件.pdf},{content:这是第二个测试文档的内容用于验证批量处理逻辑。,title:测试文档标题2,item_name:测试项目,file_title:测试文件.pdf}]})# 执行本地测试logger.info( BGE-M3向量化节点本地单元测试启动 )try:# 调用核心节点函数result_statenode_bge_embedding(test_state)# 提取测试结果result_chunksresult_state.get(chunks,[])# 打印测试结果统计logger.info(f 向量化节点本地测试完成 )logger.info(f测试任务ID{test_state.get(task_id)})logger.info(f待处理切片数2 | 实际处理切片数{len(result_chunks)})# 验证向量生成结果打印向量字段是否存在foridx,chunkinenumerate(result_chunks):has_densedense_vectorinchunk has_sparsesparse_vectorinchunk logger.info(f第{idx1}条切片稠密向量生成{ifhas_denseelse未}成功 | 稀疏向量生成{ifhas_sparseelse未}成功)exceptExceptionase:logger.error(f 向量化节点本地测试失败 f错误原因{str(e)},exc_infoTrue)# 新手友好提示给出核心排查方向logger.warning(排查提示请检查BGE-M3模型路径、显存是否充足、环境变量配置是否正确)