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

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

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

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

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

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

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

  • similarity_fn_name (strSimilarityFunction, 可选) – 要使用的相似性函数的名称。有效选项为 “cosine”、“dot”、“euclidean” 和 “manhattan”。如果未设置,当调用 similaritysimilarity_pairwisemodel.similarity_fn_name 仍为 None 时,它会自动设置为 “cosine”。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  • tokenizer_kwargs (Dict[str, Any], 可选) – 要传递给 Hugging Face Transformers 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.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]])

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

active_adapters() list[str][源代码]

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

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

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

add_adapter(*args, **kwargs) None[源代码]

为当前模型添加一个新的适配器以用于训练。如果没有传递适配器名称,则会为适配器分配一个默认名称,以遵循 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 (模块)

delete_adapter(*args, **kwargs) None[源代码]

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

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

参数:
property device: device

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

disable_adapters() None[源代码]

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

double() T

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

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

enable_adapters() None[源代码]

启用附加到模型的适配器。模型将使用 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[源代码]
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(),最好的选择是:对于有明确查询(query)和文档/段落(document/passage)区分的信息检索任务,使用 encode_query()encode_document();对于所有其他任务,使用 encode()

请注意,encode() 是最通用的方法,可用于任何任务,包括信息检索。如果模型未经预定义提示和/或任务类型训练,那么这三种方法将返回相同的嵌入。

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

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

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

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

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

  • output_value (Optional[Literal["sentence_embedding", "token_embeddings"]], 可选) – 要返回的嵌入类型:“sentence_embedding” 获取句子嵌入,“token_embeddings” 获取词元(wordpiece token)嵌入,None 获取所有输出值。默认为 “sentence_embedding”。

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

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

  • convert_to_tensor (bool, 可选) – 输出是否应为一个大的张量。此参数会覆盖 convert_to_numpy。默认为 False。

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

    用于计算的设备。

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

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

    • None,自动检测可用的设备进行单进程编码。

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

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

  • truncate_dim (int, 可选) – 将句子嵌入截断到的维度。对于 Matryoshka 模型 来说,截断尤其有意义,即那些即使嵌入维度降低也能产生有用嵌入的模型。截断后的嵌入需要更少的内存,并且执行检索的速度更快,但请注意,推理速度同样快,并且嵌入性能比完整嵌入差。如果为 None,则使用模型初始化时的 truncate_dim。默认为 None。

  • pool (Dict[Literal["input", "output", "processes"], Any], 可选) – 由 start_multi_process_pool() 创建的用于多进程编码的池。如果提供,编码任务将分布在多个进程上。建议用于大型数据集和有多个 GPU 可用时。默认为 None。

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

返回:

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

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

此方法是 encode() 的一个专门版本,仅在以下两个方面有所不同:

  1. 如果没有提供 prompt_nameprompt,它会使用预定义的 "document" 提示(如果模型 prompts 字典中存在的话)。

  2. 它将 task 设置为 "document"。如果模型有 Router 模块,它将使用 "document" 任务类型将输入路由到相应的子模块。

提示

如果你不确定应该使用 encode()encode_query() 还是 encode_document(),最好的选择是:对于有明确查询(query)和文档/段落(document/passage)区分的信息检索任务,使用 encode_query()encode_document();对于所有其他任务,使用 encode()

