SparseEncoder

SparseEncoder

class sentence_transformers.sparse_encoder.SparseEncoder(model_name_or_path: str | None = None, modules: Iterable[Module] | None = None, device: str | None = None, prompts: dict[str, str] | None = None, default_prompt_name: str | None = None, similarity_fn_name: str | SimilarityFunction | None = None, cache_folder: str | None = None, trust_remote_code: bool = False, revision: str | None = None, local_files_only: bool = False, token: bool | str | None = None, max_active_dims: int | None = None, model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None, config_kwargs: dict[str, Any] | None = None, model_card_data: SparseEncoderModelCardData | None = None, backend: Literal['torch', 'onnx', 'openvino'] = 'torch')[source]

加载或创建一个 SparseEncoder 模型，该模型可用于将句子/文本映射到稀疏嵌入。

参数:

  • model_name_or_path (str, optional) – 如果是磁盘上的文件路径，则从该路径加载模型。如果不是路径，则首先尝试下载预训练的 SparseEncoder 模型。如果失败，则尝试使用该名称从 Hugging Face Hub 构建模型。

  • modules (Iterable[nn.Module], optional) – 应按顺序调用的 torch Module 列表，可用于从头开始创建自定义 SparseEncoder 模型。

  • device (str, optional) – 应用于计算的设备（如“cuda”、“cpu”、“mps”、“npu”）。如果为 None，则检查是否可以使用 GPU。

  • prompts (Dict[str, str], optional) – 包含模型提示的字典。键是提示名称，值是提示文本。提示文本将在要编码的任何文本之前添加。例如：{“query”: “query: “, “passage”: “passage: “} 或 {“clustering”: “Identify the main category based on the titles in “}。

  • default_prompt_name (str, optional) – 默认使用的提示名称。如果未设置，则不应用任何提示。

  • similarity_fn_name (str or SimilarityFunction, optional) – 要使用的相似性函数的名称。有效选项包括“cosine”、“dot”、“euclidean”和“manhattan”。如果未设置，则当 similarity 或 similarity_pairwise 被调用且 model.similarity_fn_name 仍为 None 时，它将自动设置为“cosine”。

  • cache_folder (str, optional) – 存储模型的路径。也可以通过 SENTENCE_TRANSFORMERS_HOME 环境变量设置。

  • trust_remote_code (bool, optional) – 是否允许 Hub 上定义的自定义模型在其自己的建模文件中。此选项仅应设置为 True 的您信任并已阅读代码的存储库，因为它将在您的本地计算机上执行 Hub 上的代码。

  • revision (str, optional) – 要使用的特定模型版本。它可以是存储在 Hugging Face 上的模型的名称、标签名称或提交 ID。

  • local_files_only (bool, optional) – 是否只查看本地文件（即不尝试下载模型）。

  • token (bool or str, optional) – 用于下载私有模型的 Hugging Face 身份验证令牌。

  • max_active_dims (int, optional) – 模型输出中的最大活动（非零）维度数。默认为 None。这意味着活动维度数没有限制，如果您的模型尚未（或尚未）微调到高稀疏度，可能会导致速度缓慢或内存占用过大。

  • model_kwargs (Dict[str, Any], optional) –

    要传递给 Hugging Face Transformers 模型的其他模型配置参数。特别有用的选项是

    • torch_dtype: 覆盖默认的 torch.dtype 并以特定的 dtype 加载模型。不同的选项有

      1. torch.float16, torch.bfloat16 或 torch.float: 以指定的 dtype 加载，如果模型存在 config.torch_dtype 则忽略它。如果未指定 - 模型将以 torch.float (fp32) 加载。

      2. "auto" - 将尝试使用模型 config.json 文件中的 torch_dtype 条目。如果未找到此条目，则接下来检查检查点中第一个浮点类型权重的 dtype，并将其用作 dtype。这将使用模型训练结束时保存的 dtype 加载模型。它不能用作模型训练方式的指示器。因为它可能以一种半精度 dtype 训练，但保存为 fp32。

    • attn_implementation: 模型中使用的注意力实现（如果相关）。可以是 “eager”（注意力的手动实现）、“sdpa”（使用 F.scaled_dot_product_attention）或 “flash_attention_2”（使用 Dao-AILab/flash-attention）中的任何一个。默认情况下，如果可用，对于 torch>=2.1.1 将使用 SDPA。否则默认是手动 “eager” 实现。

    • provider: 如果 backend 是“onnx”，则这是用于推理的提供程序，例如“CPUExecutionProvider”、“CUDAExecutionProvider”等。有关所有 ONNX 执行提供程序，请参阅 https://runtime.onnx.org.cn/docs/execution-providers/。

    • file_name: 如果 backend 是“onnx”或“openvino”，这是要加载的文件名，对于加载已优化或量化的 ONNX 或 OpenVINO 模型很有用。

    • export: 如果 backend 是“onnx”或“openvino”，则这是一个布尔标志，指定此模型是否应导出到后端。如果未指定，则仅当模型存储库或目录已包含导出的模型时，才会导出模型。

    有关更多详细信息，请参阅 PreTrainedModel.from_pretrained 文档。

  • tokenizer_kwargs (Dict[str, Any], optional) – 要传递给 Hugging Face Transformers 分词器的额外分词器配置参数。有关更多详细信息，请参阅 AutoTokenizer.from_pretrained 文档。

  • config_kwargs (Dict[str, Any], optional) – 要传递给 Hugging Face Transformers 配置的其他模型配置参数。有关更多详细信息，请参阅 AutoConfig.from_pretrained 文档。

  • model_card_data (SparseEncoderModelCardData, optional) – 包含模型信息的模型卡数据对象。保存模型时用于生成模型卡。如果未设置，则会创建一个默认的模型卡数据对象。

  • backend (str) – 用于推理的后端。可以是“torch”（默认）、“onnx”或“openvino”之一。有关不同后端的基准测试信息，请参阅 https://sbert.net.cn/docs/sentence_transformer/usage/efficiency.html。

