Embeddingmodlebgelargeenv1.5

🚀 FlagEmbedding

FlagEmbedding 能够将任意文本映射到低维稠密向量，可用于检索、分类、聚类或语义搜索等任务，还能在大语言模型的向量数据库中使用。

🚀 快速开始

FlagEmbedding 提供了多种模型，涵盖不同语言和规模，可满足多样化的使用需求。你可以根据具体任务选择合适的模型，并参考以下使用示例进行操作。

✨ 主要特性

• 多语言支持：支持英语和中文等多种语言，适用于不同语言环境下的文本处理。
• 高性能表现：在 MTEB 和 C - MTEB 排行榜上取得了最先进的性能。
• 模型丰富：提供不同规模的模型，如 large、base、small 等，可根据实际需求灵活选择。
• 使用灵活：支持多种使用方式，如 FlagEmbedding、Sentence - Transformers、Langchain、Huggingface Transformers 等。

📦 安装指南

安装 FlagEmbedding

pip install -U FlagEmbedding

若安装失败，可参考 FlagEmbedding 获取更多安装方法。

安装 Sentence - Transformers

pip install -U sentence-transformers

💻 使用示例

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

# 对于短查询到长段落的检索任务，建议使用 encode_queries() 方法，它会自动为每个查询添加指令
# 检索任务中的语料库仍可使用 encode() 或 encode_corpus() 方法，因为它们不需要指令
queries = ['查询 1', '查询 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 调用嵌入模型

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)

对于短查询到长段落的检索任务，每个短查询应以指令开头（指令可参考 模型列表），但段落不需要指令。

from sentence_transformers import SentenceTransformer
queries = ['查询 1', '查询 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 调用嵌入模型

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 调用嵌入模型

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')
# 对于短查询到长段落的检索任务，为查询添加指令（段落不需要添加指令）
# encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')

# 计算标记嵌入
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)

使用 FlagEmbedding 调用重排器

from FlagEmbedding import FlagReranker
reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True) # 设置 use_fp16 为 True 可在轻微性能损失的情况下加速计算

score = reranker.compute_score(['查询', '段落'])
print(score)

scores = reranker.compute_score([['什么是熊猫？', '你好'], ['什么是熊猫？', '大熊猫（Ailuropoda melanoleuca），有时也被称为熊猫或单纯的熊猫，是中国特有的熊科物种。']])
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 = [['什么是熊猫？', '你好'], ['什么是熊猫？', '大熊猫（Ailuropoda melanoleuca），有时也被称为熊猫或单纯的熊猫，是中国特有的熊科物种。']]
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)

📚 详细文档

模型列表

bge 是 BAAI general embedding 的缩写。

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

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

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

常见问题解答

🔧 技术细节

BAAI 嵌入模型

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

BGE 重排器

交叉编码器会对输入对进行全注意力计算，比嵌入模型（即双编码器）更准确，但比嵌入模型更耗时。因此，它可用于对嵌入模型返回的前 k 个文档进行重排。我们在多语言配对数据上训练交叉编码器，数据格式与嵌入模型相同，因此你可以按照 示例 轻松微调它。更多细节请参考 ./FlagEmbedding/reranker/README.md。

📄 许可证

FlagEmbedding 采用 MIT 许可证。发布的模型可免费用于商业用途。