请注意,encode() 是最通用的方法,可用于任何任务,包括信息检索。如果模型未经预定义提示和/或任务类型训练,那么这三种方法将返回相同的嵌入。

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

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

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

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

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

  • output_value (Optional[Literal["sentence_embedding", "token_embeddings"]], 可选) – 要返回的嵌入类型:“sentence_embedding” 获取句子嵌入,“token_embeddings” 获取词元(wordpiece token)嵌入,None 获取所有输出值。默认为 “sentence_embedding”。

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

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

  • convert_to_tensor (bool, 可选) – 输出是否应为一个大的张量。此参数会覆盖 convert_to_numpy。默认为 False。

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

    用于计算的设备。

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

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

    • None,自动检测可用的设备进行单进程编码。

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

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

  • truncate_dim (int, 可选) –

    将句子嵌入截断到的维度。对于 Matryoshka 模型 来说,截断尤其有意义,即那些即使嵌入维度降低也能产生有用嵌入的模型。截断后的嵌入需要更少的内存,并且执行检索的速度更快,但请注意,推理速度同样快,并且嵌入性能比完整嵌入差。如果为 None,则使用模型初始化时的 truncate_dim。默认为 None。

  • pool (Dict[Literal["input", "output", "processes"], Any], 可选) – 由 start_multi_process_pool() 创建的用于多进程编码的池。如果提供,编码任务将分布在多个进程上。建议用于大型数据集和有多个 GPU 可用时。默认为 None。

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

返回:

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

警告

此方法已弃用。您现在可以调用 SentenceTransformer.encode 并传入相同的参数,它会自动处理使用所提供 pool 的多进程编码。

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

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

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

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

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

  • batch_size (int) – 使用此批量大小对句子进行编码。(默认值:32)

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

  • show_progress_bar (bool, 可选) – 编码句子时是否输出进度条。默认为 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, 可选) –

    将句子嵌入截断到的维度。对于 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[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]][源代码]

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

此方法是 encode() 的一个专门版本,仅在以下两个方面有所不同:

  1. 如果没有提供 prompt_nameprompt,它会使用预定义的 "query" 提示(如果模型 prompts 字典中存在的话)。

  2. 它将 task 设置为 "query"。如果模型有 Router 模块,它将使用 "query" 任务类型将输入路由到相应的子模块。

提示

如果你不确定应该使用 encode()encode_query() 还是 encode_document(),最好的选择是:对于有明确查询(query)和文档/段落(document/passage)区分的信息检索任务,使用 encode_query()encode_document();对于所有其他任务,使用 encode()

请注意,encode() 是最通用的方法,可用于任何任务,包括信息检索。如果模型未经预定义提示和/或任务类型训练,那么这三种方法将返回相同的嵌入。

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

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

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

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

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

  • output_value (Optional[Literal["sentence_embedding", "token_embeddings"]], 可选) – 要返回的嵌入类型:“sentence_embedding” 获取句子嵌入,“token_embeddings” 获取词元(wordpiece token)嵌入,None 获取所有输出值。默认为 “sentence_embedding”。

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

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

  • convert_to_tensor (bool, 可选) – 输出是否应为一个大的张量。此参数会覆盖 convert_to_numpy。默认为 False。

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

    用于计算的设备。

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

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

    • None,自动检测可用的设备进行单进程编码。

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

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

  • truncate_dim (int, 可选) –

    将句子嵌入截断到的维度。对于 Matryoshka 模型 来说,截断尤其有意义,即那些即使嵌入维度降低也能产生有用嵌入的模型。截断后的嵌入需要更少的内存,并且执行检索的速度更快,但请注意,推理速度同样快,并且嵌入性能比完整嵌入差。如果为 None,则使用模型初始化时的 truncate_dim。默认为 None。

  • pool (Dict[Literal["input", "output", "processes"], Any], 可选) – 由 start_multi_process_pool() 创建的用于多进程编码的池。如果提供,编码任务将分布在多个进程上。建议用于大型数据集和有多个 GPU 可用时。默认为 None。

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

返回:

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

将模块设置为评估模式。

这只对某些模块有影响。有关特定模块在训练/评估模式下行为的详细信息,请参阅其文档,例如 DropoutBatchNorm 等。

这等同于 self.train(False)

有关 .eval() 与其他可能与之混淆的类似机制的比较,请参见 局部禁用梯度计算

返回:

self

返回类型:

Module (模块)

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

