训练概览

为什么要进行微调？

微调稀疏编码器模型通常能显著提高模型在您的用例上的性能，因为每个任务都需要不同的相似性概念。例如，给定新闻文章：

“Apple launches the new iPad”（苹果发布新iPad）
“NVIDIA is gearing up for the next GPU generation”（英伟达正为下一代GPU做准备）

在以下用例中，我们可能对相似性有不同的理解：

用于新闻文章分类（如经济、体育、科技、政治等）的模型，应为这些文本生成相似的嵌入。
用于语义文本相似性的模型，应为这些文本生成不相似的嵌入，因为它们具有不同的含义。
用于语义搜索的模型不需要两个文档之间的相似性概念，因为它应该只比较查询和文档。

另请参阅训练示例，其中包含许多用于常见实际应用的训练脚本，您可以采用这些脚本。

训练组件

训练稀疏编码器模型涉及4到6个组件。

了解如何初始化用于训练的模型。了解如何准备用于训练的数据。了解如何准备和选择损失函数。了解哪些训练参数有用。了解如何在训练期间和之后进行评估。了解如何开始训练过程。

模型

稀疏编码器模型由一系列模块、稀疏编码器特定模块或自定义模块组成，提供了很大的灵活性。如果您想进一步微调一个SparseEncoder模型（例如，它有一个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 MLMTransformer, SparseStaticEmbedding, 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）模型在密集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
sae = 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, sae])

或者，如果您的基础模型是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维度）的密集编码器模型上最有效。

数据集

SparseEncoderTrainer使用datasets.Dataset（一个数据集）或datasets.DatasetDict实例（多个数据集，另请参阅多数据集训练）进行训练和评估。

🤗 Hugging Face Hub 上的数据

如果您想从Hugging Face Datasets加载数据，那么您应该使用datasets.load_dataset()

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="train")
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")

print(train_dataset)
"""
Dataset({
    features: ['anchor', 'positive', 'negative'],
    num_rows: 557850
})
"""

一些数据集（包括sentence-transformers/all-nli）需要您提供一个“子集”以及数据集名称。sentence-transformers/all-nli有4个子集，每个子集都有不同的数据格式：pair、pair-class、pair-score、triplet。

注意

许多Hugging Face数据集通过sentence-transformers标签与Sentence Transformers无缝集成，您可以轻松地通过浏览https://hugging-face.cn/datasets?other=sentence-transformers找到它们。我们强烈建议您浏览这些数据集，以找到对您的任务有用的训练数据集。

本地数据（CSV、JSON、Parquet、Arrow、SQL）

如果您有常用文件格式的本地数据，那么您可以使用datasets.load_dataset()轻松加载这些数据

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")

或

from datasets import load_dataset

dataset = load_dataset("json", data_files="my_file.json")

需要预处理的本地数据

如果您有需要额外预处理的本地数据，我的建议是使用datasets.Dataset.from_dict()和一个列表字典来初始化数据集，如下所示

from datasets import Dataset

anchors = []
positives = []
# Open a file, do preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "anchor": anchors,
    "positive": positives,
})

字典中的每个键都将成为结果数据集中的一列。

数据集格式

您的数据集格式与您的损失函数匹配（或者您选择与数据集格式匹配的损失函数）非常重要。验证数据集格式是否与损失函数兼容涉及两个步骤：

如果您的损失函数根据损失概述表需要一个标签（Label），那么您的数据集必须包含一个名为“label”或“score”的列。此列将自动用作标签。
根据损失概述表，所有不名为“label”或“score”的列都被视为输入（Inputs）。剩余列的数量必须与您选择的损失的有效输入数量匹配。这些列的名称无关紧要，只有顺序才重要。

例如，给定一个包含["text1", "text2", "label"]列的数据集，其中“label”列包含0到1之间的浮点相似度分数，我们可以将其与SparseCoSENTLoss、SparseAnglELoss和SparseCosineSimilarityLoss一起使用，因为它：

包含这些损失函数所需的“label”列。
有2个非标签列，正好是这些损失函数所需的数量。

如果您的列顺序不正确，请务必使用Dataset.select_columns重新排序数据集列。例如，如果您的数据集包含["good_answer", "bad_answer", "question"]作为列，那么这个数据集技术上可以与需要（anchor, positive, negative）三元组的损失一起使用，但good_answer列将被视为anchor，bad_answer视为positive，question视为negative。

此外，如果您的数据集包含无关列（例如 sample_id、metadata、source、type），您应该使用Dataset.remove_columns将其删除，否则它们将被用作输入。您也可以使用Dataset.select_columns仅保留所需的列。

损失函数

损失函数量化了模型在给定一批数据上的表现，从而允许优化器更新模型权重以产生更优（即更低）的损失值。这是训练过程的核心。

遗憾的是，没有一个损失函数能适用于所有用例。相反，使用哪个损失函数在很大程度上取决于您可用的数据和目标任务。请参阅数据集格式，了解哪些数据集对哪些损失函数有效。此外，损失概述将是您了解选项的最佳助手。

警告

