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)[源代码][源代码]

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

参数:

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

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

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

  • prompts (Dict[str, str], optional) -- 一个包含模型提示的字典。键是提示名称，值是提示文本。提示文本将附加在任何待编码文本之前。例如：{"query": "query: ", "passage": "passage: "} 或 {"clustering": "根据标题识别主要类别"}。

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

  • similarity_fn_name (str or SimilarityFunction, optional) -- 要使用的相似度函数的名称。有效选项为 "cosine"、"dot"、"euclidean" 和 "manhattan"。如果未设置，则在调用 similarity 或 similarity_pairwise 时，如果 model.similarity_fn_name 仍为 None，则会自动设置为 "cosine"。

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

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

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

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

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

  • use_auth_token (bool or str, optional) -- 已弃用的参数。请改用 token。

  • truncate_dim (int, optional) -- 要截断句子嵌入的维度。None 表示不进行截断。截断仅在调用 :meth:SentenceTransformer.encode 进行推理时适用。

  • model_kwargs (Dict[str, Any], optional) -- 要传递给 Hugging Face Transformers 模型的其他模型配置参数。特别有用的选项包括：

  • tokenizer_kwargs (Dict[str, Any], optional) -- 要传递给 Hugging Face Transformers 分词器的额外分词器配置参数。更多详情请参阅 AutoTokenizer.from_pretrained <https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained>_ 文档。

  • config_kwargs (Dict[str, Any], optional) -- 要传递给 Hugging Face Transformers 配置的附加模型配置参数。更多详情请参阅 AutoConfig.from_pretrained <https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoConfig.from_pretrained>_ 文档。

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

示例

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 共享。

bfloat16() → T

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

备注

此方法就地修改模块。

返回:

自己

返回类型:

Module

compile(*args, **kwargs)

使用 :func:torch.compile 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 :func:torch.compile。

有关此函数的参数详情，请参阅 :func:torch.compile。

cpu() → T

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

备注

此方法就地修改模块。

返回:

自己

返回类型:

Module

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

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

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

备注

此方法就地修改模块。

参数:

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

返回:

自己

返回类型:

Module

property device: device

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

double() → T

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

备注

此方法就地修改模块。

返回:

自己

返回类型:

Module

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

计算句子嵌入。

参数:

  • 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” 获取词片标记嵌入，以及 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) -- 用于计算的 :class:torch.device 。默认为 None。

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

返回:

默认情况下，返回一个形状为 [num_inputs, output_dimension] 的二维 numpy 数组。如果只提供一个字符串输入，则输出是一个形状为 [output_dimension] 的一维数组。如果 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[源代码][源代码]

通过 :meth:SentenceTransformer.encode <sentence_transformers.SentenceTransformer.encode> 使用多个进程和GPU对句子列表进行编码。句子被分块成较小的包并发送给各个进程，这些进程在不同的GPU或CPU上对它们进行编码。此方法仅适用于编码大量句子。

参数:

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

  • pool (Dict[Literal["input", "output", "processes"], Any]) -- 使用 :meth:SentenceTransformer.start_multi_process_pool <sentence_transformers.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。

返回:

