ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

FlagEmbedding BaseEmbedder 深度指南:encoder-only 稠密检索模型的查询与语料编码实战

FlagEmbedding BaseEmbedder 深度指南:encoder-only 稠密检索模型的查询与语料编码实战 FlagEmbedding BaseEmbedder 深度指南encoder-only 稠密检索模型的查询与语料编码实战【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding本篇技术指南以 FlagEmbedding 官方 API 文档 BaseEmbedder.rst 为骨架系统讲解BaseEmbedder类的全部构造参数、五个核心编码方法与底层实现原理。结合仓库内的源码、可运行示例与单元测试读者将掌握如何加载 BGE 系列等 encoder-only 稠密向量模型、如何正确区分并调用encode_queries/encode_corpus/encode完成检索任务的向量化以及查询指令拼接、设备分配、多进程并行、向量归一化与截断等内部机制从而在实际检索系统中写出高效、正确的推理代码。BaseEmbedder 在 FlagEmbedding 推理体系中的定位BaseEmbedder位于 FlagEmbedding/inference/embedder/encoder_only/base.py是 FlagEmbedding 面向encoder-only 架构即传统的 BERT 类双向编码器稠密向量模型提供的推理类。其类层次关系为AbsEmbedder抽象基类定义通用推理框架 └── BaseEmbedderencoder-only 模型的具体实现 ├── FlagModel公开别名见 encoder_only/__init__.py └── M3EmbedderBGE-M3 多向量/稀疏向量扩展见 m3.py其中 AbsEmbedder 定义了encode_queries、encode_corpus、encode的调用约定并实现了设备解析、指令拼接、多进程池等通用逻辑BaseEmbedder则补全了模型加载、tokenizer 初始化和单设备编码的具体实现。仓库在 encoder_only/init.py 中将BaseEmbedder直接导出为FlagModel这也是历史版本与现有教程中from FlagEmbedding import FlagModel的由来。从模型映射表 model_mapping.py 可以看到BGE 全系列bge-small/base/large 的 en/zh 与 v1.5 版本、E5 系列e5-large-v2、multilingual-e5-*、部分 GTE 模型以及 BCE 模型均通过FlagModel即BaseEmbedder加载且多数使用cls池化。构造函数与全部参数详解BaseEmbedder.__init__的完整签名见 base.py所有参数及默认值整理如下参数类型默认值作用model_name_or_pathstr必填本地模型路径或 HuggingFace Hub 上的模型名如BAAI/bge-small-en-v1.5normalize_embeddingsboolTrue是否对输出向量做 L2 归一化归一化后可直接用点积计算余弦相似度use_fp16boolTrue使用半精度浮点加速推理性能略有下降use_bf16boolFalse使用 bfloat16与use_fp16二选一优先级更高见 get_model_torch_dtypequery_instruction_for_retrievalOptional[str]None查询侧指令例如Represent this sentence for searching relevant passages: query_instruction_formatstr{}{}指令模板{}分别被指令与查询文本填充含\n时会被转换为换行符见 get_detailed_instructdevicesOptional[Union[str, List[str]]]None推理设备如cuda:0或[cuda:0, cuda:1]为None时按 CUDA → NPU → MUSA → MPS → CPU 顺序自动探测见 get_target_devicespooling_methodstrcls池化方式当前支持cls与mean类属性DEFAULT_POOLING_METHOD clstrust_remote_codeboolFalse加载依赖自定义代码的模型时置为Truecache_dirOptional[str]None模型缓存目录batch_sizeint256推理批大小query_max_lengthint512查询最大 token 数passage_max_lengthint512语料passage最大 token 数convert_to_numpyboolTrue输出为 NumPy 数组置False时输出 PyTorch Tensortruncate_dimOptional[int]None输出向量截断维度用于 Matryoshka 表示学习模型见 _truncate_embeddings构造时类通过AutoTokenizer.from_pretrained与AutoModel.from_pretrained加载模型并依据get_model_torch_dtype()决定加载精度base.py。其余**kwargs会被透传给 Transformers 配置或作为实例属性保存例如passage_instruction_for_retrieval/passage_instruction_format便是通过 kwargs 传入以控制语料侧指令见 AbsEmbedder.encode_corpus。五个核心编码方法BaseEmbedder对外暴露五个方法对应原文档 BaseEmbedder.rst 中的 Methods 列表它们之间存在清晰的委托关系encode_queries() ──┐ encode_corpus() ──┼──► encode() ──► encode_single_device()单设备或多进程编码 encode() ──┘ pooling() ← 在 encode_single_device 内部被调用encode_queries查询侧编码自动附加查询指令签名encode_queries(queries, batch_sizeNone, max_lengthNone, convert_to_numpyNone)。queries可以是单个字符串或字符串列表。当构造时指定了query_instruction_for_retrieval查询文本会先按query_instruction_format模板拼接指令再进入编码流程max_length缺省时使用query_max_lengthbatch_size缺省时使用构造参数AbsEmbedder.py。这一点非常关键同一模型在检索场景中query 与 passage 必须使用不同的入口方法才能保证 query 带指令、passage 不带指令从而复现训练时的不对称设定。encode_corpus语料侧编码签名encode_corpus(corpus, batch_sizeNone, max_lengthNone, convert_to_numpyNone)。语料侧默认不附加指令除非通过 kwargs 传入passage_instruction_for_retrieval与passage_instruction_formatmax_length缺省时使用passage_max_lengthAbsEmbedder.py。这正是 BGE 系列查询带指令、文档不带指令检索范式的代码落地。encode通用编码入口签名encode(sentences, batch_sizeNone, max_lengthNone, convert_to_numpyNone, instructionNone, instruction_formatNone)。这是最底层的统一入口先按需完成指令拼接然后根据输入类型与设备数量分派——单个字符串或单设备时直接调用encode_single_device多设备且输入为列表时则启动多进程池并行编码AbsEmbedder.py。注意encode本身不区分 query/corpus直接调用时max_length缺省取passage_max_length。encode_single_device单设备编码核心签名encode_single_device(sentences, batch_size256, max_length512, convert_to_numpyTrue, deviceNone)被torch.no_grad()装饰base.py。其执行管线是理解推理性能的关键设备就绪device缺省取target_devices[0]在 CPU 上推理时调用self.model.float()避免 fp16 在 CPU 上的精度/速度问题预 tokenize先按batch_size分批调用 tokenizertruncationTruemax_lengthmax_length不做 padding得到每条的准确长度按长度降序排序np.argsort([-len(x[input_ids]) for x in all_inputs])同一 batch 内文本长度相近显著减少 padding 浪费自适应批大小先用目标 batch 试跑一次捕获RuntimeError或torch.cuda.OutOfMemoryError时按batch_size batch_size * 3 // 4逐步缩减避免显存溢出导致整任务失败正式编码对排序后的批次tokenizer.pad补齐、前向得到last_hidden_state调用pooling提取向量归一化与截断依次执行truncate_dim截断与normalize_embeddingsL2 归一化顺序还原编码完成后按np.argsort(length_sorted_idx)恢复原始输入顺序单个字符串输入时返回一维向量否则返回二维数组。pooling池化函数签名pooling(last_hidden_state, attention_maskNone)base.pycls直接取last_hidden_state[:, 0]即[CLS]token 的隐状态mean对last_hidden_state按attention_mask做掩码加权平均分母为有效 token 数避免 padding 污染其他取值抛出NotImplementedError。从映射表可知BGE 系列默认cls而e5-*、gte-large/base/small等模型在FlagAutoModel自动加载时会被指定为mean与各模型官方训练设定保持一致见 model_mapping.py。多设备并行的底层机制当devices传入多个 GPU 且输入为列表时encode会走多进程路径模型先to(cpu)并share_memory()随后为每个目标设备spawn一个独立 worker 进程通过输入/输出队列分发文本块并回收结果AbsEmbedder.py。分块大小按ceil(len(sentences) / len(processes))计算各进程独立调用encode_single_device最后按 chunk 序号排序拼接。对象析构时__del__会自动调用stop_self_pool()终止进程并释放显存AbsEmbedder.py。实战示例单设备查询与语料编码仓库提供了可直接运行的示例 base_single_device.py完整代码如下import os from FlagEmbedding import FlagModel def test_base_single_device(): model FlagModel( BAAI/bge-small-en-v1.5, query_instruction_for_retrievalRepresent this sentence for searching relevant passages: , query_instruction_format{}{}, devicescuda:0, # if you dont have a GPU, you can use cpu pooling_methodcls, cache_diros.getenv(HF_HUB_CACHE, None), ) queries [ What is the capital of France?, What is the population of China?, ] * 100 passages [ Paris is the capital of France., The population of China is over 1.4 billion people. ] * 100 queries_embeddings model.encode_queries(queries) passages_embeddings model.encode_corpus(passages) cos_scores queries_embeddings passages_embeddings.T print(cos_scores[:2, :2]) if __name__ __main__: test_base_single_device() print(--------------------------------) print(Expected Output:) print([[0.7944 0.4492]\n [0.58 0.801 ]])要点解读encode_queries会自动为每个 query 拼接Represent this sentence for searching relevant passages: 前缀而encode_corpus不加指令由于normalize_embeddingsTrue默认开启两个矩阵的点积即余弦相似度对角元0.7944、0.801明显高于非对角元符合语义匹配预期示例注释明确了devicescuda:0无 GPU 时可直接换成cpu多设备版本见 base_multi_devices.py对应 BGE-M3 的版本见 m3_single_device.py 与 auto_base_single_device.py后者演示FlagAutoModel.from_finetuned的自动加载。与 M3Embedder 的关系与差异同为 encoder-only 推理类M3Embedder 继承了AbsEmbedder而非BaseEmbedder但共享同一套encode_queries/encode_corpus/encode调用骨架。差异在于M3Embedder 额外支持return_dense、return_sparse、return_colbert_vecs、colbert_dim参数可同时输出稠密向量、稀疏词权重与 ColBERT 多向量其 API 文档见 M3Embedder.rst模型加载不走AutoModel而是包装微调侧提供的EncoderOnlyEmbedderM3ModelForInference稀疏词权重由_process_token_weights过滤掉 cls/eos/pad/unk 特殊 token 后按 token id 累积权重得到ColBERT 向量则剔除 cls 位置m3.py。因此纯稠密检索场景直接使用FlagModelBaseEmbedder即可需要混合检索稠密稀疏多向量时再升级为BGEM3FlagModel。自动加载FlagAutoModel 与模型映射FlagAutoModel.from_finetuned会依据模型名的 basename 在 model_mapping.py 的AUTO_EMBEDDER_MAPPING中查找对应配置自动为FlagModel选定正确的pooling_method、trust_remote_code与query_instruction_formatauto_embedder.py。例如加载bge-small-zh-v1.5会自动使用cls池化加载e5-large-v2会自动改用mean池化若模型名不在映射表中则抛出ValueError此时可显式传入model_class如encoder-only-base手动指定加载方式。这也是BaseEmbedder参数被自动填充的标准入口详情见 FlagAutoModel。测试验证与使用建议仓库单元测试 test_infer_embedder_basic.py 以BAAI/bge-base-en-v1.5为对象验证了 BaseEmbedder 的以下行为可作为接入自检清单单字符串输入encode返回一维np.ndarray且无 NaN查询与相关语料的余弦相似度落在[0, 1]且大于 0.5批量输入返回二维数组shape[0]等于输入条数。实际使用时的几点结论检索任务务必用encode_queries编码查询、encode_corpus编码语料以保持查询指令的不对称性大规模语料离线向量化时建议调大batch_size并结合cache_dir复用模型缓存多卡环境传入设备列表即可自动并行需要 Matryoshka 截断能力如降维存储时通过truncate_dim在推理阶段直接截断输出向量输出形态可通过convert_to_numpy在 NumPy 数组与 Torch Tensor 之间切换use_bf16在非 CPU 设备上会自动安全转 float32 后再转 NumPy见 _convert_to_numpy。综上BaseEmbedder是 FlagEmbedding 中 encoder-only 稠密向量模型推理的基石上层有FlagAutoModel自动装配参数下层有AbsEmbedder提供多设备与指令处理框架自身则以预 tokenize → 长度排序 → 池化 → 归一化的高效管线完成向量化是搭建 BGE 系列检索系统的首选入口。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表