计算嵌入¶

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

from sentence_transformers import SentenceTransformer

# 1. 加载预训练的 Sentence Transformer 模型
model = SentenceTransformer("all-MiniLM-L6-v2")

# 要编码的句子
sentences = [
    "今天天气真好。",
    "外面阳光明媚！",
    "他开车去了体育场。",
]

# 2. 通过调用 model.encode() 计算嵌入
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 384]

# 3. 计算嵌入的相似度
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")
# 或者，你可以传递本地模型目录的路径：
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": "对以下文本进行分类：",
        "retrieval": "检索语义相似的文本：",
        "clustering": "根据文本识别主题或主题：",
    },
)
# 或者
model.prompts = {
    "classification": "对以下文本进行分类：",
    "retrieval": "检索语义相似的文本：",
    "clustering": "根据文本识别主题或主题：",
}

default_prompt_name 是一个可选参数，用于确定默认使用的提示。它必须对应于 prompts 中的提示名称。如果为 None，则默认不使用提示。例如:

model = SentenceTransformer(
    "intfloat/multilingual-e5-large",
    prompts={
        "classification": "对以下文本进行分类：",
        "retrieval": "检索语义相似的文本：",
        "clustering": "根据文本识别主题或主题：",
    },
    default_prompt_name="retrieval",
)
# 或者
model.default_prompt_name="retrieval"

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

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

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

embeddings = model.encode("如何制作草莓蛋糕", prompt="检索语义相似的文本：")

在 SentenceTransformer.encode 中显式使用 prompt_name 选项，依赖于从 a) 初始化或 b) 模型配置加载的提示:
```
embeddings = model.encode("如何制作草莓蛋糕", prompt_name="retrieval")
```
如果在 SentenceTransformer.encode 中未指定 prompt 或 prompt_name，则将应用由 default_prompt_name 指定的提示。如果为 None，则不应用提示:
```
embeddings = model.encode("如何制作草莓蛋糕")
```

输入序列长度¶

对于像 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("最大序列长度:", model.max_seq_length)
# => 最大序列长度: 256

# 将长度更改为 200
model.max_seq_length = 200

print("最大序列长度:", model.max_seq_length)
# => 最大序列长度: 200

备注

你不能将长度增加到超过相应 transformer 模型支持的最大值。另请注意，如果模型是在短文本上训练的，长文本的表示可能不会那么好。

多进程 / 多 GPU 编码¶

你可以使用多个 GPU（或在 CPU 机器上使用多个进程）对输入文本进行编码。例如，请参见：computing_embeddings_multi_gpu.py。

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