我们已更新了新版重排序模型，支持更长的输入长度、更多语言，并实现了更好的性能。

<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="#citation">引用</a> | <a href="#license">许可证</a> <p> </h4>

更多详细信息请访问我们的GitHub：FlagEmbedding。

English | 中文

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

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

新闻

• 2024/3/18：发布新版重排序模型，基于强大的M3和LLM（GEMMA和MiniCPM，实际上没那么大）骨干构建，支持多语言处理和更大输入，在BEIR、C-MTEB/Retrieval、MIRACL、LlamaIndex Evaluation上实现了大幅排名性能提升。
• 2024/3/18：发布Visualized-BGE，为BGE赋予视觉能力。Visualized-BGE可用于为混合图像-文本数据生成嵌入。
• 2024/1/30：发布BGE-M3，BGE系列的新成员！M3代表Multi-linguality（100+语言）、Multi-granularities（输入长度最高8192）、Multi-Functionality（统一密集检索、稀疏检索、多向量/colbert检索）。 它是首个支持三种检索方法的嵌入模型，在多语言（MIRACL）和跨语言（MKQA）基准测试上取得了新的SOTA成绩。技术报告和代码。🔥
• 2024/1/9：发布Activation-Beacon，一种有效、高效、兼容且低成本（训练）的方法来扩展LLM的上下文长度。技术报告 🔥
• 2023/12/24：发布LLaRA，基于LLaMA-7B的密集检索器，在MS MARCO和BEIR上实现了最先进性能。模型和代码将开源，敬请期待。技术报告
• 2023/11/23：发布LM-Cocktail，一种在微调过程中保持通用能力的方法，通过合并多个语言模型来实现。技术报告 🔥
• 2023/10/12：发布LLM-Embedder，一个统一的嵌入模型，支持LLM的多样化检索增强需求。技术报告
• 2023/9/15：BGE的技术报告已发布
• 2023/9/15：BGE的大规模训练数据已发布
• 2023/9/12：新模型发布：
  • 新版重排序模型：发布交叉编码器模型BAAI/bge-reranker-base和BAAI/bge-reranker-large，比嵌入模型更强大。我们建议使用/微调它们来对嵌入模型返回的top-k文档进行重排序。
  • 更新嵌入模型：发布bge-*-v1.5嵌入模型，以缓解相似度分布问题，并在无需指令的情况下增强检索能力。

<details> <summary>更多</summary>

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

</details>

模型列表

bge是BAAI general embedding的缩写。

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

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

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

常见问题

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

按照此示例准备数据并微调您的模型。 一些建议：

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

</details>

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

建议使用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，我们提升了不使用指令时的检索能力。 不使用指令仅比使用指令有轻微的性能下降。 因此，为了方便起见，您可以在所有情况下都不使用指令生成嵌入。

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

</details>

使用说明

嵌入模型的使用示例

以下是使用FlagEmbedding、Sentence-Transformers、Langchain或Huggingface Transformers使用bge模型的示例。

使用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

您也可以使用sentence-transformers来使用bge模型：

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]）的最后一层隐藏状态作为句子嵌入。

from transformers import AutoTokenizer, AutoModel
import torch
# 我们想要获取句子嵌入的句子
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嵌入
with torch.no_grad():
    model_output = model(**encoded_input)
    # 执行池化。此处为cls池化。
    sentence_embeddings = model_output[0][:, 0]
# 归一化嵌入
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
print("Sentence embeddings:", sentence_embeddings)

重排序器的使用

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

使用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)

使用ONNX文件运行重排序器

from optimum.onnxruntime import ORTModelForSequenceClassification  # type: ignore

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-large')
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
model_ort = ORTModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base', file_name="onnx/model.onnx")

# 我们想要获取句子嵌入的句子
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.']]

# 对句子进行分词
encoded_input = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt')

scores_ort = model_ort(**encoded_input, return_dict=True).logits.view(-1, ).float()
# 计算token嵌入
with torch.inference_mode():
    scores = model_ort(**encoded_input, return_dict=True).logits.view(-1, ).float()

# scores和scores_ort是相同的

使用infinity运行重排序器

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

import asyncio
from infinity_emb import AsyncEmbeddingEngine, EngineArgs

query='what is a panda?'
docs = ['The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear', "Paris is in France."]

engine = AsyncEmbeddingEngine.from_args(
    EngineArgs(model_name_or_path = "BAAI/bge-reranker-base", device="cpu", engine="torch" # 或使用engine="optimum"以使用onnx
))

async def main(): 
    async with engine:
        ranking, usage = await engine.rerank(query=query, docs=docs)
        print(list(zip(ranking, docs)))
asyncio.run(main())

评估

baai-general-embedding模型在MTEB和C-MTEB排行榜上都实现了最先进的性能！ 更多详细信息和评估工具请参阅我们的脚本。

• MTEB：

• C-MTEB： 我们创建了用于中文文本嵌入的基准测试C-MTEB，包含来自6个任务的31个数据集。 请参阅C_MTEB获取详细简介。

• 重排序： 请参阅C_MTEB获取评估脚本。

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

训练

BAAI Embedding

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

BGE 重排序器

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

引用

如果您发现此仓库有用，请考虑给一个星标⭐并引用

@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许可证授权。发布的模型可免费用于商业目的。