模块
sentence_transformers.sparse_encoder.models
定义了不同的构建块,可用于从头开始创建 SparseEncoder 网络。更多详情,请参阅训练概览。请注意,sentence_transformers.models
中的模块也可用于稀疏模型,例如SentenceTransformer > 模块中的 sentence_transformers.models.Transformer
。
SPLADE 池化
- class sentence_transformers.sparse_encoder.models.SpladePooling(pooling_strategy: Literal['max', 'sum'] = 'max', activation_function: Literal['relu', 'log1p_relu'] = 'relu', word_embedding_dimension: int | None = None, chunk_size: int | None = None)[源代码]
用于创建稀疏嵌入的 SPLADE 池化模块。
此模块实现了 SPLADE 池化机制,该机制
从掩码语言模型(MLM)中获取词元 logits。
通过一个激活函数后接 log1p(即 log(1 + activation(MLM_logits)))来应用稀疏变换。
应用 max 或 sum 池化策略来产生稀疏嵌入。
生成的嵌入是高度稀疏的,并捕获了词汇信息,使其适用于高效的信息检索。
- 参数:
pooling_strategy (str) –
跨词元维度的池化方法。可选值
sum:求和池化(用于原始 SPLADE,见 https://arxiv.org/pdf/2107.05720)。
max:最大池化(用于 SPLADEv2 及更高版本模型,见 https://arxiv.org/pdf/2109.10086 或 https://arxiv.org/pdf/2205.04733)。
activation_function (str) –
在 log1p 变换之前应用的激活函数。可选值
relu:ReLU 激活函数(所有 Splade 模型的标准配置)。
log1p_relu:log(1 + ReLU(x)) 变体,用于 Opensearch Splade 模型,见 https://arxiv.org/pdf/2504.14839。
word_embedding_dimension (int, 可选) – 输出嵌入的维度(如果需要)。
chunk_size (int, 可选) – 沿序列长度维度的块大小(即每个块的词元数)。如果为 None,则一次性处理整个序列。使用较小的块可以减少内存使用,但可能会降低训练和推理速度。默认为 None。
MLM 转换器
-
class sentence_transformers.sparse_encoder.models.MLMTransformer(model_name_or_path: str, max_seq_length: int | None = None, model_args: dict[str, Any] | None = None, tokenizer_args: dict[str, Any] | None = None, config_args: dict[str, Any] | None = None, cache_dir: str | None = None, do_lower_case: bool = False, tokenizer_name_or_path: str | None = None, backend
: str = 'torch'
)[源代码] MLMTransformer 调整一个掩码语言模型(MLM)以适应稀疏编码应用。
该类扩展了 Transformer 类,专门用于处理带有 MLM 头的模型(如 BERT、RoBERTa 等),并设计与 SpladePooling 一起使用以创建 SPLADE 稀疏表示。
MLMTransformer 访问 MLM 预测头以获取每个词元的词汇 logits,这些 logits 随后被 SpladePooling 用于创建稀疏词汇表示。
- 参数:
model_name_or_path – Hugging Face 模型名称 (https://huggingface.co/models)
max_seq_length – 截断任何超过 max_seq_length 的输入
model_args – 传递给 Hugging Face MLMTransformers 模型的关键字参数
tokenizer_args – 传递给 Hugging Face MLMTransformers 分词器的关键字参数
config_args – 传递给 Hugging Face MLMTransformers 配置的关键字参数
cache_dir – Hugging Face MLMTransformers 存储/加载模型的缓存目录
do_lower_case – 如果为 true,则将输入转换为小写(无论模型是否区分大小写)
tokenizer_name_or_path – 分词器的名称或路径。如果为 None,则使用 model_name_or_path
backend – 用于模型推理的后端。目前对于此类只能是 torch。
稀疏自编码器
- class sentence_transformers.sparse_encoder.models.SparseAutoEncoder(input_dim: int, hidden_dim: int = 512, k: int = 8, k_aux: int = 512, normalize: bool = False, dead_threshold: int = 30)[源代码]
该模块实现了基于论文《Beyond Matryoshka: Revisiting Sparse Coding for Adaptive Representation》的稀疏自编码器架构,见 https://arxiv.org/abs/2503.01776
该模块通过以下方式将密集嵌入转换为稀疏表示
应用多层前馈网络
应用 top-k 稀疏化以仅保留最大值
支持用于训练稳定性的辅助损失(通过 k_aux 参数)
- 参数:
input_dim – 输入嵌入的维度。
hidden_dim – 隐藏层的维度。默认为 512。
k – 在最终稀疏表示中保留的最大值的数量。默认为 8。
k_aux – 为辅助损失计算保留的最大值的数量。默认为 512。
normalize – 是否对输入嵌入应用层归一化。默认为 False。
dead_threshold – 死神经元的阈值。非零激活值低于此阈值的神经元被认为是死的。默认为 30。
稀疏静态嵌入
- class sentence_transformers.sparse_encoder.models.SparseStaticEmbedding(tokenizer: PreTrainedTokenizer, weight: torch.Tensor | None = None, frozen: bool = False)[源代码]
用于高效稀疏表示的 SparseStaticEmbedding 模块。
这个轻量级模块通过将输入词元映射到静态权重(如 IDF,即逆文档频率权重)来计算稀疏表示。它被设计用于根据输入中词元的出现情况,将查询或文档编码为固定大小的嵌入。
一个常见的场景是使用此模块对查询进行编码,而对文档编码则使用更重量级的模块,如 SPLADE (MLMTransformer + SpladePooling)。
- 参数:
tokenizer (PreTrainedTokenizer) – 用于将输入文本分词为输入 ID 的 PreTrainedTokenizer。
weight (torch.Tensor | None) – 词汇表中词元的静态权重(例如 IDF 权重),形状应为 (vocab_size,)。如果为 None,则将权重初始化为全一向量。默认为 None。
frozen (bool) – 权重是否应被冻结(不可训练)。默认为 False。