ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

FlagEmbedding 实战:使用 BGE-Code-v1 构建面向代码检索与多语言检索的 LLM 嵌入系统

FlagEmbedding 实战:使用 BGE-Code-v1 构建面向代码检索与多语言检索的 LLM 嵌入系统 FlagEmbedding 实战使用 BGE-Code-v1 构建面向代码检索与多语言检索的 LLM 嵌入系统【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbeddingBGE-Code-v1代码名 CodeR是 BAAI 发布的首个基于 LLMdecoder-only架构的代码嵌入模型在 FlagEmbedding 项目中通过FlagLLMModel统一封装。本文以官方文档 bge_code.rst 为核心骨架结合 FlagLLMModel 源码 与 BGE_Coder 研究与评测仓库完整讲解该模型的三种检索能力、参数语义、底层 last-token 池化原理、多设备推理、任务指令模板以及 CoIR / CodeRAG 评测复现帮助你直接落地代码检索、Text2SQL、代码补全检索等真实场景。模型概览BGE-Code-v1 是什么BGE-Code-v1 是一个LLM-based 代码嵌入模型参数量 1.5B模型文件约 6.18 GB多语言官方文档归纳其三大核心能力代码检索能力支持英文与中文的自然语言查询覆盖20 种编程语言的代码检索文本检索能力保持与同规模文本嵌入模型相当的文字检索性能多语言检索能力在英语、中文、日语、法语等语言上均有出色的检索表现。模型语言参数量模型大小说明BAAI/bge-code-v1多语言1.5B6.18 GB官方定位为 SOTA 级代码检索模型同时具备多语言文本检索能力在 FlagEmbedding 的模型映射表 model_mapping.py 中bge-code-v1被注册到BGE_MAPPING对应配置为bge-code-v1, EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_codeTrue, query_instruction_formatinstruct{}\nquery{})这说明它是一个decoder-only 基座模型由FlagLLMModel加载、采用last_token 池化、需要trust_remote_codeTrue模型仓库自带自定义代码并且查询侧必须套用instruct{}\nquery{}指令模板。从源码结构可以推断其与bge-en-icl、bge-multilingual-gemma2、Qwen3-Embedding、E5、GTE 等共享同一套 LLM 嵌入推理框架属于 FlagEmbedding 中 decoder-only 嵌入模型家族的通用范式。环境准备与安装BGE-Code-v1 基于 decoder-only 大模型需要依赖 FlagEmbedding 的 LLM 推理路径安装方式与项目整体一致git clone https://github.com/FlagOpen/FlagEmbedding.git cd FlagEmbedding pip install -e .安装完成后即可导入FlagLLMModel其别名定义见 decoder_only/init.py 与 embedder/init.py。由于模型约 6.18 GB建议具备支持 CUDA 的 GPU 环境无 GPU 时也可通过devicescpu运行但推理速度会明显下降。快速上手用 FlagLLMModel 完成代码检索官方文档给出了一个非常直观的 Text2SQL 检索示例用自然语言查询去检索对应的 SQL 语句。from FlagEmbedding import FlagLLMModel queries [ Delete the record with ID 4 from the Staff table., Delete all records in the Livestock table where age is greater than 5 ] documents [ DELETE FROM Staff WHERE StaffID 4;, DELETE FROM Livestock WHERE age 5; ] model FlagLLMModel(BAAI/bge-code-v1, query_instruction_formatinstruct{}\nquery{}, query_instruction_for_retrievalGiven a question in text, retrieve SQL queries that are appropriate responses to the question., trust_remote_codeTrue, use_fp16True) # Setting use_fp16 to True speeds up computation with a slight performance degradation embeddings_1 model.encode_queries(queries) embeddings_2 model.encode_corpus(documents) similarity embeddings_1 embeddings_2.T print(similarity)要点拆解encode_queries负责对查询侧做编码encode_corpus负责对文档/代码侧做编码。两者在 base.py 中最终都汇入统一的encode→encode_single_device管线查询与文档使用不同的处理路径查询会套上instruct{指令}\nquery{查询}模板由query_instruction_format与query_instruction_for_retrieval拼装而文档侧不附加指令这是检索模型的标准不对称编码设计相似度直接用向量内积embeddings_1 embeddings_2.T计算。由于模型默认会对向量做 L2 归一化normalize_embeddingsTrue内积等价于余弦相似度数值区间约在 [0, 1]。官方给出的任务指令示例对应 Text2SQL 场景Given a question in text, retrieve SQL queries that are appropriate responses to the question.结合query_instruction_formatinstruct{}\nquery{}实际送入模型的查询文本会被构造成instructGiven a question in text, retrieve SQL queries that are appropriate responses to the question. queryDelete the record with ID 4 from the Staff table.构造参数与默认值全解FlagLLMModel的核心参数定义在 BaseLLMEmbedder 构造函数 中结合源码注释整理如下参数默认值说明model_name_or_path必填模型名如BAAI/bge-code-v1或本地模型路径normalize_embeddingsTrue是否对嵌入向量做 L2 归一化为True时内积即余弦相似度use_fp16True半精度推理显著加速仅带来轻微性能损失use_bf16False是否使用 bfloat16 精度query_instruction_for_retrievalNone检索任务的查询指令与query_instruction_format配合使用query_instruction_formatInstruct: {}\nQuery: {}指令模板BGE-Code-v1 必须改为instruct{}\nquery{}devicesNone推理设备如cuda:0或[cuda:0, cuda:1]None时默认使用全部可用 GPUtrust_remote_codeFalse是否信任 HuggingFace Hub 上模型的远程代码BGE-Code-v1 必须设为Truecache_dirNone模型缓存目录可用环境变量HF_HUB_CACHE指定batch_size256推理批大小显存不足时会自动按3/4比例回退query_max_length512查询最大 token 长度passage_max_length512文档/代码最大 token 长度convert_to_numpyTrue输出为 numpy 数组否则返回 torch Tensortruncate_dimNone截断嵌入维度如做 Matryoshka 降维两个需要特别注意的约束池化方式锁定源码在初始化末尾强制校验——if self.kwargs.get(pooling_method, last_token) ! last_token: raise ValueError(Pooling method must be last_token for LLM-based models.)。即 LLM 嵌入模型只允许 last_token 池化传入其他池化方式会直接报错指令模板必须匹配BGE-Code-v1 的映射表中模板为instruct{}\nquery{}如果沿用默认的Instruct: {}\nQuery: {}检索效果会显著下降因为模型微调时就是按前者组织输入。底层原理last-token 池化与推理管线LLM 嵌入模型与 BERT 类编码器模型的本质区别在于池化策略。BGE-Code-v1 取最后一个 token 的隐藏状态作为整段文本的向量对应实现 last_token_pooldef last_token_pool(last_hidden_states, attention_mask): left_padding (attention_mask[:, -1].sum() attention_mask.shape[0]) if left_padding: return last_hidden_states[:, -1] else: sequence_lengths attention_mask.sum(dim1) - 1 batch_size last_hidden_states.shape[0] return last_hidden_states[torch.arange(batch_size), sequence_lengths]该函数同时处理左填充与右填充两种情况若批内全部为左填充attention_mask最后一列全为 1直接取[:, -1]否则按每个样本的真实序列长度attention_mask.sum(dim1) - 1取对应位置的隐藏状态保证无论填充方向如何都拿到“真正最后一个 token”的向量。在 encode_single_device 的完整推理管线中还有几个值得关注的工程细节预分词 按长度降序排序先对全部句子分词truncationTrue、max_lengthmax_length再按长度排序分批减少 padding 浪费显存自适应回退首轮尝试batch_size批量前向若抛出RuntimeError或torch.cuda.OutOfMemoryError自动将 batch 缩减为原来的 3/4 后重试推理结束恢复顺序输出向量在拼接后按np.argsort(length_sorted_idx)恢复原始输入顺序调用方无需关心内部排序归一化与类型转换normalize_embeddings时执行torch.nn.functional.normalize(embeddings, dim-1)convert_to_numpyTrue时最终np.concatenate为单个数组。多设备与显存控制FlagLLMModel默认会使用所有可用 GPU做并行推理。官方建议通过环境变量精确控制可见设备import os os.environ[CUDA_VISIBLE_DEVICES] 0,1 # 只使用 0、1 号 GPU # os.environ[CUDA_VISIBLE_DEVICES] # 强制全部使用 CPU也可以直接在构造函数中显式指定devices参考官方示例 base_single_device.pymodel FlagLLMModel( BAAI/bge-code-v1, query_instruction_for_retrievalGiven a question in text, retrieve SQL queries that are appropriate responses to the question., query_instruction_formatinstruct{}\nquery{}, devicescuda:0, # 单卡无 GPU 时可改为 cpu cache_diros.getenv(HF_HUB_CACHE, None), )同目录下的base_multi_devices.py、auto_base_single_device.py等示例进一步展示了多卡与 Auto 模型的写法。值得说明的是BGE-Code-v1 属于 1.5B 参数量级的模型若显存不足除了调小batch_size开启use_fp16True也是官方文档明确推荐的第一优化手段。其他加载方式Sentence Transformers 与原生 Transformers除了 FlagEmbeddingresearch/BGE_Coder/README.md还提供了两种等价用法供已依赖其他生态的团队参考。方式一Sentence Transformersfrom sentence_transformers import SentenceTransformer import torch model SentenceTransformer( BAAI/bge-code-v1, trust_remote_codeTrue, model_kwargs{torch_dtype: torch.float16}, ) instruction Given a question in text, retrieve SQL queries that are appropriate responses to the question. prompt finstruct{instruction}\nquery queries [Delete the record with ID 4 from the Staff table.] documents [DELETE FROM Staff WHERE StaffID 4;] query_embeddings model.encode(queries, promptprompt) # 查询带 prompt document_embeddings model.encode(documents) # 文档不带 prompt similarities model.similarity(query_embeddings, document_embeddings)方式二原生 HuggingFace Transformersimport torch import torch.nn.functional as F from torch import Tensor from transformers import AutoTokenizer, AutoModel def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) - Tensor: left_padding (attention_mask[:, -1].sum() attention_mask.shape[0]) if left_padding: return last_hidden_states[:, -1] else: sequence_lengths attention_mask.sum(dim1) - 1 batch_size last_hidden_states.shape[0] return last_hidden_states[torch.arange(batch_size), sequence_lengths] tokenizer AutoTokenizer.from_pretrained(BAAI/bge-code-v1, trust_remote_codeTrue) model AutoModel.from_pretrained(BAAI/bge-code-v1, trust_remote_codeTrue) model.eval() max_length 4096 # 长代码场景可放大 max_length input_texts queries documents batch_dict tokenizer(input_texts, max_lengthmax_length, paddingTrue, truncationTrue, return_tensorspt, pad_to_multiple_of8) with torch.no_grad(): outputs model(**batch_dict) embeddings last_token_pool(outputs.last_hidden_state, batch_dict[attention_mask]) embeddings F.normalize(embeddings, p2, dim1) scores (embeddings[:2] embeddings[2:].T) * 100 # 放大 100 倍展示注意两种方式都必须设置trust_remote_codeTrue且查询侧要手动拼接instruct{指令}\nquery前缀原生写法中pad_to_multiple_of8与max_length4096是官方评测脚本的习惯配置用于长代码片段对齐。任务指令模板让检索对齐目标任务BGE-Code-v1 的能力高度依赖query_instruction_for_retrieval的任务指令。research/BGE_Coder/README.md给出了完整指令集合可直接复制使用任务指令query_instruction_for_retrievalAppsGiven a code contest problem description, retrieve relevant code that can help solve the problem.CosQAGiven a web search query, retrieve relevant code that can help answer the query.Text2SQLGiven a question in text, retrieve SQL queries that are appropriate responses to the question.CSNGiven a piece of code, retrieve the document string that summarizes the code.CSN-CCRGiven a piece of code segment, retrieve the code segment that is the latter part of the code.CodeTrans-DLGiven a piece of code, retrieve code that is semantically equivalent to the input code.CodeTrans-ContestGiven a piece of Python code, retrieve C code that is semantically equivalent to the input code.StackOverFlow-QAGiven a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.CodeFeedBack-STGiven a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.CodeFeedBack-MTGiven a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.HummanEvalGiven a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.MBPPGiven a textual explanation of code functionality, retrieve the corresponding code implementation.DS-1000Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.ODEXGiven a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.RepoEvalGiven a piece of code segment, retrieve the code segment that is the latter part of the code.SWE-bench-LiteGiven a code snippet containing a bug and a natural language description of the bug or error, retrieve code snippets that demonstrate solutions or fixes for similar bugs or errors (the desired documents).这些指令覆盖了题目检索、Web 查询、Text2SQL、代码-文档互检索、跨语言代码翻译、缺陷修复检索等典型代码检索场景。在实际业务中你可以仿照该格式为自定义任务编写指令例如针对你仓库的 API 文档检索然后通过query_instruction_for_retrieval传入。评测基准与复现CoIR 与 CodeRAGresearch/BGE_Coder/README.md中报告了 BGE-Code-v1 在两大代码检索基准上的成绩数据出自仓库可直接复现核验。CoIR 基准nDCG10在 Apps、CosQA、Text2SQL、CSN、CSN-CCR、CodeTrans-Contest、CodeTrans-DL、StackOverFlow-QA、CodeFeedBack-ST、CodeFeedBack-MT 十个子任务上平均81.77其中 Apps98.08、CSN-CCR98.30、CodeFeedBack-MT94.38等子任务表现突出。CodeRAG 基准Recall覆盖 HummanEval、MBPP、DS-1000、ODEX、RepoEval、SWE-bench-Lite 六个任务平均72.8其中 DS-1000 达到 40.9、SWE-bench-Lite 达到 67.4。评测脚本位于 research/BGE_Coder/evaluationCoIR进入evaluation/coir_eval目录克隆 CoIR 评测工具包后执行bash eval.sh入口为 coir_eval/main.pyCodeRAG进入evaluation/coderag_eval目录克隆 CodeRAG 评测工具包将test/下的任务构造脚本复制进工具包依次执行bash prepare_data.sh与bash eval.sh任务构造脚本位于 coderag_eval/test/create覆盖 humaneval、mbpp、ds1000、odex、repoeval、swebench 等数据集。需要提醒的是以上数据与结论均为仓库research/BGE_Coder目录内自行评测报告的结果复现时应以相同评测脚本与数据划分为准。延伸模型的数据生成与训练管线research/BGE_Coder目录还完整开放了 CodeR / BGE-Code-v1 的数据生成管线位于 data_generationconstant.py任务与指令常量定义corpus_generator.py代码语料生成triplet_generator.py检索三元组查询、正例、负例构造llm.py与run_generation.py调用大模型批量生成合成数据search.py负样本挖掘hard negative mining相关检索format_generated_examples.py生成样本格式化。如果你想基于 BGE-Code-v1 微调出面向特定语言或特定业务代码的嵌入模型可以参考该管线构造训练数据再接入 FlagEmbedding 的 finetune/embedder/decoder_only 微调框架。小结BGE-Code-v1 是 FlagEmbedding 生态中首个以 LLM 为基座的代码嵌入模型其“last-token 池化 指令模板”的组合是 decoder-only 嵌入模型的通用范式。本文从官方文档 bge_code.rst 出发依次覆盖了模型能力概览、FlagLLMModel快速上手、全部构造参数、底层池化与推理管线、多设备控制、Sentence Transformers / 原生 Transformers 等价用法、任务指令模板以及 CoIR / CodeRAG 评测复现路径。对想快速落地的开发者直接复制“快速上手”一节即可跑通 Text2SQL 检索对想做深度定制的团队research/BGE_Coder的数据生成与评测代码则提供了完整的参考实现。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表