CrossEncoder

CrossEncoder

有关 Cross-Encoders 的介绍，请参阅 Cross-Encoders。

class sentence_transformers.cross_encoder.CrossEncoder(model_name_or_path: str, num_labels: int = None, max_length: int = None, activation_fn: Callable | None = None, device: str | None = None, cache_folder: str = None, trust_remote_code: bool = False, revision: str | None = None, local_files_only: bool = False, token: bool | str | None = None, model_kwargs: dict = None, tokenizer_kwargs: dict = None, config_kwargs: dict = None, model_card_data: CrossEncoderModelCardData | None = None, backend: Literal['torch', 'onnx', 'openvino'] = 'torch')[source]

CrossEncoder 接受正好两个句子/文本作为输入，并为这个句子对预测一个分数或标签。例如，它可以预测句子对的相似度，范围为 0 … 1。

它不产生句子嵌入，也不适用于单个句子。

参数:

  • model_name_or_path (str) – Hugging Face Hub 中的模型名称，可以使用 AutoModel 加载，或者本地模型的路径。我们提供了几个预训练的 CrossEncoder 模型，可用于常见任务。

  • num_labels (int, 可选) – 分类器的标签数量。如果为 1，则 CrossEncoder 是一个回归模型，输出 0…1 的连续分数。如果 > 1，它会输出多个分数，可以进行 soft-max 运算以获得不同类别的概率分数。默认为 None。

  • max_length (int, 可选) – 输入序列的最大长度。更长的序列将被截断。如果为 None，将使用模型的最大长度。默认为 None。

  • activation_fn (Callable, 可选) – 关于应在 model.predict() 之上使用的默认激活函数的可调用对象（如 nn.Sigmoid）。如果为 None。如果 num_labels=1，将使用 nn.Sigmoid()，否则使用 nn.Identity()。默认为 None。

  • device (str, 可选) – 用于计算的设备（如 “cuda”、“cpu”、“mps”、“npu”）。如果为 None，则检查是否可以使用 GPU。

  • cache_folder (str, Path, optional) – 缓存文件存储在的文件夹路径。

  • trust_remote_code (bool, 可选) – 是否允许在 Hub 上自定义模型，这些模型在其自己的建模文件中定义。此选项仅应针对您信任且已阅读代码的存储库设置为 True，因为它将在您的本地计算机上执行 Hub 上存在的代码。默认为 False。

  • revision (str, 可选) – 要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID，对于 Hugging Face 上的存储模型。默认为 None。

  • local_files_only (bool, 可选) – 是否只查看本地文件（即，不尝试下载模型）。

  • token (bool or str, 可选) – Hugging Face 身份验证令牌，用于下载私有模型。

  • model_kwargs (Dict[str, Any], 可选) –

    要传递给 Hugging Face Transformers 模型的其他模型配置参数。特别有用的选项是

    • torch_dtype：覆盖默认的 torch.dtype，并以特定的 dtype 加载模型。不同的选项是

      1. torch.float16, torch.bfloat16 或 torch.float：以指定的 dtype 加载，忽略模型的 config.torch_dtype（如果存在）。如果未指定 - 模型将以 torch.float (fp32) 加载。

      2. "auto" - 将尝试使用模型的 config.json 文件中的 torch_dtype 条目。如果未找到此条目，则接下来检查检查点中第一个浮点类型的权重的 dtype，并将其用作 dtype。这将使用模型在训练结束时保存的 dtype 加载模型。它不能用作模型如何训练的指标。因为它可能以半精度 dtype 训练，但以 fp32 保存。

    • attn_implementation：模型中要使用的注意力实现（如果相关）。可以是 “eager”（注意力的手动实现）、“sdpa”（使用 F.scaled_dot_product_attention）或 “flash_attention_2”（使用 Dao-AILab/flash-attention）中的任何一个。默认情况下，如果可用，SDPA 将用于 torch>=2.1.1。否则默认为手动 “eager” 实现。

    有关更多详细信息，请参阅 AutoModelForSequenceClassification.from_pretrained 文档。

  • tokenizer_kwargs (Dict[str, Any], 可选) – 要传递给 Hugging Face Transformers 分词器的其他分词器配置参数。有关更多详细信息，请参阅 AutoTokenizer.from_pretrained 文档。

  • config_kwargs (Dict[str, Any], 可选) – 要传递给 Hugging Face Transformers 配置的其他模型配置参数。有关更多详细信息，请参阅 AutoConfig.from_pretrained 文档。例如，您可以通过此参数设置 classifier_dropout。

  • model_card_data (SentenceTransformerModelCardData, optional) – 包含有关模型信息的模型卡数据对象。这用于在保存模型时生成模型卡。如果未设置，则会创建默认模型卡数据对象。

  • backend (str) – 用于推理的后端。可以是 “torch”（默认）、“onnx” 或 “openvino” 之一。有关不同后端的基准测试信息，请参阅 https://sbert.hugging-face.cn/docs/cross_encoder/usage/efficiency.html。

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

