计算嵌入

安装 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 模型。您可以使用预训练模型中的任何模型或本地模型。有关参数信息，另请参阅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()

可以在调用SentenceTransformer.start_multi_process_pool()并传入设备列表（例如["cuda:0", "cuda:1"]或["cpu", "cpu", "cpu", "cpu"]）后，提供pool参数。这样做的好处是，该池可以重复用于多次调用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 个文本的批次进行嵌入。这有助于内存管理和性能，特别是对于大型数据集。