更多详情请参考我们的 GitHub 仓库:https://github.com/FlagOpen/FlagEmbedding
在本项目中,我们推出了 BGE-M3,它以其在多功能性、多语言性和多粒度方面的卓越表现而著称。
RAG 中检索 pipeline 的一些建议
我们建议使用以下 pipeline:混合检索 + 重排序。
混合检索利用了各种方法的优势,提供了更高的准确性和更强的泛化能力。 一个经典示例:同时使用嵌入检索和 BM25 算法。 现在,您可以尝试使用 BGE-M3,它同时支持嵌入检索和稀疏检索。 这使您在生成稠密嵌入时,能够零额外成本地获得 token 权重(类似于 BM25)。 要使用混合检索,您可以参考 Vespa 和 Milvus。
作为交叉编码器模型,重排序器(re-ranker)比双编码器嵌入模型具有更高的准确性。 在检索之后使用重排序模型(例如 bge-reranker、bge-reranker-v2)可以进一步筛选所选文本。
| 模型名称 | 维度 | 序列长度 | 介绍 |
|---|---|---|---|
| BAAI/bge-m3 | 1024 | 8192 | 多语言;基于bge-m3-unsupervised进行统一微调(dense、sparse及colbert) |
| BAAI/bge-m3-unsupervised | 1024 | 8192 | 多语言;基于bge-m3-retromae进行对比学习 |
| BAAI/bge-m3-retromae | -- | 8192 | 多语言;将xlm-roberta的最大长度扩展至8192,并通过retromae进一步预训练 |
| BAAI/bge-large-en-v1.5 | 1024 | 512 | 英文模型 |
| BAAI/bge-base-en-v1.5 | 768 | 512 | 英文模型 |
| BAAI/bge-small-en-v1.5 | 384 | 512 | 英文模型 |
| 数据集 | 介绍 |
|---|---|
| MLDR | 文档检索数据集,涵盖13种语言 |
| bge-m3-data | bge-m3所使用的微调数据 |
1. 不同检索方法介绍
2. 如何在其他项目中使用BGE-M3?
对于嵌入检索,您可以采用与BGE相同的方法使用BGE-M3模型。
唯一的区别是BGE-M3模型不再需要为查询添加指令。
3. 如何微调bge-M3模型?
您可以按照此示例中的常规方法来微调稠密嵌入。
如果您想微调m3的所有嵌入功能(dense、sparse和colbert),可以参考统一微调示例
安装:
git clone https://github.com/FlagOpen/FlagEmbedding.git
cd FlagEmbedding
pip install -e .或:
pip install -U FlagEmbeddingfrom FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('BAAI/bge-m3',
use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
embeddings_1 = model.encode(sentences_1,
batch_size=12,
max_length=8192, # If you don't need such a long length, you can set a smaller value to speed up the encoding process.
)['dense_vecs']
embeddings_2 = model.encode(sentences_2)['dense_vecs']
similarity = embeddings_1 @ embeddings_2.T
print(similarity)
# [[0.6265, 0.3477], [0.3499, 0.678 ]]您也可以使用 sentence-transformers 和 huggingface transformers 来生成 dense embeddings。 详情请参考 baai_general_embedding。
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
output_1 = model.encode(sentences_1, return_dense=True, return_sparse=True, return_colbert_vecs=False)
output_2 = model.encode(sentences_2, return_dense=True, return_sparse=True, return_colbert_vecs=False)
# you can see the weight for each token:
print(model.convert_id_to_token(output_1['lexical_weights']))
# [{'What': 0.08356, 'is': 0.0814, 'B': 0.1296, 'GE': 0.252, 'M': 0.1702, '3': 0.2695, '?': 0.04092},
# {'De': 0.05005, 'fin': 0.1368, 'ation': 0.04498, 'of': 0.0633, 'BM': 0.2515, '25': 0.3335}]
# compute the scores via lexical mathcing
lexical_scores = model.compute_lexical_matching_score(output_1['lexical_weights'][0], output_2['lexical_weights'][0])
print(lexical_scores)
# 0.19554901123046875
print(model.compute_lexical_matching_score(output_1['lexical_weights'][0], output_1['lexical_weights'][1]))
# 0.0from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
output_1 = model.encode(sentences_1, return_dense=True, return_sparse=True, return_colbert_vecs=True)
output_2 = model.encode(sentences_2, return_dense=True, return_sparse=True, return_colbert_vecs=True)
print(model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][0]))
print(model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][1]))
# 0.7797
# 0.4620from openmind import AutoTokenizer, AutoModel, is_torch_npu_available
from openmind_hub import snapshot_download
import torch.nn.functional as F
from torch import Tensor
import openmind
import torch
import argparse
# Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
token_embeddings = model_output[0] # First element of model_output contains all token embeddings
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to model",
default="jeffding/bge-m3-openmind",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
model_path = args.model_name_or_path
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path).to(device)
sentences = ['如何更换花呗绑定银行卡', 'How to replace the Huabei bundled bank card']
# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt').to(device)
# Compute token embeddings
with torch.no_grad():
model_output = model(**encoded_input)
# Perform pooling. In this case, mean pooling.
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
print("Sentence embeddings:")
print(sentence_embeddings)
if __name__ == "__main__":
main()输入文本对列表,即可获得通过不同方法计算出的得分。
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
sentence_pairs = [[i,j] for i in sentences_1 for j in sentences_2]
print(model.compute_score(sentence_pairs,
max_passage_length=128, # a smaller max length leads to a lower latency
weights_for_different_modes=[0.4, 0.2, 0.4])) # weights_for_different_modes(w) is used to do weighted sum: w[0]*dense_score + w[1]*sparse_score + w[2]*colbert_score
# {
# 'colbert': [0.7796499729156494, 0.4621465802192688, 0.4523794651031494, 0.7898575067520142],
# 'sparse': [0.195556640625, 0.00879669189453125, 0.0, 0.1802978515625],
# 'dense': [0.6259765625, 0.347412109375, 0.349853515625, 0.67822265625],
# 'sparse+dense': [0.482503205537796, 0.23454029858112335, 0.2332356721162796, 0.5122477412223816],
# 'colbert+sparse+dense': [0.6013619303703308, 0.3255828022956848, 0.32089319825172424, 0.6232916116714478]
# }
BGE-M3模型在该基准测试中表现最佳(OAI是OpenAI的缩写)。
更多详情,请参考文章和Github仓库。


