计算 Embedding

一旦您安装了 Sentence Transformers，您可以轻松使用 Sentence Transformer 模型

from sentence_transformers import SentenceTransformer

# 1. Load a pretrained Sentence Transformer model
model = SentenceTransformer("all-MiniLM-L6-v2")

# The sentences to encode
sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]

# 2. Calculate embeddings by calling model.encode()
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 384]

# 3. Calculate the embedding similarities
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[1.0000, 0.6660, 0.1046],
#         [0.6660, 1.0000, 0.1411],
#         [0.1046, 0.1411, 1.0000]])

注意

即使我们谈论句子 Embedding，您也可以将 Sentence Transformers 用于较短的短语以及包含多个句子的较长文本。有关较长文本的 Embedding 的说明，请参阅输入序列长度。

初始化句子 Transformer 模型

第一步是加载预训练的 Sentence Transformer 模型。您可以使用预训练模型中的任何模型或本地模型。另请参阅 SentenceTransformer 以获取有关参数的信息。

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-mpnet-base-v2")
# Alternatively, you can pass a path to a local model directory:
model = SentenceTransformer("output/models/mpnet-base-finetuned-all-nli")

该模型将自动放置在性能最佳的可用设备上，例如，如果可用，则为 cuda 或 mps。您也可以显式指定设备

model = SentenceTransformer("all-mpnet-base-v2", device="cuda")

计算 Embedding

计算 Embedding 的方法是 SentenceTransformer.encode。

Prompt 模板

某些模型需要使用特定的文本prompts才能实现最佳性能。例如，对于 intfloat/multilingual-e5-large，您应该为所有查询添加前缀 "query: "，为所有段落添加前缀 "passage: "。另一个例子是 BAAI/bge-large-en-v1.5，当输入文本前缀为 "Represent this sentence for searching relevant passages: " 时，它在检索方面表现最佳。

句子 Transformer 模型可以使用 prompts 和 default_prompt_name 参数进行初始化

prompts 是一个可选参数，它接受 prompts 字典，其中包含 prompt 名称到 prompt 文本的映射。prompt 将在推理期间添加到输入文本的前面。例如

model = SentenceTransformer(
    "intfloat/multilingual-e5-large",
    prompts={
        "classification": "Classify the following text: ",
        "retrieval": "Retrieve semantically similar text: ",
        "clustering": "Identify the topic or theme based on the text: ",
    },
)
# or
model.prompts = {
    "classification": "Classify the following text: ",
    "retrieval": "Retrieve semantically similar text: ",
    "clustering": "Identify the topic or theme based on the text: ",
}

default_prompt_name 是一个可选参数，用于确定要使用的默认 prompt。它必须与 prompts 中的 prompt 名称相对应。如果为 None，则默认不使用任何 prompt。例如

model = SentenceTransformer(
    "intfloat/multilingual-e5-large",
    prompts={
        "classification": "Classify the following text: ",
        "retrieval": "Retrieve semantically similar text: ",
        "clustering": "Identify the topic or theme based on the text: ",
    },
    default_prompt_name="retrieval",
)
# or
model.default_prompt_name="retrieval"

这两个参数也可以在已保存模型的 config_sentence_transformers.json 文件中指定。这样，您在加载时不必手动指定这些选项。当您保存 Sentence Transformer 模型时，这些选项也将自动保存。

在推理期间，可以以几种不同的方式应用 prompts。所有这些场景都会导致嵌入相同的文本

在 SentenceTransformer.encode 中显式使用 prompt 选项

embeddings = model.encode("How to bake a strawberry cake", prompt="Retrieve semantically similar text: ")

通过依赖于从 a) 初始化或 b) 模型配置加载的 prompts，在 SentenceTransformer.encode 中显式使用 prompt_name 选项
```
embeddings = model.encode("How to bake a strawberry cake", prompt_name="retrieval")
```
如果在 SentenceTransformer.encode 中未指定 prompt 也未指定 prompt_name，则将应用由 default_prompt_name 指定的 prompt。如果为 None，则不会应用任何 prompt
```
embeddings = model.encode("How to bake a strawberry cake")
```

输入序列长度

对于像 BERT、RoBERTa、DistilBERT 等 Transformer 模型，运行时和内存需求随输入长度呈二次方增长。这限制了 Transformer 的输入长度。基于 BERT 的模型的常用值是 512 个 token，这相当于大约 300-400 个单词（对于英语）。

每个模型都有一个最大序列长度，在 model.max_seq_length 下，这是可以处理的最大 token 数。较长的文本将被截断为前 model.max_seq_length 个 token

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
print("Max Sequence Length:", model.max_seq_length)
# => Max Sequence Length: 256

# Change the length to 200
model.max_seq_length = 200

print("Max Sequence Length:", model.max_seq_length)
# => Max Sequence Length: 200

注意

您不能将长度增加到高于各个 Transformer 模型最大支持的长度。另请注意，如果模型是在短文本上训练的，则长文本的表示可能不是那么好。

多进程 / 多 GPU 编码

您可以使用多个 GPU（或 CPU 计算机上的多个进程）对输入文本进行编码。有关示例，请参见：computing_embeddings_multi_gpu.py。

相关方法是 start_multi_process_pool()，它启动用于编码的多个进程。