SentenceTransformer

SentenceTransformer

class sentence_transformers.SentenceTransformer(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, use_auth_token: bool | str | None = None, truncate_dim: 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: SentenceTransformerModelCardData | None = None, backend: Literal['torch', 'onnx', 'openvino'] = 'torch')[source]

加载或创建 SentenceTransformer 模型，该模型可用于将句子/文本映射为嵌入。

参数:

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

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

  • 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”。如果未设置，当调用 similarity 或 similarity_pairwise 且 model.similarity_fn_name 仍为 None 时，它将自动设置为“cosine”。

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

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

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

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

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

  • use_auth_token (bool 或 str, 可选) – 已弃用的参数。请使用 token 代替。

  • truncate_dim (int, 可选) – 截断句子嵌入的维度。默认为 None。

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

    要传递给 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: 如果后端是“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 tokenizer 的额外 tokenizer 配置参数。有关更多详细信息，请参阅 AutoTokenizer.from_pretrained 文档。

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

  • model_card_data (SentenceTransformerModelCardData, 可选) – 一个模型卡数据对象，包含有关模型的信息。此对象用于在保存模型时生成模型卡。如果未设置，将创建默认的模型卡数据对象。

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

示例

from sentence_transformers import SentenceTransformer

# Load a pre-trained SentenceTransformer model
model = SentenceTransformer('all-mpnet-base-v2')

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

# Get the similarity scores between all sentences
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[1.0000, 0.6817, 0.0492],
#         [0.6817, 1.0000, 0.0421],
#         [0.0492, 0.0421, 1.0000]])

初始化内部模块状态，由 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 数据类型。

注意

此方法会原地修改模块。

返回:

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, 可选) – 如果指定，所有参数都将复制到该设备。

返回:

self

返回类型:

Module

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 数据类型。

注意

此方法会原地修改模块。

返回:

self

返回类型:

Module

enable_adapters() → None[source]

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

