使用Sentence Transformers v5訓練和微調稀疏嵌入模型

Sentence Transformers是一個Python庫，用於使用和訓練嵌入模型和重排序模型，適用於廣泛的應用，例如檢索增強生成、語義搜尋、語義文字相似度、釋義挖掘等。最近幾個主要版本對訓練進行了重大改進

• v3.0: （改進的）Sentence Transformer（密集嵌入）模型訓練
• v4.0: （改進的）Cross Encoder（重排序器）模型訓練
• v5.0: （新增的）稀疏嵌入模型訓練

在這篇博文中，我將向您展示如何使用它來微調稀疏編碼器/嵌入模型，並解釋為什麼您可能想要這樣做。這得到了sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq，一個廉價的模型，在混合搜尋或檢索和重排序場景中表現出色。

微調稀疏嵌入模型涉及幾個元件：模型、資料集、損失函式、訓練引數、評估器和訓練器類。我將詳細介紹每個元件，並附上實際示例，說明它們如何用於微調強大的稀疏嵌入模型。

除了訓練您自己的模型外，您還可以從Hugging Face Hub上提供的各種預訓練稀疏編碼器中進行選擇。為了幫助您在這個不斷增長的領域中導航，我們策劃了一個SPLADE模型集合，重點介紹了一些最相關的模型。
我們在文件的預訓練模型中列出了最著名的模型及其基準測試結果。

目錄

• 什麼是稀疏嵌入模型？
  • 查詢和文件擴充套件
  • 為什麼要使用稀疏嵌入模型？
• 為什麼要微調？
• 訓練元件
• 模型
  • Splade
  • 免推理Splade
  • 對比稀疏表示（CSR）
  • 架構選擇指南
• 資料集
  • Hugging Face Hub上的資料
  • 本地資料（CSV、JSON、Parquet、Arrow、SQL）
  • 需要預處理的本地資料
  • 資料集格式
• 損失函式
• 訓練引數
• 評估器
  • SparseNanoBEIREvaluator
  • 使用STSb的SparseEmbeddingSimilarityEvaluator
  • 使用AllNLI的SparseTripletEvaluator
• 訓練器
  • 回撥
  • 多資料集訓練
• 評估
• 訓練技巧
• 向量資料庫整合
• 其他資源
  • 訓練示例
  • 文件

什麼是稀疏嵌入模型？

更廣義的“嵌入模型”是指將某種輸入（通常是文字）轉換為向量表示（嵌入），以捕獲輸入語義的模型。與原始輸入不同，您可以對這些嵌入執行數學運算，從而得到可用於各種任務（例如搜尋、聚類或分類）的相似性分數。

對於密集嵌入模型（即常見型別），嵌入通常是低維向量（例如384、768或1024維），其中大多數值是非零的。另一方面，稀疏嵌入模型生成高維向量（例如30,000+維），其中大多數值是零。通常，稀疏嵌入中的每個活躍維度（即具有非零值的維度）對應於模型詞彙表中的特定標記，從而實現可解釋性。

讓我們以最先進的稀疏嵌入模型naver/splade-v3為例

from sentence_transformers import SparseEncoder

# Download from the 🤗 Hub
model = SparseEncoder("naver/splade-v3")

# Run inference
sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# (3, 30522)

# Get the similarity scores for the embeddings
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[   32.4323,     5.8528,     0.0258],
#         [    5.8528,    26.6649,     0.0302],
#         [    0.0258,     0.0302,    24.0839]])

# Let's decode our embeddings to be able to interpret them
decoded = model.decode(embeddings, top_k=10)
for decoded, sentence in zip(decoded, sentences):
    print(f"Sentence: {sentence}")
    print(f"Decoded: {decoded}")
    print()

Sentence: The weather is lovely today.
Decoded: [('weather', 2.754288673400879), ('today', 2.610959529876709), ('lovely', 2.431990623474121), ('currently', 1.5520408153533936), ('beautiful', 1.5046082735061646), ('cool', 1.4664798974990845), ('pretty', 0.8986214995384216), ('yesterday', 0.8603134155273438), ('nice', 0.8322536945343018), ('summer', 0.7702118158340454)]

Sentence: It's so sunny outside!
Decoded: [('outside', 2.6939032077789307), ('sunny', 2.535827398300171), ('so', 2.0600898265838623), ('out', 1.5397940874099731), ('weather', 1.1198079586029053), ('very', 0.9873268604278564), ('cool', 0.9406591057777405), ('it', 0.9026399254798889), ('summer', 0.684999406337738), ('sun', 0.6520509123802185)]

Sentence: He drove to the stadium.
Decoded: [('stadium', 2.7872302532196045), ('drove', 1.8208855390548706), ('driving', 1.6665740013122559), ('drive', 1.5565159320831299), ('he', 1.4721972942352295), ('stadiums', 1.449463129043579), ('to', 1.0441515445709229), ('car', 0.7002660632133484), ('visit', 0.5118278861045837), ('football', 0.502326250076294)]

