FlagEmbedding项目教程:BGE Auto Embedder使用指南
概述
FlagEmbedding项目提供了一个名为FlagAutoModel
的高级封装类,它统一了多种流行开源嵌入模型的推理接口。除了BGE系列模型外,还支持E5、GTE、SFR等其他优秀模型。本文将详细介绍如何使用这个强大的工具。
安装准备
首先需要安装FlagEmbedding库:
pip install FlagEmbedding
基础使用方法
模型初始化
从FlagEmbedding导入FlagAutoModel
类,并使用from_finetuned()
函数初始化模型:
from FlagEmbedding import FlagAutoModel
model = FlagAutoModel.from_finetuned(
'BAAI/bge-base-en-v1.5',
query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
devices="cuda:0" # 可指定使用GPU或CPU
)
关键参数说明:
- 第一个参数指定模型名称或路径
query_instruction_for_retrieval
设置查询指令前缀devices
参数控制计算设备
编码与相似度计算
初始化后,使用方法与FlagModel
完全一致:
queries = ["query 1", "query 2"]
corpus = ["passage 1", "passage 2"]
# 编码查询和文本
q_embeddings = model.encode_queries(queries)
p_embeddings = model.encode_corpus(corpus)
# 计算相似度分数
scores = q_embeddings @ p_embeddings.T
print(scores)
支持模型列表
FlagAutoModel
内部维护了一个有序字典MODEL_MAPPING
,存储所有支持的模型配置。当前支持的模型包括:
- BGE系列:bge-m3、bge-large-en-v1.5、bge-base-en-v1.5等
- E5系列:e5-mistral-7b-instruct、e5-large-v2等
- GTE系列:gte-Qwen2-7B-instruct、gte-large-en-v1.5等
- SFR系列:SFR-Embedding-2_R、SFR-Embedding-Mistral等
- 其他模型:Linq-Embed-Mistral等
模型配置解析
每个模型对应一个EmbedderConfig
配置对象,包含四个关键属性:
model_class
:模型类,决定是编码器还是解码器pooling_method
:池化方法(CLS token、平均池化或最后token)trust_remote_code
:是否需要远程代码信任query_instruction_format
:查询指令格式
例如,BGE-ICL模型的配置如下:
EmbedderConfig(
model_class=ICLLLMEmbedder,
pooling_method=PoolingMethod.LAST_TOKEN,
trust_remote_code=False,
query_instruction_format='<instruct>{}\n<query>{}'
)
自定义模型支持
如果需要将自己的模型接入FlagAutoModel
系统,可以按照以下步骤操作:
- 确定模型类型:区分是编码器还是解码器架构
- 选择池化方法:CLS token、平均池化或最后token
- 设置远程代码信任:根据模型需求决定
- 定义查询指令格式:如有检索需求则需设置
确认这四个属性后,将模型名称和对应的EmbedderConfig
添加到MODEL_MAPPING
中即可。
最佳实践建议
- 设备选择:对于大型模型,建议使用GPU加速
- 批处理:对大量文本进行编码时,合理设置batch_size
- 模型选择:根据任务需求选择合适大小的模型
- 指令优化:针对特定任务调整查询指令格式
通过FlagAutoModel
的统一接口,开发者可以轻松切换不同嵌入模型,大大提高了实验效率和代码复用性。
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考