要训练一个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
})
"""

训练参数

SparseEncoderTrainingArguments类可用于指定影响训练性能的参数以及定义跟踪/调试参数。尽管它是可选的，但强烈建议尝试各种有用的参数。

eval_strategy eval_steps save_strategy save_steps save_total_limit report_to run_name log_level logging_steps push_to_hub hub_model_id hub_strategy hub_private_repo

以下是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 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,  # losses that use "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",  # Will be used in W&B if `wandb` is installed
)

评估器

您可以为SparseEncoderTrainer提供一个eval_dataset以在训练期间获得评估损失，但在训练期间获得更具体的指标也可能很有用。为此，您可以使用评估器在训练之前、期间或之后使用有用的指标评估模型的性能。您可以同时使用eval_dataset和评估器，或两者之一，或两者都不用。它们根据eval_strategy和eval_steps训练参数进行评估。

以下是Sentence Transformers为稀疏编码器模型实现的评估器：

评估器	所需数据
`SparseBinaryClassificationEvaluator`	带有类别标签的对。
`SparseEmbeddingSimilarityEvaluator`	带相似度分数的句子对。
`SparseInformationRetrievalEvaluator`	查询 (qid => 问题)、语料库 (cid => 文档) 和相关文档 (qid => set[cid])。
`SparseNanoBEIREvaluator`	无需数据。
`SparseMSEEvaluator`	用于用教师模型嵌入的源句子和用于用学生模型嵌入的目标句子。可以是相同的文本。
`SparseRerankingEvaluator`	字典列表 `{'query': '...', 'positive': [...], 'negative': [...]}`。
`SparseTranslationEvaluator`	两种不同语言的句子对。
`SparseTripletEvaluator`	（锚点，正例，负例）对。

此外，SequentialEvaluator应用于将多个评估器组合成一个评估器，该评估器可以传递给SparseEncoderTrainer。

有时您没有所需的评估数据来自己准备这些评估器之一，但您仍然希望跟踪模型在某些常见基准测试上的表现。在这种情况下，您可以使用这些带有 Hugging Face 数据的评估器。

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

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

# Load the STSB dataset (https://hugging-face.cn/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

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://hugging-face.cn/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将所有先前的组件整合在一起。我们只需使用模型、训练参数（可选）、训练数据集、评估数据集（可选）、损失函数、评估器（可选）来指定训练器，然后就可以开始训练了。让我们看一个所有这些组件都整合在一起的脚本：

SPLADE

import logging

from datasets import load_dataset

from sentence_transformers import (
    SparseEncoder,
    SparseEncoderModelCardData,
    SparseEncoderTrainer,
    SparseEncoderTrainingArguments,
)
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
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
model = SparseEncoder(
    "distilbert/distilbert-base-uncased",
    model_card_data=SparseEncoderModelCardData(
        language="en",
        license="apache-2.0",
        model_name="DistilBERT base 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"]

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

# 5. (Optional) Specify training arguments
run_name = "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,
    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
    # 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)

无推理SPLADE

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 MLMTransformer, SparseStaticEmbedding, 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-4,
)

# 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)

回调

此稀疏编码器训练器集成了对各种transformers.TrainerCallback子类的支持，例如：

SpladeRegularizerWeightSchedulerCallback用于在训练期间调度SpladeLoss损失的lambda参数。
WandbCallback，如果安装了wandb，则自动将训练指标记录到W&B
TensorBoardCallback，如果tensorboard可访问，则将训练指标记录到TensorBoard。
CodeCarbonCallback，如果安装了codecarbon，则在训练期间跟踪模型的碳排放。
- 注意：这些碳排放量将被包含在您自动生成的模型卡片中。

有关集成回调以及如何编写您自己的回调的更多信息，请参阅 Transformers 回调文档。

多数据集训练

性能最佳的模型是同时使用多个数据集训练的。通常，这相当棘手，因为每个数据集都有不同的格式。然而，SparseEncoderTrainer可以训练多个数据集，而无需将每个数据集转换为相同的格式。它甚至可以对每个数据集应用不同的损失函数。使用多个数据集进行训练的步骤是：

使用Dataset实例的字典（或DatasetDict）作为train_dataset（可选地，也作为eval_dataset）。
（可选）使用一个损失函数字典，将数据集名称映射到损失。仅当您希望为不同的数据集使用不同的损失函数时才需要。

每个训练/评估批次将仅包含来自一个数据集的样本。从多个数据集中采样批次的顺序由MultiDatasetBatchSamplers枚举定义，该枚举可以通过multi_dataset_batch_sampler传递给SparseEncoderTrainingArguments。有效选项包括：

MultiDatasetBatchSamplers.ROUND_ROBIN: 从每个数据集中轮流采样，直到其中一个耗尽。使用此策略，可能不会使用每个数据集中的所有样本，但每个数据集都被同等地采样。
MultiDatasetBatchSamplers.PROPORTIONAL（默认）：按其大小比例从每个数据集采样。使用此策略，每个数据集中的所有样本都会被使用，并且更大的数据集会更频繁地被采样。

训练技巧

训练稀疏编码器模型时，您应该注意一些特性：

稀疏编码器模型不应仅通过评估分数进行评估，还应通过嵌入的稀疏性进行评估。毕竟，低稀疏性意味着模型嵌入存储成本高昂且检索缓慢。这也意味着决定稀疏性的参数（例如SpladeLoss中的query_regularizer_weight、document_regularizer_weight以及CSRLoss中的beta和gamma）应进行调整，以在性能和稀疏性之间实现良好平衡。每个评估器都会输出active_dims和sparsity_ratio指标，可用于评估嵌入的稀疏性。
不建议在训练前对未训练的模型使用Evaluator，因为稀疏度会非常低，内存使用量可能会出乎意料地高。
更强的稀疏编码器模型几乎完全通过从更强的教师模型（例如CrossEncoder模型）进行蒸馏来训练，而不是直接从文本对或三元组进行训练。例如，参见SPLADE-v3论文，它使用SparseDistillKLDivLoss和SparseMarginMSELoss进行蒸馏。
虽然大多数密集嵌入模型训练时都与余弦相似度一起使用，但SparseEncoder模型通常训练时与点积一起用于计算相似度。有些损失函数要求您提供相似度函数，这时使用点积可能会更好。请注意，您通常可以为损失函数提供model.similarity或model.similarity_pairwise。