<h1 align="center">FlagEmbedding</h1>

<h4 align="center"> <p> <a href=#model-list>模型列表</a> | <a href=#frequently-asked-questions>常见问题</a> | <a href=#usage>使用</a> | <a href="#evaluation">评测</a> | <a href="#train">训练</a> | <a href="#contact">联系我们</a> | <a href="#citation">引用</a> | <a href="#license">许可证</a> <p> </h4>

更多详情请参阅我们的 GitHub：FlagEmbedding。

如果您需要支持更多语言、更长文本和其他检索方式的模型，可以尝试使用 bge-m3。

English | 中文

FlagEmbedding 专注于检索增强的 LLM，目前包含以下项目：

• 长上下文 LLM：Activation Beacon
• LM 微调：LM-Cocktail
• 稠密检索：BGE-M3、LLM Embedder、BGE Embedding
• 重排序模型：BGE Reranker
• 基准测试：C-MTEB

最新消息

• 2024/1/30：发布 BGE-M3，BGE 模型系列的新成员！M3 代表 Multi-linguality（多语言，支持 100+ 种语言）、Multi-granularities（多粒度，输入长度可达 8192）、Multi-Functionality（多功能，统一稠密检索、词汇检索、多向量/colbert 检索）。 这是首个支持三种检索方法的Embedding模型，在多语言（MIRACL）和跨语言（MKQA）基准测试中取得了新的 SOTA。 技术报告 和 代码。:fire:
• 2024/1/9：发布 Activation-Beacon，一种有效、高效、兼容且低成本（训练）的 LLM 上下文长度扩展方法。技术报告 :fire:
• 2023/12/24：发布 LLaRA，基于 LLaMA-7B 的稠密检索器，在 MS MARCO 和 BEIR 上取得了最先进的性能。模型和代码将开源，敬请期待。技术报告 :fire:
• 2023/11/23：发布 LM-Cocktail，一种在微调时通过合并多个语言模型来保持通用能力的方法。技术报告 :fire:
• 2023/10/12：发布 LLM-Embedder，一个统一的 Embedding 模型，支持 LLM 的多样化检索增强需求。技术报告
• 2023/9/15：BGE 的 技术报告 和 大规模训练数据 已发布
• 2023/9/12：新模型发布：
  • 新重排序模型：发布跨编码器模型 BAAI/bge-reranker-base 和 BAAI/bge-reranker-large，比 Embedding 模型更强大。建议使用/微调它们来对 Embedding 模型返回的前 k 个文档进行重排序。
  • 更新 Embedding 模型：发布 bge-*-v1.5 Embedding 模型，以改善相似度分布问题，并在没有指令的情况下增强检索能力。

<details> <summary>更多</summary> <!-- ### 更多 -->

• 2023/9/7：更新 微调代码：添加了挖掘困难负样本的脚本，并支持在微调时添加指令。
• 2023/8/9：BGE 模型已集成到 Langchain，您可以这样使用 #using-langchain；C-MTEB 排行榜 已上线 available。
• 2023/8/5：发布 base 和 small 规模模型，同规模模型中最佳性能 🤗
• 2023/8/2：发布 bge-large-*（BAAI General Embedding 的缩写）模型，在 MTEB 和 C-MTEB 基准测试中排名第一！ :tada: :tada:
• 2023/8/1：发布 中文大规模文本 Embedding 基准（C-MTEB），包含 31 个测试数据集。

</details>

模型列表

bge 是 BAAI general embedding 的缩写。

[1]：如果您需要根据查询检索相关段落，建议将指令添加到查询中；在其他情况下，不需要指令，直接使用原始查询即可。在所有情况下，都不需要为段落添加指令。

[2]：与 Embedding 模型不同，重排序器使用问题和文档作为输入，直接输出相似度而不是 Embedding。为了平衡准确性和时间成本，跨编码器被广泛用于对其他简单模型检索出的前 k 个文档进行重排序。 例如，使用 bge Embedding 模型检索前 100 个相关文档，然后使用 bge 重排序器对这 100 个文档进行重排序，得到最终的前 3 个结果。

所有模型都已上传到 Hugging Face Hub，您可以在 https://huggingface.co/BAAI 查看。 如果您无法访问 Hugging Face Hub，也可以从 https://model.baai.ac.cn/models 下载模型。

常见问题

<details> <summary>1. 如何微调 bge Embedding 模型？</summary>

<!-- ### 如何微调 bge Embedding 模型？ --> 按照此 示例 准备数据并微调您的模型。 一些建议：

• 按照此 示例 挖掘困难负样本，这可以提高检索性能。
• 如果您在自己的数据上预训练 bge，预训练模型不能直接用于计算相似度，必须先用对比学习进行微调才能计算相似度。
• 如果微调后模型的准确率仍然不高，建议使用/微调跨编码器模型（bge-reranker）对 top-k 结果进行重排序。微调重排序器也需要困难负样本。

