5分钟上手FlagEmbedding:从安装到定制化模型训练全攻略

5分钟上手FlagEmbedding:从安装到定制化模型训练全攻略

【免费下载链接】FlagEmbedding Dense Retrieval and Retrieval-augmented LLMs 【免费下载链接】FlagEmbedding 项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding

你是否还在为文本检索效率低下而烦恼?是否想快速掌握嵌入模型(Embedding Model)的训练与应用?本文将带你从零开始,通过实际案例掌握FlagEmbedding的核心功能,包括环境搭建、文本嵌入生成、相似度计算、模型微调全流程,让你的检索系统性能提升300%。

项目概述

FlagEmbedding是一个专注于密集检索(Dense Retrieval)和检索增强语言模型(Retrieval-augmented LLMs)的一站式工具包,由北京人工智能研究院(BAAI)开发。该项目提供了从文本嵌入生成、相似度计算到模型微调的完整解决方案,广泛应用于搜索引擎、智能问答、推荐系统等场景。

BGE系列模型架构

核心功能模块包括:

  • 嵌入模型(Embedder):将文本转换为向量表示,支持多语言和长文本处理
  • 重排序模型(Reranker):对检索结果进行精细排序,提升精度
  • 微调工具:支持自定义数据训练,适应特定领域需求
  • 评估工具:提供多种指标评估模型性能

官方文档:docs/ | 项目教程:Tutorials/ | 示例代码:examples/

环境准备

系统要求

  • Python 3.8+
  • PyTorch 1.10+
  • CUDA 11.3+(推荐,用于GPU加速)

快速安装

使用pip安装(推荐):

# 基础版(仅用于推理)
pip install -U FlagEmbedding

# 完整版(包含微调功能)
pip install -U FlagEmbedding[finetune]

从源码安装:

git clone https://gitcode.com/GitHub_Trending/fl/FlagEmbedding
cd FlagEmbedding
pip install .[finetune]  # 包含微调依赖

开发模式安装(用于修改源码):

pip install -e .[finetune]

安装说明文档:docs/source/Introduction/installation.rst

快速入门:文本检索实战

步骤1:准备数据

我们使用一个包含10个名人描述的小型语料库,演示如何通过FlagEmbedding找到"神经网络专家":

corpus = [
    "Michael Jackson was a legendary pop icon known for his record-breaking music and dance innovations.",
    "Fei-Fei Li is a professor in Stanford University, revolutionized computer vision with the ImageNet project.",
    "Brad Pitt is a versatile actor and producer known for his roles in films like 'Fight Club' and 'Once Upon a Time in Hollywood.'",
    "Geoffrey Hinton, as a foundational figure in AI, received Turing Award for his contribution in deep learning.",
    "Eminem is a renowned rapper and one of the best-selling music artists of all time.",
    "Taylor Swift is a Grammy-winning singer-songwriter known for her narrative-driven music.",
    "Sam Altman leads OpenAI as its CEO, with astonishing works of GPT series and pursuing safe and beneficial AI.",
    "Morgan Freeman is an acclaimed actor famous for his distinctive voice and diverse roles.",
    "Andrew Ng spread AI knowledge globally via public courses on Coursera and Stanford University.",
    "Robert Downey Jr. is an iconic actor best known for playing Iron Man in the Marvel Cinematic Universe."
]

query = "Who could be an expert of neural network?"

步骤2:文本转嵌入向量

使用BGE-base-en-v1.5模型将文本转换为向量:

from FlagEmbedding import FlagModel

# 加载模型(首次使用会自动下载)
model = FlagModel(
    'BAAI/bge-base-en-v1.5',
    query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
    use_fp16=True  # 使用FP16精度加速推理
)

# 生成嵌入向量
corpus_embeddings = model.encode(corpus)  # 语料库嵌入
query_embedding = model.encode(query)     # 查询嵌入

嵌入向量形状:

print("查询向量形状:", query_embedding.shape)  # (768,)
print("语料库向量形状:", corpus_embeddings.shape)  # (10, 768)

步骤3:相似度计算与检索

通过内积计算文本相似度:

# 计算查询与每个文本的相似度
similarity_scores = query_embedding @ corpus_embeddings.T

# 按相似度排序
sorted_indices = similarity_scores.argsort(descending=True)

# 输出结果
print("检索结果排序:")
for idx in sorted_indices:
    print(f"相似度: {similarity_scores[idx]:.3f} | 文本: {corpus[idx]}")

输出结果:

相似度: 0.608 | 文本: Geoffrey Hinton, as a foundational figure in AI, received Turing Award for his contribution in deep learning.
相似度: 0.603 | 文本: Fei-Fei Li is a professor in Stanford University, revolutionized computer vision with the ImageNet project.
相似度: 0.528 | 文本: Andrew Ng spread AI knowledge globally via public courses on Coursera and Stanford University.
...

完整示例代码:Tutorials/quick_start.ipynb

进阶应用:模型微调

数据准备

数据格式

微调数据需为JSONL格式(每行一个JSON对象),包含以下字段:

{
    "query": "查询文本",
    "pos": ["相关文本1", "相关文本2"],  // 正例
    "neg": ["不相关文本1", "不相关文本2"],  // 负例
    "type": "retrieval"  // 任务类型
}

