Modules

sentence_transformers.models 定义了不同的构建块，也称为 Modules，可用于从头开始创建 SentenceTransformer 模型。有关更多详细信息，请参阅 创建自定义模型。

Main Modules

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://hugging-face.cn/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 – 用于模型推理的后端。可以是 torch、onnx 或 openvino。默认为 torch。

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（如果底层词嵌入模型返回了它）。您可以将多个池化层连接在一起。

参数:

  • 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 – 在每个维度上取所有 tokens 的最大值。

  • pooling_mode_mean_tokens – 执行平均池化

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

  • pooling_mode_weightedmean_tokens – 执行（位置）加权平均池化。请参阅 SGPT: GPT Sentence Embeddings for Semantic Search。

  • pooling_mode_lasttoken –

    执行 last token 池化。请参阅 SGPT: GPT Sentence Embeddings for Semantic Search 和 Text and Code Embeddings by Contrastive Pre-Training。

  • include_prompt – 如果设置为 false，则不将 prompt tokens 包含在池化中。这对于重现不包含 prompt tokens 的工作（如 INSTRUCTOR）很有用，但否则不推荐。

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 – 线性层偏置的初始值

class sentence_transformers.models.Normalize[source]

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

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

此模型允许创建非对称 SentenceTransformer 模型，这些模型根据指定的路由（例如，“query”或“document”）应用不同的模块。对于具有不同查询和文档编码器的模型特别有用。

值得注意的是，model.encode 的 task 参数可用于指定要使用的路由，而 model.encode_query 和 model.encode_document 分别是 task="query" 和 task="document" 的简写。这些方法还可以选择性地应用特定于查询或文档的 prompts。

注意

在使用 Router 模块训练模型时，您必须在 SentenceTransformerTrainingArguments 或 SparseEncoderTrainingArguments 中使用 router_mapping 参数，将训练数据集列映射到正确的路由（“query”或“document”）。例如，如果您的训练数据集包含 ["question", "positive", "negative"] 列，则可以使用以下映射：

args = SparseEncoderTrainingArguments(
    ...,
    router_mapping={
        "question": "query",
        "positive": "document",
        "negative": "document",
    }
)

此外，通常为不同的路由使用不同的学习率。为此，您应该在 SentenceTransformerTrainingArguments 或 SparseEncoderTrainingArguments 中使用 learning_rate_mapping 参数，将参数模式映射到其学习率。例如，如果您想为 SparseStaticEmbedding 模块使用 1e-3 的学习率，为模型的其余部分使用 2e-5 的学习率，您可以这样做：

args = SparseEncoderTrainingArguments(
    ...,
    learning_rate=2e-5,
    learning_rate_mapping={
        r"SparseStaticEmbedding\.*": 1e-3,
    }
)

在下面的示例中，Router 模型用于创建具有不同查询和文档编码器的非对称模型。在这些示例中，“query”路由是高效的（例如，使用 SparseStaticEmbedding），而“document”路由使用更复杂的模型（例如 Transformers 模块）。这使得查询编码高效，同时仍使用强大的文档编码器，但组合不限于此。

示例

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import Router, Normalize

# Use a regular SentenceTransformer for the document embeddings, and a static embedding model for the query embeddings
document_embedder = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
query_embedder = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1")
router = Router.for_query_document(
    query_modules=list(query_embedder.children()),
    document_modules=list(document_embedder.children()),
)
normalize = Normalize()

# Create an asymmetric model with different encoders for queries and documents
model = SentenceTransformer(
    modules=[router, normalize],
)

# ... requires more training to align the vector spaces

# Use the query & document routes
query_embedding = model.encode_query("What is the capital of France?")
document_embedding = model.encode_document("Paris is the capital of France.")

from sentence_transformers.models import Router
from sentence_transformers.sparse_encoder import SparseEncoder
from sentence_transformers.sparse_encoder.models import MLMTransformer, SparseStaticEmbedding, SpladePooling

# Load an asymmetric model with different encoders for queries and documents
doc_encoder = MLMTransformer("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill")
router = Router.for_query_document(
    query_modules=[
        SparseStaticEmbedding.from_json(
            "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill",
            tokenizer=doc_encoder.tokenizer,
            frozen=True,
        ),
    ],
    document_modules=[
        doc_encoder,
        SpladePooling(pooling_strategy="max", activation_function="log1p_relu"),
    ],
)

model = SparseEncoder(modules=[router], similarity_fn_name="dot")