</details>

<details> <summary>2. 两个不相似句子的相似度分数高于 0.5</summary>

<!-- ### 两个不相似句子的相似度分数高于 0.5 --> 建议使用 bge v1.5，它改善了相似度分布问题。

由于我们使用温度为 0.01 的对比学习对模型进行微调， 当前 BGE 模型的相似度分布大约在 [0.6, 1] 区间内。 因此，相似度分数大于 0.5 并不表示两个句子相似。

对于下游任务，如段落检索或语义相似度，重要的是分数的相对顺序，而不是绝对值。 如果您需要根据相似度阈值筛选相似句子，请根据您数据上的相似度分布选择合适的阈值（如 0.8、0.85，甚至 0.9）。

</details>

<details> <summary>3. 何时需要使用查询指令</summary>

<!-- ### 何时需要使用查询指令 -->

对于 bge-*-v1.5，我们改进了不使用指令时的检索能力。 不使用指令相比使用指令只有轻微的性能下降。 因此，为了方便，您可以在所有情况下都不使用指令生成 Embedding。

对于使用短查询查找长相关文档的检索任务，建议为这些短查询添加指令。 决定是否为查询添加指令的最佳方法是在您的任务上选择性能更好的设置。 在所有情况下，文档/段落都不需要添加指令。

</details>

使用

Embedding 模型使用

以下是使用 bge 模型的示例，支持 FlagEmbedding、Sentence-Transformers、Langchain 和 Huggingface Transformers。

使用 FlagEmbedding

pip install -U FlagEmbedding

如果安装不成功，请参阅 FlagEmbedding 获取更多安装方法。

from FlagEmbedding import FlagModel
sentences_1 = ["样例数据-1", "样例数据-2"]
sentences_2 = ["样例数据-3", "样例数据-4"]
model = FlagModel('BAAI/bge-large-zh-v1.5', 
                  query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                  use_fp16=True) # 设置 use_fp16 为 True 可以加速计算，但会有轻微的性能下降
embeddings_1 = model.encode(sentences_1)
embeddings_2 = model.encode(sentences_2)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

# 对于 s2p（短查询到长段落）检索任务，建议使用 encode_queries()，它会自动为每个查询添加指令
# 检索任务中的语料库仍然可以使用 encode() 或 encode_corpus()，因为它们不需要指令
queries = ['query_1', 'query_2']
passages = ["样例文档-1", "样例文档-2"]
q_embeddings = model.encode_queries(queries)
p_embeddings = model.encode(passages)
scores = q_embeddings @ p_embeddings.T

关于参数 query_instruction_for_retrieval 的取值，请参阅 模型列表。

默认情况下，FlagModel 会在编码时使用所有可用的 GPU。请设置 os.environ["CUDA_VISIBLE_DEVICES"] 来选择特定的 GPU。 您也可以设置 os.environ["CUDA_VISIBLE_DEVICES"]="" 来使所有 GPU 不可用。

使用 Sentence-Transformers

您也可以将 bge 模型与 sentence-transformers 一起使用：

pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer
sentences_1 = ["样例数据-1", "样例数据-2"]
sentences_2 = ["样例数据-3", "样例数据-4"]
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
embeddings_1 = model.encode(sentences_1, normalize_embeddings=True)
embeddings_2 = model.encode(sentences_2, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

对于 s2p（短查询到长段落）检索任务， 每个短查询应以指令开头（指令见 模型列表）。 但段落不需要指令。

from sentence_transformers import SentenceTransformer
queries = ['query_1', 'query_2']
passages = ["样例文档-1", "样例文档-2"]
instruction = "为这个句子生成表示以用于检索相关文章："

model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
q_embeddings = model.encode([instruction+q for q in queries], normalize_embeddings=True)
p_embeddings = model.encode(passages, normalize_embeddings=True)
scores = q_embeddings @ p_embeddings.T

使用 Langchain

您可以像这样在 langchain 中使用 bge：

from langchain.embeddings import HuggingFaceBgeEmbeddings
model_name = "BAAI/bge-large-en-v1.5"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True} # 设置为 True 以计算余弦相似度
model = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章："
)
model.query_instruction = "为这个句子生成表示以用于检索相关文章："

使用 HuggingFace Transformers

使用 transformers 包，您可以像这样使用模型：首先，将输入通过 transformer 模型，然后选择第一个 token（即 [CLS]）的最后一个隐藏状态作为句子 Embedding。

from transformers import AutoTokenizer, AutoModel
import torch
# 我们需要句子 Embedding 的句子
sentences = ["样例数据-1", "样例数据-2"]

