CrossEncoder

CrossEncoder

有关交叉编码器的介绍,请参阅交叉编码器

class sentence_transformers.cross_encoder.CrossEncoder(model_name_or_path: str, num_labels: int | None = None, max_length: int | None = None, activation_fn: Callable | None = None, device: str | None = None, cache_folder: str | None = None, trust_remote_code: bool = False, revision: str | None = None, local_files_only: bool = False, token: bool | str | None = None, model_kwargs: dict | None = None, tokenizer_kwargs: dict | None = None, config_kwargs: dict | None = 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, 可选) – 用于模型 .predict() 之上的默认激活函数的 Callable (如 nn.Sigmoid)。如果为 None,当 num_labels=1 时将使用 nn.Sigmoid(),否则使用 nn.Identity()。默认为 None。

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

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

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

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

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

  • token (boolstr, 可选) – 用于下载私有模型的 Hugging Face 认证令牌。

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

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

    • torch_dtype: 覆盖默认的 torch.dtype 并在特定 dtype 下加载模型。不同的选项是

      1. torch.float16, torch.bfloat16torch.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)中的任何一种。默认情况下,如果可用,对于 torch>=2.1.1 将使用 SDPA。否则默认是手动 “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, 可选) – 包含模型信息的模型卡数据对象。此对象用于在保存模型时生成模型卡。如果未设置,将创建默认的模型卡数据对象。

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

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

bfloat16() T

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

注意

此方法会就地修改模块。

返回:

自身

返回类型:

Module

compile(*args, **kwargs)

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

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

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

cpu() T

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

注意

此方法会就地修改模块。

返回:

自身

返回类型:

Module

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

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

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

注意

此方法会就地修改模块。

参数:

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

返回:

自身

返回类型:

Module

double() T

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

注意

此方法会就地修改模块。

返回:

自身

返回类型:

Module

eval() T

将模块设置为评估模式。

这仅对某些模块有影响。有关特定模块在训练/评估模式下行为的详细信息(如果受影响),请参阅其文档,例如 DropoutBatchNorm 等。

这等同于 self.train(False)

有关 .eval() 与几个可能混淆的类似机制之间的比较,请参阅 局部禁用梯度计算

返回:

自身

返回类型:

Module

fit(train_dataloader: DataLoader, evaluator: SentenceEvaluator | None = 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 = 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]

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

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

此方法在 v4.0 中的结果应与 v4.0 之前的结果相同,但如果您现有训练脚本遇到任何问题,则可以考虑改用 CrossEncoder.old_fit。后者使用 v4.0 之前的旧训练方法。

参数:
  • train_dataloader – 包含 InputExample 实例的 DataLoader

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

  • epochs – 训练轮数

  • 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 – 在每次评估后调用的回调函数。它必须按此顺序接受以下三个参数:scoreepochsteps

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

float() T

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

注意

此方法会就地修改模块。

返回:

自身

返回类型:

Module

half() T

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

注意

此方法会就地修改模块。

返回:

自身

返回类型:

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 和损失函数来训练模型。

参数:
  • train_dataloader – 包含 InputExample 实例的 DataLoader

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

  • epochs – 训练轮数

  • 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 – 在每次评估后调用的回调函数。它必须按此顺序接受以下三个参数:scoreepochsteps

  • 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]]) – 句子对列表,例如 [(句子1, 句子2), (句子3, 句子4)];或单个句子对 (句子1, 句子2)。

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

  • show_progress_bar (bool, 可选) – 是否输出进度条。默认为 None。

  • activation_fn (可调用对象, 可选) – 应用于 CrossEncoder 的 logits 输出的激活函数。如果为 None,则将使用 model.activation_fn,当 num_labels=1 时,默认为 torch.nn.Sigmoid,否则为 torch.nn.Identity。默认为 None。

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

  • apply_softmax (bool, 可选) – 如果设置为 True 且 model.num_labels > 1,则对 logits 输出应用 softmax,使每个样本中各类的分数总和为 1。默认为 False。

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

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

返回:

对传入的句子对进行预测。返回类型取决于 convert_to_numpyconvert_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, 可选) – 用于 Hugging Face Hub API 的身份验证令牌。如果未提供,将使用通过 Hugging Face CLI 存储的令牌。

  • private (bool, 可选) – 是否创建私有存储库。如果未指定,存储库将是公开的。

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

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

  • exist_ok (bool, 可选) – 如果为 True,则如果存储库已存在,则不引发错误。如果 create_pr=True,则忽略此参数。默认为 False。

  • revision (str, 可选) – 要提交到的 Git 分支。默认为“main”分支的头部。

  • create_pr (bool, 可选) – 是通过上传创建拉取请求还是直接提交。默认为 False。

  • tags (list[str], 可选) – 要添加到模型卡(model card)的标签列表。默认为 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 = 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], 可选) – 返回前 k 个文档。如果为 None,则返回所有文档。默认为 None。

  • return_documents (bool, 可选) – 如果为 True,也返回文档。如果为 False,则只返回索引和分数。默认为 False。

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

  • show_progress_bar (bool, 可选) – 是否输出进度条。默认为 None。

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

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

  • apply_softmax (bool, 可选) – 如果维度大于 2 且 apply_softmax=True,则对 logits 输出应用 softmax。默认为 False。

  • convert_to_tensor (bool, 可选) – 是否将输出转换为张量。默认为 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]

将模型和分词器保存到指定路径。

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

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

  • **kwargs – 传递给模型和分词器底层保存方法的额外参数。

返回:

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 和 device 是此模块中所有参数和缓冲区的所需 dtype 和 device

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

返回:

自身

返回类型:

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

将模块设置为训练模式。

这仅对某些模块有影响。有关特定模块在训练/评估模式下行为的详细信息(如果受影响),请参阅其文档,例如 DropoutBatchNorm 等。

参数:

mode (bool) – 是设置为训练模式 (True) 还是评估模式 (False)。默认值:True

返回:

自身

返回类型:

Module

属性 transformers_model: PreTrainedModel | None

获取底层 transformers PreTrainedModel 实例的属性。

返回:

底层 transformers 模型,如果未找到则为 None。

返回类型:

PreTrainedModel 或 None

示例

from sentence_transformers import CrossEncoder

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

# You can now access the underlying transformers model
transformers_model = model.transformers_model
print(type(transformers_model))
# => <class 'transformers.models.bert.modeling_bert.BertForSequenceClassification'>

CrossEncoderModelCardData

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 = None, tags: list[str] | None = <factory>, local_files_only: bool = False, generate_widget_examples: bool = True)[source]

一个用于存储模型卡(model card)中数据的 dataclass。

参数:
  • 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”]。

  • local_files_only (bool) – 如果为 True,则不尝试在 Hub 上查找数据集或基础模型信息。默认为 False。

提示

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

示例

>>> 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",
...     ),
... )