query = "What's the weather in ny now?"
document = "Currently New York is rainy."

query_embed = model.encode_query(query)
document_embed = model.encode_document(document)

sim = model.similarity(query_embed, document_embed)
print(f"Similarity: {sim}")

# Visualize top tokens for each text
top_k = 10
print(f"Top tokens {top_k} for each text:")

decoded_query = model.decode(query_embed, top_k=top_k)
decoded_document = model.decode(document_embed)

for i in range(min(top_k, len(decoded_query))):
    query_token, query_score = decoded_query[i]
    doc_score = next((score for token, score in decoded_document if token == query_token), 0)
    if doc_score != 0:
        print(f"Token: {query_token}, Query score: {query_score:.4f}, Document score: {doc_score:.4f}")

'''
Similarity: tensor([[11.1105]], device='cuda:0')
Top tokens 10 for each text:
Token: ny, Query score: 5.7729, Document score: 0.8049
Token: weather, Query score: 4.5684, Document score: 0.9710
Token: now, Query score: 3.5895, Document score: 0.4720
Token: ?, Query score: 3.3313, Document score: 0.0286
Token: what, Query score: 2.7699, Document score: 0.0787
Token: in, Query score: 0.4989, Document score: 0.0417
'''

注意

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

参数:

  • sub_modules – 路由键到模块列表的映射。每个键对应一个特定的任务类型，通常是“query”或“document”，列表包含应用于该任务类型的模块。

  • default_route – 如果未指定任务类型，则使用的默认路由。如果为 None，则在未指定任务类型时将抛出异常。如果 allow_empty_key 为 True，则 sub_modules 中的第一个键将用作默认路由。默认为 None。

  • allow_empty_key – 如果为 True，则允许在 default_route 为 None 时将默认路由设置为 sub_modules 中的第一个键。默认为 True。

classmethod for_query_document(query_modules: list[Module], document_modules: list[Module], default_route: str | None = None, allow_empty_key: bool = True) → Self[source]

创建了一个专门用于查询和文档模块的 Router 模型，可以通过 model.encode_query 和 model.encode_document 方便使用。

参数:

  • query_modules – 应用于“query”任务类型的模块列表。

  • document_modules – 应用于“document”任务类型的模块列表。

  • default_route – 如果未指定任务类型，则使用的默认路由。如果为 None，则在未指定任务类型时将抛出异常。如果 allow_empty_key 为 True，则 sub_modules 中的第一个键将用作默认路由。默认为 None。

  • allow_empty_key – 如果为 True，则允许在 default_route 为 None 时将默认路由设置为 sub_modules 中的第一个键。默认为 True。

返回:

一个 Router 模型实例，包含指定的查询和文档模块。

返回类型:

Router (路由器)

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) – 要使用的分词器。必须是 transformers 或 tokenizers 中的快速分词器。

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

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

提示

由于此模块架构的效率极高，将输入移到 GPU 的开销可能大于实际计算时间。因此，请考虑使用 CPU 设备进行推理和训练。

示例

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)
引发:

  • 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, sif_coefficient: float | None = 0.0001, token_remove_pattern: str | None = '\\[unused\\d+\\]', quantize_to: str = 'float32', use_subword: bool = True, **kwargs: Any) → 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) – 在蒸馏过程中是否应用 Zipf 定律。默认为 True。

  • sif_coefficient (float | None, optional) – SIF 加权的系数。默认为 1e-4。

  • token_remove_pattern (str | None, optional) – 用于从词汇表中删除 token 的正则表达式模式。默认为 r”[unusedd+]”。

  • quantize_to (str) – 要量化权重的目标数据类型。默认为 ‘float32’。

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

返回:

一个 StaticEmbedding 实例，使用蒸馏模型的

分词器和嵌入权重进行初始化。

返回类型:

StaticEmbedding

引发:

ImportError – 如果未安装 model2vec 包。

classmethod from_model2vec(model_id_or_path: str) → StaticEmbedding[source]

从 model2vec 模型创建 StaticEmbedding 实例。此方法加载预训练的 model2vec 模型，并提取嵌入权重和分词器以创建 StaticEmbedding 实例。

参数:

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

返回:

一个 StaticEmbedding 实例，使用 model2vec 模型的

分词器和嵌入权重进行初始化。

返回类型:

StaticEmbedding

引发:

ImportError – 如果未安装 model2vec 包。

Further Modules

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 向量。输出向量的大小是词汇表的大小。

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