示例

from sentence_transformers import SparseEncoder

# Load a pre-trained SparseEncoder model
model = SparseEncoder('naver/splade-cocondenser-ensembledistil')

# Encode some texts
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 between all sentences
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[   35.629,     9.154,     0.098],
#         [    9.154,    27.478,     0.019],
#         [    0.098,     0.019,    29.553]])

初始化内部 Module 状态，由 nn.Module 和 ScriptModule 共享。

active_adapters() → list[str][source]

如果您不熟悉适配器和 PEFT 方法，我们邀请您阅读 PEFT 官方文档中的更多内容：https://hugging-face.cn/docs/peft

获取模型的当前活动适配器。对于多适配器推理（组合多个适配器进行推理），返回所有活动适配器的列表，以便用户可以相应地处理它们。

对于之前的 PEFT 版本（不支持多适配器推理），module.active_adapter 将返回单个字符串。

add_adapter(*args, **kwargs) → None[source]

为模型添加一个全新的适配器以供训练。如果未传入适配器名称，则会分配一个默认名称以遵循 PEFT 库的约定（在 PEFT 中，我们使用“default”作为默认适配器名称）。

要求 peft 作为后端来加载适配器权重，并且底层模型必须与 PEFT 兼容。

参数:
bfloat16() → T

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的前向。

此模块的 __call__ 方法已编译，所有参数按原样传递给 torch.compile()。

有关此函数参数的详细信息，请参阅 torch.compile()。

cpu() → T

将所有模型参数和缓冲区移动到 CPU。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

cuda(device: int | device | None = None) → T

将所有模型参数和缓冲区移至 GPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 GPU 上，则应在构建优化器之前调用此方法。

注意

此方法就地修改模块。

参数:

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回:

self

返回类型:

Module (模块)

decode(embeddings: Tensor, top_k: int | None = None) → list[tuple[str, float]] | list[list[tuple[str, float]]][source]

从稀疏嵌入中解码前 K 个 token 和权重。如果为空，则仅返回所有 token 和权重。

参数:

  • embeddings (torch.Tensor) – 稀疏嵌入张量（batch, vocab）或（vocab）。

  • top_k (int, optional) – 每个样本返回的 top token 数量。如果为 None，则返回所有非零 token。

返回:

每个嵌入的（token，权重）元组列表。如果输入是批处理的，则返回元组列表的列表。

返回类型:

list[tuple[str, float]] | list[list[tuple[str, float]]]

delete_adapter(*args, **kwargs) → None[source]

如果您不熟悉适配器和 PEFT 方法，我们邀请您阅读 PEFT 官方文档中的更多内容：https://hugging-face.cn/docs/peft

从底层模型中删除适配器的 LoRA 层。

参数:
property device: device

从模块获取 torch.device，假设整个模块都在一个设备上。如果没有 PyTorch 参数，则回退到 CPU。

disable_adapters() → None[source]

禁用附加到模型的所有适配器。这将仅使用基础模型进行推理。

double() → T

将所有浮点参数和缓冲区转换为 double 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

enable_adapters() → None[source]

启用附加到模型的适配器。模型将使用 self.active_adapter()

encode(sentences: str | list[str] | ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, convert_to_tensor: bool = True, convert_to_sparse_tensor: bool = True, save_to_cpu: bool = False, device: str | list[str | device] | None = None, max_active_dims: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs: Any) → list[Tensor] | ndarray | Tensor | dict[str, Tensor] | list[dict[str, Tensor]][source]

计算稀疏句子嵌入。

提示

如果您不确定是否应该使用 encode()、encode_query() 或 encode_document()，最好为信息检索任务使用 encode_query() 和 encode_document()，并为所有其他任务使用 encode()。

请注意，encode() 是最通用的方法，可用于任何任务，包括信息检索，并且如果模型没有使用预定义的提示和/或任务类型进行训练，则所有三个方法将返回相同的嵌入。

