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

<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
• Reranker模型：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。 技术报告和代码。:fire:
• 2024/1/9：发布Activation-Beacon，一种有效、高效、兼容且低成本（训练）的LLM上下文长度扩展方法。技术报告:fire:
• 2023/12/24：发布LLaRA，基于LLaMA-7B的稠密检索器，在MS MARCO和BEIR上达到最先进的性能。模型和代码将开源，敬请期待。技术报告
• 2023/11/23：发布LM-Cocktail，一种在微调时通过合并多个语言模型来保持通用能力的方法。技术报告:fire:
• 2023/10/12：发布LLM-Embedder，统一嵌入模型，支持LLM的多样化检索增强需求。技术报告
• 2023/9/15：BGE的技术报告已发布
• 2023/9/15：BGE的大规模训练数据已发布
• 2023/9/12：新模型：
  • 新版重排序模型：发布cross-encoder模型BAAI/bge-reranker-base和BAAI/bge-reranker-large，比embedding模型更强大。我们推荐使用/微调它们来对embedding模型返回的top-k文档进行重排序。
  • 更新embedding模型：发布bge-*-v1.5embedding模型，缓解相似度分布问题，在不添加指令的情况下增强检索能力。

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

• 2023/9/7：更新微调代码：添加挖掘困难负样本的脚本，支持在微调时添加指令。
• 2023/8/9：BGE模型已集成到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：我们发布了中文大规模文本嵌入基准测试（C-MTEB），包含31个测试数据集。

</details>

模型列表

bge是BAAI general embedding的缩写。

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

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

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

常见问题

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

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

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

</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>

使用说明

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

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)

Reranker使用方法

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

使用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文件的reranker

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部署reranker

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

BGE Reranker

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

引用

如果您觉得这个仓库有用，请考虑给一个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许可证。 发布的模型可用于商业用途，免费使用。