示例数据:examples/finetune/embedder/example_data/

难负例挖掘

难负例(Hard Negatives)是提升模型性能的关键,可通过以下命令生成:

cd scripts
python hn_mine.py \
    --input_file toy_finetune_data.jsonl \
    --output_file toy_finetune_data_minedHN.jsonl \
    --range_for_sampling 2-200 \
    --negative_number 15 \
    --use_gpu_for_searching \
    --embedder_name_or_path BAAI/bge-base-en-v1.5

参数说明:

  • range_for_sampling: 负例采样范围(如2-200表示从排名2-200的结果中采样)
  • negative_number: 每个查询采样的负例数量
  • embedder_name_or_path: 用于生成嵌入的模型

微调训练

以标准编码器模型为例:

torchrun --nproc_per_node 2 \
    -m FlagEmbedding.finetune.embedder.encoder_only.base \
    --model_name_or_path BAAI/bge-large-en-v1.5 \
    --cache_dir ./cache/model \
    --train_data ./example_data/retrieval \
    --cache_path ./cache/data \
    --train_group_size 8 \
    --query_max_len 512 \
    --passage_max_len 512 \
    --query_instruction_for_retrieval 'Represent this sentence for searching relevant passages: ' \
    --output_dir ./custom_embedder \
    --learning_rate 1e-5 \
    --fp16 \
    --num_train_epochs 2 \
    --per_device_train_batch_size 2 \
    --negatives_cross_device \
    --temperature 0.02 \
    --normalize_embeddings True

关键参数说明:

  • --model_name_or_path: 预训练模型名称或路径
  • --train_group_size: 每组训练样本数(包含1个查询+多个正负例)
  • --negatives_cross_device: 在多GPU间共享负例
  • --temperature: 温度参数,控制相似度分布

微调详细教程:examples/finetune/embedder/README.md

多模型支持

FlagEmbedding支持多种模型微调,包括:

  1. 标准模型(如bge-base-en-v1.5)
  2. 多语言模型(如bge-m3)
  3. 指令微调模型(如bge-en-icl)

以bge-m3(多语言、多粒度模型)为例:

torchrun --nproc_per_node 2 \
    -m FlagEmbedding.finetune.embedder.encoder_only.m3 \
    --model_name_or_path BAAI/bge-m3 \
    --train_data ./example_data/retrieval \
    --knowledge_distillation True \
    --unified_finetuning True \
    --use_self_distill True \
    --output_dir ./custom_m3_model

模型评估

评估指标

FlagEmbedding提供多种评估指标:

  • MRR (Mean Reciprocal Rank):平均倒数排名
  • NDCG (Normalized Discounted Cumulative Gain):归一化折损累积增益
  • MAP (Mean Average Precision):平均精度均值

快速评估

使用MTEB (Massive Text Embedding Benchmark)评估模型:

cd examples/evaluation/mteb
python eval_mteb.py \
    --model_name_or_path ./custom_embedder \
    --task_types Retrieval \
    --batch_size 32 \
    --output_dir ./evaluation_results

评估结果将保存为JSON文件,包含各任务的指标得分。

评估工具使用说明:examples/evaluation/

实际案例:构建智能问答系统

系统架构

RAG系统架构

  1. 文档处理:将文档分割为段落,生成嵌入向量
  2. 检索模块:根据用户问题检索相关段落
  3. 重排序:使用Reranker优化检索结果
  4. 生成回答:结合检索结果生成最终回答

核心代码示例

from FlagEmbedding import FlagModel, FlagReranker

# 加载嵌入模型和重排序模型
embedder = FlagModel('BAAI/bge-base-en-v1.5', use_fp16=True)
reranker = FlagReranker('BAAI/bge-reranker-base', use_fp16=True)

# 文档嵌入
documents = ["文档段落1", "文档段落2", ...]
doc_embeddings = embedder.encode(documents)

# 检索相关文档
query = "用户问题"
query_embedding = embedder.encode(query)
similarities = query_embedding @ doc_embeddings.T
top_k = similarities.argsort(descending=True)[:5]
candidates = [documents[i] for i in top_k]

# 重排序
pairs = [(query, doc) for doc in candidates]
scores = reranker.compute_score(pairs)
reranked_candidates = [candidates[i] for i in scores.argsort(descending=True)]

# 生成回答(需结合LLM)

RAG应用教程:Tutorials/6_RAG/

总结与展望

通过本文学习,你已掌握FlagEmbedding的核心功能:

  1. 快速安装与环境配置
  2. 文本嵌入生成与相似度计算
  3. 自定义数据微调模型
  4. 模型性能评估方法

FlagEmbedding持续更新中,即将支持:

  • 多模态嵌入(文本+图像)
  • 更大规模的预训练模型
  • 更高效的微调算法

资源获取

下一步行动

  1. 尝试微调自己的数据集,提升特定领域性能
  2. 探索BGE-M3模型的多语言能力
  3. 结合LLM构建端到端的RAG系统

立即开始你的FlagEmbedding之旅,让文本检索变得前所未有的高效!

【免费下载链接】FlagEmbedding Dense Retrieval and Retrieval-augmented LLMs 【免费下载链接】FlagEmbedding 项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值