CNN 层，具有在词嵌入上使用的多个核尺寸

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

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

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

Token 嵌入是其不同隐藏层表示的加权平均值

class sentence_transformers.models.WordEmbeddings(tokenizer: WordTokenizer | PreTrainedTokenizerBase, embedding_weights, update_embeddings: bool = False, max_seq_length: int = 1000000)[source]
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]) – token 到 float 权重的映射。词嵌入将乘以该 float 值。word_weights 中的 token 不能等于 vocab（可以包含更多或更少的值）。

  • unknown_word_weight (float, optional) – 不在 word_weights 查找中的 vocab 中词的权重。这些可以是例如 vocab 中不存在权重的稀有词。默认为 1。

Base Modules

class sentence_transformers.models.Module(*args, **kwargs)[source]

Sentence Transformers 库中所有模块的基类。

此类为所有模块提供了一个通用接口，包括加载和保存模块配置和权重的方​​法。它还提供了一个执行模块前向传递的方法。

在此类中定义了两个抽象方法，子类必须实现它们：

可选地，您可能还需要覆盖

为了协助加载和保存模块，提供了几个实用方法：

并且定义了几个类变量以协助加载和保存模块：

config_file_name: str = 'config.json'

用于保存模块配置的配置文件名称。加载预训练模型时，此文件用于初始化模块。

config_keys: list[str] = []

用于保存模块配置的键列表。将模块保存到磁盘时，这些键用于保存模块的配置。

abstract forward(features: dict[str, Tensor | Any], **kwargs) → dict[str, Tensor | Any][source]

模块的前向传递。子类应覆盖此方法以实现模块的特定行为。

forward 方法以一个特征字典作为输入，并返回一个特征字典作为输出。features 字典中的键取决于模块在模型管道中的位置，因为 features 字典会从一个模块传递到下一个模块。features 字典中的常用键有：

  • input_ids: 输入文本中 token 的输入 ID。

  • attention_mask: 输入 token 的注意力掩码。

  • token_type_ids: 输入 token 的 token 类型 ID。

  • token_embeddings: 输入 token 的 token 嵌入。

  • sentence_embedding: 输入文本的句子嵌入，即池化的 token 嵌入。

可选地，forward 方法可以接受额外的关键字参数（**kwargs），这些参数可用于将额外信息从 model.encode 传递给此模块。

参数:

  • features (dict[str, torch.Tensor | Any]) – 要由模块处理的特征字典。

  • **kwargs – 可用于将额外信息从 model.encode 传递的其他关键字参数。

返回:

经过模块处理后的特征字典。

返回类型:

dict[str, torch.Tensor | Any]

get_config_dict() → dict[str, Any][source]

返回模块配置参数的字典。

这些参数用于将模块配置保存到磁盘，并再次用于从预训练模型加载模块时初始化模块。字典中使用的键在 config_keys 类变量中定义。

返回:

模块配置参数的字典。

返回类型:

dict[str, Any]

classmethod load(model_name_or_path: str, subfolder: str = '', token: bool | str | None = None, cache_folder: str | None = None, revision: str | None = None, local_files_only: bool = False, **kwargs) → Self[source]

从模型检查点加载此模块。检查点可以是本地目录，也可以是 Hugging Face 上的模型 ID。

参数:

  • model_name_or_path (str) – 模型目录的路径或 Hugging Face 上的模型名称。

  • subfolder (str, optional) – 要从中加载的模型目录中的子文件夹，例如 "1_Pooling"。默认为 ""。

  • token (bool | str | None, optional) – 加载 Hugging Face 时用于身份验证的 token。如果为 None，则尝试使用使用 huggingface-cli login 或 HF_TOKEN 环境变量保存的 token。默认为 None。

  • cache_folder (str | None, optional) – 用于缓存模型文件的文件夹。如果为 None，则使用 Hugging Face 的默认缓存文件夹 ~/.cache/huggingface。默认为 None。

  • revision (str | None, optional) – 要加载的模型版本。如果为 None，则使用最新版本。默认为 None。

  • local_files_only (bool, optional) – 是否仅加载本地文件。默认为 False。

  • **kwargs – 在重写的 load 方法中使用的其他模块特定参数，例如 trust_remote_code、model_kwargs、tokenizer_kwargs、config_kwargs、backend 等。

返回:

加载的模块。

返回类型:

Self

