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, 可选) – 如果是磁盘上的文件路径,则从该路径加载模型。如果不是路径,则首先尝试下载预训练的 SparseEncoder 模型。如果失败,则尝试从 Hugging Face Hub 使用该名称构建模型。

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

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

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

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

  • similarity_fn_name (str 或 SimilarityFunction, 可选) – 要使用的相似度函数的名称。有效选项包括“cosine”、“dot”、“euclidean”和“manhattan”。如果未设置,则在调用 similaritysimilarity_pairwisemodel.similarity_fn_name 仍为 None 时,会自动设置为“cosine”。

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

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

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

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

  • token (bool 或 str, 可选) – 用于下载私有模型的 Hugging Face 认证令牌。

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

  • model_kwargs (Dict[str, Any], 可选) –

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

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

      1. torch.float16, torch.bfloat16torch.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: 如果后端是“onnx”,这是用于推理的提供者,例如“CPUExecutionProvider”、“CUDAExecutionProvider”等。有关所有 ONNX 执行提供者,请参阅 https://runtime.onnx.org.cn/docs/execution-providers/

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

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

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

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

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

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

  • backend (str) – 用于推理的后端。可以是“torch”(默认)、“onnx”或“openvino”之一。有关不同后端上的基准测试信息,请参阅 https://sbert.hugging-face.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]])

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

active_adapters() list[str][source]

如果您不熟悉适配器和 PEFT 方法,我们邀请您在 PEFT 官方文档中阅读更多相关信息:https://huggingface.co/docs/peft

获取模型的当前活动适配器。在多适配器推理(组合多个适配器进行推理)的情况下,返回所有活动适配器列表,以便用户可以相应地处理它们。

对于以前的 PEFT 版本(不支持多适配器推理),module.active_adapter 将返回一个字符串。

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

为当前模型添加一个新的适配器以进行训练。如果未传递适配器名称,则会为适配器分配一个默认名称,以遵循 PEFT 库的约定(在 PEFT 中,我们使用“default”作为默认适配器名称)。

需要 peft 作为后端来加载适配器权重,并且基础模型要与 PEFT 兼容。

参数:
bfloat16() T

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

注意

此方法会原地修改模块。

返回:

自身

返回类型:

Module

compile(*args, **kwargs)

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

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

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

cpu() T

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

注意

此方法会原地修改模块。

返回:

自身

返回类型:

Module

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

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

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

注意

此方法会原地修改模块。

参数:

device (int, 可选) – 如果指定,所有参数都将复制到该设备

返回:

自身

返回类型:

Module

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

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

参数:
  • embeddings (torch.Tensor) – 稀疏嵌入张量 (batch, vocab) 或 (vocab)。

  • top_k (int, 可选) – 每个样本要返回的前 K 个 token 数量。如果为 None,则返回所有非零 token。

返回:

每个嵌入的 (token, weight) 元组列表。如果是批量输入,则返回元组列表的列表。

返回类型:

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

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

如果您不熟悉适配器和 PEFT 方法,我们邀请您在 PEFT 官方文档中阅读更多相关信息:https://huggingface.co/docs/peft

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

参数:
property device: device

从模块获取 torch.device,假设整个模块只有一个设备。如果不存在 PyTorch 参数,则回退到 CPU。

disable_adapters() None[source]

禁用连接到模型的所有适配器。这使得推理仅使用基础模型。

double() T

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

注意

此方法会原地修改模块。

返回:

自身

返回类型:

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], 可选) – 用于编码的提示名称。必须是 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], 可选) – 用于编码的提示。例如,如果提示是“query: ”,那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”,因为句子被附加到提示中。如果 prompt 已设置,则 prompt_name 将被忽略。默认为 None。

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

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

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

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

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

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

    用于计算的设备。可以是

    • 单个设备字符串(例如,“cuda:0”、“cpu”)用于单进程编码

    • 设备字符串列表(例如,[“cuda:0”, “cuda:1”],[“cpu”, “cpu”, “cpu”, “cpu”])用于在多个进程中分发编码

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

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

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

  • pool (Dict[Literal[“input”, “output”, “processes”], Any], 可选) – 由 start_multi_process_pool() 创建的用于多进程编码的池。如果提供,编码将在多个进程中分发。对于大型数据集和可用多个 GPU 的情况,建议使用此选项。默认为 None。

  • chunk_size (int, 可选) – 多进程编码的块大小。仅在多进程处理时使用,即当 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_nameprompt,它将使用预定义的“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], 可选) – 用于编码的提示名称。必须是 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], 可选) – 用于编码的提示。例如,如果提示是“query: ”,那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”,因为句子被附加到提示中。如果 prompt 已设置,则 prompt_name 将被忽略。默认为 None。

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

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

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

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

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

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

    用于计算的设备。可以是

    • 单个设备字符串(例如,“cuda:0”、“cpu”)用于单进程编码

    • 设备字符串列表(例如,[“cuda:0”, “cuda:1”],[“cpu”, “cpu”, “cpu”, “cpu”])用于在多个进程中分发编码

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

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

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

  • pool (Dict[Literal[“input”, “output”, “processes”], Any], 可选) – 由 start_multi_process_pool() 创建的用于多进程编码的池。如果提供,编码将在多个进程中分发。对于大型数据集和可用多个 GPU 的情况,建议使用此选项。默认为 None。

  • chunk_size (int, 可选) – 多进程编码的块大小。仅在多进程处理时使用,即当 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_nameprompt,它将使用预定义的“query”提示(如果模型 prompts 字典中可用)。

  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], 可选) – 用于编码的提示名称。必须是 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], 可选) – 用于编码的提示。例如,如果提示是“query: ”,那么句子“What is the capital of France?”将被编码为“query: What is the capital of France?”,因为句子被附加到提示中。如果 prompt 已设置,则 prompt_name 将被忽略。默认为 None。

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

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

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

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

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

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

    用于计算的设备。可以是

    • 单个设备字符串(例如,“cuda:0”、“cpu”)用于单进程编码

    • 设备字符串列表(例如,[“cuda:0”, “cuda:1”],[“cpu”, “cpu”, “cpu”, “cpu”])用于在多个进程中分发编码

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

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

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

  • pool (Dict[Literal[“input”, “output”, “processes”], Any], 可选) – 由 start_multi_process_pool() 创建的用于多进程编码的池。如果提供,编码将在多个进程中分发。对于大型数据集和可用多个 GPU 的情况,建议使用此选项。默认为 None。

  • chunk_size (int, 可选) – 多进程编码的块大小。仅在多进程处理时使用,即当 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

