FlagEmbedding: 先进的文本嵌入和检索增强大语言模型框架

FlagEmbedding

FlagEmbedding简介

FlagEmbedding是由北京智源人工智能研究院(BAAI)开发的开源框架,专注于检索增强大语言模型(Retrieval-augmented LLMs)相关技术。该项目提供了一系列先进的文本嵌入模型和工具,可用于信息检索、文本分类、聚类等多种自然语言处理任务。

FlagEmbedding的核心特点包括:

提供多种高性能的文本嵌入模型,如BGE系列模型。
支持长文本处理的大语言模型技术。
提供模型微调和融合的方法。
包含多语言、多功能、多粒度的嵌入模型。
提供重排序模型以提高检索精度。
包含用于评估模型性能的基准测试。

主要项目

FlagEmbedding包含以下几个主要项目:

BGE-M3

BGE-M3是一个创新的嵌入模型,具有以下特点:

多功能性:支持密集检索、多向量检索和稀疏检索三种常见的检索功能。
多语言性:支持100多种工作语言。
多粒度:能够处理从短句到长达8192个token的长文档等不同粒度的输入。

BGE-M3 Model

Visualized-BGE

Visualized-BGE将图像token嵌入整合到BGE文本嵌入框架中,可用于多模态知识检索、组合图像检索等混合模态检索任务。该模型在多个零样本混合模态检索任务中表现出色。

LongLLM QLoRA

该项目通过QLoRA微调将Llama-3-8B-Instruct的上下文长度从8K扩展到80K。整个训练周期非常高效,在一台8xA800 (80G) GPU机器上仅需8小时。生成的模型在NIHS、主题检索和长上下文语言理解等多项评估任务中表现优异。

Activation Beacon

Activation Beacon是一种有效、高效、兼容且低成本(训练)的方法,用于扩展LLM的上下文长度。它将LLM的原始激活压缩成更紧凑的形式,使其能够在有限的上下文窗口内感知更长的上下文。

LM-Cocktail

LM-Cocktail使用简单的函数自动合并微调模型和基础模型,计算合并权重。它可以用于提高目标领域的性能,同时不降低目标领域之外的通用能力,也可以在不进行微调的情况下为新任务生成模型。

LLM Embedder

LLM Embedder基于LLM的反馈进行微调,支持大型语言模型的检索增强需求,包括知识检索、记忆检索、示例检索和工具检索。它在6个任务上进行了微调:问答、对话搜索、长对话、长程语言建模、上下文学习和工具学习。

BGE Reranker

BGE Reranker是一个交叉编码器模型,可以对输入对执行全注意力操作,比嵌入模型(即双编码器)更准确但耗时更长。因此,它可以用于对嵌入模型返回的top-k文档进行重新排序。

安装和快速入门

可以通过pip安装FlagEmbedding:

pip install -U FlagEmbedding

使用示例:

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)

sentences_1 = ["I love NLP", "I love machine learning"]
sentences_2 = ["I love BGE", "I love text retrieval"]
embeddings_1 = model.encode(sentences_1)
embeddings_2 = model.encode(sentences_2)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

模型列表

FlagEmbedding提供了多种预训练模型,包括:

BGE-M3:多语言、多功能、多粒度的嵌入模型
BGE系列:包括large、base、small等不同规模的英文和中文嵌入模型
LLM Embedder:统一的嵌入模型,支持LLM的多样化检索增强需求
BGE Reranker:用于重新排序的交叉编码器模型

这些模型在MTEB和C-MTEB等基准测试中都取得了优异的成绩。

总结

FlagEmbedding为检索增强大语言模型提供了一套全面的解决方案,包括高性能的文本嵌入模型、长文本处理技术、模型融合方法等。它不仅在多个基准测试中取得了领先成绩,还提供了丰富的工具和示例,方便研究人员和开发者进行进一步的探索和应用。无论是进行学术研究还是开发实际应用,FlagEmbedding都是一个值得关注的开源框架。