在此示例中，嵌入是30,522維向量，其中每個維度對應於模型詞彙表中的一個標記。`decode`方法返回了嵌入中值最高的10個標記，使我們能夠解釋哪些標記對嵌入的貢獻最大。

我們甚至可以確定嵌入之間的交集或重疊，這對於確定為什麼兩個文字被認為是相似或不相似非常有用。

# Let's also compute the intersection/overlap of the first two embeddings
intersection_embedding = model.intersection(embeddings[0], embeddings[1])
decoded_intersection = model.decode(intersection_embedding)
print(decoded_intersection)

Decoded: [('weather', 3.0842742919921875), ('cool', 1.379457712173462), ('summer', 0.5275946259498596), ('comfort', 0.3239051103591919), ('sally', 0.22571465373039246), ('julian', 0.14787325263023376), ('nature', 0.08582140505313873), ('beauty', 0.0588383711874485), ('mood', 0.018594780936837196), ('nathan', 0.000752730411477387)]

查詢和文件擴充套件

神經稀疏嵌入模型的一個關鍵組成部分是**查詢/文件擴充套件**。與BM25等傳統詞彙方法僅匹配精確標記不同，神經稀疏模型通常會自動用語義相關的術語擴充套件原始文字。

• 傳統詞彙方法（例如BM25）：僅匹配文字中的精確標記。
• 神經稀疏模型：自動擴充套件相關術語。

例如，在上面的程式碼輸出中，句子“The weather is lovely today”被擴充套件為包含“beautiful”、“cool”、“pretty”和“nice”等原始文字中沒有的詞語。類似地，“It's so sunny outside!”被擴充套件為包含“weather”、“summer”和“sun”。

這種擴充套件使得神經稀疏模型即使沒有精確的詞元匹配也能匹配語義相關的內容或同義詞，處理拼寫錯誤，並克服詞彙不匹配問題。這就是為什麼像SPLADE這樣的神經稀疏模型在保持稀疏表示效率優勢的同時，通常優於傳統的詞彙搜尋方法。

然而，擴充套件也有其風險。例如，對“星期二天氣如何？”的查詢擴充套件很可能也會擴充套件到“星期一”、“星期三”等，這可能不是所希望的。

為什麼要使用稀疏嵌入模型？

簡而言之，神經稀疏嵌入模型在傳統詞彙方法（如BM25）和密集嵌入模型（如Sentence Transformers）之間佔據了一個重要的利基市場。它們具有以下優點：

• 混合潛力：可與密集模型高效結合，密集模型在詞彙匹配重要的搜尋中可能表現不佳。
• 可解釋性：您可以準確地看到哪些標記促成了匹配。
• 效能：在許多檢索任務中與密集模型相比具有競爭力或更優。

在本部落格中，我將交替使用“稀疏嵌入模型”和“稀疏編碼器模型”。

為什麼要微調？

大多數（神經）稀疏嵌入模型都採用上述查詢/文件擴充套件，這樣您就可以匹配含義幾乎相同的文字，即使它們不共享任何單詞。簡而言之，模型必須識別同義詞，以便這些標記可以放置在最終嵌入中。

大多數開箱即用的稀疏嵌入模型都能輕易識別出“supermarket”、“food”和“market”是包含“grocery”的文字的有用擴充套件，但例如

• “The patient complained of severe cephalalgia.”（患者抱怨嚴重的頭痛。）

擴充套件為

'##lal', 'severe', '##pha', 'ce', '##gia', 'patient', 'complaint', 'patients', 'complained', 'warning', 'suffered', 'had', 'disease', 'complain', 'diagnosis', 'syndrome', 'mild', 'pain', 'hospital', 'injury'

然而，我們希望它擴充套件為“headache”（頭痛），這是“cephalalgia”的常用詞。這個例子擴充套件到許多領域，例如無法識別“Java”是一種程式語言，“Audi”製造汽車，或者“NVIDIA”是一家制造顯示卡的的公司。

透過微調，模型可以學習專門關注對您重要的領域和/或語言。

訓練元件

訓練Sentence Transformer模型涉及以下元件

1. 模型：要訓練或微調的模型，可以是預訓練的稀疏編碼器模型或基礎模型。
2. 資料集：用於訓練和評估的資料。
3. 損失函式：量化模型效能並指導最佳化過程的函式。
4. 訓練引數（可選）：影響訓練效能和跟蹤/除錯的引數。
5. 評估器（可選）：用於在訓練前、訓練中或訓練後評估模型的工具。
6. 訓練器：將模型、資料集、損失函式和其他元件整合在一起進行訓練。