参数:

  • sentences (Union[str, List[str]]) – 要嵌入的句子。

  • prompt_name (Optional[str], optional) – 用于编码的提示的名称。必须是构造函数中设置或从模型配置加载的 prompts 字典中的一个键。例如，如果 prompt_name 是“query”，并且 prompts 是 {“query”: “query: “, …}，那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”，因为句子被追加到提示之后。如果也设置了 prompt，则忽略此参数。默认为 None。

  • prompt (Optional[str], optional) – 用于编码的提示。例如，如果提示是“query: ”，那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”，因为句子被追加到提示之后。如果设置了 prompt，则忽略 prompt_name。默认为 None。

  • batch_size (int, optional) – 用于计算的批大小。默认为 32。

  • show_progress_bar (bool, optional) – 在编码句子时是否输出进度条。默认为 None。

  • convert_to_tensor (bool, optional) – 输出应该是单个堆叠张量（True）还是单个张量列表（False）。稀疏张量可能难以切片，因此此选项允许您输出张量列表。默认为 True。

  • convert_to_sparse_tensor (bool, optional) – 输出格式是否应为稀疏（COO）张量。默认为 True。

  • save_to_cpu (bool, optional) – 输出是应移动到 CPU 还是保留在计算它的设备上。默认为 False。

  • device (Union[str, List[str], None], optional) –

    用于计算的设备。可以是

    • 单个设备字符串（例如，“cuda:0”、“cpu”）用于单进程编码

    • 设备字符串列表（例如，[“cuda:0”, “cuda:1”], [“cpu”, “cpu”, “cpu”, “cpu”]）用于在多个进程之间分发编码

    • None 以自动检测可用于单进程编码的设备

    如果提供了列表，则使用多进程编码。默认为 None。

  • max_active_dims (int, optional) – 模型输出中的最大活动（非零）维度数。None 表示将使用模型配置中的值。默认为 None。如果模型配置中为 None，则表示活动维度数没有限制，如果您的模型尚未（或尚未）微调到高稀疏度，可能会导致速度缓慢或内存占用过大。

  • pool (Dict[Literal["input", "output", "processes"], Any], optional) – 由start_multi_process_pool()创建的用于多进程编码的池。如果提供，编码将被分发到多个进程。这对于大型数据集和多个 GPU 可用时推荐使用。默认为 None。

  • chunk_size (int, optional) – 多进程编码的块大小。仅在 multiprocessing 时使用，即当 pool 不为 None 或 device 是列表时。如果为 None，则计算一个合理的默认值。默认为 None。

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 torch 稀疏张量。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 save_to_cpu 为 True，则嵌入将移至 CPU。

返回类型:

Union[List[Tensor], ndarray, Tensor]

示例

from sentence_transformers import SparseEncoder

# Load a pre-trained SparseEncoder model
model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

# Encode some texts
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)
encode_document(sentences: str | list[str] | ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, convert_to_tensor: bool = True, convert_to_sparse_tensor: bool = True, save_to_cpu: bool = False, device: str | list[str | device] | None = None, max_active_dims: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs: Any) → list[Tensor] | ndarray | Tensor | dict[str, Tensor] | list[dict[str, Tensor]][source]

计算专门为文档/段落表示优化的句子嵌入。

此方法是 encode() 的专门版本，仅在两个方面有所不同：

  1. 如果未提供 prompt_name 或 prompt，它将使用预定义的“document”提示（如果该提示存在于模型的 prompts 字典中）。

  2. 它将 task 设置为“document”。如果模型具有 Router 模块，它将使用“document”任务类型将输入路由到适当的子模块。

提示

如果您不确定是否应该使用 encode()、encode_query() 或 encode_document()，最好为信息检索任务使用 encode_query() 和 encode_document()，并为所有其他任务使用 encode()。

请注意，encode() 是最通用的方法，可用于任何任务，包括信息检索，并且如果模型没有使用预定义的提示和/或任务类型进行训练，则所有三个方法将返回相同的嵌入。

参数:

  • sentences (Union[str, List[str]]) – 要嵌入的句子。

  • prompt_name (Optional[str], optional) – 用于编码的提示的名称。必须是构造函数中设置或从模型配置加载的 prompts 字典中的一个键。例如，如果 prompt_name 是“query”，并且 prompts 是 {“query”: “query: “, …}，那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”，因为句子被追加到提示之后。如果也设置了 prompt，则忽略此参数。默认为 None。

  • prompt (Optional[str], optional) – 用于编码的提示。例如，如果提示是“query: ”，那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”，因为句子被追加到提示之后。如果设置了 prompt，则忽略 prompt_name。默认为 None。

  • batch_size (int, optional) – 用于计算的批大小。默认为 32。

  • show_progress_bar (bool, optional) – 在编码句子时是否输出进度条。默认为 None。

  • convert_to_tensor (bool, optional) – 输出应该是单个堆叠张量（True）还是单个张量列表（False）。稀疏张量可能难以切片，因此此选项允许您输出张量列表。默认为 True。

  • convert_to_sparse_tensor (bool, optional) – 输出格式是否应为稀疏（COO）张量。默认为 True。

  • save_to_cpu (bool, optional) – 输出是应移动到 CPU 还是保留在计算它的设备上。默认为 False。

  • device (Union[str, List[str], None], optional) –

    用于计算的设备。可以是

    • 单个设备字符串（例如，“cuda:0”、“cpu”）用于单进程编码

    • 设备字符串列表（例如，[“cuda:0”, “cuda:1”], [“cpu”, “cpu”, “cpu”, “cpu”]）用于在多个进程之间分发编码

    • None 以自动检测可用于单进程编码的设备

    如果提供了列表，则使用多进程编码。默认为 None。

  • max_active_dims (int, optional) – 模型输出中的最大活动（非零）维度数。None 表示将使用模型配置中的值。默认为 None。如果模型配置中为 None，则表示活动维度数没有限制，如果您的模型尚未（或尚未）微调到高稀疏度，可能会导致速度缓慢或内存占用过大。

  • pool (Dict[Literal["input", "output", "processes"], Any], optional) – 由start_multi_process_pool()创建的用于多进程编码的池。如果提供，编码将被分发到多个进程。这对于大型数据集和多个 GPU 可用时推荐使用。默认为 None。

  • chunk_size (int, optional) – 多进程编码的块大小。仅在 multiprocessing 时使用，即当 pool 不为 None 或 device 是列表时。如果为 None，则计算一个合理的默认值。默认为 None。

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 torch 稀疏张量。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 save_to_cpu 为 True，则嵌入将移至 CPU。