bfloat16() → T

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法已编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数的详细信息，请参阅 torch.compile()。

cpu() → T

将所有模型参数和缓冲区移动到 CPU。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module

cuda(device: int | device | None = None) → T

将所有模型参数和缓冲区移动到 GPU。

这也使得关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 GPU 上，则应在构造优化器之前调用它。

注意

此方法就地修改模块。

参数:

device (int, 可选) – 如果指定，所有参数都将复制到该设备

返回:

self

返回类型:

Module

double() → T

将所有浮点参数和缓冲区转换为 double 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module

eval() → T

将模块设置为评估模式。

这仅对某些模块有效。有关它们在训练/评估模式下的行为的详细信息，请参阅特定模块的文档（如果它们受到影响），例如 Dropout、BatchNorm 等。

这等效于 self.train(False)。

有关 .eval() 与可能与之混淆的几种类似机制之间的比较，请参阅 在本地禁用梯度计算。

返回:

self

返回类型:

Module

fit(train_dataloader: DataLoader, evaluator: SentenceEvaluator = None, epochs: int = 1, loss_fct=None, activation_fct=Identity(), scheduler: str = 'WarmupLinear', warmup_steps: int = 10000, optimizer_class: type[Optimizer] = <class 'torch.optim.adamw.AdamW'>, optimizer_params: dict[str, object] = {'lr': 2e-05}, weight_decay: float = 0.01, evaluation_steps: int = 0, output_path: str = None, save_best_model: bool = True, max_grad_norm: float = 1, use_amp: bool = False, callback: Callable[[float, int, int], None] = None, show_progress_bar: bool = True) → None[source]

v4.0 之前的已弃用训练方法，建议改用 CrossEncoderTrainer。此方法在后台使用 CrossEncoderTrainer，但不如 Trainer 本身灵活。

此训练方法使用 DataLoader 和 Loss 函数来训练模型。

此方法应在 v4.0 中产生与 v4.0 之前等效的结果，但如果您在现有训练脚本中遇到任何问题，则可能希望改用 CrossEncoder.old_fit。它使用 v4.0 之前的旧训练方法。

参数:

  • train_dataloader – 带有 InputExample 实例的 DataLoader

  • evaluator – 评估器 (sentence_transformers.cross_encoder.evaluation) 在训练期间在保留的开发数据上评估模型性能。它用于确定保存到磁盘的最佳模型。

  • epochs – 训练的 epoch 数

  • loss_fct – 用于训练的损失函数。如果为 None，如果 self.config.num_labels == 1，将使用 BinaryCrossEntropy()，否则使用 CrossEntropyLoss()。默认为 None。

  • activation_fct – 应用于模型 logits 输出之上的激活函数。

  • scheduler – 学习率调度器。可用的调度器：constantlr, warmupconstant, warmuplinear, warmupcosine, warmupcosinewithhardrestarts

  • warmup_steps – 行为取决于调度器。对于 WarmupLinear (默认)，学习率从 0 增加到最大学习率。经过这么多训练步数后，学习率线性地降低回零。

  • optimizer_class – 优化器

  • optimizer_params – 优化器参数

  • weight_decay – 模型参数的权重衰减

  • evaluation_steps – 如果 > 0，则在每次训练步数后使用评估器评估模型

  • output_path – 模型和评估文件的存储路径

  • save_best_model – 如果为 true，则最佳模型（根据评估器）将存储在 output_path

  • max_grad_norm – 用于梯度归一化。

  • use_amp – 使用自动混合精度 (AMP)。仅适用于 Pytorch >= 1.6.0

  • callback – 在每次评估后调用的回调函数。它必须按顺序接受以下三个参数：score、epoch、steps

  • show_progress_bar – 如果为 True，则输出 tqdm 进度条