classmethod load_config(model_name_or_path: str, subfolder: str = '', config_filename: str | None = None, token: bool | str | None = None, cache_folder: str | None = None, revision: str | None = None, local_files_only: bool = False) → dict[str, Any][source]

从模型检查点加载模块的配置。检查点可以是本地目录，也可以是 Hugging Face 上的模型 ID。配置从 JSON 文件加载，其中包含用于初始化模块的参数。

参数:

  • model_name_or_path (str) – 模型目录的路径或 Hugging Face 上的模型名称。

  • subfolder (str, optional) – 要从中加载的模型目录中的子文件夹，例如 "1_Pooling"。默认为 ""。

  • config_filename (str | None, optional) – 要加载的配置文件名称。如果为 None，则使用 config_file_name 类变量中定义的默认配置文件名称。默认为 None。

  • token (bool | str | None, optional) – 加载 Hugging Face 时用于身份验证的 token。如果为 None，则尝试使用使用 huggingface-cli login 或 HF_TOKEN 环境变量保存的 token。默认为 None。

  • cache_folder (str | None, optional) – 用于缓存模型文件的文件夹。如果为 None，则使用 Hugging Face 的默认缓存文件夹 ~/.cache/huggingface。默认为 None。

  • revision (str | None, optional) – 要加载的模型版本。如果为 None，则使用最新版本。默认为 None。

  • local_files_only (bool, optional) – 是否仅加载本地文件。默认为 False。

返回:

模块配置参数的字典。

返回类型:

dict[str, Any]

static load_dir_path(model_name_or_path: str, subfolder: str = '', token: bool | str | None = None, cache_folder: str | None = None, revision: str | None = None, local_files_only: bool = False) → str[source]

一个实用函数，用于从模型检查点加载目录。检查点可以是本地目录，也可以是 Hugging Face 上的模型 ID。文件将从模型目录中的指定子文件夹加载。

参数:

  • model_name_or_path (str) – 模型目录的路径或 Hugging Face 上的模型名称。

  • subfolder (str, optional) – 要从中加载的模型目录中的子文件夹，例如 "1_Pooling"。默认为 ""。

  • token (bool | str | None, optional) – 加载 Hugging Face 时用于身份验证的 token。如果为 None，则尝试使用使用 huggingface-cli login 或 HF_TOKEN 环境变量保存的 token。默认为 None。

  • cache_folder (str | None, optional) – 用于缓存模型文件的文件夹。如果为 None，则使用 Hugging Face 的默认缓存文件夹 ~/.cache/huggingface。默认为 None。

  • revision (str | None, optional) – 要加载的模型版本。如果为 None，则使用最新版本。默认为 None。

  • local_files_only (bool, optional) – 是否仅加载本地文件。默认为 False。

返回:

加载的目录的路径。

返回类型:

str

static load_file_path(model_name_or_path: str, filename: str, subfolder: str = '', token: bool | str | None = None, cache_folder: str | None = None, revision: str | None = None, local_files_only: bool = False) → str | None[source]

一个实用函数，用于从模型检查点加载文件。检查点可以是本地目录，也可以是 Hugging Face 上的模型 ID。文件将从模型目录中的指定子文件夹加载。

参数:

  • model_name_or_path (str) – 模型目录的路径或 Hugging Face 上的模型名称。

  • filename (str) – 要加载的文件名。

  • subfolder (str, optional) – 要从中加载的模型目录中的子文件夹，例如 "1_Pooling"。默认为 ""。

  • token (bool | str | None, optional) – 加载 Hugging Face 时用于身份验证的 token。如果为 None，则尝试使用使用 huggingface-cli login 或 HF_TOKEN 环境变量保存的 token。默认为 None。

  • cache_folder (str | None, optional) – 用于缓存模型文件的文件夹。如果为 None，则使用 Hugging Face 的默认缓存文件夹 ~/.cache/huggingface。默认为 None。

  • revision (str | None, optional) – 要加载的模型版本。如果为 None，则使用最新版本。默认为 None。

  • local_files_only (bool, optional) – 是否仅加载本地文件。默认为 False。

返回:

加载的文件的路径，如果文件未找到则为 None。

返回类型:

str | None

classmethod load_torch_weights(model_name_or_path: str, subfolder: str = '', token: bool | str | None = None, cache_folder: str | None = None, revision: str | None = None, local_files_only: bool = False, model: Self | None = None)[source]