长文档检索
MLDR:
请注意,MLDR是我们通过LLM构建的文档检索数据集,涵盖13种语言,包括测试集、验证集和训练集。我们利用MLDR的训练集来增强模型的长文档检索能力。因此,将基线与“Dense w.o.long”(不使用长文档数据集进行微调)进行比较更为公平。此外,该长文档检索数据集将开源,以解决当前开源多语言长文本检索数据集匮乏的问题。我们相信这些数据将对开源社区训练文档检索模型有所帮助。
NarritiveQA:

与BM25的比较
我们使用Pyserini实现BM25,测试结果可通过此脚本复现。我们使用两种不同的分词器对BM25进行了测试:一种使用Lucene Analyzer,另一种使用与M3相同的分词器(即xlm-roberta的分词器)。结果表明,BM25仍然是一个具有竞争力的基线,尤其在长文档检索方面。

更多详情请参考我们的报告。
感谢开源数据集的作者,包括 Miracl、MKQA、NarritiveQA 等。感谢开源库,如 Tevatron、Pyserini。
如果您觉得本仓库对您有帮助,请考虑给予星级 :star: 并引用。
@misc{bge-m3,
title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
year={2024},
eprint={2402.03216},
archivePrefix={arXiv},
primaryClass={cs.CL}
}