将模块设置为评估模式。

这只对某些模块有效。有关特定模块在训练/评估模式下的行为详情,如果它们受到影响,例如 DropoutBatchNorm 等,请参阅其文档。

这等同于 self.train(False)

有关 .eval() 与可能与它混淆的几种类似机制之间的比较,请参阅 局部禁用梯度计算

返回:

自身

返回类型:

Module

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

根据评估器评估模型

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

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

返回:

评估结果。

float() T

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

注意

此方法会原地修改模块。

返回:

自身

返回类型:

Module

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

如果您不熟悉适配器和 PEFT 方法,我们邀请您在 PEFT 官方文档中阅读更多相关信息:https://huggingface.co/docs/peft

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

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

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

返回:

用于推理的后端。

返回类型:

str

get_max_seq_length() int | None[source]

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

返回:

模型接受的最大序列长度,如果未定义则为 None。

返回类型:

Optional[int]

get_sentence_embedding_dimension() int | None[source]

返回 SparseEncoder.encode 输出中的维度数量。我们覆盖了此函数,没有更新关于截断维度,因为对于稀疏模型,输出的维度是相同的,只改变活动维度的数量。

返回:

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

返回类型:

Optional[int]

half() T

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

注意

此方法会原地修改模块。

返回:

自身

返回类型:

Module

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

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

参数:
  • embeddings_1 (torch.Tensor) – 第一个嵌入张量,(词汇表)。

  • embeddings_2 (torch.Tensor) – 第二个嵌入张量,(词汇表) 或 (batch_size, 词汇表)。

返回:

两个嵌入的交集。

返回类型:

torch.Tensor

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

从文件或远程 Hub 文件夹加载适配器权重。” 如果您不熟悉适配器和 PEFT 方法,我们邀请您阅读 PEFT 官方文档以了解更多信息:https://huggingface.co/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]

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]

将此稀疏编码器的所有元素上传到新的 HuggingFace Hub 仓库。

参数:
  • repo_id (str) – 模型在 Hub 中的仓库名称,包括用户或组织。

  • token (str, 可选) – 身份验证令牌(参见 https://huggingface.co/settings/token

  • private (bool, 可选) – 设置为 true,用于托管私有模型

  • safe_serialization (bool, 可选) – 如果为 true,则使用 safetensors 保存模型。如果为 false,则以传统的 PyTorch 方式保存模型。

  • commit_message (str, 可选) – 推送时提交的消息。

  • local_model_path (str, 可选) – 模型本地路径。如果设置,此文件路径将被上传。否则,当前模型将被上传。

  • exist_ok (bool, 可选) – 如果为 true,则可以保存到现有仓库。如果为 false,则只能保存到新仓库。

  • replace_model_card (bool, 可选) – 如果为 true,则用自动创建的模型卡替换 Hub 中现有的模型卡。

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

  • revision (str, 可选) – 上传文件推送到的分支。

  • create_pr (bool, 可选) – 如果为 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, 可选) – 可选模型名称。

  • create_model_card (bool, 可选) – 如果为 True,则创建包含此模型基本信息的 README.md。

  • train_datasets (List[str], 可选) – 可选列表,包含用于训练模型的数据集名称。

  • safe_serialization (bool, 可选) – 如果为 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) – 是否将提示包含在池化层中。

返回:

None

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 张量,包含相似度分数。

返回类型:

Tensor

示例

>>> 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 张量,包含成对相似度分数。

返回类型:

Tensor

示例

>>> 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 中的批次转换为模型张量批次。此处,批次是 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 模块的块大小(如果存在)。此块大小沿序列长度维度(即,每个块的令牌数)。如果为 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], 可选) – 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]) – 包含输入队列、输出队列和进程列表的字典。

返回:

None

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 时,它会尝试在可能的情况下异步转换/移动(相对于主机),例如,将具有固定内存的 CPU 张量移动到 CUDA 设备。

请参阅下文了解示例。

注意

此方法会原地修改模块。

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

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

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

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

返回:

自身

返回类型:

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

将模块设置为训练模式。

这只对某些模块有效。有关特定模块在训练/评估模式下的行为详情,如果它们受到影响,例如 DropoutBatchNorm 等,请参阅其文档。

参数:

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

返回:

自身

返回类型:

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]) – 模型的友好名称,例如 “基于 answerdotai/ModernBERT-base 的 SparseEncoder”。

  • 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) – 模型训练所针对的易读任务,例如“语义搜索和稀疏检索”。

  • 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])