尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

FlagEmbedding BGE-M3 推理完全指南:M3Embedder 三路编码与多路打分详解

FlagEmbedding BGE-M3 推理完全指南:M3Embedder 三路编码与多路打分详解 FlagEmbedding BGE-M3 推理完全指南M3Embedder 三路编码与多路打分详解【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding本指南以 FlagEmbedding 仓库中 M3Embedder 的 API 文档 为核心结合其完整源码与官方示例系统讲解 BGE-M3 推理阶段的全部能力如何构造模型、如何用 dense稠密向量/ sparse稀疏词权重/ colbert多向量三种模式编码查询与语料以及如何用compute_score完成五类相关度打分。读完本文你将能直接用BGEM3FlagModel即M3Embedder的公开别名搭建一套多路召回 加权融合的检索管线。一、M3Embedder 是什么M3Embedder是 FlagEmbedding 中 BGE-M3 模型BAAI/bge-m3的推理封装类定义于 FlagEmbedding/inference/embedder/encoder_only/m3.py并被公开导出为BGEM3FlagModel别名见 encoder_only/init.py 中的from .m3 import M3Embedder as BGEM3FlagModel因此你可以直接通过from FlagEmbedding import BGEM3FlagModel使用。它的核心能力是三路并出一次前向同时得到三种检索信号——dense_vecs稠密向量用于向量相似度检索lexical_weights稀疏词权重类似 BM25 的词粒度匹配colbert_vecs多向量token 级向量用于 ColBERT 式精细交互打分。从源码结构看M3Embedder继承自抽象基类 AbsEmbedder后者统一提供了设备发现、多进程池、指令拼接、numpy 转换、维度截断等基础设施而M3Embedder负责 M3 特有的三路编码与打分逻辑最终前向计算由 EncoderOnlyEmbedderM3ModelForInference 完成。二、构造参数详解M3Embedder.__init__的参数分为三组基类通用参数、M3 特有模型参数、推理输出控制参数。from FlagEmbedding import BGEM3FlagModel model BGEM3FlagModel( model_name_or_pathBAAI/bge-m3, # 本地路径或 HuggingFace 模型名 normalize_embeddingsTrue, # 是否归一化稠密向量 use_fp16True, # 半精度推理 use_bf16False, # bf16 推理与 use_fp16 二选一 query_instruction_for_retrievalNone, # 查询侧指令 query_instruction_format{}{}, # 指令拼接模板 devicesNone, # 如 cuda:0 或 [cuda:0, cuda:1] pooling_methodcls, # 池化方式默认 cls trust_remote_codeFalse, cache_dirNone, colbert_dim-1, # colbert 线性层维度-1 表示使用 hidden_size batch_size256, query_max_length512, passage_max_length512, return_denseTrue, return_sparseFalse, return_colbert_vecsFalse, truncate_dimNone, # 截断维度Matryoshka 模型用 )各参数含义与默认值如下依据 m3.py 与基类 AbsEmbedder.py 的 docstring参数类型默认值说明model_name_or_pathstr必填本地模型目录或 HuggingFace Hub 上的模型名会自动下载normalize_embeddingsboolTrue归一化稠密向量使内积即余弦相似度use_fp16boolTrue半精度推理加速但有轻微精度损失use_bf16boolFalsebf16 推理基类get_model_torch_dtype中优先级高于 fp16query_instruction_for_retrievalOptional[str]None查询侧指令仅用于encode_queriesquery_instruction_formatstr{}{}指令模板如{}{}或含\n的模板devicesOptional[Union[str, List[str]]]None推理设备None时自动探测 cuda/npu/musa/mps/cpupooling_methodstrcls池化方式M3Embedder.DEFAULT_POOLING_METHOD为clstrust_remote_codeboolFalse是否信任远程自定义代码cache_dirOptional[str]NoneHF 模型缓存目录colbert_dimint-1ColBERT 线性投影维度-1时使用模型 hidden_sizebatch_sizeint256推理批大小query_max_lengthint512查询最大 token 长度passage_max_lengthint512段落最大 token 长度return_denseboolTrue是否返回稠密向量return_sparseboolFalse是否返回稀疏词权重return_colbert_vecsboolFalse是否返回 colbert 多向量truncate_dimOptional[int]None截断输出向量到指定维度Matryoshka 表示学习模型适用关于devices的解析细节基类静态方法get_target_devices支持字符串如cuda:0、整数如0映射为cuda:0或列表传None时按cuda → npu → musa → mps → cpu顺序自动探测可用设备。关于query_instruction_for_retrieval的生效路径encode_queries会把指令经query_instruction_format模板拼接到每个查询上AbsEmbedder.get_detailed_instruct且会处理模板中的\n转义而encode_corpus不会使用查询指令如需语料侧指令可通过kwargs传入passage_instruction_for_retrieval见AbsEmbedder.encode_corpus实现。三、三路编码输出格式encode_queries、encode_corpus、encode三个方法签名一致均返回一个字典{ dense_vecs: np.ndarray | None, # shape (N, hidden) 或 (N, truncate_dim) lexical_weights: List[Dict[str, float]] | None, # 每条文本一个 {token_id: weight} colbert_vecs: List[np.ndarray] | None, # 每条文本一个 token 级向量矩阵 }dense_vecs归一化后的稠密向量numpy 数组lexical_weights每个 token id 对应一个权重分数的字典列表。需要注意encode输出的lexical_weights的键是token id字符串形式如需还原成可读的 token 文本用convert_id_to_token转换colbert_vecs去掉[CLS]与 padding 后的 token 向量序列供后续colbert_score/compute_score使用。输出控制遵循方法级参数优先于构造参数的规则三个编码方法都接受batch_size、max_length、return_dense、return_sparse、return_colbert_vecs可选参数传None时回落到构造时设置的对应属性encode_queries默认长度上限用query_max_length而encode_corpus与encode默认用passage_max_length见 m3.py。四、encode_queries / encode_corpus / encode 的调用链这三个方法的行为差异在于指令与长度默认值encode_queries(queries)自动拼接query_instruction_for_retrieval指令默认max_lengthquery_max_length适合对检索查询编码encode_corpus(corpus)不拼查询指令默认max_lengthpassage_max_length适合对文档语料编码encode(sentences)通用入口默认max_lengthpassage_max_length可手动传instruction/instruction_format。调用链为encode_queries → AbsEmbedder.encode → encode_single_device单设备时若配置了多个设备则走start_multi_process_pool encode_multi_process由多个子进程并行处理再汇总汇总逻辑见 AbsEmbedder.py 的_concatenate_results_from_multi_process。五、encode_single_device 内部流程encode_single_device是单设备编码的核心实现m3.py其流程可以概括为四步无 padding 预分词先按batch_size分批调用 tokenizer只截断、不 padding得到每条的 input_ids按长度降序排序np.argsort([-len(x[input_ids]) ...])让同批文本长度接近、减少 padding 浪费自适应批大小先试探性前向一次若抛RuntimeError/torch.cuda.OutOfMemoryError则batch_size batch_size * 3 // 4缩小后重试循环直到成功批量前向并后处理每批调用self.model(inputs_batch, return_dense..., return_sparse..., return_colbert_vecs..., truncate_dimself.truncate_dim)随后dense 直接收集sparse 经_process_token_weights过滤掉cls/eos/pad/unk特殊 token 且权重大于 0 的项按 token id 聚合权重同一个 token 取最大权重colbert 经_process_colbert_vecs按attention_mask去掉 padding 向量并去掉[CLS]位置的向量取前tokens_num - 1个全部批次完成后按步骤 2 的排序索引np.argsort(length_sorted_idx)恢复原始输入顺序若输入是单个字符串还会把结果从单元素列表解包为标量。整体方法以torch.no_grad()装饰推理过程不计算梯度。六、稀疏词权重的还原与词粒度打分convert_id_to_tokenid 到 token 的还原convert_id_to_token(lexical_weights)把encode*返回的{token_id: weight}字典转为{token: weight}可读形式。实现上对每个 id 调用self.tokenizer.decode([int(id)])。输入可以是单个字典或字典列表若输入为单个字典返回时也会解包为单个字典而非列表。compute_lexical_matching_score词粒度匹配分compute_lexical_matching_score(lexical_weights_1, lexical_weights_2)计算两组词权重的匹配得分核心逻辑是词级点积对lw1中的每个 token若它也出现在lw2中累加weight1 * weight2。两个输入同时为dict时返回标量float同时为list时返回np.ndarray第[i][j]个元素是lexical_weights_1[i]与lexical_weights_2[j]的匹配分类型不匹配时抛出ValueError。官方单卡示例m3_single_device.py正是用它做稀疏召回sparse_scores model.compute_lexical_matching_score( queries_embeddings[lexical_weights], passages_embeddings[lexical_weights], )colbert_score多向量交互分colbert_score(q_reps, p_reps)接收 numpy 格式的多向量表示内部转 torch 后计算einsum(in,jn-ij)得到 token 级相似度矩阵沿最后一维取最大max-token 交互再对 batch 取平均返回torch.Tensor。七、compute_score五类相关度打分compute_score(sentence_pairs, batch_sizeNone, max_query_lengthNone, max_passage_lengthNone, weights_for_different_modesNone)输入(query, passage)对的列表一次性输出五种分数返回字典{ colbert: [...], # 多向量交互分 sparse: [...], # 稀疏匹配分 dense: [...], # 稠密相似度分 sparsedense: [...], # 稠密 稀疏 融合 colbertsparsedense: [...], # 三路融合 }关键行为见 m3.py 的compute_score_single_device查询与语料分别用max_query_length、max_passage_length截断默认 512前向时强制return_denseTrue, return_sparseTrue, return_colbert_vecsTrue, return_sparse_embeddingTrue三种表示全部计算单个分数分别来自底层模型的compute_dense_score、compute_sparse_score内积相似度除以温度见 modeling.py、compute_colbert_scoretoken 级einsum(qin,pjn-qipj)取 max 后按有效 token 数平均再除以温度融合权重weights_for_different_modes是一个长度为 3 的列表依次对应[dense, sparse, colbert]的权重传None时默认[1., 1., 1.]并打印日志sparsedense为(sparse*w1 dense*w0) / (w1 w0)colbertsparsedense为三路加权后除以权重总和weight_sum输入为单个 pair即列表首元素是字符串时返回的每个键是标量而非列表。compute_score对多设备场景会自动分发单设备时直接调compute_score_single_device多设备时启动_compute_score_multi_process_worker进程池经compute_score_multi_process按 chunk 切分chunk_size ceil(len(pairs) / 进程数)分发到各设备再从输出队列收集并按 chunk id 排序、由_concatenate_compute_score_results_from_multi_process拼接。该多进程实现参考自 sentence-transformers 的对应实现源码中有明确注释。八、单卡与多卡完整可运行示例仓库 examples/inference/embedder/encoder_only 下提供了可直接运行的四个 M3 示例。以下为单卡检索示例m3_single_device.pyimport os from FlagEmbedding import BGEM3FlagModel model BGEM3FlagModel( BAAI/bge-m3, devicescuda:0, # 无 GPU 时改用 cpu pooling_methodcls, cache_diros.getenv(HF_HUB_CACHE, None), ) queries [ What is BGE M3?, Defination of BM25 ] * 100 passages [ BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction., BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document ] * 100 queries_embeddings model.encode_queries( queries, return_denseTrue, return_sparseTrue, return_colbert_vecsFalse, ) passages_embeddings model.encode_corpus( passages, return_denseTrue, return_sparseTrue, return_colbert_vecsFalse, ) # 稠密打分向量内积向量已归一化 dense_scores queries_embeddings[dense_vecs] passages_embeddings[dense_vecs].T # 稀疏打分词权重匹配 sparse_scores model.compute_lexical_matching_score( queries_embeddings[lexical_weights], passages_embeddings[lexical_weights], )多卡版本m3_multi_devices.py只需把devices改为[cuda:0, cuda:1]无 GPU 时可用[cpu, cpu]其余代码不变。直接打分示例m3_single_device_compute_score.py演示了三路融合打分注意这里把 sparse 权重压低为 0.3sentence_pairs list(zip(queries, passages)) scores_dict model.compute_score( sentence_pairs, weights_for_different_modes[1., 0.3, 1.] # dense, sparse, colbert 权重 ) print(scores_dict)示例注释中给出了BAAI/bge-m3在这些样例输入上的期望输出例如compute_score第一对输入约为{colbert: 0.7798, sparse: 0.1956, dense: 0.6260, sparsedense: 0.5266, colbertsparsedense: 0.6367}可供本地运行后对照验证实现一致性。九、其他实用说明精度与设备use_fp16True会在 GPU 上以半精度推理CPU 上encode_single_device会把模型转回float()。基类的_convert_to_numpy对 bf16 推理做了兼容——NumPy 不支持 bfloat16非 CPU 设备上 bf16 张量会先升到 float32 再转 numpyAbsEmbedder.py。维度截断truncate_dim支持 Matryoshka 式截断编码时会透传给模型_truncate_embeddings负责按[..., :truncate_dim]切分。多进程资源释放基类提供stop_multi_process_pool终止并回收进程对象析构时__del__会调用stop_self_pool释放模型到 CPU 并清空显存。源码地图核心实现位于 m3.py底层前向与三种打分见 modeling.py抽象基类见 AbsEmbedder.py可复现示例见 examples/inference/embedder/encoder_only测试用例见 tests/test_infer_embedder_basic.py。掌握以上 API 之后你就可以基于BGEM3FlagModel自由组合 dense / sparse / colbert 三种信号离线对语料做三路索引在线对查询做三路编码再用compute_score的加权融合公式做最终排序构建自己的混合检索系统。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表