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 許可證。發佈的模型可免費用於商業用途。