<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检索）。 这是首个支持三种检索方法的嵌入模型，在多语言（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，统一嵌入模型，支持LLM的多样化检索增强需求。技术报告
• 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基准测试中排名第一！ :tada: :tada:
• 2023/8/1：我们发布了 中文大规模文本嵌入基准测试（C-MTEB），包含31个测试数据集。

</details>

模型列表

bge 是 BAAI general embedding 的缩写。

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

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

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

常见问题

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

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

• 按照这个 示例 挖掘硬负样本，这可以提高检索性能。
• 如果您在自己的数据上预训练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，我们改进了其不使用指令时的检索能力。 不使用指令与使用指令相比，检索性能仅有轻微下降。 因此，为了方便，您可以在所有情况下都不带指令生成嵌入。

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

</details>

使用方法

嵌入模型使用方法

以下是使用 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

您也可以使用 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)

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")

# 我们需要句子嵌入的句子
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嵌入
with torch.no_grad():
    model_output = model(**encoded_input)

# model_output 和 model_output_ort 是相同的

通过 infinity 使用

也可以使用 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())

重排序模型使用方法

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

使用 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：
  我们创建了中文文本嵌入基准测试C-MTEB，包含来自6个任务的31个数据集。 请参阅 C_MTEB 了解详细说明。

• 重排序： 评估脚本请参见 C_MTEB。

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

训练

BAAI 嵌入模型

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

BGE 重排序模型

跨编码器将对输入对执行全注意力计算， 这比嵌入模型（即双编码器）更准确，但比嵌入模型更耗时。 因此，它可以用于对嵌入模型返回的top-k文档进行重排序。 我们在多语言配对数据上训练跨编码器， 数据格式与嵌入模型相同，因此您可以按照我们的 示例 轻松进行微调。 更多详情请参阅 ./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许可证。发布的模型可以免费用于商业目的。