一个形状为 [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

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 :class:Dropout, :class:BatchNorm 等。

这等同于 :meth:self.train(False) <torch.nn.Module.train>。

请参阅 :ref:locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回:

自己

返回类型:

Module

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

基于评估者评估模型

参数:

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

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

返回:

评估结果。

fit(train_objectives: ~typing.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[源代码]

在 Sentence Transformers v3.0 之前的弃用训练方法，建议使用 :class:~sentence_transformers.trainer.SentenceTransformerTrainer 替代。此方法在后台使用 :class:~sentence_transformers.trainer.SentenceTransformerTrainer，但提供的灵活性不如 Trainer 本身。

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

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

参数:

  • train_objectives -- 元组 (DataLoader, LossFunction)。传递多个以进行多任务学习

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

  • epochs -- 训练的轮数

  • 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 -- 如果为真，最佳模型（根据评估器）将存储在 output_path 中

  • max_grad_norm -- 用于梯度归一化。

  • use_amp -- 使用自动混合精度（AMP）。仅适用于 Pytorch >= 1.6.0。

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

  • show_progress_bar -- 如果为真，输出一个 tqdm 进度条

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

  • checkpoint_save_steps -- 将在这么多步之后保存一个检查点

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

float() → T

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

备注

此方法就地修改模块。

返回:

自己

返回类型:

Module

get_max_seq_length() → int | None[源代码][源代码]

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

返回:

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

返回类型:

Optional[int]

get_sentence_embedding_dimension() → int | None[源代码][源代码]

返回 :meth:SentenceTransformer.encode <sentence_transformers.SentenceTransformer.encode> 输出的维度数。

返回:

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

返回类型:

Optional[int]

half() → T

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

备注

此方法就地修改模块。

返回:

自己

返回类型:

Module

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: ~typing.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[源代码]

在 Sentence Transformers v3.0 之前使用的已弃用训练方法，建议使用 :class:sentence_transformers.trainer.SentenceTransformerTrainer 代替。只有在升级到 v3.0+ 后遇到现有训练脚本问题时，才应使用此方法。

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

参数:

  • train_objectives -- 元组 (DataLoader, LossFunction)。传递多个以进行多任务学习

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

  • epochs -- 训练的轮数

  • 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 -- 如果为真，最佳模型（根据评估器）将存储在 output_path 中

  • max_grad_norm -- 用于梯度归一化。

  • use_amp -- 使用自动混合精度（AMP）。仅适用于 Pytorch >= 1.6.0。

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

  • show_progress_bar -- 如果为真，输出一个 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 = 'Add new SentenceTransformer model.', 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) → str[源代码][源代码]

将此 Sentence Transformer 的所有元素上传到一个新的 HuggingFace Hub 仓库。

参数:

  • repo_id (str) -- Hub 中模型的仓库名称，包括用户或组织。

  • token (str, optional) -- 一个认证令牌（参见 https://huggingface.co/settings/token）

  • private (bool, optional) -- 设置为 true，用于托管私有模型

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

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

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

  • exist_ok (bool, optional) -- 如果为真，保存到现有仓库是可以的。如果为假，只能保存到新仓库。

  • replace_model_card (bool, optional) -- 如果为真，则用自动创建的模型卡替换hub中的现有模型卡

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

  • revision (str, optional) -- 分支以推送上传的文件到

返回:

你的模型在 Hugging Face Hub 仓库中的提交的 URL。

返回类型:

str

save(path: str, model_name: str | None = None, create_model_card: bool = True, train_datasets: list[str] | None = None, safe_serialization: bool = True) → None[源代码][源代码]

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

参数:

  • path (str) -- 模型将被保存的磁盘路径。

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

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

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

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

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[源代码][源代码]

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

参数:

  • path (str) -- 模型将被保存的磁盘路径。

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

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

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

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

set_pooling_include_prompt(include_prompt: bool) → None[源代码][源代码]

在模型中设置池化层的 include_prompt 属性（如果有的话）。

这对于INSTRUCTOR模型很有用，因为对于这些模型，提示应从池化策略中排除。

参数:

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

返回:

无

property similarity: Callable[[Tensor | ndarray, Tensor | ndarray], Tensor]

计算两个嵌入集合之间的相似度。输出将是一个矩阵，其中包含第一个参数中的所有嵌入与第二个参数中的所有嵌入之间的相似度分数。这与 similarity_pairwise 不同，后者计算每对嵌入之间的相似度。

参数:

  • 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: str | None

返回由 :meth:SentenceTransformer.similarity 和 :meth:SentenceTransformer.similarity_pairwise 使用的相似度函数的名称。

返回:

相似度函数的名称。如果未设置，则可以为 None，在这种情况下，:meth:SentenceTransformer.similarity 和 :meth:SentenceTransformer.similarity_pairwise 的任何使用将默认使用 "cosine"。

返回类型:

Optional[str]

示例

>>> model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")
>>> model.similarity_fn_name
'dot'
property similarity_pairwise: Callable[[Tensor | ndarray, Tensor | ndarray], Tensor]

计算两个嵌入集合之间的相似度。输出将是一个包含每对嵌入之间相似度分数的向量。

参数:

  • 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][源代码]

将来自 SmartBatchingDataset 的一批数据转换为模型的一批张量。这里，batch 是一个 InputExample 实例的列表：[InputExample(...), ...]

参数:

batch -- 来自 SmartBatchingDataset 的一个批次

返回:

模型的一批张量

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

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

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

参数:

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

返回:

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

返回类型:

Dict[str, Any]

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

停止所有使用 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)

它的签名类似于 :meth:torch.Tensor.to ，但只接受浮点数或复数的 :attr:dtype 。此外，此方法只会将浮点数或复数参数和缓冲区转换为 :attr:dtype （如果给出）。整数参数和缓冲区将被移动到 :attr:device ，如果给出，但数据类型保持不变。当 :attr:non_blocking 被设置时，它会尝试尽可能地异步转换/移动，例如，将具有固定内存的CPU张量移动到CUDA设备。

请参见以下示例。

备注

此方法就地修改模块。

参数:

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

  • dtype (torch.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]]) → dict[str, Tensor][源代码][源代码]

对文本进行分词。

参数:

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

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 :class:Dropout, :class:BatchNorm 等。

参数:

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

返回:

自己

返回类型:

Module

truncate_sentence_embeddings(truncate_dim: int | None) → Iterator[None][源代码][源代码]

在这种情况下，:meth:SentenceTransformer.encode <sentence_transformers.SentenceTransformer.encode> 输出在 truncate_dim 维度上截断的句子嵌入。

当你在不同的应用中使用同一个模型，但需要不同的维度时，这可能会很有用。

参数:

truncate_dim (int, optional) -- 要截断句子嵌入的维度。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')[源代码][源代码]

一个用于存储模型卡片中使用的数据的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]]) -- 模型的标签列表，例如：["句子转换器", "句子相似性", "特征提取"]。

小技巧

安装 codecarbon <https://github.com/mlco2/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",
...     ),
... )

相似度函数

class sentence_transformers.SimilarityFunction(value)[源代码][源代码]

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

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

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

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

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

static possible_values() → list[str][源代码][源代码]

返回 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][源代码][源代码]

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

参数:

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][源代码][源代码]

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

成对相似度函数返回相似度矩阵的对角向量，即它仅计算每个 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])