現在，讓我們更詳細地探討這些元件。

模型

稀疏編碼器模型由一系列模組、稀疏編碼器特定模組或自定義模組組成，具有很大的靈活性。如果您想進一步微調稀疏編碼器模型（例如，它有一個modules.json檔案），那麼您無需擔心使用哪些模組

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

但如果您想從其他檢查點或從頭開始訓練，那麼這些是最常見的架構

Splade

Splade模型使用MLMTransformer模組，後接SpladePooling模組。前者載入預訓練的掩碼語言建模Transformer模型（例如BERT、RoBERTa、DistilBERT、ModernBERT等），後者將MLMHead的輸出池化以生成詞彙表大小的單一稀疏嵌入。

from sentence_transformers import models, SparseEncoder
from sentence_transformers.sparse_encoder.models import MLMTransformer, SpladePooling

# Initialize MLM Transformer (use a fill-mask model)
mlm_transformer = MLMTransformer("google-bert/bert-base-uncased")

# Initialize SpladePooling module
splade_pooling = SpladePooling(pooling_strategy="max")

# Create the Splade model
model = SparseEncoder(modules=[mlm_transformer, splade_pooling])

如果您向SparseEncoder提供一個fill-mask模型架構，此架構是預設的，因此使用快捷方式會更容易

from sentence_transformers import SparseEncoder

model = SparseEncoder("google-bert/bert-base-uncased")
# SparseEncoder(
#   (0): MLMTransformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertForMaskedLM'})
#   (1): SpladePooling({'pooling_strategy': 'max', 'activation_function': 'relu', 'word_embedding_dimension': None})
# )

免推理Splade

免推理Splade 使用Router模組，其中查詢和文件使用不同的模組。通常，對於這種型別的架構，文件部分是傳統的Splade架構（一個MLMTransformer後接一個SpladePooling模組），而查詢部分是SparseStaticEmbedding模組，它只返回查詢中每個標記的預計算分數。

from sentence_transformers import SparseEncoder
from sentence_transformers.models import Router
from sentence_transformers.sparse_encoder.models import SparseStaticEmbedding, MLMTransformer, SpladePooling

# Initialize MLM Transformer for document encoding
doc_encoder = MLMTransformer("google-bert/bert-base-uncased")

# Create a router model with different paths for queries and documents
router = Router.for_query_document(
    query_modules=[SparseStaticEmbedding(tokenizer=doc_encoder.tokenizer, frozen=False)],
    # Document path: full MLM transformer + pooling
    document_modules=[doc_encoder, SpladePooling("max")],
)

# Create the inference-free model
model = SparseEncoder(modules=[router], similarity_fn_name="dot")
# SparseEncoder(
#   (0): Router(
#     (query_0_SparseStaticEmbedding): SparseStaticEmbedding ({'frozen': False}, dim:30522, tokenizer: BertTokenizerFast)
#     (document_0_MLMTransformer): MLMTransformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertForMaskedLM'})
#     (document_1_SpladePooling): SpladePooling({'pooling_strategy': 'max', 'activation_function': 'relu', 'word_embedding_dimension': None})
#   )
# )

這種架構允許使用輕量級的SparseStaticEmbedding方法進行快速查詢時處理，該方法可以訓練並視為線性權重，而文件則透過完整的MLM transformer和SpladePooling進行處理。

免推理Splade特別適用於查詢延遲至關重要的搜尋應用，因為它將計算複雜度轉移到可以離線完成的文件索引階段。

訓練帶有`Router`模組的模型時，您必須在`SparseEncoderTrainingArguments`中使用`router_mapping`引數將訓練資料集列對映到正確的路由（“query”或“document”）。例如，如果您的資料集有`["question", "answer"]`列，那麼您可以使用以下對映

args = SparseEncoderTrainingArguments(
    ...,
    router_mapping={
        "question": "query",
        "answer": "document",
    }
)

此外，建議對SparseStaticEmbedding模組使用比模型其餘部分更高的學習率。為此，您應該在`SparseEncoderTrainingArguments`中使用`learning_rate_mapping`引數將引數模式對映到其學習率。例如，如果您希望SparseStaticEmbedding模組的學習率為`1e-3`，模型其餘部分為`2e-5`，您可以這樣做

args = SparseEncoderTrainingArguments(
    ...,
    learning_rate=2e-5,
    learning_rate_mapping={
        r"SparseStaticEmbedding\.*": 1e-3,
    }
)

對比稀疏表示 (CSR)

對比稀疏表示（CSR）模型，在Beyond Matryoshka: Revisiting Sparse Coding for Adaptive Representation中引入，在密集Sentence Transformer模型之上應用了SparseAutoEncoder模組，後者通常由Transformer後接Pooling模組組成。您可以像這樣從頭開始初始化一個：