返回类型:

Union[List[Tensor], ndarray, Tensor]

示例

from sentence_transformers import SparseEncoder

# Load a pre-trained SparseEncoder model
model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

# Encode some texts
sentences = [
    "This research paper discusses the effects of climate change on marine life.",
    "The article explores the history of artificial intelligence development.",
    "This document contains technical specifications for the new product line.",
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# (3, 30522)
encode_query(sentences: str | list[str] | ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, convert_to_tensor: bool = True, convert_to_sparse_tensor: bool = True, save_to_cpu: bool = False, device: str | list[str | device] | None = None, max_active_dims: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs: Any) → list[Tensor] | ndarray | Tensor | dict[str, Tensor] | list[dict[str, Tensor]][source]

计算专门为查询表示优化的句子嵌入。

此方法是 encode() 的专门版本，仅在两个方面有所不同：

  1. 如果未提供 prompt_name 或 prompt，它会使用模型 prompts 字典中预定义的“query”提示（如果可用）。

  2. 它将 task 设置为“query”。如果模型具有 Router 模块，它将使用“query”任务类型将输入路由到适当的子模块。

提示

如果您不确定是否应该使用 encode()、encode_query() 或 encode_document()，最好为信息检索任务使用 encode_query() 和 encode_document()，并为所有其他任务使用 encode()。

请注意，encode() 是最通用的方法，可用于任何任务，包括信息检索，并且如果模型没有使用预定义的提示和/或任务类型进行训练，则所有三个方法将返回相同的嵌入。

参数:

  • sentences (Union[str, List[str]]) – 要嵌入的句子。

  • prompt_name (Optional[str], optional) – 用于编码的提示的名称。必须是构造函数中设置或从模型配置加载的 prompts 字典中的一个键。例如，如果 prompt_name 是“query”，并且 prompts 是 {“query”: “query: “, …}，那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”，因为句子被追加到提示之后。如果也设置了 prompt，则忽略此参数。默认为 None。

  • prompt (Optional[str], optional) – 用于编码的提示。例如，如果提示是“query: ”，那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”，因为句子被追加到提示之后。如果设置了 prompt，则忽略 prompt_name。默认为 None。

  • batch_size (int, optional) – 用于计算的批大小。默认为 32。

  • show_progress_bar (bool, optional) – 在编码句子时是否输出进度条。默认为 None。

  • convert_to_tensor (bool, optional) – 输出应该是单个堆叠张量（True）还是单个张量列表（False）。稀疏张量可能难以切片，因此此选项允许您输出张量列表。默认为 True。

  • convert_to_sparse_tensor (bool, optional) – 输出格式是否应为稀疏（COO）张量。默认为 True。

  • save_to_cpu (bool, optional) – 输出是应移动到 CPU 还是保留在计算它的设备上。默认为 False。

  • device (Union[str, List[str], None], optional) –

    用于计算的设备。可以是

    • 单个设备字符串（例如，“cuda:0”、“cpu”）用于单进程编码

    • 设备字符串列表（例如，[“cuda:0”, “cuda:1”], [“cpu”, “cpu”, “cpu”, “cpu”]）用于在多个进程之间分发编码

    • None 以自动检测可用于单进程编码的设备

    如果提供了列表，则使用多进程编码。默认为 None。

  • max_active_dims (int, optional) – 模型输出中的最大活动（非零）维度数。None 表示将使用模型配置中的值。默认为 None。如果模型配置中为 None，则表示活动维度数没有限制，如果您的模型尚未（或尚未）微调到高稀疏度，可能会导致速度缓慢或内存占用过大。

  • pool (Dict[Literal["input", "output", "processes"], Any], optional) – 由start_multi_process_pool()创建的用于多进程编码的池。如果提供，编码将被分发到多个进程。这对于大型数据集和多个 GPU 可用时推荐使用。默认为 None。

  • chunk_size (int, optional) – 多进程编码的块大小。仅在 multiprocessing 时使用，即当 pool 不为 None 或 device 是列表时。如果为 None，则计算一个合理的默认值。默认为 None。

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 torch 稀疏张量。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 save_to_cpu 为 True，则嵌入将移至 CPU。

返回类型:

Union[List[Tensor], ndarray, Tensor]

示例

from sentence_transformers import SparseEncoder

# Load a pre-trained SparseEncoder model
model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

# Encode some texts
queries = [
    "What are the effects of climate change?",
    "History of artificial intelligence",
    "Technical specifications product XYZ",
]
embeddings = model.encode_query(queries)
print(embeddings.shape)
# (3, 30522)
eval() → T

将模块设置为评估模式。

这仅对某些模块有影响。有关它们在训练/评估模式下的行为细节，请参阅特定模块的文档，例如 Dropout、BatchNorm 等。

这相当于 self.train(False)。

有关 .eval() 和其他可能混淆的类似机制的比较，请参阅 局部禁用梯度计算。

返回:

self

返回类型:

Module (模块)

evaluate(evaluator: SentenceEvaluator, output_path: str | None = None) → dict[str, float] | float[source]

根据评估器评估模型

参数:

  • evaluator (SentenceEvaluator) – 用于评估模型的评估器。

  • output_path (str, optional) – 评估器可以写入结果的路径。默认为 None。

返回:

评估结果。

float() → T

将所有浮点参数和缓冲区转换为 float 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

get_adapter_state_dict(*args, **kwargs) → dict[source]

如果您不熟悉适配器和 PEFT 方法，我们邀请您阅读 PEFT 官方文档中的更多内容：https://hugging-face.cn/docs/peft

获取仅包含指定适配器适配器名称的权重张量的适配器状态字典。如果未传入适配器名称，则使用活动适配器。

参数:
get_backend() → Literal['torch', 'onnx', 'openvino'][source]

返回用于推理的后端，可以是“torch”、“onnx”或“openvino”之一。

返回:

用于推理的后端。

返回类型:

str

get_max_seq_length() → int | None[source]

返回模型可接受的最大序列长度。更长的输入将被截断。

返回:

模型可接受的最大序列长度，如果未知则为 None。

返回类型:

Optional[int]

get_model_kwargs() → list[str][source]

获取此模型特定于 encode、encode_query 或 encode_document 方法的关键字参数。

示例

>>> from sentence_transformers import SentenceTransformer, SparseEncoder
>>> SentenceTransformer("all-MiniLM-L6-v2").get_model_kwargs()
[]
>>> SentenceTransformer("jinaai/jina-embeddings-v4", trust_remote_code=True).get_model_kwargs()
['task', 'truncate_dim']
>>> SparseEncoder("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill").get_model_kwargs()
['task']
返回:

用于前向传递的关键字参数列表。

返回类型:

list[str]

get_sentence_embedding_dimension() → int | None[source]

返回 SparseEncoder.encode 输出中的维度数。我们覆盖了该函数，但未根据截断维度进行更新，因为对于稀疏模型，输出的维度与输入的维度相同，只有活动维度的数量会发生变化。

返回:

encode 输出中的维度数。如果未知，则为 None。

返回类型:

Optional[int]

half() → T

将所有浮点参数和缓冲区转换为 half 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

intersection(embeddings_1: Tensor, embeddings_2: Tensor) → Tensor[source]

计算两个稀疏嵌入的交集。

参数:

  • embeddings_1 (torch.Tensor) – 第一个嵌入张量，（vocab）。

  • embeddings_2 (torch.Tensor) – 第二个嵌入张量，（vocab）或（batch_size, vocab）。

返回:

两个嵌入的交集。

返回类型:

torch.Tensor

load_adapter(*args, **kwargs) → None[source]

从文件或远程 Hub 文件夹加载适配器权重。“如果您不熟悉适配器和 PEFT 方法，我们邀请您阅读 PEFT 官方文档中的更多内容：https://hugging-face.cn/docs/peft

要求 peft 作为后端来加载适配器权重，并且底层模型必须与 PEFT 兼容。

参数:
property max_seq_length: int

返回模型可接受的最大输入序列长度。更长的输入将被截断。

返回:

最大输入序列长度。

返回类型:

int

示例

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-cocondenser-ensembledistil")
print(model.max_seq_length)
# => 512
model_card_data_class[source]

alias of SparseEncoderModelCardData

push_to_hub(repo_id: str, token: str | None = None, private: bool | None = None, safe_serialization: bool = True, commit_message: str | None = None, local_model_path: str | None = None, exist_ok: bool = False, replace_model_card: bool = False, train_datasets: list[str] | None = None, revision: str | None = None, create_pr: bool = False) → str[source]

将此SparseEncoder的所有元素上传到新的HuggingFace Hub存储库。

参数:

  • repo_id (str) – 在Hub上存储你的模型的存储库名称，包括用户或组织。

  • token (str, optional) – 认证令牌（参见 https://hugging-face.cn/settings/token）

  • private (bool, optional) – 如果设置为True，则托管私有模型

  • safe_serialization (bool, optional) – 如果为True，则使用safetensors保存模型。如果为False，则以传统的PyTorch方式保存模型

  • commit_message (str, optional) – 推送时使用的提交消息。

  • local_model_path (str, optional) – 模型在本地的路径。如果设置了，将上传此文件路径。否则，将上传当前模型

  • exist_ok (bool, optional) – 如果为True，则保存到现有存储库是OK的。如果为False，则只能保存到新存储库

  • replace_model_card (bool, optional) – 如果为True，则用自动创建的模型卡替换Hub中已有的模型卡

  • train_datasets (List[str], optional) – 用于训练模型的数据集。如果设置了，这些数据集将被添加到Hub的模型卡中。

  • revision (str, optional) – 要将上传的文件推送到哪个分支

  • create_pr (bool, optional) – 如果为True，则创建拉取请求而不是直接推送到主分支

返回:

你的模型在Hugging Face Hub上的存储库中的提交URL。

返回类型:

str

save_pretrained(path: str, model_name: str | None = None, create_model_card: bool = True, train_datasets: list[str] | None = None, safe_serialization: bool = True) → None[source]

将模型及其配置文件保存到目录中，以便可以使用 SparseEncoder(path) 再次加载。

参数:

  • path (str) – 模型将要保存到的磁盘路径。

  • model_name (str, optional) – 可选的模型名称。

  • create_model_card (bool, optional) – 如果为True，则创建一个包含此模型基本信息的README.md文件。

  • train_datasets (List[str], optional) – 用于训练模型的名称列表。

  • safe_serialization (bool, optional) – 如果为True，则使用safetensors保存模型。如果为False，则以传统的（但不安全的）PyTorch方式保存模型。

set_adapter(*args, **kwargs) → None[source]

通过强制模型使用特定适配器并禁用其他适配器来设置特定的适配器。

参数:
set_pooling_include_prompt(include_prompt: bool) → None[source]

设置模型中池化层的 include_prompt 属性，如果存在的话。

这对于INSTRUCTOR模型非常有用，因为对于这些模型，不应将提示包含在池化策略中。

参数:

include_prompt (bool) – 是否在池化层中包含提示。

返回:

无

property similarity: Callable[[Tensor | ndarray[Any, dtype[float32]], Tensor | ndarray[Any, dtype[float32]]], Tensor]

计算两个嵌入集合之间的相似度。输出将是一个矩阵，其中包含第一个参数中的所有嵌入与第二个参数中的所有嵌入之间的相似度得分。这与 similarity_pairwise 不同，后者计算每个嵌入对之间的相似度。此方法仅支持 fp32 精度的嵌入，不支持量化嵌入。

参数:

  • embeddings1 (Union[Tensor, ndarray]) – [num_embeddings_1, embedding_dim] 或 [embedding_dim] 形状的 numpy 数组或 torch 张量。

  • embeddings2 (Union[Tensor, ndarray]) – [num_embeddings_2, embedding_dim] 或 [embedding_dim] 形状的 numpy 数组或 torch 张量。

返回:

一个 [num_embeddings_1, num_embeddings_2] 形状的 torch 张量，包含相似度得分。

返回类型:

张量

示例

>>> model = SparseEncoder("naver/splade-cocondenser-ensembledistil")
>>> sentences = [
...     "The weather is so nice!",
...     "It's so sunny outside.",
...     "He's driving to the movie theater.",
...     "She's going to the cinema.",
... ]
>>> embeddings = model.encode(sentences, normalize_embeddings=True)
>>> model.similarity(embeddings, embeddings)
tensor([[   30.953,    12.871,     0.000,     0.011],
        [   12.871,    27.505,     0.580,     0.578],
        [    0.000,     0.580,    36.068,    15.301],
        [    0.011,     0.578,    15.301,    39.466]])
>>> model.similarity_fn_name
"dot"
>>> model.similarity_fn_name = "cosine"
>>> model.similarity(embeddings, embeddings)
tensor([[    1.000,     0.441,     0.000,     0.000],
        [    0.441,     1.000,     0.018,     0.018],
        [    0.000,     0.018,     1.000,     0.406],
        [    0.000,     0.018,     0.406,     1.000]])
property similarity_fn_name: Literal['cosine', 'dot', 'euclidean', 'manhattan']

返回 SparseEncoder.similarity() 和 SparseEncoder.similarity_pairwise() 使用的相似度函数的名称。

返回:

相似度函数的名称。如果未设置，则可能为 None，在这种情况下，它将在第一次调用时

默认为“cosine”。

返回类型:

Optional[str]

示例

>>> model = SparseEncoder("naver/splade-cocondenser-ensembledistil")
>>> model.similarity_fn_name
'dot'
property similarity_pairwise: Callable[[Tensor | ndarray[Any, dtype[float32]], Tensor | ndarray[Any, dtype[float32]]], Tensor]

计算两个嵌入集合之间的相似度。输出将是一个向量，其中包含每个嵌入对的相似度得分。此方法仅支持 fp32 精度的嵌入，不支持量化嵌入。

参数:

  • embeddings1 (Union[Tensor, ndarray]) – [num_embeddings, embedding_dim] 或 [embedding_dim] 形状的 numpy 数组或 torch 张量。

  • embeddings2 (Union[Tensor, ndarray]) – [num_embeddings, embedding_dim] 或 [embedding_dim] 形状的 numpy 数组或 torch 张量。

返回:

一个 [num_embeddings] 形状的 torch 张量，包含成对相似度得分。

返回类型:

张量

示例

>>> model = SparseEncoder("naver/splade-cocondenser-ensembledistil")
>>> sentences = [
...     "The weather is so nice!",
...     "It's so sunny outside.",
...     "He's driving to the movie theater.",
...     "She's going to the cinema.",
... ]
>>> embeddings = model.encode(sentences, convert_to_sparse_tensor=False)
>>> model.similarity_pairwise(embeddings[::2], embeddings[1::2])
tensor([12.871, 15.301])
>>> model.similarity_fn_name
"dot"
>>> model.similarity_fn_name = "cosine"
>>> model.similarity_pairwise(embeddings[::2], embeddings[1::2])
tensor([0.441, 0.406])
smart_batching_collate(batch: list[InputExample]) → tuple[list[dict[str, Tensor]], Tensor][source]

将 SmartBatchingDataset 的批次转换为模型张量的批次。这里，batch 是 InputExample 实例的列表：[InputExample(…), …]

参数:

batch – 来自 SmartBatchingDataset 的批次

返回:

模型张量的批次

static sparsity(embeddings: Tensor) → dict[str, float][source]

计算给定嵌入的稀疏性统计信息，包括激活维度的平均数量和平均稀疏率。

参数:

embeddings (torch.Tensor) – 要分析的嵌入。

返回:

包含平均激活维度和平均稀疏率的字典。

返回类型:

dict[str, float]

示例
from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-cocondenser-ensembledistil")
embeddings = model.encode(["The weather is so nice!", "It's so sunny outside."])
stats = model.sparsity(embeddings)
print(stats)
# => {'active_dims': 44.0, 'sparsity_ratio': 0.9985584020614624}
property splade_pooling_chunk_size: int | None

返回 SpladePooling 模块的块大小（如果存在）。此块大小沿序列长度维度（即每个块的 token 数）。如果为 None，则一次处理整个序列。使用较小的块可以减少内存使用，但可能会降低训练和推理速度。默认为 None。

返回:

块大小，如果未找到 SpladePooling 或未设置 chunk_size，则为 None。

返回类型:

Optional[int]

start_multi_process_pool(target_devices: list[str] | None = None) → dict[Literal['input', 'output', 'processes'], Any][source]

启动一个多进程池，通过 SentenceTransformer.encode_multi_process 使用多个独立进程来处理编码。

如果你想在多个 GPU 或 CPU 上进行编码，推荐使用此方法。建议每个 GPU 只启动一个进程。此方法与 encode_multi_process 和 stop_multi_process_pool 一起工作。

参数:

target_devices (List[str], optional) – PyTorch 目标设备，例如 [“cuda:0”, “cuda:1”, …], [“npu:0”, “npu:1”, …] 或 [“cpu”, “cpu”, “cpu”, “cpu”]。如果 target_devices 为 None 且 CUDA/NPU 可用，则将使用所有可用的 CUDA/NPU 设备。如果 target_devices 为 None 且 CUDA/NPU 不可用，则将使用 4 个 CPU 设备。

返回:

一个字典，包含目标进程、输入队列和输出队列。

返回类型:

Dict[str, Any]

static stop_multi_process_pool(pool: dict[Literal['input', 'output', 'processes'], Any]) → None[source]

停止所有通过 start_multi_process_pool 启动的进程。

参数:

pool (Dict[str, object]) – 包含输入队列、输出队列和进程列表的字典。

返回:

无

to(*args, **kwargs)

移动和/或转换参数和缓冲区。

可以这样调用

to(device=None, dtype=None, non_blocking=False)
to(dtype, non_blocking=False)
to(tensor, non_blocking=False)
to(memory_format=torch.channels_last)

其签名与 torch.Tensor.to() 类似，但只接受浮点数或复数 dtype。此外，此方法只会将浮点数或复数参数和缓冲区转换为 dtype（如果已给出）。整数参数和缓冲区将被移动到 device（如果已给出），但 dtype 不会改变。当 non_blocking 设置为 True 时，它会尝试与主机异步进行转换/移动（如果可能），例如将具有固定内存的 CPU 张量移动到 CUDA 设备。

请参阅下面的示例。

注意

此方法就地修改模块。

参数:

  • device (torch.device) – 此模块中参数和缓冲区的所需设备

  • dtype (torch.dtype) – 此模块中参数和缓冲区的所需浮点数或复数 dtype

  • tensor (torch.Tensor) – 张量，其 dtype 和设备是此模块中所有参数和缓冲区的所需 dtype 和设备

  • memory_format (torch.memory_format) – 此模块中 4D 参数和缓冲区的所需内存格式（仅限关键字参数）

返回:

self

返回类型:

Module (模块)

示例

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
tokenize(texts: list[str] | list[dict] | list[tuple[str, str]], **kwargs) → dict[str, Tensor][source]

对文本进行分词。

参数:

texts (Union[List[str], List[Dict], List[Tuple[str, str]]]) – 要分词的文本列表。

返回:

包含分词后文本的张量字典。常用键包括 “input_ids” 、

”attention_mask” 和 “token_type_ids”。

返回类型:

Dict[str, Tensor]

property tokenizer: Any

获取此模型使用的分词器的属性

train(mode: bool = True) → T

将模块设置为训练模式。

这仅对某些模块有影响。有关它们在训练/评估模式下的行为细节，请参阅特定模块的文档，例如 Dropout、BatchNorm 等。

参数:

mode (bool) – 是否设置训练模式 (True) 或评估模式 (False)。默认值：True。

返回:

self

返回类型:

Module (模块)

property transformers_model: PreTrainedModel | None

获取底层 transformers PreTrainedModel 实例的属性（如果存在）。请注意，一个模型可能拥有多个底层 transformers 模型，但此属性将返回它在模块层次结构中找到的第一个模型。

返回:

底层的 transformers 模型，如果未找到则为 None。

返回类型:

PreTrainedModel or None

示例

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-v3")

# You can now access the underlying transformers model
transformers_model = model.transformers_model
print(type(transformers_model))
# => <class 'transformers.models.bert.modeling_bert.BertForMaskedLM'>

SparseEncoderModelCardData

class sentence_transformers.sparse_encoder.model_card.SparseEncoderModelCardData(language: str | list[str] | None = <factory>, license: str | None = None, model_name: str | None = None, model_id: str | None = None, train_datasets: list[dict[str, str]] = <factory>, eval_datasets: list[dict[str, str]] = <factory>, task_name: str | None = None, tags: list[str] | None = <factory>, local_files_only: bool = False, generate_widget_examples: bool = True)[source]

用于存储模型卡片中数据的数据类。

参数:

  • language (Optional[Union[str, List[str]]]) – 模型语言，可以是字符串或列表，例如“en”或 [“en”, “de”, “nl”]

  • license (Optional[str]) – 模型的许可证，例如“apache-2.0”、“mit”或“cc-by-nc-sa-4.0”

  • model_name (Optional[str]) – 模型的漂亮名称，例如“SparseEncoder based on answerdotai/ModernBERT-base”。

  • model_id (Optional[str]) – 将模型推送到Hub时的模型ID，例如“tomaarsen/se-mpnet-base-ms-marco”。

  • train_datasets (List[Dict[str, str]]) – 训练数据集的名称和/或 Hugging Face 数据集 ID 列表。例如 [{“name”: “SNLI”, “id”: “stanfordnlp/snli”}, {“name”: “MultiNLI”, “id”: “nyu-mll/multi_nli”}, {“name”: “STSB”}]

  • eval_datasets (List[Dict[str, str]]) – 评估数据集的名称和/或 Hugging Face 数据集 ID 列表。例如 [{“name”: “SNLI”, “id”: “stanfordnlp/snli”}, {“id”: “mteb/stsbenchmark-sts”}]

  • task_name (str) – 模型训练的面向人类的任务的名称，例如“semantic search and sparse retrieval”。

  • tags (Optional[List[str]]) – 模型的标签列表，例如 [“sentence-transformers”, “sparse-encoder”]。

  • local_files_only (bool) – 如果为True，则不尝试在Hub上查找数据集或基础模型信息。添加注释更多操作 默认为 False。

  • generate_widget_examples (bool) – 如果为True，则从评估或训练数据集中生成小部件示例，并计算它们的相似度。默认为 True。

提示

安装 codecarbon 以自动跟踪碳排放使用情况并将其包含在您的模型卡片中。

示例

>>> model = SparseEncoder(
...     "microsoft/mpnet-base",
...     model_card_data=SparseEncoderModelCardData(
...         model_id="tomaarsen/se-mpnet-base-allnli",
...         train_datasets=[{"name": "SNLI", "id": "stanfordnlp/snli"}, {"name": "MultiNLI", "id": "nyu-mll/multi_nli"}],
...         eval_datasets=[{"name": "SNLI", "id": "stanfordnlp/snli"}, {"name": "MultiNLI", "id": "nyu-mll/multi_nli"}],
...         license="apache-2.0",
...         language="en",
...     ),
... )
pipeline_tag: str = None
task_name: str = None

SimilarityFunction

class sentence_transformers.SimilarityFunction(value)[source]

支持的相似度函数的枚举类。支持以下函数：

  • SimilarityFunction.COSINE ("cosine"): 余弦相似度

  • SimilarityFunction.DOT_PRODUCT ("dot", dot_product): 点积相似度

  • SimilarityFunction.EUCLIDEAN ("euclidean"): 欧几里得距离

  • SimilarityFunction.MANHATTAN ("manhattan"): 曼哈顿距离

static possible_values() → list[str][source]

返回 SimilarityFunction 枚举的所有可能值列表。

返回:

SimilarityFunction 枚举的所有可能值列表。

返回类型:

list

示例

>>> possible_values = SimilarityFunction.possible_values()
>>> possible_values
['cosine', 'dot', 'euclidean', 'manhattan']
static to_similarity_fn(similarity_function: str | SimilarityFunction) → Callable[[Tensor | ndarray, Tensor | ndarray], Tensor][source]

将相似度函数名称或枚举值转换为相应的相似度函数。

参数:

similarity_function (Union[str, SimilarityFunction]) – 相似度函数的名称或枚举值。

返回:

相应的相似度函数。

返回类型:

Callable[[Union[Tensor, ndarray], Union[Tensor, ndarray]], Tensor]

引发:

ValueError – 如果提供的函数不受支持。

示例

>>> similarity_fn = SimilarityFunction.to_similarity_fn("cosine")
>>> similarity_scores = similarity_fn(embeddings1, embeddings2)
>>> similarity_scores
tensor([[0.3952, 0.0554],
        [0.0992, 0.1570]])
static to_similarity_pairwise_fn(similarity_function: str | SimilarityFunction) → Callable[[Tensor | ndarray, Tensor | ndarray], Tensor][source]

将相似度函数转换为成对相似度函数。

成对相似度函数返回相似度矩阵的对角线向量，即它只计算每个 i（在输入张量的范围内）的 similarity(a[i], b[i])，而不是计算 a 和 b 之间所有对的相似度。

参数:

similarity_function (Union[str, SimilarityFunction]) – 相似度函数的名称或枚举值。

返回:

成对相似度函数。

返回类型:

Callable[[Union[Tensor, ndarray], Union[Tensor, ndarray]], Tensor]

引发:

ValueError – 如果提供的相似度函数不受支持。

示例

>>> pairwise_fn = SimilarityFunction.to_similarity_pairwise_fn("cosine")
>>> similarity_scores = pairwise_fn(embeddings1, embeddings2)
>>> similarity_scores
tensor([0.3952, 0.1570])