基于评估器评估模型

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

  • output_path (str, 可选) – 评估器可写入结果的路径。默认为 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[源代码]

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

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

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

参数:
  • train_objectives – (DataLoader, LossFunction) 的元组。传递多个以进行多任务学习。

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

  • epochs – 训练的轮数

  • steps_per_epoch – 每轮的训练步数。如果设置为 None(默认值),则一轮等于 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 – 每次评估后调用的回调函数。它必须按此顺序接受以下三个参数:scoreepochsteps

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

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

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

参数:
get_backend() Literal['torch', 'onnx', 'openvino'][源代码]

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

返回:

用于推理的后端。

返回类型:

str

get_max_seq_length() int | None[源代码]

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

返回:

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

返回类型:

Optional[int]

get_sentence_embedding_dimension() int | None[源代码]

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

返回:

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

返回类型:

Optional[int]

half() T

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

注意

此方法就地修改模块。

返回:

self

返回类型:

Module (模块)

load_adapter(*args, **kwargs) None[源代码]

从文件或远程 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
model_card_data_class[源代码]

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

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

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

参数:
  • train_objectives – (DataLoader, LossFunction) 的元组。传递多个以进行多任务学习。

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

  • epochs – 训练的轮数

  • steps_per_epoch – 每轮的训练步数。如果设置为 None(默认值),则一轮等于 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 – 每次评估后调用的回调函数。它必须按此顺序接受以下三个参数:scoreepochsteps

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

将此 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,则创建一个拉取请求 (pull request),而不是直接推送到主分支

返回:

模型在 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[源码]

将模型及其配置文件保存到目录中,以便可以使用 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[源码]

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

参数:
set_pooling_include_prompt(include_prompt: bool) None[源码]

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

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

参数:

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

返回:

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

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

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

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

返回:

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

返回类型:

张量

示例

>>> model = 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]])
属性 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'
属性 similarity_pairwise: Callable[[Tensor | ndarray[Any, dtype[float32]], Tensor | ndarray[Any, dtype[float32]]], Tensor]

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

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

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

返回:

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

返回类型:

张量

示例

>>> model = 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 中的批次转换为模型的张量批次。此处,批次是 InputExample 实例的列表:[InputExample(…), …]

参数:

batch – 来自 SmartBatchingDataset 的批次

返回:

模型的张量批次

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

启动一个多进程池,通过 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]

静态 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)

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

请参阅下面的示例。

注意

此方法就地修改模块。

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

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

  • tensor (torch.Tensor) – 其 dtype 和 device 是此模块中所有参数和缓冲区的目标 dtype 和 device

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

对文本进行分词。

参数:

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

返回:

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

“attention_mask”和“token_type_ids”。

返回类型:

Dict[str, Tensor]

属性 tokenizer: Any

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

train(mode: bool = True) T

将模块设置为训练模式。

这只对某些模块有影响。有关特定模块在训练/评估模式下行为的详细信息,请参阅其文档,例如 DropoutBatchNorm 等。

参数:

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

返回:

self

返回类型:

Module (模块)

属性 transformers_model: PreTrainedModel | None

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

返回:

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

返回类型:

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

在这种情况下,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)[源码]

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

参数:
  • 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) – 模型训练所针对的人类可读任务,例如“语义文本相似度、语义搜索、释义挖掘、文本分类、聚类等”。

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

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

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

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

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

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

静态 possible_values() list[str][源码]

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

返回:

SimilarityFunction 枚举的可能值列表。

返回类型:

list

示例

>>> possible_values = SimilarityFunction.possible_values()
>>> possible_values
['cosine', 'dot', 'euclidean', 'manhattan']
静态 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]])
静态 to_similarity_pairwise_fn(similarity_function: str | SimilarityFunction) Callable[[Tensor | ndarray, Tensor | ndarray], Tensor][源码]

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

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

参数:

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

返回:

成对相似度函数。

返回类型:

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

引发:

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

示例

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