FlagEmbedding

Github： https://github.com/FlagOpen/FlagEmbedding

智源发布最强开源可商用中英文语义向量模型BGE（BAAI General Embedding），在中英文语义检索精度与整体语义表征能力均超越了社区所有同类模型，如OpenAI 的text embedding 002等。此外，BGE 保持了同等参数量级模型中的最小向量维度，使用成本更低。

BGE模型将任意文本映射为低维稠密向量，以用于检索、分类、聚类或语义匹配等任务，并可支持为大模型调用外部知识。

模型列表

*: 如果您需要为一个简短的查询搜索相关的文档，您需要在查询中添加指令；在其他情况下，不需要指令，直接使用原始查询即可。在任何情况下，您都不需要为候选文档增加指令。

**: 不同于向量模型输出向量，reranker交叉编码器使用问题和文档作为输入，直接输出相似度。为了平衡准确率和时间成本，交叉编码器一般用于对其他简单模型检索到的top-k文档进行重排序。例如，使用bge向量模型检索前100个相关文档，然后使用bge reranker对前100个文档重新排序，得到最终的top-3结果。

常见问题

1. 如何微调bge模型 遵循这个示例 来准备数据并微调模型。 一些建议：

• 按照这个命令 挖掘难负样本，这可以明显提高检索性能。
• train_group_size 参数在我们实验中（使用了难负样本）默认设为8。可以根据数据中的平均负样本数据进行设置：train_group_size=负样本数量+1。
• query_max_len 和passage_max_len参数应该按照实际数据长度进行设置，数据都很长的话应该增大，但不能超过512。
• 如果微调向量模型的准确率仍然不高，建议使用或者微调交叉编码器模型(bge-reranker)对top-k结果进行重新排序。 交叉编码器模型的数据格式与向量模型一致，同时也建议使用难负样本。在正确微调的情况下，交叉编码器模型的准确度会高于向量模型。
• 如果进行了预训练，预训练后的模型无法直接用于计算相似度，必须经过微调才能进行相似度。

1. 不相似句子之间的相似度分数很高 建议使用bge v1.5，它缓解了相似度分布的问题。

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

对于下游任务，如段落检索或语义相似性， 重要的是分数的相对顺序，而不是绝对值。 如果你需要根据相似度阈值过滤相似句子， 请根据数据的相似度分布(如0.8,0.85，甚至0.9)选择合适的相似度阈值。

1. 什么时候需要添加查询指令 对于bge-*-v1.5，我们提高了它在不使用指令时的检索能力。无指令检索性能仅比使用指令检索性能略有下降。 因此，如果想要更加方便的话，您在所有情况下都可以在没有指令的情况下生成向量。

对于一个使用短查询寻找相关长文档的检索任务，查询与文档之间长度非常不一致，推荐为短查询添加指令。其他任务，推荐不添加指令。 最好的选择方式，是根据实际情况选择其中表现最好的方式。 在所有情况下，文档端都不用添加指令，只是查询端可以选择是否添加指令。

例子

模型使用例子

这里展示了一些通过 FlagEmbedding, Sentence-Transformers, Langchain, or Huggingface Transformers. 使用bge模型的方法。

基于 FlagEmbedding 库例子

pip install -U FlagEmbedding

如果您使用了镜像，可能无法找到最新版的FlagEmbedding。 可以参考FlagEmbedding 下载改项目进行安装。

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

# 对于短查询到长文档的检索任务，请对查询使用 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

Instruction参数 query_instruction_for_retrieval 请参照： Model List. 当加载你微调后的模型时，如果你没有在训练的json文件中为query添加指令，则将其设置为空字符串""; 如果你在训练数据中为query添加了指令，更改为你新设置的指令。

FlagModel支持GPU也支持CPU推理。如果GPU可用，其默认优先使用GPU。如果想禁止其使用GPU，设置os.environ["CUDA_VISIBLE_DEVICES"]="" 为提高效率，FlagModel默认会使用所有的GPU进行推理。如果想要使用具体的GPU，请设置os.environ["CUDA_VISIBLE_DEVICES"]。

基于 Sentence-Transformers 例子

安装 sentence-transformers:

pip install -U sentence-transformers

基于Sentence-Transformers的使用方法:

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

对于短查询到长文档的检索任务， 每个查询都应该以一条指令开始(指令参考 Model List). 但对于文档，不需要添加任何指令。

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

如果想使用sentence_transformers加载你微调后的模型，可以参考这里 。 同时，对于instruction, 如果你没有在训练的json文件中为query添加指令，则将其设置为空字符串""; 如果你在训练数据中为query添加了指令，更改为你新设置的指令。

基于 Langchain 例子

在Langchian中使用bge模型：

from langchain.embeddings import HuggingFaceBgeEmbeddings
model_name = "BAAI/bge-large-en-v1.5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
model = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

基于 HuggingFace Transformers 例子

使用transformers库时，您可以这样使用模型:首先，将输入传递给transformer模型，然后选择第一个标记的最后一个隐藏状态(即[CLS])作为句子嵌入。

from transformers import AutoTokenizer, AutoModel
import torch
# Sentences we want sentence embeddings for
sentences = ["样例数据-1", "样例数据-2"]

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-large-zh-v1.5')
model = AutoModel.from_pretrained('BAAI/bge-large-zh-v1.5')

# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# 对于短查询到长文档的检索任务, 为查询加上指令
# encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')

# Compute embeddings
with torch.no_grad():
    model_output = model(**encoded_input)
    # Perform pooling. In this case, cls pooling.
    sentence_embeddings = model_output[0][:, 0]
# normalize embeddings
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
print("Sentence embeddings:", sentence_embeddings)

Usage for Reranker

不同于向量模型，reranker无法对单个文本输出向量，其需要输入一个文本对直接计算分数。 你可以通过在reranker中输入query和passage来获得相关度分数，分数越高代表越相关。 该重排序器基于交叉熵损失进行优化，因此相关性分数没有一个特定的数值范围。

Using FlagEmbedding

pip install -U FlagEmbedding

计算相关分数，越高表示越相关:

from FlagEmbedding import FlagReranker
reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True) #设置 fp16 为True可以加快推理速度，效果会有可以忽略的下降

score = reranker.compute_score(['query', 'passage']) # 计算 query 和 passage的相似度
print(score)

scores = reranker.compute_score([['query 1', 'passage 1'], ['query 2', 'passage 2']])
print(scores)

Using 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排行榜上都实现了最先进的性能! 更多细节和评估脚本请参见 C_MTEB.

• MTEB:

• C-MTEB:

我们建立了一个中文文本嵌入的基准测试集合C-MTEB，其包括6个任务的31个数据集。 请参阅C_MTEB获取详细介绍。

所有的模型文件都已上传到huggingface上： https://huggingface.co/BAAI. 如果你无法连接到huggingface,可以通过智源网站进行下载： https://model.baai.ac.cn/models .

• Reranking: 评估脚本参考 C_MTEB.

* : T2RerankingZh2En 是跨语言检索数据集，使用中文检索英文， T2RerankingEn2Zh是使用英文检索中文。

License

FlagEmbedding基于MIT License开源协议。发布的模型权重可商用。