Modules

sentence_transformers.sparse_encoder.models 定义了不同的构建块,可用于从头开始创建 SparseEncoder 网络。更多详情,请参阅 训练概述。请注意,来自 sentence_transformers.models 的模块也可用于稀疏模型,例如 SentenceTransformer > Modules 中的 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)[source]

SPLADE 池化模块,用于创建稀疏嵌入。

该模块实现了 SPLADE 池化机制,该机制

  1. 从掩码语言模型 (MLM) 中获取 token logits。

  2. 应用稀疏转换,使用激活函数后接 log1p(即 log(1 + activation(MLM_logits)))。

  3. 应用池化策略 maxsum 来生成稀疏嵌入。

生成的嵌入具有高度稀疏性,并捕获词汇信息,使其适用于高效的信息检索。

参数:
  • pooling_strategy (str) –

    跨 token 维度的池化方法。选项包括:

  • activation_function (str) –

    在 log1p 转换之前应用的激活函数。选项包括:

  • word_embedding_dimension (int, optional) – 输出嵌入的维度(如果需要)。

  • chunk_size (int, optional) – 沿序列长度维度的块大小(即每个块的 token 数)。如果为 None,则一次处理整个序列。使用较小的块会降低内存使用量,但可能会降低训练和推理速度。默认为 None。

MLM Transformer

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')[source]

MLMTransformer 适配了掩码语言模型 (MLM),用于稀疏编码应用。

此类扩展了 Transformer 类,专门用于处理具有 MLM 头的模型(如 BERT、RoBERTa 等),并且设计用于与 SpladePooling 一起创建 SPLADE 稀疏表示。

MLMTransformer 访问 MLM 预测头以获取每个 token 的词汇 logits,这些 logits 稍后由 SpladePooling 用于创建稀疏词汇表示。

参数:
  • model_name_or_path – Hugging Face 模型名称(https://hugging-face.cn/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

SparseAutoEncoder

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)[source]

该模块实现了基于以下论文的稀疏自编码器架构:Beyond Matryoshka: Revisiting Sparse Coding for Adaptive Representation, https://hugging-face.cn/papers/2503.01776

该模块通过以下方式将密集嵌入转换为稀疏表示:

  1. 应用多层前馈网络

  2. 应用 top-k 稀疏化以仅保留最大的值

  3. 支持用于训练稳定性的辅助损失(通过 k_aux 参数)

参数:
  • input_dim – 输入嵌入的维度。

  • hidden_dim – 隐藏层的维度。默认为 512。

  • k – 在最终稀疏表示中要保留的最大值的数量。默认为 8。

  • k_aux – 用于计算辅助损失的最大值数量。默认为 512。

  • normalize – 是否对输入嵌入应用层归一化。默认为 False。

  • dead_threshold – 死神经元的阈值。非零激活值低于此阈值的神经元被视为死神经元。默认为 30。

SparseStaticEmbedding

class sentence_transformers.sparse_encoder.models.SparseStaticEmbedding(tokenizer: PreTrainedTokenizer, weight: torch.Tensor | None = None, frozen: bool = False)[source]

SparseStaticEmbedding 模块,用于高效的稀疏表示。

这个轻量级模块通过将输入 token 映射到静态权重(例如 IDF(逆文档频率)权重)来计算稀疏表示。它被设计用于对查询或文档进行编码,生成基于输入中 token 存在情况的固定大小的嵌入。

一个常见的场景是使用此模块对查询进行编码,并使用像 SPLADE(MLMTransformer + SpladePooling)这样的更复杂的模块来对文档进行编码。

参数:
  • tokenizer (PreTrainedTokenizer) – 用于将输入文本分词为输入 ID 的 PreTrainedTokenizer。

  • weight (torch.Tensor | None) – 词汇表 token 的静态权重(例如 IDF 权重),形状应为 (vocab_size,)。如果为 None,则将权重初始化为全 1 向量。默认为 None。

  • frozen (bool) – 权重是否应该被冻结(不可训练)。默认为 False。