# 从 HuggingFace Hub 加载模型
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-large-zh-v1.5')
model = AutoModel.from_pretrained('BAAI/bge-large-zh-v1.5')
model.eval()

# 对句子进行分词
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# 对于 s2p（短查询到长段落）检索任务，为查询添加指令（不要为段落添加指令）
# encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')

# 计算 token Embedding
with torch.no_grad():
    model_output = model(**encoded_input)
    # 执行池化。在这种情况下，使用 CLS 池化。
    sentence_embeddings = model_output[0][:, 0]
# 归一化 Embedding
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
print("Sentence embeddings:", sentence_embeddings)

ONNX 文件的使用

from optimum.onnxruntime import ORTModelForFeatureExtraction  # type: ignore

import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-large-en-v1.5')
model = AutoModel.from_pretrained('BAAI/bge-large-en-v1.5', revision="refs/pr/13")
model_ort = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-large-en-v1.5', revision="refs/pr/13",file_name="onnx/model.onnx")

# 我们需要句子 Embedding 的句子
sentences = ["样例数据-1", "样例数据-2"]

# 对句子进行分词
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# 对于 s2p（短查询到长段落）检索任务，为查询添加指令（不要为段落添加指令）
# encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')

model_output_ort = model_ort(**encoded_input)
# 计算 token Embedding
with torch.no_grad():
    model_output = model(**encoded_input)

# model_output 和 model_output_ort 是相同的

也可以使用 infinity_emb pip 包来部署 onnx 文件。

import asyncio
from infinity_emb import AsyncEmbeddingEngine, EngineArgs

sentences = ["Embed this is sentence via Infinity.", "Paris is in France."]
engine = AsyncEmbeddingEngine.from_args(
    EngineArgs(model_name_or_path = "BAAI/bge-large-en-v1.5", device="cpu", engine="optimum" # 或 engine="torch"
))

async def main(): 
    async with engine:
        embeddings, usage = await engine.embed(sentences=sentences)
asyncio.run(main())

重排序器使用

与 Embedding 模型不同，重排序器使用问题和文档作为输入，直接输出相似度而不是 Embedding。 您可以通过将查询和段落输入重排序器来获得相关性分数。 重排序器基于交叉熵损失优化，因此相关性分数不受特定范围的限制。

使用 FlagEmbedding

pip install -U FlagEmbedding

获取相关性分数（分数越高表示越相关）：

from FlagEmbedding import FlagReranker
reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True) # 设置 use_fp16 为 True 可以加速计算，但会有轻微的性能下降

score = reranker.compute_score(['query', 'passage'])
print(score)

scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
print(scores)

使用 Huggingface transformers

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-large')
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-large')
model.eval()

pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
    print(scores)

评测

baai-general-embedding 模型在 MTEB 和 C-MTEB 排行榜上都取得了最先进的性能！ 更多详情和评测工具请参阅我们的 脚本。

• MTEB：

• C-MTEB：
  我们创建了中文文本 Embedding 基准 C-MTEB，包含来自 6 个任务的 31 个数据集。 请参阅 C_MTEB 获取详细说明。

• 重排序： 评测脚本请参阅 C_MTEB。

*：T2RerankingZh2En 和 T2RerankingEn2Zh 是跨语言检索任务

训练

BAAI Embedding

我们使用 retromae 预训练模型，并使用对比学习在大规模配对数据上训练。 您可以按照我们的 示例 在您自己的数据上微调 Embedding 模型。 我们还提供了 预训练示例。 请注意，预训练的目的是重建文本，预训练模型不能直接用于相似度计算，需要进行微调。 更多 bge 的训练细节请参阅 baai_general_embedding。

BGE 重排序器

跨编码器将对输入对执行全注意力，比 Embedding 模型（即双编码器）更准确，但比 Embedding 模型更耗时。 因此，它可以用于对 Embedding 模型返回的前 k 个文档进行重排序。 我们在多语言配对数据上训练跨编码器， 数据格式与 Embedding 模型相同，因此您可以按照我们的 示例 轻松进行微调。 更多详情请参阅 ./FlagEmbedding/reranker/README.md

联系我们

如果您有任何与此项目相关的问题或建议，请随时提交 issue 或 pull request。 您也可以发送邮件至 Shitao Xiao(stxiao@baai.ac.cn) 和 Zheng Liu(liuzheng@baai.ac.cn)。

引用

如果您觉得这个仓库有用，请考虑给一个 star :star: 并引用

@misc{bge_embedding,
      title={C-Pack: Packaged Resources To Advance General Chinese Embedding}, 
      author={Shitao Xiao and Zheng Liu and Peitian Zhang and Niklas Muennighoff},
      year={2023},
      eprint={2309.07597},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

许可证

FlagEmbedding 采用 MIT 许可证。发布的模型可以免费用于商业目的。