计算嵌入

一旦你安装了 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]])

注意

尽管我们谈论的是句子嵌入，但你也可以将 Sentence Transformers 用于较短的短语以及包含多个句子的较长文本。有关较长文本嵌入的注意事项，请参阅输入序列长度。

初始化一个 Sentence 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")

计算嵌入

计算嵌入的方法是 SentenceTransformer.encode。

提示模板

有些模型需要使用特定的文本提示以达到最佳性能。例如，对于 intfloat/multilingual-e5-large 模型，你应该在所有查询前加上 "query: "，在所有段落前加上 "passage: "。另一个例子是 BAAI/bge-large-en-v1.5，当输入文本前缀为 "Represent this sentence for searching relevant passages: " 时，其检索性能最佳。

Sentence Transformer 模型可以用 prompts 和 default_prompt_name 参数进行初始化。

prompts 是一个可选参数，接受一个从提示名称到提示文本的字典。在推理过程中，提示将被前置到输入文本。例如：

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 是一个可选参数，用于确定要使用的默认提示。它必须与 prompts 中的一个提示名称相对应。如果为 None，则默认不使用任何提示。例如：

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 模型时，这些选项也会被自动保存。

在推理过程中，提示可以以几种不同的方式应用。所有这些情况都会导致嵌入相同的文本：

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

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

在 SentenceTransformer.encode 中显式使用 prompt_name 选项，依赖于从 a) 初始化或 b) 模型配置中加载的提示：
```
embeddings = model.encode("How to bake a strawberry cake", prompt_name="retrieval")
```
如果在 SentenceTransformer.encode 中既未指定 prompt 也未指定 prompt_name，则将应用由 default_prompt_name 指定的提示。如果为 None，则不应用任何提示。
```
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。

你可以将 SentenceTransformer.encode() （或 SentenceTransformer.encode_query() 或 SentenceTransformer.encode_document()) 与以下任一方式结合使用：

device 参数，可以设置为例如 "cuda:0" 或 "cpu" 用于单进程计算，也可以设置为设备列表用于多进程或多 GPU 计算，例如 ["cuda:0", "cuda:1"] 或 ["cpu", "cpu", "cpu", "cpu"]。

from sentence_transformers import SentenceTransformer

def main():
    model = SentenceTransformer("all-MiniLM-L6-v2")
    # Encode with multiple GPUs
    embeddings = model.encode(
        inputs,
        device=["cuda:0", "cuda:1"]  # or ["cpu", "cpu", "cpu", "cpu"]
    )

if __name__ == "__main__":
    main()

pool 参数可以在调用 SentenceTransformer.start_multi_process_pool() 后提供，并附带一个设备列表，例如 ["cuda:0", "cuda:1"] 或 ["cpu", "cpu", "cpu", "cpu"]。这样做的好处是，该进程池可以被多次调用 SentenceTransformer.encode() 重复使用，这比为每次调用都启动一个新池要高效得多。

from sentence_transformers import SentenceTransformer

def main():
    model = SentenceTransformer("all-MiniLM-L6-v2")
    # Start a multi-process pool with multiple GPUs
    pool = model.start_multi_process_pool(devices=["cuda:0", "cuda:1"])
    # Encode with multiple GPUs
    embeddings = model.encode(inputs, pool=pool)
    # Don't forget to stop the pool after usage
    model.stop_multi_process_pool(pool)

if __name__ == "__main__":
    main()

此外，你可以使用 chunk_size 参数来控制发送到每个进程的块的大小。这与 batch_size 参数不同。例如，如果 chunk_size=1000 且 batch_size=32，输入文本将被分割成大小为 1000 的块，每个块将被发送到一个进程，并以 32 个文本为一批进行嵌入。这有助于内存管理和性能，尤其是在处理大型数据集时。