from sentence_transformers import models, SparseEncoder
from sentence_transformers.sparse_encoder.models import SparseAutoEncoder

# Initialize transformer (can be any dense encoder model)
transformer = models.Transformer("google-bert/bert-base-uncased")

# Initialize pooling
pooling = models.Pooling(transformer.get_word_embedding_dimension(), pooling_mode="mean")

# Initialize SparseAutoEncoder module
sparse_auto_encoder = SparseAutoEncoder(
    input_dim=transformer.get_word_embedding_dimension(),
    hidden_dim=4 * transformer.get_word_embedding_dimension(),
    k=256,  # Number of top values to keep
    k_aux=512,  # Number of top values for auxiliary loss
)
# Create the CSR model
model = SparseEncoder(modules=[transformer, pooling, sparse_auto_encoder])

或者，如果您的基礎模型是 1) 一個密集 Sentence Transformer 模型，或者 2) 一個非 MLM Transformer 模型（預設情況下這些模型會作為 Splade 模型載入），那麼這個快捷方式將自動為您初始化 CSR 模型

from sentence_transformers import SparseEncoder

model = SparseEncoder("mixedbread-ai/mxbai-embed-large-v1")
# SparseEncoder(
#   (0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertModel'})
#   (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
#   (2): SparseAutoEncoder({'input_dim': 1024, 'hidden_dim': 4096, 'k': 256, 'k_aux': 512, 'normalize': False, 'dead_threshold': 30})
# )

與（免推理）Splade模型不同，CSR模型生成的稀疏嵌入與基礎模型的詞彙表大小不同。這意味著您無法像Splade模型那樣直接解釋嵌入中激活了哪些單詞，因為Splade模型中每個維度都對應於詞彙表中的特定標記。

此外，CSR模型在那些使用高維表示（例如1024-4096維）的密集編碼器模型上最有效。

架構選擇指南

如果您不確定使用哪種架構，這裡有一個快速指南

• 您想稀疏化現有的密集嵌入模型嗎？如果想，請使用CSR。
• 您希望查詢推理即時完成，但代價是效能略有下降嗎？如果是，請使用免推理SPLADE。
• 否則，請使用SPLADE。

資料集

SparseEncoderTrainer使用datasets.Dataset或datasets.DatasetDict例項進行訓練和評估。您可以從Hugging Face資料集中心載入資料，或使用各種格式的本地資料，如CSV、JSON、Parquet、Arrow或SQL。

注意：許多與Sentence Transformers開箱即用的公共資料集都已在Hugging Face Hub上標記為sentence-transformers，因此您可以在https://huggingface.co/datasets?other=sentence-transformers上輕鬆找到它們。考慮瀏覽這些資料集，以找到可能對您的任務、領域或語言有用的現成資料集。

Hugging Face Hub上的資料

您可以使用load_dataset函式從Hugging Face Hub中的資料集載入資料。

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/natural-questions", split="train")

print(train_dataset)
"""
Dataset({
    features: ['query', 'answer'],
    num_rows: 100231
})
"""

一些資料集，例如nthakur/swim-ir-monolingual，具有多個不同資料格式的子集。您需要指定子集名稱以及資料集名稱，例如dataset = load_dataset("nthakur/swim-ir-monolingual", "de", split="train")。

本地資料（CSV、JSON、Parquet、Arrow、SQL）

您也可以使用load_dataset載入某些檔案格式的本地資料

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
# or
dataset = load_dataset("json", data_files="my_file.json")

需要預處理的本地資料

如果您的本地資料需要預處理，您可以使用datasets.Dataset.from_dict。這允許您使用字典列表初始化資料集。

from datasets import Dataset

queries = []
documents = []
# Open a file, perform preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "query": queries,
    "document": documents,
})

字典中的每個鍵都將成為結果資料集中的一列。

資料集格式

確保資料集格式與您選擇的損失函式匹配至關重要。這涉及檢查兩件事

1. 如果您的損失函式需要一個標籤（如損失概述表中所示），則您的資料集必須有一個名為"label"或"score"的列。
2. 除"label"或"score"之外的所有列都被視為輸入（如損失概述表中所示）。這些列的數量必須與您所選損失函式的有效輸入數量相匹配。列的名稱不重要，只有它們的順序重要。

例如，如果您的損失函式接受(anchor, positive, negative)三元組，那麼您的第一個、第二個和第三個資料集列分別對應anchor、positive和negative。這意味著您的第一列和第二列必須包含應該緊密嵌入的文字，而您的第一列和第三列必須包含應該相距遙遠的文字。這就是為什麼根據您的損失函式，您的資料集列順序很重要。

考慮一個包含列["text1", "text2", "label"]的資料集，其中"label"列包含浮點相似度分數。此資料集可與SparseCoSENTLoss、SparseAnglELoss和SparseCosineSimilarityLoss一起使用，因為