一个用于从检查点加载 PyTorch 权重的实用函数。检查点可以是本地目录，也可以是 Hugging Face 上的模型 ID。权重将根据哪个文件可用，从 model.safetensors 文件或 pytorch_model.bin 文件加载。此方法将权重加载到模型中，或将权重作为状态字典返回。

参数:

  • model_name_or_path (str) – 模型目录的路径或 Hugging Face 上的模型名称。

  • subfolder (str, optional) – 要从中加载的模型的子文件夹，例如 "2_Dense"。默认为 ""。

  • token (bool | str | None, optional) – 加载 Hugging Face 时用于身份验证的 token。如果为 None，则尝试使用使用 huggingface-cli login 或 HF_TOKEN 环境变量保存的 token。默认为 None。

  • cache_folder (str | None, optional) – 用于缓存模型文件的文件夹。如果为 None，则使用 Hugging Face 的默认缓存文件夹 ~/.cache/huggingface。默认为 None。

  • revision (str | None, optional) – 要加载的模型版本。如果为 None，则使用最新版本。默认为 None。

  • local_files_only (bool, optional) – 是否仅加载本地文件。默认为 False。

  • model (Self | None, optional) – 要将权重加载到的模型。如果为 None，则将权重作为状态字典返回。默认为 None。

引发:

ValueError – 如果在 subfolder 中的模型检查点中找不到 model.safetensors 文件或 pytorch_model.bin 文件。

返回:

加载了权重的模型或作为状态字典的权重，

具体取决于 model 参数的值。

返回类型:

Self | dict[str, torch.Tensor]

abstract save(output_path: str, *args, safe_serialization: bool = True, **kwargs) → None[source]

将模块保存到磁盘。子类应覆盖此方法以实现模块的特定行为。

参数:

  • output_path (str) – 应保存模块的目录路径。

  • *args – 可用于将附加信息传递给 save 方法的附加参数。

  • safe_serialization (bool, optional) – 是否使用 safetensors 格式保存模型权重。默认为 True。

  • **kwargs – 可用于将附加信息传递给 save 方法的附加关键字参数。

save_config(output_path: str, filename: str | None = None) → None[source]

将模块的配置保存到 JSON 文件。

参数:

  • output_path (str) – 应保存配置文件所在的目录路径。

  • filename (str | None, optional) – 配置文件的名称。如果为 None，则使用 config_file_name 类变量中定义的默认配置文件名。默认为 None。

返回:

无

save_in_root: bool = False

是否将模块的配置保存在模型的根目录中，还是保存在以模块命名的子目录中。

save_torch_weights(output_path: str, safe_serialization: bool = True) → None[source]

将模块的 PyTorch 权重保存到磁盘。

参数:

  • output_path (str) – 应保存权重的目录路径。

  • safe_serialization (bool, optional) – 是否使用 safetensors 格式保存模型权重。默认为 True。

返回:

无

class sentence_transformers.models.InputModule(*args, **kwargs)[source]

Sentence Transformers 库中所有输入模块的基类，即用于处理输入并可能在前向传播中执行处理的模块。它是 sentence_transformers.models.Module 的子类。

此类为所有输入模块提供了一个通用接口，包括加载和保存模块配置和权重的方​​法，以及输入处理。它还提供了一个用于执行模块前向传播的方法。

此类定义了三个抽象方法，子类必须实现这些方法：

可选地，您可能还需要覆盖

为了协助加载和保存模块，提供了几个实用方法：

并且定义了几个类变量以协助加载和保存模块：

save_in_root: bool = True

是否将模块的配置保存在模型的根目录中，还是保存在以模块命名的子目录中。

save_tokenizer(output_path: str, **kwargs) → None[source]

将分词器保存到指定输出路径。

参数:

  • output_path (str) – 保存分词器的路径。

  • **kwargs – 用于保存分词器的附加关键字参数。

返回:

无

abstract tokenize(texts: list[str], **kwargs) → dict[str, Tensor | Any][source]

标记输入文本并返回标记化特征的字典。

参数:

  • texts (list[str]) – 要标记化的输入文本列表。

  • **kwargs – 标记化的附加关键字参数，例如 task。

返回:

包含标记化特征的字典，例如

{"input_ids": ..., "attention_mask": ...}

返回类型:

dict[str, torch.Tensor | Any]

tokenizer: PreTrainedTokenizerBase | Tokenizer

用于标记化输入文本的分词器。它可以是 transformers.PreTrainedTokenizerBase 子类，也可以是 tokenizers 库中的 Tokenizer。