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, 可选) – 要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID，用于 Hugging Face 上的存储模型。

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

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

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

  • truncate_dim (int, 可选) – 句子嵌入要截断到的维度。None 不进行截断。截断仅在调用 SentenceTransformer.encode() 时在推理期间适用。

  • 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）中的任何一个。默认情况下，如果可用，SDPA 将用于 torch>=2.1.1。否则，默认设置为手动 “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 (SentenceTransformerModelCardData, 可选) – 包含有关模型信息的模型卡数据对象。这用于在保存模型时生成模型卡。如果未设置，则会创建默认模型卡数据对象。

  • backend (str) – 用于推理的后端。可以是“torch”（默认）、“onnx”或“openvino”之一。有关不同后端的基准测试信息，请参阅 https://sbert.net.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://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() 编译此模块的 forward。

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

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

cpu() → T

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

注意

此方法会就地修改模块。

返回:

self

返回类型:

Module

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

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

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

注意

此方法会就地修改模块。

参数:

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

返回:

self

返回类型:

Module

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 | None = None, normalize_embeddings: bool = False, **kwargs) → Tensor[source]
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, 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 | None = None, normalize_embeddings: bool = False, **kwargs) → ndarray
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, 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 | None = None, normalize_embeddings: bool = False, **kwargs) → Tensor
encode(sentences: 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'] = 'sentence_embedding', precision: Literal['float32', 'int8', 'uint8', 'binary', 'ubinary'] = 'float32', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: str | None = None, normalize_embeddings: bool = False, **kwargs) → list[Tensor]
encode(sentences: list[str] | 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 | None = None, normalize_embeddings: bool = False, **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 | None = None, normalize_embeddings: bool = False, **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 | None = None, normalize_embeddings: bool = False, **kwargs) → Tensor

计算句子嵌入。

参数:

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

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

  • prompt (Optional[str], optional) – 用于编码的 prompt。例如，如果 prompt 是 “query: “，则句子 “What is the capital of France?” 将被编码为 “query: What is the capital of France?”，因为句子被附加到 prompt。如果设置了 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 词元嵌入，以及 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 (str, optional) – 用于计算的 torch.device。默认为 None。

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

返回:

默认情况下，返回形状为 [num_inputs, output_dimension] 的 2d numpy 数组。如果仅提供一个字符串输入，则输出为形状为 [output_dimension] 的 1d 数组。如果 convert_to_tensor，则返回 torch Tensor。如果 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_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) → ndarray[source]

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

参数:

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

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

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

  • prompt (Optional[str], optional) – 用于编码的 prompt。例如，如果 prompt 是 “query: “，则句子 “What is the capital of France?” 将被编码为 “query: What is the capital of France?”，因为句子被附加到 prompt。如果设置了 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。

返回:

形状为 [num_inputs, output_dimension] 的 2D 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()
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 和 Loss 函数列表来训练模型。每个 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（默认值），学习率从 0 增加到最大学习率。在这些训练步骤之后，学习率线性降回零。

  • 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://hugging-face.cn/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://hugging-face.cn/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
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 和 Loss 函数列表来训练模型。每个 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（默认值），学习率从 0 增加到最大学习率。在这些训练步骤之后，学习率线性降回零。

  • 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://hugging-face.cn/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 模型很有用，因为对于这些模型，应从池化策略中排除 prompt。

参数:

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

返回:

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

返回:

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

返回:

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

获取此模型使用的 tokenizer 的属性

train(mode: bool = True) → T

将模块设置为训练模式。

这仅对某些模块有效。有关其在训练/评估模式下的行为的详细信息，请参阅特定模块的文档（如果它们受到影响），例如 Dropout、BatchNorm 等。

参数:

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

返回:

self

返回类型:

Module

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>, generate_widget_examples: ~typing.Literal['deprecated'] = 'deprecated')[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]) – 模型的漂亮名称，例如 “SentenceTransformer based on microsoft/mpnet-base”。

  • 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) – 模型训练的人类可读任务名称，例如 “semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more”。

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

提示

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

返回类型:

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]

Raises:

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]

Raises:

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

示例

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