1. 資料集有一個“label”列，這是這些損失函式所必需的。
2. 資料集有2個非標籤列，與這些損失函式所需的輸入數量相匹配。

如果您的資料集中的列順序不正確，請使用Dataset.select_columns重新排序。此外，使用Dataset.remove_columns刪除任何多餘的列（例如，sample_id、metadata、source、type），否則它們將被視為輸入。

損失函式

損失函式衡量模型在給定批次資料上的表現，並指導最佳化過程。損失函式的選擇取決於您可用的資料和目標任務。有關選項的全面列表，請參閱損失概述。

要訓練SparseEncoder，您需要一個SpladeLoss或CSRLoss，具體取決於架構。這些是包裝損失，它們在主損失函式之上新增稀疏性正則化，主損失函式必須作為引數提供。唯一可以獨立使用的損失是SparseMSELoss，因為它執行嵌入級別蒸餾，透過直接複製教師的稀疏嵌入來確保稀疏性。

大多數損失函式可以透過僅用您正在訓練的SparseEncoder以及一些可選引數進行初始化，例如

from datasets import load_dataset
from sentence_transformers import SparseEncoder
from sentence_transformers.sparse_encoder.losses import SpladeLoss, SparseMultipleNegativesRankingLoss

# Load a model to train/finetune
model = SparseEncoder("distilbert/distilbert-base-uncased")

# Initialize the SpladeLoss with a SparseMultipleNegativesRankingLoss
# This loss requires pairs of related texts or triplets
loss = SpladeLoss(
    model=model,
    loss=SparseMultipleNegativesRankingLoss(model=model),
    query_regularizer_weight=5e-5,  # Weight for query loss
    document_regularizer_weight=3e-5,
) 

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/natural-questions", split="train")
print(train_dataset)
"""
Dataset({
    features: ['query', 'answer'],
    num_rows: 100231
})
"""

文件

• SpladeLoss
• CSRLoss
• 損失API參考

訓練引數

SparseEncoderTrainingArguments類允許您指定影響訓練效能和跟蹤/除錯的引數。雖然是可選的，但嘗試這些引數有助於提高訓練效率並提供對訓練過程的洞察。

在Sentence Transformers文件中，我概述了一些最有用的訓練引數。我建議您閱讀訓練概述 > 訓練引數。

以下是初始化SparseEncoderTrainingArguments的示例

from sentence_transformers import SparseEncoderTrainingArguments

args = SparseEncoderTrainingArguments(
    # Required parameter:
    output_dir="models/splade-distilbert-base-uncased-nq",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if your GPU can't handle FP16
    bf16=False,  # Set to True if your GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Losses using "in-batch negatives" benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="splade-distilbert-base-uncased-nq",  # Used in W&B if `wandb` is installed
)

請注意，`eval_strategy`是在`transformers`版本`4.41.0`中引入的。之前的版本應使用`evaluation_strategy`。

評估器

您可以向SparseEncoderTrainer提供eval_dataset以在訓練期間獲取評估損失，但在訓練期間獲取更具體的指標也可能很有用。為此，您可以使用評估器在訓練前、訓練中或訓練後使用有用的指標評估模型的效能。您可以同時使用eval_dataset和評估器，或其中一個，或都不使用。它們根據eval_strategy和eval_steps訓練引數進行評估。

以下是Sentence Transformers為稀疏編碼器模型實現的評估器：

此外，應使用SequentialEvaluator將多個評估器組合成一個評估器，該評估器可以傳遞給SparseEncoderTrainer。

有時，您沒有所需的評估資料來自行準備這些評估器，但您仍然希望跟蹤模型在某些常見基準上的表現。在這種情況下，您可以將這些評估器與Hugging Face的資料一起使用。

SparseNanoBEIREvaluator

文件

• sentence_transformers.sparse_encoder.evaluation.SparseNanoBEIREvaluator

from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator

# Initialize the evaluator. Unlike most other evaluators, this one loads the relevant datasets
# directly from Hugging Face, so there's no mandatory arguments
dev_evaluator = SparseNanoBEIREvaluator()
# You can run evaluation like so:
# results = dev_evaluator(model)

使用STSb的SparseEmbeddingSimilarityEvaluator

文件

• sentence-transformers/stsb
• sentence_transformers.sparse_encoder.evaluation.SparseEmbeddingSimilarityEvaluator
• sentence_transformers.SimilarityFunction

from datasets import load_dataset
from sentence_transformers.evaluation import SimilarityFunction
from sentence_transformers.sparse_encoder.evaluation import SparseEmbeddingSimilarityEvaluator

# Load the STSB dataset (https://huggingface.co/datasets/sentence-transformers/stsb)
eval_dataset = load_dataset("sentence-transformers/stsb", split="validation")

