模型

sentence_transformers.models 定义了不同的构建块,可用于从头开始创建 SentenceTransformer 网络。有关更多详细信息,请参阅 训练概述

主要类

class sentence_transformers.models.Transformer(model_name_or_path: str, max_seq_length: int | None = None, model_args: dict[str, Any] | None = None, tokenizer_args: dict[str, Any] | None = None, config_args: dict[str, Any] | None = None, cache_dir: str | None = None, do_lower_case: bool = False, tokenizer_name_or_path: str | None = None, backend: str = 'torch')[source]

Hugging Face AutoModel 用于生成 token 嵌入。加载正确的类,例如 BERT / RoBERTa 等。

参数:
  • model_name_or_path – Hugging Face 模型名称 (https://huggingface.co/models)

  • max_seq_length – 截断任何长于 max_seq_length 的输入

  • model_args – 传递给 Hugging Face Transformers 模型的关键字参数

  • tokenizer_args – 传递给 Hugging Face Transformers 分词器的关键字参数

  • config_args – 传递给 Hugging Face Transformers 配置的关键字参数

  • cache_dir – Hugging Face Transformers 存储/加载模型的缓存目录

  • do_lower_case – 如果为 true,则将输入转换为小写(与模型是否区分大小写无关)

  • tokenizer_name_or_path – 分词器的名称或路径。如果为 None,则使用 model_name_or_path

  • backend – 用于模型推理的后端。可以是 torchonnxopenvino。默认为 torch

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.Pooling(word_embedding_dimension: int, pooling_mode: str | None = None, pooling_mode_cls_token: bool = False, pooling_mode_max_tokens: bool = False, pooling_mode_mean_tokens: bool = True, pooling_mode_mean_sqrt_len_tokens: bool = False, pooling_mode_weightedmean_tokens: bool = False, pooling_mode_lasttoken: bool = False, include_prompt: bool = True)[source]

对 token 嵌入执行池化(最大值或平均值)。

使用池化,它可以从可变大小的句子生成固定大小的句子嵌入。如果底层词嵌入模型返回 CLS token,则此层还允许使用 CLS token。您可以将多个池化层连接在一起。

参数:
  • word_embedding_dimension – 词嵌入的维度

  • pooling_mode – 可以是 “cls”、“lasttoken”、“max”、“mean”、“mean_sqrt_len_tokens” 或 “weightedmean”。如果设置,则覆盖其他 pooling_mode_* 设置

  • pooling_mode_cls_token – 使用第一个 token (CLS token) 作为文本表示

  • pooling_mode_max_tokens – 在所有 token 上,在每个维度中使用最大值。

  • pooling_mode_mean_tokens – 执行均值池化

  • pooling_mode_mean_sqrt_len_tokens – 执行均值池化,但除以 sqrt(input_length)。

  • pooling_mode_weightedmean_tokens – 执行(位置)加权均值池化。参见 SGPT:用于语义搜索的 GPT 句子嵌入

  • pooling_mode_lasttoken

    执行最后一个 token 池化。参见 SGPT:用于语义搜索的 GPT 句子嵌入通过对比预训练的文本和代码嵌入

  • include_prompt – 如果设置为 false,则提示 token 不包含在池化中。这对于复现不包含提示 token 的工作(如 INSTRUCTOR)很有用,但在其他情况下不建议使用。

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.Dense(in_features: int, out_features: int, bias: bool = True, activation_function: Callable[[Tensor], Tensor] | None = Tanh(), init_weight: Tensor | None = None, init_bias: Tensor | None = None)[source]

带有激活函数的前馈函数。

此层接受固定大小的句子嵌入,并将其传递到前馈层。可用于生成深度平均网络 (DAN)。

参数:
  • in_features – 输入维度的大小

  • out_features – 输出大小

  • bias – 添加偏置向量

  • activation_function – 应用于输出的 Pytorch 激活函数

  • init_weight – 线性层矩阵的初始值

  • init_bias – 线性层偏置的初始值

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

更多类

class sentence_transformers.models.Asym(sub_modules: dict[str, list[Module]], allow_empty_key: bool = True)[source]

此模型允许创建非对称 SentenceTransformer 模型,这些模型根据指定的输入键应用不同的模型。

在下面的示例中,我们为 ‘query’ 和 ‘doc’ 创建了两个不同的 Dense 模型。作为 {‘query’: ‘My query’} 传递的文本将沿着第一个 Dense 模型传递,而作为 {‘doc’: ‘My document’} 传递的文本将使用另一个 Dense 模型。

请注意,当您调用 encode() 时,只能对相同类型的输入进行编码。无法对混合类型进行编码。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

# Load a SentenceTransformer model (pretrained or not), and add an Asym module
model = SentenceTransformer("microsoft/mpnet-base")
dim = model.get_sentence_embedding_dimension()
asym_model = models.Asym({
    'query': [models.Dense(dim, dim)],
    'doc': [models.Dense(dim, dim)]
})
model.add_module("asym", asym_model)

train_dataset = Dataset.from_dict({
    "query": ["is toprol xl the same as metoprolol?", "are eyes always the same size?"],
    "answer": ["Metoprolol succinate is also known by the brand name Toprol XL.", "The eyes are always the same size from birth to death."],
})

# This mapper turns normal texts into a dictionary mapping Asym keys to the text
def mapper(sample):
    return {
        "question": {"query": sample["question"]},
        "answer": {"doc": sample["answer"]},
    }

train_dataset = train_dataset.map(mapper)
loss = losses.MultipleNegativesRankingLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

# For inference, you can pass dictionaries with the Asym keys:
model.encode([
    {'query': 'how long do you have to wait to apply for cerb?'},
    {'query': '<3 what does this symbol mean?'},
    {'doc': 'The definition of <3 is "Love".'}]
)

注意

这些模型不一定比非非对称模型更强大。初步实验表明,在大多数情况下,非 Asym 模型表现更好。

参数:
  • sub_modules – 格式为 str -> List[models] 的字典。指定列表中的模型将应用于标记有相应键的输入。

  • allow_empty_key – 如果为 true,则可以处理没有键的输入。如果为 false,则在未指定键时将抛出异常。

class sentence_transformers.models.BoW(vocab: list[str], word_weights: dict[str, float] = {}, unknown_word_weight: float = 1, cumulative_term_frequency: bool = True)[source]

实现词袋 (BoW) 模型以导出句子嵌入。

可以添加权重以允许生成 tf-idf 向量。输出向量的大小为词汇表的大小。

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.CNN(in_word_embedding_dimension: int, out_channels: intkernel_sizes: list[int] = [1, 3, 5], stride_sizes: list[int] | None = None)[source]

CNN 层,具有多个内核大小,作用于词嵌入

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.LSTM(word_embedding_dimension: int, hidden_dim: int, num_layers: int = 1, dropout: float = 0, bidirectional: bool = True)[source]

在词嵌入上运行的双向 LSTM。

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.Normalize[source]

此层将嵌入归一化为单位长度

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.StaticEmbedding(tokenizer: Tokenizer | PreTrainedTokenizerFast, embedding_weights: ndarray | Tensor | None = None, embedding_dim: int | None = None, **kwargs)[source]

初始化给定分词器的 StaticEmbedding 模型。该模型是一个简单的嵌入包模型,它取训练后的每个 token 嵌入的均值来计算文本嵌入。

参数:
  • tokenizer (Tokenizer | PreTrainedTokenizerFast) – 要使用的分词器。必须是来自 transformerstokenizers 的快速分词器。

  • embedding_weights (np.ndarray | torch.Tensor | None, optional) – 预训练的嵌入权重。默认为 None。

  • embedding_dim (int | None, optional) – 嵌入的维度。如果未提供 embedding_weights,则为必需项。默认为 None。

示例

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding
from tokenizers import Tokenizer

# Pre-distilled embeddings:
static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-8M")
# or distill your own embeddings:
static_embedding = StaticEmbedding.from_distillation("BAAI/bge-base-en-v1.5", device="cuda")
# or start with randomized embeddings:
tokenizer = Tokenizer.from_pretrained("FacebookAI/xlm-roberta-base")
static_embedding = StaticEmbedding(tokenizer, embedding_dim=512)

model = SentenceTransformer(modules=[static_embedding])

embeddings = model.encode(["What are Pandas?", "The giant panda, also known as the panda bear or simply the panda, is a bear native to south central China."])
similarity = model.similarity(embeddings[0], embeddings[1])
# tensor([[0.8093]]) (If you use potion-base-8M)
# tensor([[0.6234]]) (If you use the distillation method)
# tensor([[-0.0693]]) (For example, if you use randomized embeddings)
Raises:
  • ValueError – 如果分词器不是快速分词器。

  • ValueError – 如果既未提供 embedding_weights 也未提供 embedding_dim

classmethod from_distillation(model_name: str, vocabulary: list[str] | None = None, device: str | None = None, pca_dims: int | None = 256, apply_zipf: bool = True, use_subword: bool = True) StaticEmbedding[source]

使用 model2vec 包,通过蒸馏过程创建一个 StaticEmbedding 实例。

参数:
  • model_name (str) – 要蒸馏的模型名称。

  • vocabulary (list[str] | None, optional) – 要使用的词汇表列表。默认为 None。

  • device (str) – 运行蒸馏的设备(例如,‘cpu’, ‘cuda’)。如果未指定,则自动检测最强的设备。默认为 None。

  • pca_dims (int | None, optional) – PCA 降维的维度数。默认为 256。

  • apply_zipf (bool) – 在蒸馏期间是否应用齐普夫定律。默认为 True。

  • use_subword (bool) – 是否使用子词分词。默认为 True。

返回值:

使用蒸馏模型的 tokenizer 和 embedding weights 初始化的 StaticEmbedding 实例

tokenizer 和 embedding weights。

返回类型:

StaticEmbedding

Raises:

ImportError – 如果未安装 model2vec 包。

classmethod from_model2vec(model_id_or_path: str) StaticEmbedding[source]

从 model2vec 模型创建一个 StaticEmbedding 实例。此方法加载预训练的 model2vec 模型,并提取 embedding weights 和 tokenizer 以创建 StaticEmbedding 实例。

参数:

model_id_or_path (str) – 预训练 model2vec 模型的标识符或路径。

返回值:

使用 tokenizer 和 embedding weights 初始化的 StaticEmbedding 实例

model2vec 模型。

返回类型:

StaticEmbedding

Raises:

ImportError – 如果未安装 model2vec 包。

class sentence_transformers.models.WeightedLayerPooling(word_embedding_dimension, num_hidden_layers: int = 12, layer_start: int = 4, layer_weights=None)[source]

Token embeddings 是它们不同隐藏层表示的加权平均值

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.WordEmbeddings(tokenizer: WordTokenizer, embedding_weights, update_embeddings: bool = False, max_seq_length: int = 1000000)[source]

初始化内部模块状态,由 nn.Module 和 ScriptModule 共享。

class sentence_transformers.models.WordWeights(vocab: list[str], word_weights: dict[str, float], unknown_word_weight: float = 1)[source]

此模型可以对词嵌入进行加权,例如使用 idf 值。

初始化 WordWeights 类。

参数:
  • vocab (List[str]) – 分词器的词汇表。

  • word_weights (Dict[str, float]) – 令牌到浮点权重值的映射。词嵌入乘以这个浮点值。word_weights 中的令牌不能等于 vocab(可以包含更多或更少的值)。

  • unknown_word_weight (float, optional) – vocab 中未出现在 word_weights 查找中的单词的权重。这些可能是 vocab 中不存在权重的罕见词。默认为 1。