创建自定义模型¶
句子转换器模型的结构¶
一个 Sentence Transformer 模型由一系列模块 (文档 ) 组成。最常见的架构是 Transformer 模块、Pooling 模块的组合，以及可选的 Dense 模块和/或 Normalize 模块。
Transformer: 该模块负责处理输入文本并生成上下文化的嵌入。Pooling: 该模块通过聚合嵌入来减少Transformer模块输出的维度。常见的池化策略包括均值池化和CLS池化。Dense: 该模块包含一个线性层，用于后处理从池化模块输出的嵌入。Normalize: 该模块对前一层的嵌入进行归一化。
例如，流行的 all-MiniLM-L6-v2 模型也可以通过初始化构成该模型的3个特定模块来加载：
from sentence_transformers import models, SentenceTransformer
transformer = models.Transformer("sentence-transformers/all-MiniLM-L6-v2", max_seq_length=256)
pooling = models.Pooling(transformer.get_word_embedding_dimension(), pooling_mode="mean")
normalize = models.Normalize()
model = SentenceTransformer(modules=[transformer, pooling, normalize])
保存句子转换器模型¶
每当保存一个 Sentence Transformer 模型时，会生成三种类型的文件：
modules.json: 这个文件包含模块名称、路径和类型的列表，用于重建模型。config_sentence_transformers.json: 这个文件包含 Sentence Transformer 模型的一些配置选项，包括保存的提示、模型的相似度函数以及模型作者使用的 Sentence Transformer 包版本。特定模块文件：每个模块都保存在一个单独的文件夹中，第一个模块保存在根文件夹中，所有后续模块保存在以模块索引和模型名称命名的子文件夹中（例如，
1_Pooling，2_Normalize）。大多数模块文件夹包含一个config.json``（或 ``sentence_bert_config.json用于Transformer模块）文件，该文件存储传递给该模块的关键字参数的默认值。因此，一个sentence_bert_config.json的内容如下：{ "max_seq_length": 4096, "do_lower_case": false }
意味着
Transformer模块将以max_seq_length=4096和do_lower_case=False进行初始化。
因此，如果我在前面的代码片段中对 model 调用 SentenceTransformer.save_pretrained("local-all-MiniLM-L6-v2") ，将生成以下文件：
local-all-MiniLM-L6-v2/
├── 1_Pooling
│ └── config.json
├── 2_Normalize
├── README.md
├── config.json
├── config_sentence_transformers.json
├── model.safetensors
├── modules.json
├── sentence_bert_config.json
├── special_tokens_map.json
├── tokenizer.json
├── tokenizer_config.json
└── vocab.txt
这包含一个 modules.json ，其内容如下：
[
{
"idx": 0,
"name": "0",
"path": "",
"type": "sentence_transformers.models.Transformer"
},
{
"idx": 1,
"name": "1",
"path": "1_Pooling",
"type": "sentence_transformers.models.Pooling"
},
{
"idx": 2,
"name": "2",
"path": "2_Normalize",
"type": "sentence_transformers.models.Normalize"
}
]
以及一个包含以下内容的 config_sentence_transformers.json：
{
"__version__": {
"sentence_transformers": "3.0.1",
"transformers": "4.43.4",
"pytorch": "2.5.0"
},
"prompts": {},
"default_prompt_name": null,
"similarity_fn_name": null
}
此外，1_Pooling 目录包含 Pooling 模块的配置文件，而 2_Normalize 目录是空的，因为 Normalize 模块不需要任何配置。sentence_bert_config.json 文件包含 Transformer 模块的配置，该模块还在根目录中保存了许多与分词器和模型本身相关的文件。
加载 Sentence Transformer 模型¶
要从保存的模型目录加载 Sentence Transformer 模型，会读取 modules.json 以确定构成模型的模块。每个模块都使用相应模块目录中存储的配置进行初始化，之后使用加载的模块实例化 SentenceTransformer 类。
来自 Transformers 模型的句子转换器模型¶
当你使用纯 Transformers 模型（例如 BERT、RoBERTa、DistilBERT、T5）初始化一个 Sentence Transformer 模型时，Sentence Transformers 默认会创建一个 Transformer 模块和一个均值池化模块。这提供了一种简单的方式来利用预训练语言模型进行句子嵌入。
具体来说，这两个片段是相同的:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("bert-base-uncased")
from sentence_transformers import models, SentenceTransformer
transformer = models.Transformer("bert-base-uncased")
pooling = models.Pooling(transformer.get_word_embedding_dimension(), pooling_mode="mean")
model = SentenceTransformer(modules=[transformer, pooling])
高级：自定义模块¶
要创建自定义的 Sentence Transformer 模型，您可以通过子类化 PyTorch 的 torch.nn.Module 类并实现以下方法来实现自己的模块：
一个接受
features字典的torch.nn.Module.forward()方法，字典的键可能包括input_ids、attention_mask、token_type_ids、token_embeddings和sentence_embedding，具体取决于模块在模型流水线中的位置。一个
save方法，接受一个save_dir参数，并将模块的配置保存到该目录中。一个
load静态方法，它接受一个load_dir参数，并根据该目录中的模块配置初始化模块。(如果第一个模块) 一个
get_max_seq_length方法，返回模块可以处理的最大序列长度。仅在模块处理输入文本时需要。(如果第1个模块) 一个
tokenize方法，接受输入列表并返回一个字典，字典的键如input_ids、attention_mask、token_type_ids、pixel_values等。这个字典将被传递给模块的forward方法。(可选) 一个
get_sentence_embedding_dimension方法，返回该模块生成的句子嵌入的维度。仅在模块生成嵌入或更新嵌入维度时需要。(可选) 一个
get_config_dict方法，返回一个包含模块配置的字典。此方法可用于将模块配置保存到磁盘，并将模块配置保存到模型卡中。
例如，我们可以通过实现一个自定义模块来创建一个自定义池化方法。
# decay_pooling.py
import json
import os
import torch
import torch.nn as nn
class DecayMeanPooling(nn.Module):
def __init__(self, dimension: int, decay: float = 0.95) -> None:
super(DecayMeanPooling, self).__init__()
self.dimension = dimension
self.decay = decay
def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict [str, torch.Tensor]:
token_embeddings = features["token_embeddings"]
attention_mask = features["attention_mask"].unsqueeze(-1)
# Apply the attention mask to filter away padding tokens
token_embeddings = token_embeddings * attention_mask
# Calculate mean of token embeddings
sentence_embeddings = token_embeddings.sum(1) / attention_mask.sum(1)
# Apply exponential decay
importance_per_dim = self.decay ** torch.arange(sentence_embeddings. size(1), device=sentence_embeddings.device)
features["sentence_embedding"] = sentence_embeddings * importance_per_dim
return features
def get_config_dict(self) -> dict[str, float]:
return {"dimension": self.dimension, "decay": self.decay}
def get_sentence_embedding_dimension(self) -> int:
return self.dimension
def save(self, save_dir: str, **kwargs) -> None:
with open(os.path.join(save_dir, "config.json"), "w") as fOut:
json.dump(self.get_config_dict(), fOut, indent=4)
def load(load_dir: str, **kwargs) -> "DecayMeanPooling":
with open(os.path.join(load_dir, "config.json")) as fIn:
config = json.load(fIn)
return DecayMeanPooling(**config)
备注
建议在 __init__、forward、save、load 和 tokenize 方法中添加 **kwargs，以确保这些方法与 Sentence Transformers 库的未来更新兼容。
这现在可以作为一个模块在 Sentence Transformer 模型中使用:
from sentence_transformers import models, SentenceTransformer
from decay_pooling import DecayMeanPooling
transformer = models.Transformer("bert-base-uncased", max_seq_length=256)
decay_mean_pooling = DecayMeanPooling(transformer.get_word_embedding_dimension(), decay=0.99)
normalize = models.Normalize()
model = SentenceTransformer(modules=[transformer, decay_mean_pooling, normalize])
print(model)
"""
SentenceTransformer(
(0): Transformer({'max_seq_length': 256, 'do_lower_case': False}) with Transformer model: BertModel
(1): DecayMeanPooling()
(2): Normalize()
)
"""
texts = [
"Hello, World!",
"The quick brown fox jumps over the lazy dog.",
"I am a sentence that is used for testing purposes.",
"This is a test sentence.",
"This is another test sentence.",
]
embeddings = model.encode(texts)
print(embeddings.shape)
# [5, 384]
您可以使用 SentenceTransformer.save_pretrained 保存此模型，结果将生成一个 modules.json 文件，内容如下：
[
{
"idx": 0,
"name": "0",
"path": "",
"type": "sentence_transformers.models.Transformer"
},
{
"idx": 1,
"name": "1",
"path": "1_DecayMeanPooling",
"type": "decay_pooling.DecayMeanPooling"
},
{
"idx": 2,
"name": "2",
"path": "2_Normalize",
"type": "sentence_transformers.models.Normalize"
}
]
为了确保可以导入 decay_pooling.DecayMeanPooling，你应该将 decay_pooling.py 文件复制到你保存模型的目录中。如果你将模型推送到 Hugging Face Hub，那么你也应该将 decay_pooling.py 文件上传到模型的仓库中。这样，每个人都可以通过调用 SentenceTransformer("your-username/your-model-id", trust_remote_code=True) 来使用你的自定义模块。
备注
使用存储在 Hugging Face Hub 上的远程代码的自定义模块时，要求您的用户在加载模型时将 trust_remote_code 指定为 True。这是一项安全措施，以防止远程代码执行攻击。
如果你在 Hugging Face Hub 上有你的模型和自定义建模代码，那么将你的自定义模块分离到一个单独的仓库中可能是有意义的。这样，你只需要维护一个自定义模块的实现，并且可以在多个模型中重用它。你可以通过更新 modules.json 文件中的 type 来实现这一点，将其路径包含到存储自定义模块的仓库中，例如 {repository_id}--{dot_path_to_module}。例如，如果 decay_pooling.py 文件存储在一个名为 my-user/my-model-implementation 的仓库中，并且模块名为 DecayMeanPooling，那么 modules.json 文件可能看起来像这样:
[
{
"idx": 0,
"name": "0",
"path": "",
"type": "sentence_transformers.models.Transformer"
},
{
"idx": 1,
"name": "1",
"path": "1_DecayMeanPooling",
"type": "my-user/my-model-implementation--decay_pooling.DecayMeanPooling"
},
{
"idx": 2,
"name": "2",
"path": "2_Normalize",
"type": "sentence_transformers.models.Normalize"
}
]
高级：自定义模块中的关键字参数传递¶
如果你想让你的用户能够通过 SentenceTransformer.encode 方法指定自定义的关键字参数，那么你可以将它们的名称添加到 modules.json 文件中。例如，如果我的模块在用户指定 task_type 关键字参数时应表现不同，那么你的 modules.json 可能看起来像:
[
{
"idx": 0,
"name": "0",
"path": "",
"type": "custom_transformer.CustomTransformer",
"kwargs": ["task_type"]
},
{
"idx": 1,
"name": "1",
"path": "1_Pooling",
"type": "sentence_transformers.models.Pooling"
},
{
"idx": 2,
"name": "2",
"path": "2_Normalize",
"type": "sentence_transformers.models.Normalize"
}
]
然后，你可以在自定义模块的 forward 方法中访问 task_type 关键字参数:
from sentence_transformers.models import Transformer
class CustomTransformer(Transformer):
def forward(self, features: dict[str, torch.Tensor], task_type: Optional[str] = None) -> dict[str, torch.Tensor]:
if task_type == "default":
# Do something
else:
# Do something else
return features
这样，用户在调用 SentenceTransformer.encode 时可以指定 task_type 关键字参数:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("your-username/your-model-id", trust_remote_code=True)
texts = [...]
model.encode(texts, task_type="default")