# Initialize the evaluator
dev_evaluator = SparseEmbeddingSimilarityEvaluator(
    sentences1=eval_dataset["sentence1"],
    sentences2=eval_dataset["sentence2"],
    scores=eval_dataset["score"],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-dev",
)
# You can run evaluation like so:
# results = dev_evaluator(model)

使用AllNLI的SparseTripletEvaluator

文件

• sentence-transformers/all-nli
• sentence_transformers.sparse_encoder.evaluation.SparseTripletEvaluator
• sentence_transformers.SimilarityFunction

from datasets import load_dataset
from sentence_transformers.evaluation import SimilarityFunction
from sentence_transformers.sparse_encoder.evaluation import SparseTripletEvaluator

# Load triplets from the AllNLI dataset (https://huggingface.co/datasets/sentence-transformers/all-nli)
max_samples = 1000
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split=f"dev[:{max_samples}]")

# Initialize the evaluator
dev_evaluator = SparseTripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    main_distance_function=SimilarityFunction.DOT,
    name="all-nli-dev",
)
# You can run evaluation like so:
# results = dev_evaluator(model)

如果在訓練期間使用較小的eval_steps頻繁評估，請考慮使用微小的eval_dataset以最大程度地減少評估開銷。如果您擔心評估集大小，90-1-9的訓練-評估-測試劃分可以提供一個平衡，為最終評估保留一個合理大小的測試集。訓練結束後，您可以使用trainer.evaluate(test_dataset)評估模型的測試損失，或使用test_evaluator(model)初始化測試評估器以獲取詳細的測試指標。

如果您在訓練後但在儲存模型之前進行評估，則自動生成的模型卡仍將包含測試結果。

使用分散式訓練時，評估器僅在第一個裝置上執行，與訓練和評估資料集不同，後者在所有裝置上共享。

訓練器

SparseEncoderTrainer是所有先前元件的集合。我們只需指定訓練器、模型、訓練引數（可選）、訓練資料集、評估資料集（可選）、損失函式、評估器（可選），然後就可以開始訓練了。讓我們看看一個將所有這些元件結合在一起的指令碼

import logging

from datasets import load_dataset

from sentence_transformers import (
    SparseEncoder,
    SparseEncoderModelCardData,
    SparseEncoderTrainer,
    SparseEncoderTrainingArguments,
)
from sentence_transformers.models import Router
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
from sentence_transformers.sparse_encoder.models import SparseStaticEmbedding, MLMTransformer, SpladePooling
from sentence_transformers.training_args import BatchSamplers

logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

# 1. Load a model to finetune with 2. (Optional) model card data
mlm_transformer = MLMTransformer("distilbert/distilbert-base-uncased", tokenizer_args={"model_max_length": 512})
splade_pooling = SpladePooling(
    pooling_strategy="max", word_embedding_dimension=mlm_transformer.get_sentence_embedding_dimension()
)
router = Router.for_query_document(
    query_modules=[SparseStaticEmbedding(tokenizer=mlm_transformer.tokenizer, frozen=False)],
    document_modules=[mlm_transformer, splade_pooling],
)

model = SparseEncoder(
    modules=[router],
    model_card_data=SparseEncoderModelCardData(
        language="en",
        license="apache-2.0",
        model_name="Inference-free SPLADE distilbert-base-uncased trained on Natural-Questions tuples",
    ),
)

# 3. Load a dataset to finetune on
full_dataset = load_dataset("sentence-transformers/natural-questions", split="train").select(range(100_000))
dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"]
eval_dataset = dataset_dict["test"]
print(train_dataset)
print(train_dataset[0])

# 4. Define a loss function
loss = SpladeLoss(
    model=model,
    loss=SparseMultipleNegativesRankingLoss(model=model),
    query_regularizer_weight=0,
    document_regularizer_weight=3e-3,
)