float() → T

将所有浮点参数和缓冲区转换为 float 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module

half() → T

将所有浮点参数和缓冲区转换为 half 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

Module

old_fit(train_dataloader: ~torch.utils.data.dataloader.DataLoader, evaluator: ~sentence_transformers.evaluation.SentenceEvaluator.SentenceEvaluator | None = None, epochs: int = 1, loss_fct=None, activation_fct=Identity(), scheduler: str = 'WarmupLinear', warmup_steps: int = 10000, optimizer_class: type[~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adamw.AdamW'>, optimizer_params: dict[str, object] = {'lr': 2e-05}, weight_decay: float = 0.01, evaluation_steps: int = 0, output_path: str | None = None, save_best_model: bool = True, max_grad_norm: float = 1, use_amp: bool = False, callback: ~typing.Callable[[float, int, int], None] | None = None, show_progress_bar: bool = True) → None[source]

不推荐使用的训练方法，源于 Sentence Transformers v4.0 之前，建议使用 CrossEncoderTrainer 代替。 仅当您在升级到 v4.0 后遇到现有训练脚本问题时，才应使用此方法。

此训练方法使用 DataLoader 和 Loss 函数来训练模型。

参数:

  • train_dataloader – 带有 InputExample 实例的 DataLoader

  • evaluator – 评估器 (sentence_transformers.cross_encoder.evaluation) 在训练期间在保留的开发数据上评估模型性能。它用于确定保存到磁盘的最佳模型。

  • epochs – 训练的 epoch 数

  • loss_fct – 用于训练的损失函数。如果为 None，如果 self.config.num_labels == 1，将使用 BinaryCrossEntropy()，否则使用 CrossEntropyLoss()。默认为 None。

  • activation_fct – 应用于模型 logits 输出之上的激活函数。

  • scheduler – 学习率调度器。可用的调度器：constantlr, warmupconstant, warmuplinear, warmupcosine, warmupcosinewithhardrestarts

  • warmup_steps – 行为取决于调度器。对于 WarmupLinear (默认)，学习率从 0 增加到最大学习率。经过这么多训练步数后，学习率线性地降低回零。

  • optimizer_class – 优化器

  • optimizer_params – 优化器参数

  • weight_decay – 模型参数的权重衰减

  • evaluation_steps – 如果 > 0，则在每次训练步数后使用评估器评估模型

  • output_path – 模型和评估文件的存储路径

  • save_best_model – 如果为 true，则最佳模型（根据评估器）将存储在 output_path

  • max_grad_norm – 用于梯度归一化。

  • use_amp – 使用自动混合精度 (AMP)。仅适用于 Pytorch >= 1.6.0

  • callback – 在每次评估后调用的回调函数。它必须按顺序接受以下三个参数：score、epoch、steps

  • show_progress_bar – 如果为 True，则输出 tqdm 进度条

predict(sentences: tuple[str, str] | list[str], batch_size: int = 32, show_progress_bar: bool | None = None, activation_fn: Callable | None = None, apply_softmax: bool | None = False, convert_to_numpy: Literal[False] = True, convert_to_tensor: Literal[False] = False) → torch.Tensor[source]
predict(sentences: list[tuple[str, str]] | list[list[str]] | tuple[str, str] | list[str], batch_size: int = 32, show_progress_bar: bool | None = None, activation_fn: Callable | None = None, apply_softmax: bool | None = False, convert_to_numpy: Literal[True] = True, convert_to_tensor: Literal[False] = False) → np.ndarray
predict(sentences: list[tuple[str, str]] | list[list[str]] | tuple[str, str] | list[str], batch_size: int = 32, show_progress_bar: bool | None = None, activation_fn: Callable | None = None, apply_softmax: bool | None = False, convert_to_numpy: bool = True, convert_to_tensor: Literal[True] = False) → torch.Tensor
predict(sentences: list[tuple[str, str]] | list[list[str]], batch_size: int = 32, show_progress_bar: bool | None = None, activation_fn: Callable | None = None, apply_softmax: bool | None = False, convert_to_numpy: Literal[False] = True, convert_to_tensor: Literal[False] = False) → list[torch.Tensor]

对给定的句子对执行 CrossEncoder 预测。

参数:

  • sentences (Union[List[Tuple[str, str]], Tuple[str, str]]) – 句子对列表，例如 [(Sent1, Sent2), (Sent3, Sent4)] 或一个句子对 (Sent1, Sent2)。

  • batch_size (int, optional) – 编码的批大小。默认为 32。

  • show_progress_bar (bool, optional) – 输出进度条。默认为 None。

  • activation_fn (callable, optional) – 应用于 CrossEncoder 的 logits 输出的激活函数。如果为 None，将使用 model.activation_fn，如果 num_labels=1，则默认为 torch.nn.Sigmoid，否则为 torch.nn.Identity。默认为 None。

  • convert_to_numpy (bool, optional) – 将输出转换为 numpy 矩阵。默认为 True。

  • apply_softmax (bool, optional) – 如果设置为 True 且 model.num_labels > 1，则对 logits 输出应用 softmax，以便对于每个样本，每个类别的分数总和为 1。默认为 False。

  • convert_to_numpy – 输出是否应为 numpy 向量列表。如果为 False，则输出 PyTorch 张量列表。默认为 True。

  • convert_to_tensor (bool, optional) – 输出是否应为一个大的张量。覆盖 convert_to_numpy。默认为 False。

返回:

传递的句子对的预测。返回类型取决于 convert_to_numpy 和 convert_to_tensor 参数。如果 convert_to_tensor 为 True，则输出将为 torch.Tensor。如果 convert_to_numpy 为 True，则输出将为 numpy.ndarray。否则，输出将是 torch.Tensor 值的列表。

返回类型:

Union[List[torch.Tensor], np.ndarray, torch.Tensor]

示例

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/stsb-roberta-base")
sentences = [["I love cats", "Cats are amazing"], ["I prefer dogs", "Dogs are loyal"]]
model.predict(sentences)
# => array([0.6912767, 0.4303499], dtype=float32)
push_to_hub(repo_id: str, *, token: str | None = None, private: bool | None = None, safe_serialization: bool = True, commit_message: str | None = None, exist_ok: bool = False, revision: str | None = None, create_pr: bool = False, tags: list[str] | None = None) → str[source]

将 CrossEncoder 模型上传到 Hugging Face Hub。

示例

from sentence_transformers import CrossEncoder

model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
model.push_to_hub("username/my-crossencoder-model")
# => "https://huggingface.co/username/my-crossencoder-model"
参数:

  • repo_id (str) – Hugging Face Hub 上的仓库名称，例如 “username/repo_name”、“organization/repo_name” 或仅 “repo_name”。

  • token (str, optional) – 用于 Hugging Face Hub API 的身份验证令牌。如果未提供，将使用通过 Hugging Face CLI 存储的令牌。

  • private (bool, optional) – 是否创建私有仓库。如果未指定，仓库将是公开的。

  • safe_serialization (bool, optional) – 是否将模型权重转换为 safetensors 格式以实现更安全的序列化。默认为 True。

  • commit_message (str, optional) – 用于推送的提交消息。默认为 “Add new CrossEncoder model”。

  • exist_ok (bool, optional) – 如果为 True，则如果仓库已存在，则不引发错误。如果 create_pr=True，则忽略。默认为 False。

  • revision (str, optional) – 要提交到的 git 分支。默认为 ‘main’ 分支的头部。

  • create_pr (bool, optional) – 是否创建包含上传内容的拉取请求或直接提交。默认为 False。

  • tags (list[str], optional) – 要添加到模型卡的标签列表。默认为 None。

返回:

提交或拉取请求的 URL (如果 create_pr=True)

返回类型:

str

rank(query: str, documents: list[str], top_k: int | None = None, return_documents: bool = False, batch_size: int = 32, show_progress_bar: bool = None, activation_fn: Callable | None = None, apply_softmax=False, convert_to_numpy: bool = True, convert_to_tensor: bool = False) → list[dict[Literal['corpus_id', 'score', 'text'], int | float | str]][source]

对给定的查询和文档执行 CrossEncoder 排序。返回按文档索引和分数排序的列表。

参数:

  • query (str) – 单个查询。

  • documents (List[str]) – 文档列表。

  • top_k (Optional[int], optional) – 返回前 k 个文档。如果为 None，则返回所有文档。默认为 None。

  • return_documents (bool, optional) – 如果为 True，则同时返回文档。如果为 False，则仅返回索引和分数。默认为 False。

  • batch_size (int, optional) – 编码的批大小。默认为 32。

  • show_progress_bar (bool, optional) – 输出进度条。默认为 None。

  • activation_fn ([type], optional) – 应用于 CrossEncoder 的 logits 输出的激活函数。如果为 None，如果 num_labels=1，则使用 nn.Sigmoid()，否则使用 nn.Identity。默认为 None。

  • convert_to_numpy (bool, optional) – 将输出转换为 numpy 矩阵。默认为 True。

  • apply_softmax (bool, optional) – 如果维度超过 2 且 apply_softmax=True，则对 logits 输出应用 softmax。默认为 False。

  • convert_to_tensor (bool, optional) – 将输出转换为张量。默认为 False。

返回:

包含文档的 “corpus_id”、“score” 和可选的 “text” 的排序列表。

返回类型:

List[Dict[Literal[“corpus_id”, “score”, “text”], Union[int, float, str]]]

示例

from sentence_transformers import CrossEncoder
model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")

query = "Who wrote 'To Kill a Mockingbird'?"
documents = [
    "'To Kill a Mockingbird' is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.",
    "The novel 'Moby-Dick' was written by Herman Melville and first published in 1851. It is considered a masterpiece of American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.",
    "Harper Lee, an American novelist widely known for her novel 'To Kill a Mockingbird', was born in 1926 in Monroeville, Alabama. She received the Pulitzer Prize for Fiction in 1961.",
    "Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment upon the British landed gentry at the end of the 18th century.",
    "The 'Harry Potter' series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the most popular and critically acclaimed books of the modern era.",
    "'The Great Gatsby', a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan."
]

model.rank(query, documents, return_documents=True)

[{'corpus_id': 0,
'score': 10.67858,
'text': "'To Kill a Mockingbird' is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature."},
{'corpus_id': 2,
'score': 9.761677,
'text': "Harper Lee, an American novelist widely known for her novel 'To Kill a Mockingbird', was born in 1926 in Monroeville, Alabama. She received the Pulitzer Prize for Fiction in 1961."},
{'corpus_id': 1,
'score': -3.3099542,
'text': "The novel 'Moby-Dick' was written by Herman Melville and first published in 1851. It is considered a masterpiece of American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil."},
{'corpus_id': 5,
'score': -4.8989105,
'text': "'The Great Gatsby', a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan."},
{'corpus_id': 4,
'score': -5.082967,
'text': "The 'Harry Potter' series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the most popular and critically acclaimed books of the modern era."}]
save_pretrained(path: str, *, safe_serialization: bool = True, **kwargs) → None[source]

将模型和 tokenizer 保存到指定路径。

参数:

  • path (str) – 模型应保存到的目录

  • safe_serialization (bool, optional) – 是否使用 safetensors 或传统的 PyTorch 方式保存。默认为 True。

  • **kwargs – 传递给模型和 tokenizer 的底层保存方法的其他参数。

返回:

None

set_config_value(key: str, value) → None[source]

在底层模型的配置中设置一个值。

参数:

  • key (str) – 要设置的键。

  • value – 要设置的值。

to(*args, **kwargs)

移动和/或转换参数和缓冲区。

可以这样调用

to(device=None, dtype=None, non_blocking=False)
to(dtype, non_blocking=False)
to(tensor, non_blocking=False)
to(memory_format=torch.channels_last)

其签名与 torch.Tensor.to() 类似，但仅接受浮点型或复数型 dtype。此外，此方法将仅转换浮点型或复数型参数和缓冲区为 dtype（如果给定）。整型参数和缓冲区将移动到 device（如果给定），但 dtype 保持不变。当设置 non_blocking 时，如果可能，它会尝试相对于主机异步转换/移动，例如，将具有固定内存的 CPU 张量移动到 CUDA 设备。

请参见以下示例。

注意

此方法就地修改模块。

参数:

  • device (torch.device) – 此模块中参数和期望的设备

  • dtype (torch.dtype) – 此模块中参数和缓冲区的期望的浮点型或复数型 dtype

  • tensor (torch.Tensor) – 张量，其 dtype 和设备是此模块中所有参数和缓冲区的期望 dtype 和设备

  • memory_format (torch.memory_format) – 此模块中 4D 参数和缓冲区的期望内存格式（仅关键字参数）

返回:

self

返回类型:

Module

示例

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
train(mode: bool = True) → T

将模块设置为训练模式。

这仅对某些模块有效。有关它们在训练/评估模式下的行为的详细信息，请参阅特定模块的文档（如果它们受到影响），例如 Dropout、BatchNorm 等。

参数:

mode (bool) – 是否设置训练模式 (True) 或评估模式 (False)。默认值：True。

返回:

self

返回类型:

Module

CrossEncoderModelCardData

class sentence_transformers.cross_encoder.model_card.CrossEncoderModelCardData(language: str | list[str] | None = <factory>, license: str | None = None, model_name: str | None = None, model_id: str | None = None, train_datasets: list[dict[str, str]] = <factory>, eval_datasets: list[dict[str, str]] = <factory>, task_name: str = None, tags: list[str] | None = <factory>, generate_widget_examples: Literal['deprecated'] = 'deprecated')[source]

用于存储模型卡中使用的数据的数据类。

参数:

  • language (Optional[Union[str, List[str]]]) – 模型语言，可以是字符串或列表，例如 “en” 或 [“en”, “de”, “nl”]

  • license (Optional[str]) – 模型的许可证，例如 “apache-2.0”、“mit” 或 “cc-by-nc-sa-4.0”

  • model_name (Optional[str]) – 模型的漂亮名称，例如 “基于 answerdotai/ModernBERT-base 的 CrossEncoder”。

  • model_id (Optional[str]) – 将模型推送到 Hub 时的模型 ID，例如 “tomaarsen/ce-mpnet-base-ms-marco”。

  • train_datasets (List[Dict[str, str]]) – 训练数据集的名称和/或 Hugging Face 数据集 ID 列表。例如 [{“name”: “SNLI”, “id”: “stanfordnlp/snli”}, {“name”: “MultiNLI”, “id”: “nyu-mll/multi_nli”}, {“name”: “STSB”}]

  • eval_datasets (List[Dict[str, str]]) – 评估数据集的名称和/或 Hugging Face 数据集 ID 列表。例如 [{“name”: “SNLI”, “id”: “stanfordnlp/snli”}, {“id”: “mteb/stsbenchmark-sts”}]

  • task_name (str) – 模型训练的人类可读的任务名称，例如 “语义搜索和释义挖掘”。

  • tags (Optional[List[str]]) – 模型的标签列表，例如 [“sentence-transformers”, “cross-encoder”]。

提示

安装 codecarbon 以自动跟踪碳排放使用情况并将其包含在您的模型卡中。

示例

>>> model = CrossEncoder(
...     "microsoft/mpnet-base",
...     model_card_data=CrossEncoderModelCardData(
...         model_id="tomaarsen/ce-mpnet-base-allnli",
...         train_datasets=[{"name": "SNLI", "id": "stanfordnlp/snli"}, {"name": "MultiNLI", "id": "nyu-mll/multi_nli"}],
...         eval_datasets=[{"name": "SNLI", "id": "stanfordnlp/snli"}, {"name": "MultiNLI", "id": "nyu-mll/multi_nli"}],
...         license="apache-2.0",
...         language="en",
...     ),
... )