encode(sentences: str, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: Literal['sentence_embedding', 'token_embeddings'] = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: Literal[False] = True, convert_to_tensor: bool = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → Tensor[source]
encode(sentences: str | list[str] | np.ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: Literal['sentence_embedding'] = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: Literal[True] = True, convert_to_tensor: Literal[False] = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → np.ndarray
encode(sentences: str | list[str] | np.ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: Literal['sentence_embedding'] = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: Literal[True] = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → Tensor
encode(sentences: list[str] | np.ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: Literal['sentence_embedding', 'token_embeddings'] = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → list[Tensor]
encode(sentences: list[str] | np.ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: None = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → list[dict[str, Tensor]]
encode(sentences: str, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: None = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → dict[str, Tensor]
encode(sentences: str, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: Literal['token_embeddings'] = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | list[str | torch.device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → Tensor

计算句子嵌入。

提示

如果您不确定应该使用 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。

  • output_value (Optional[Literal["sentence_embedding", "token_embeddings"]], optional) – 要返回的嵌入类型：“sentence_embedding”用于获取句子嵌入，“token_embeddings”用于获取词片（wordpiece）token嵌入，以及 None 用于获取所有输出值。默认为“sentence_embedding”。

  • precision (Literal["float32", "int8", "uint8", "binary", "ubinary"], optional) – 用于嵌入的精度。可以是“float32”、“int8”、“uint8”、“binary”或“ubinary”。所有非 float32 精度都是量化嵌入。量化嵌入尺寸更小，计算速度更快，但准确性可能较低。它们有助于减小语料库嵌入的尺寸以进行语义搜索等任务。默认为“float32”。

  • convert_to_numpy (bool, optional) – 输出是否应为 numpy 向量列表。如果为 False，则为 PyTorch 张量列表。默认为 True。

  • convert_to_tensor (bool, optional) – 输出是否应为一个大型张量。覆盖 convert_to_numpy。默认为 False。

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

    用于计算的设备。可以是

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

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

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

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

  • normalize_embeddings (bool, optional) – 是否将返回的向量归一化为长度 1。在这种情况下，可以使用更快的点积 (util.dot_score) 而不是余弦相似度。默认为 False。

  • truncate_dim (int, optional) – 截断句子嵌入到的维度。截断对于 Matryoshka 模型特别有用，即即使嵌入维度减小，也能生成有用嵌入的模型。截断后的嵌入需要更少的内存，并且检索速度更快，但请注意推理速度相同，且嵌入性能不如完整嵌入。如果为 None，则使用模型初始化中的 truncate_dim。默认为 None。

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

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

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 numpy 数组。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 convert_to_tensor 为 True，则返回一个 torch 张量。如果 self.truncate_dim <= output_dimension，则 output_dimension 为 self.truncate_dim。

返回类型:

Union[List[Tensor], ndarray, Tensor]

示例

from sentence_transformers import SentenceTransformer

# Load a pre-trained SentenceTransformer model
model = SentenceTransformer("all-mpnet-base-v2")

# 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, 768)
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, output_value: Literal['sentence_embedding', 'token_embeddings'] | None = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | list[str | device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → 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。

  • output_value (Optional[Literal["sentence_embedding", "token_embeddings"]], optional) – 要返回的嵌入类型：“sentence_embedding”用于获取句子嵌入，“token_embeddings”用于获取词片（wordpiece）token嵌入，以及 None 用于获取所有输出值。默认为“sentence_embedding”。

  • precision (Literal["float32", "int8", "uint8", "binary", "ubinary"], optional) – 用于嵌入的精度。可以是“float32”、“int8”、“uint8”、“binary”或“ubinary”。所有非 float32 精度都是量化嵌入。量化嵌入尺寸更小，计算速度更快，但准确性可能较低。它们有助于减小语料库嵌入的尺寸以进行语义搜索等任务。默认为“float32”。

  • convert_to_numpy (bool, optional) – 输出是否应为 numpy 向量列表。如果为 False，则为 PyTorch 张量列表。默认为 True。

  • convert_to_tensor (bool, optional) – 输出是否应为一个大型张量。覆盖 convert_to_numpy。默认为 False。

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

    用于计算的设备。可以是

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

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

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

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

  • normalize_embeddings (bool, optional) – 是否将返回的向量归一化为长度 1。在这种情况下，可以使用更快的点积 (util.dot_score) 而不是余弦相似度。默认为 False。

  • truncate_dim (int, optional) –

    截断句子嵌入到的维度。截断对于 Matryoshka 模型特别有用，即即使嵌入维度减小，也能生成有用嵌入的模型。截断后的嵌入需要更少的内存，并且检索速度更快，但请注意推理速度相同，且嵌入性能不如完整嵌入。如果为 None，则使用模型初始化中的 truncate_dim。默认为 None。

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

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

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 numpy 数组。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 convert_to_tensor 为 True，则返回一个 torch 张量。如果 self.truncate_dim <= output_dimension，则 output_dimension 为 self.truncate_dim。

返回类型:

Union[List[Tensor], ndarray, Tensor]

示例

from sentence_transformers import SentenceTransformer

# Load a pre-trained SentenceTransformer model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# Encode some documents
documents = [
    "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.",
]

# Using document-specific encoding
embeddings = model.encode_document(documents)
print(embeddings.shape)
# (3, 768)
encode_multi_process(sentences: list[str], pool: dict[Literal['input', 'output', 'processes'], Any], prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, chunk_size: int | None = None, show_progress_bar: bool | None = None, precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', normalize_embeddings: bool = False, truncate_dim: int | None = None) → ndarray[source]

警告

此方法已弃用。您现在可以直接使用相同的参数调用 SentenceTransformer.encode，它将自动使用提供的 pool 处理多进程编码。

使用多个进程和 GPU，通过 SentenceTransformer.encode 编码句子列表。句子被分成更小的块并发送到单独的进程，这些进程在不同的 GPU 或 CPU 上对它们进行编码。此方法仅适用于编码大量句子。

参数:

  • sentences (List[str]) – 要编码的句子列表。

  • pool (Dict[Literal["input", "output", "processes"], Any]) – 使用 SentenceTransformer.start_multi_process_pool 启动的工作池。

  • 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) – 以批处理大小编码句子。（默认值：32）

  • chunk_size (int) – 句子被分块并发送到各个进程。如果为 None，则确定一个合理的大小。默认为 None。

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

  • precision (Literal["float32", "int8", "uint8", "binary", "ubinary"]) – 用于嵌入的精度。可以是“float32”、“int8”、“uint8”、“binary”或“ubinary”。所有非 float32 精度都是量化嵌入。量化嵌入尺寸更小，计算速度更快，但准确性可能较低。它们有助于减小语料库嵌入的尺寸以进行语义搜索等任务。默认为“float32”。

  • normalize_embeddings (bool) – 是否将返回的向量归一化为长度 1。在这种情况下，可以使用更快的点积 (util.dot_score) 而不是余弦相似度。默认为 False。

  • truncate_dim (int, optional) –

    截断句子嵌入到的维度。截断对于 Matryoshka 模型特别有用，即即使嵌入维度减小，也能生成有用嵌入的模型。截断后的嵌入需要更少的内存，并且检索速度更快，但请注意推理速度相同，且嵌入性能不如完整嵌入。如果为 None，则使用模型初始化中的 truncate_dim。默认为 None。

返回:

一个形状为 [num_inputs, output_dimension] 的二维 numpy 数组。

返回类型:

np.ndarray

示例

from sentence_transformers import SentenceTransformer

def main():
    model = SentenceTransformer("all-mpnet-base-v2")
    sentences = ["The weather is so nice!", "It's so sunny outside.", "He's driving to the movie theater.", "She's going to the cinema."] * 1000

    pool = model.start_multi_process_pool()
    embeddings = model.encode_multi_process(sentences, pool)
    model.stop_multi_process_pool(pool)

    print(embeddings.shape)
    # => (4000, 768)

if __name__ == "__main__":
    main()
encode_query(sentences: str | list] | ndarray, prompt_name: str | None = None, prompt: str | None = None, batch_size: int = 32, show_progress_bar: bool | None = None, output_value: Literal['sentence_embedding', 'token_embeddings'] | None = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | list[str | device] | None = None, normalize_embeddings: bool = False, truncate_dim: int | None = None, pool: dict[Literal['input', 'output', 'processes'], Any] | None = None, chunk_size: int | None = None, **kwargs) → list[Tensor] | ndarray | Tensor | dict[str,Tensor] | list[dict[str, Tensor]][source]

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

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

  1. 如果未提供 prompt_name 或 prompt，它将使用预定义的“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], 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。

  • output_value (Optional[Literal["sentence_embedding", "token_embeddings"]], optional) – 要返回的嵌入类型：“sentence_embedding”用于获取句子嵌入，“token_embeddings”用于获取词片（wordpiece）token嵌入，以及 None 用于获取所有输出值。默认为“sentence_embedding”。

  • precision (Literal["float32", "int8", "uint8", "binary", "ubinary"], optional) – 用于嵌入的精度。可以是“float32”、“int8”、“uint8”、“binary”或“ubinary”。所有非 float32 精度都是量化嵌入。量化嵌入尺寸更小，计算速度更快，但准确性可能较低。它们有助于减小语料库嵌入的尺寸以进行语义搜索等任务。默认为“float32”。

  • convert_to_numpy (bool, optional) – 输出是否应为 numpy 向量列表。如果为 False，则为 PyTorch 张量列表。默认为 True。

  • convert_to_tensor (bool, optional) – 输出是否应为一个大型张量。覆盖 convert_to_numpy。默认为 False。

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

    用于计算的设备。可以是

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

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

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

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

  • normalize_embeddings (bool, optional) – 是否将返回的向量归一化为长度 1。在这种情况下，可以使用更快的点积 (util.dot_score) 而不是余弦相似度。默认为 False。

  • truncate_dim (int, optional) –

    截断句子嵌入到的维度。截断对于 Matryoshka 模型特别有用，即即使嵌入维度减小，也能生成有用嵌入的模型。截断后的嵌入需要更少的内存，并且检索速度更快，但请注意推理速度相同，且嵌入性能不如完整嵌入。如果为 None，则使用模型初始化中的 truncate_dim。默认为 None。

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

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

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 numpy 数组。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 convert_to_tensor 为 True，则返回一个 torch 张量。如果 self.truncate_dim <= output_dimension，则 output_dimension 为 self.truncate_dim。

返回类型:

Union[List[Tensor], ndarray, Tensor]

示例

from sentence_transformers import SentenceTransformer

# Load a pre-trained SentenceTransformer model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# Encode some queries
queries = [
    "What are the effects of climate change?",
    "History of artificial intelligence",
    "Technical specifications product XYZ",
]

# Using query-specific encoding
embeddings = model.encode_query(queries)
print(embeddings.shape)
# (3, 768)
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。

返回:

评估结果。

fit(train_objectives: ~collections.abc.Iterable[tuple[~torch.utils.data.dataloader.DataLoader, ~torch.nn.modules.module.Module]], evaluator: ~sentence_transformers.evaluation.SentenceEvaluator.SentenceEvaluator | None = None, epochs: int = 1, steps_per_epoch=None, scheduler: str = 'WarmupLinear', warmup_steps: int = 10000, optimizer_class: type[~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adamw.AdamW'>, optimizer_params: dict[str, object] = {'lr': 2e-05}, weight_decay: float = 0.01, evaluation_steps: int = 0, output_path: str | None = None, save_best_model: bool = True, max_grad_norm: float = 1, use_amp: bool = False, callback: ~typing.Callable[[float, int, int], None] | None = None, show_progress_bar: bool = True, checkpoint_path: str | None = None, checkpoint_save_steps: int = 500, checkpoint_save_total_limit: int = 0, resume_from_checkpoint: bool = False) → None[source]

Sentence Transformers v3.0 之前的已弃用训练方法，建议改用 SentenceTransformerTrainer。此方法在后台使用 SentenceTransformerTrainer，但不如 Trainer 本身灵活。

此训练方法使用 DataLoader 和损失函数列表来训练模型。每个 DataLoader 轮流采样一个批次。我们从每个 DataLoader 中采样的批次数量与最小的 DataLoader 中的批次数量相同，以确保每个数据集的训练量相等，即轮询采样。

此方法在 v3.0+ 中应产生与 v3.0 之前相同的结果，但如果您现有训练脚本遇到任何问题，则可能希望改用 SentenceTransformer.old_fit。它使用 v3.0 之前的旧训练方法。

参数:

  • train_objectives – (DataLoader, LossFunction) 元组。对于多任务学习，可传入多个。

  • evaluator – 评估器 (sentence_transformers.evaluation) 用于在训练期间在保留的开发数据上评估模型性能。它用于确定并保存到磁盘的最佳模型。

  • epochs – 训练的 epoch 数量

  • steps_per_epoch – 每个 epoch 的训练步数。如果设置为 None（默认），一个 epoch 等于 train_objectives 中 DataLoader 的大小。

  • scheduler – 学习率调度器。可用调度器：constantlr, warmupconstant, warmuplinear, warmupcosine, warmupcosinewithhardrestarts

  • warmup_steps – 行为取决于调度器。对于 WarmupLinear（默认），学习率从零增加到最大学习率。经过这么多训练步骤后，学习率线性递减至零。

  • optimizer_class – 优化器

  • optimizer_params – 优化器参数

  • weight_decay – 模型参数的权重衰减

  • evaluation_steps – 如果 > 0，在每经过一定数量的训练步数后，使用评估器评估模型

  • output_path – 模型和评估文件的存储路径

  • save_best_model – 如果为 true，最佳模型（根据评估器）将存储在 output_path。

  • max_grad_norm – 用于梯度归一化。

  • use_amp – 使用自动混合精度 (AMP)。仅适用于 PyTorch >= 1.6.0

  • callback – 在每次评估后调用的回调函数。它必须按以下顺序接受这三个参数：score、epoch、steps

  • show_progress_bar – 如果为 True，则输出 tqdm 进度条

  • checkpoint_path – 训练期间保存检查点的文件夹

  • checkpoint_save_steps – 在经过这么多步后保存一个检查点

  • checkpoint_save_total_limit – 要存储的检查点总数

  • resume_from_checkpoint – 如果为 true，则搜索检查点以继续训练。

float() → T

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

注意

此方法会原地修改模块。

返回:

self

返回类型:

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]

返回 SentenceTransformer.encode 输出中的维度数量。

返回:

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

返回类型:

Optional[int]

half() → T

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

注意

此方法会原地修改模块。

返回:

self

返回类型:

Module

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 SentenceTransformer

model = SentenceTransformer("all-mpnet-base-v2")
print(model.max_seq_length)
# => 384
model_card_data_class[source]

SentenceTransformerModelCardData 的别名

old_fit(train_objectives: ~collections.abc.Iterable[tuple[~torch.utils.data.dataloader.DataLoader, ~torch.nn.modules.module.Module]], evaluator: ~sentence_transformers.evaluation.SentenceEvaluator.SentenceEvaluator | None = None, epochs: int = 1, steps_per_epoch=None, scheduler: str = 'WarmupLinear', warmup_steps: int = 10000, optimizer_class: type[~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adamw.AdamW'>, optimizer_params: dict[str, object] = {'lr': 2e-05}, weight_decay: float = 0.01, evaluation_steps: int = 0, output_path: str | None = None, save_best_model: bool = True, max_grad_norm: float = 1, use_amp: bool = False, callback: ~typing.Callable[[float, int, int], None] | None = None, show_progress_bar: bool = True, checkpoint_path: str | None = None, checkpoint_save_steps: int = 500, checkpoint_save_total_limit: int = 0) → None[source]

这是Sentence Transformers v3.0之前的已弃用训练方法，建议改用 sentence_transformers.trainer.SentenceTransformerTrainer。仅当您在升级到v3.0+后，现有训练脚本遇到问题时才应使用此方法。

此训练方法使用 DataLoader 和损失函数列表来训练模型。每个 DataLoader 轮流采样一个批次。我们从每个 DataLoader 中采样的批次数量与最小的 DataLoader 中的批次数量相同，以确保每个数据集的训练量相等，即轮询采样。

参数:

  • train_objectives – (DataLoader, LossFunction) 元组。对于多任务学习，可传入多个。

  • evaluator – 评估器 (sentence_transformers.evaluation) 用于在训练期间评估模型在保留的开发数据上的性能。它用于确定要保存到磁盘的最佳模型。

  • epochs – 训练的 epoch 数量

  • steps_per_epoch – 每个 epoch 的训练步数。如果设置为 None（默认），一个 epoch 等于 train_objectives 中 DataLoader 的大小。

  • scheduler – 学习率调度器。可用调度器：constantlr, warmupconstant, warmuplinear, warmupcosine, warmupcosinewithhardrestarts

  • warmup_steps – 行为取决于调度器。对于 WarmupLinear（默认），学习率从零增加到最大学习率。经过这么多训练步骤后，学习率线性递减至零。

  • optimizer_class – 优化器

  • optimizer_params – 优化器参数

  • weight_decay – 模型参数的权重衰减

  • evaluation_steps – 如果 > 0，在每经过一定数量的训练步数后，使用评估器评估模型

  • output_path – 模型和评估文件的存储路径

  • save_best_model – 如果为 true，最佳模型（根据评估器）将存储在 output_path。

  • max_grad_norm – 用于梯度归一化。

  • use_amp – 使用自动混合精度 (AMP)。仅适用于 PyTorch >= 1.6.0

  • callback – 在每次评估后调用的回调函数。它必须按以下顺序接受这三个参数：score、epoch、steps

  • show_progress_bar – 如果为 True，则输出 tqdm 进度条

  • checkpoint_path – 训练期间保存检查点的文件夹

  • checkpoint_save_steps – 在经过这么多步后保存一个检查点

  • checkpoint_save_total_limit – 要存储的检查点总数

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]

将此Sentence Transformer的所有元素上传到新的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]

将模型及其配置文件保存到目录，以便可以使用 SentenceTransformer(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) – 是否在池化层中包含提示。

返回:

无

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 = SentenceTransformer("all-mpnet-base-v2")
>>> 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([[1.0000, 0.7235, 0.0290, 0.1309],
        [0.7235, 1.0000, 0.0613, 0.1129],
        [0.0290, 0.0613, 1.0000, 0.5027],
        [0.1309, 0.1129, 0.5027, 1.0000]])
>>> model.similarity_fn_name
"cosine"
>>> model.similarity_fn_name = "euclidean"
>>> model.similarity(embeddings, embeddings)
tensor([[-0.0000, -0.7437, -1.3935, -1.3184],
        [-0.7437, -0.0000, -1.3702, -1.3320],
        [-1.3935, -1.3702, -0.0000, -0.9973],
        [-1.3184, -1.3320, -0.9973, -0.0000]])
property similarity_fn_name: Literal['cosine', 'dot', 'euclidean', 'manhattan']

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

返回:

相似度函数名称。如果未设置，则可以为None，在这种情况下，它将

首次调用时默认为“cosine”。

返回类型:

Optional[str]

示例

>>> model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")
>>> 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 = SentenceTransformer("all-mpnet-base-v2")
>>> 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_pairwise(embeddings[::2], embeddings[1::2])
tensor([0.7235, 0.5027])
>>> model.similarity_fn_name
"cosine"
>>> model.similarity_fn_name = "euclidean"
>>> model.similarity_pairwise(embeddings[::2], embeddings[1::2])
tensor([-0.7437, -0.9973])
smart_batching_collate(batch: list[InputExample]) → tuple[list[dict[str, Tensor]], Tensor][source]

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

参数:

batch – 来自SmartBatchingDataset的批次

返回:

模型的张量批次

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]) – 包含输入队列、输出队列和进程列表的字典。

返回:

无

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参数和缓冲区的目标内存格式（仅关键字参数）

返回:

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 SentenceTransformer

model = SentenceTransformer("all-mpnet-base-v2")

# You can now access the underlying transformers model
transformers_model = model.transformers_model
print(type(transformers_model))
# => <class 'transformers.models.mpnet.modeling_mpnet.MPNetModel'>
truncate_sentence_embeddings(truncate_dim: int | None) → Iterator[None][source]

在此上下文中，SentenceTransformer.encode 输出在维度 truncate_dim 处截断的句子嵌入。

当您在不同应用中使用相同模型但需要不同维度时，这可能很有用。

参数:

truncate_dim (int, 可选) – 要截断句子嵌入到的维度。None 表示不进行截断。

示例

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-mpnet-base-v2")

with model.truncate_sentence_embeddings(truncate_dim=16):
    embeddings_truncated = model.encode(["hello there", "hiya"])
assert embeddings_truncated.shape[-1] == 16

SentenceTransformerModelCardData

class sentence_transformers.model_card.SentenceTransformerModelCardData(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 = 'semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more', tags: list[str] | None = <factory>, local_files_only: bool = False, generate_widget_examples: bool = True)[source]

一个用于存储模型卡中数据的dataclass。

参数:

  • 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]) – 模型的美观名称，例如“基于microsoft/mpnet-base的SentenceTransformer”。

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

  • 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”, “sentence-similarity”, “feature-extraction”]。

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

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

提示

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

示例

>>> model = SentenceTransformer(
...     "microsoft/mpnet-base",
...     model_card_data=SentenceTransformerModelCardData(
...         model_id="tomaarsen/sbert-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",
...     ),
... )

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枚举的可能值列表。

返回类型:

列表

示例

>>> 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的相似度(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])