# 5. (Optional) Specify training arguments
run_name = "inference-free-splade-distilbert-base-uncased-nq"
args = SparseEncoderTrainingArguments(
    # Required parameter:
    output_dir=f"models/{run_name}",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    learning_rate=2e-5,
    learning_rate_mapping={r"SparseStaticEmbedding\.weight": 1e-3},  # Set a higher learning rate for the SparseStaticEmbedding module
    warmup_ratio=0.1,
    fp16=True,  # Set to False if you get an error that your GPU can't run on FP16
    bf16=False,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
    router_mapping={"query": "query", "answer": "document"},  # Map the column names to the routes
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=1000,
    save_strategy="steps",
    save_steps=1000,
    save_total_limit=2,
    logging_steps=200,
    run_name=run_name,  # Will be used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"], batch_size=16)

# 7. Create a trainer & train
trainer = SparseEncoderTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# 8. Evaluate the model performance again after training
dev_evaluator(model)

# 9. Save the trained model
model.save_pretrained(f"models/{run_name}/final")

# 10. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(run_name)

在這個例子中，我正在從distilbert/distilbert-base-uncased微調，這是一個尚未成為稀疏編碼器模型的基礎模型。這比微調現有稀疏編碼器模型（如naver/splade-cocondenser-ensembledistil）需要更多的訓練資料。

執行此指令碼後，sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq模型已為我上傳。該模型在NanoMSMARCO上得分0.5241 NDCG@10，在NanoNFCorpus上得分0.3299 NDCG@10，在NanoNQ上得分0.5357 NDCG@10，對於一個基於distilbert的免推理模型，僅在自然語言問答資料集的10萬對資料上進行訓練，這是一個不錯的結果。

該模型在文件的稀疏嵌入中平均使用184個活躍維度，而查詢的活躍維度為7.7個（即查詢中的平均標記數）。這分別對應於99.39%和99.97%的稀疏度。

所有這些資訊都儲存在自動生成的模型卡中，包括基礎模型、語言、許可證、評估結果、訓練和評估資料集資訊、超引數、訓練日誌等等。無需任何努力，您上傳的模型將包含您的潛在使用者確定模型是否適合他們所需的所有資訊。

回撥

Sentence Transformers訓練器支援各種transformers.TrainerCallback子類，包括

• 如果安裝了wandb，則使用WandbCallback將訓練指標記錄到W&B。
• 如果可以訪問tensorboard，則使用TensorBoardCallback將訓練指標記錄到TensorBoard。
• 如果安裝了codecarbon，則使用CodeCarbonCallback跟蹤訓練期間的碳排放。

只要安裝了所需的依賴項，這些功能就會自動使用，您無需進行任何指定。

有關這些回撥以及如何建立自己的回撥的更多資訊，請參閱Transformers回撥文件。

多資料集訓練

表現出色的模型通常透過同時使用多個數據集進行訓練。SparseEncoderTrainer透過允許您使用多個數據集進行訓練而無需將它們轉換為相同格式來簡化此過程。您甚至可以對每個資料集應用不同的損失函式。以下是多資料集訓練的步驟

1. 使用datasets.Dataset例項（或datasets.DatasetDict）的字典作為train_dataset和eval_dataset。
2. （可選）如果想對不同資料集使用不同的損失函式，請使用將資料集名稱對映到損失函式的字典。

每個訓練/評估批次將僅包含來自一個數據集的樣本。從多個數據集中取樣批次的順序由MultiDatasetBatchSamplers列舉決定，該列舉可以透過multi_dataset_batch_sampler傳遞給SparseEncoderTrainingArguments。有效選項包括

• MultiDatasetBatchSamplers.ROUND_ROBIN: 以迴圈方式從每個資料集中取樣，直到其中一個耗盡。此策略可能不會使用每個資料集中的所有樣本，但它確保從每個資料集中進行相同的取樣。
• MultiDatasetBatchSamplers.PROPORTIONAL（預設）：根據每個資料集的大小按比例從每個資料集中取樣。此策略確保使用每個資料集中的所有樣本，並且更頻繁地從更大的資料集中取樣。

評估

讓我們使用NanoMSMARCO資料集評估我們新訓練的免推理SPLADE模型，看看它與密集檢索方法相比如何。我們還將探討結合稀疏和密集向量的混合檢索方法，以及重排序以進一步提高搜尋質量。

執行我們略作修改的hybrid_search.py指令碼後，我們獲得了NanoMSMARCO資料集的以下結果，使用了這些模型：

• 稀疏：sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq（我們剛剛訓練的模型）
• 密集：sentence-transformers/all-MiniLM-L6-v2
• 重排序器：cross-encoder/ms-marco-MiniLM-L6-v2

稀疏和密集排名可以使用倒數排名融合（RRF）進行組合，這是一種結合多個排名結果的簡單方法。如果應用了重排序器，它將對先前檢索步驟的結果進行重排序。

結果表明，對於此資料集，結合密集和稀疏排名表現非常出色，與密集和稀疏基線相比，分別提高了12.3%和18.7%。簡而言之，結合稀疏和密集檢索方法是提高搜尋效能的非常有效的方法。

此外，對任何排名應用重排序器都將效能提高到大約66.3 NDCG@10，這表明無論是稀疏、密集還是混合（密集+稀疏）都在其前100名中找到了相關文件，然後重排序器將其排到前10名。因此，將密集 -> 重排序器管道替換為稀疏 -> 重排序器管道可能會同時提高延遲和成本

• 稀疏嵌入儲存成本更低，例如我們的模型對MS MARCO文件只使用約180個活躍維度，而不是密集模型常用的1024個維度。
• 一些稀疏編碼器支援免推理查詢處理，實現近乎即時的第一階段檢索，類似於BM25等詞彙解決方案。

訓練技巧

稀疏編碼器模型在訓練時有一些您應該注意的特點：

1. 稀疏編碼器模型不應僅透過評估分數進行評估，還應透過嵌入的稀疏性進行評估。畢竟，低稀疏性意味著模型嵌入儲存成本高昂且檢索速度慢。
2. 更強的稀疏編碼器模型幾乎完全透過從更強的教師模型（例如CrossEncoder模型）進行蒸餾來訓練，而不是直接從文字對或三元組進行訓練。例如，參見SPLADE-v3論文，其中使用SparseDistillKLDivLoss和SparseMarginMSELoss進行蒸餾。我們在此部落格中沒有詳細介紹，因為它需要更多的資料準備，但蒸餾設定應該被認真考慮。

向量資料庫整合

訓練稀疏嵌入模型後，下一個關鍵步驟是將其有效地部署到生產環境中。向量資料庫提供了儲存、索引和大規模檢索稀疏嵌入所需的基礎設施。流行的選項包括Qdrant、OpenSearch、Elasticsearch和Seismic等。

有關上述向量資料庫的全面示例，請參閱向量資料庫語義搜尋文件或下方的Qdrant示例。

Qdrant整合示例

Qdrant為稀疏向量提供了出色的支援，具有高效的儲存和快速檢索功能。以下是一個全面的實現示例

先決條件：

• Qdrant本地執行（或可訪問），更多詳細資訊請參閱Qdrant快速入門。
• 安裝了Python Qdrant客戶端

  pip install qdrant-client
  

此示例演示瞭如何設定 Qdrant 進行稀疏向量搜尋，透過展示如何使用稀疏編碼器高效編碼和索引文件，使用稀疏向量構建搜尋查詢，並提供互動式查詢介面。請參閱下文

import time

from datasets import load_dataset
from sentence_transformers import SparseEncoder
from sentence_transformers.sparse_encoder.search_engines import semantic_search_qdrant

# 1. Load the natural-questions dataset with 100K answers
dataset = load_dataset("sentence-transformers/natural-questions", split="train")
num_docs = 10_000
corpus = dataset["answer"][:num_docs]

# 2. Come up with some queries
queries = dataset["query"][:2]

# 3. Load the model
sparse_model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

# 4. Encode the corpus
corpus_embeddings = sparse_model.encode_document(
    corpus, convert_to_sparse_tensor=True, batch_size=16, show_progress_bar=True
)

# Initially, we don't have a qdrant index yet
corpus_index = None
while True:
    # 5. Encode the queries using the full precision
    start_time = time.time()
    query_embeddings = sparse_model.encode_query(queries, convert_to_sparse_tensor=True)
    print(f"Encoding time: {time.time() - start_time:.6f} seconds")

    # 6. Perform semantic search using qdrant
    results, search_time, corpus_index = semantic_search_qdrant(
        query_embeddings,
        corpus_index=corpus_index,
        corpus_embeddings=corpus_embeddings if corpus_index is None else None,
        top_k=5,
        output_index=True,
    )

    # 7. Output the results
    print(f"Search time: {search_time:.6f} seconds")
    for query, result in zip(queries, results):
        print(f"Query: {query}")
        for entry in result:
            print(f"(Score: {entry['score']:.4f}) {corpus[entry['corpus_id']]}, corpus_id: {entry['corpus_id']}")
        print("")

    # 8. Prompt for more queries
    queries = [input("Please enter a question: ")]

附加資源

訓練示例

以下頁面包含帶有解釋和程式碼連結的訓練示例。我們建議您瀏覽這些示例，熟悉訓練迴圈

• 模型蒸餾 - 使模型更小、更快、更輕的示例。
• MS MARCO - 在 MS MARCO 資訊檢索資料集上進行訓練的示例訓練指令碼。
• 檢索器 - 在通用資訊檢索資料集上進行訓練的示例訓練指令碼。
• 自然語言推理 - 自然語言推理 (NLI) 資料對於預訓練和微調模型以建立有意義的稀疏嵌入非常有幫助。
• Quora 重複問題 - Quora 重複問題是一個大型語料庫，包含來自 Quora 社群的重複問題。該資料夾包含如何訓練模型用於重複問題挖掘和語義搜尋的示例。
• STS - 訓練模型最基本的方法是使用語義文字相似度 (STS) 資料。在這裡，我們使用句子對和表示語義相似度的分數。

文件

此外，以下頁面可能有助於您瞭解更多關於 Sentence Transformers 的資訊

• 安裝
• 快速入門
• 用法
• 預訓練模型
• 訓練概述 （這篇部落格文章是訓練概述文件的精髓）
• 資料集概述
• 損失概述
• API參考

最後，這裡有一些您可能感興趣的高階頁面

• 超引數最佳化
• 分散式訓練