Transformers documentation

Reformer

Transformers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

该模型于 2020-01-13 在 HF papers 上发布，并于 2020-11-16 贡献至 Hugging Face Transformers。

Reformer

概述

Reformer 模型由 Nikita Kitaev、Łukasz Kaiser 和 Anselm Levskaya 在论文 Reformer: The Efficient Transformer 中提出。

该论文的摘要如下：

大型 Transformer 模型在许多任务上都能取得最先进的结果，但训练这些模型的成本可能高得令人望而却步，尤其是在处理长序列时。我们提出了两种技术来提高 Transformer 的效率。其一，我们用基于局部敏感哈希（locality-sensitive hashing）的注意力机制替代了点积注意力，将其复杂度从 O(L^2) 降低到 O(Llog(L))，其中 L 为序列长度。此外，我们使用可逆残差层（reversible residual layers）替代标准残差层，这使得在训练过程中只需存储一次激活值，而不是 N 次，其中 N 为层数。由此得到的模型 Reformer 在性能上与 Transformer 模型相当，同时在长序列上具有更高的内存效率和更快的速度。

该模型由 patrickvonplaten 贡献。作者的代码可在此处找到。

使用技巧

由于 PyTorch 中的一个 bug，Reformer 无法与 torch.nn.DataParallel 配合使用，请参阅 issue #36035。
使用轴向位置编码（更多详情请参见下文）。这是一种通过将位置编码矩阵分解为较小的矩阵，从而避免在序列长度非常大时生成巨大的位置编码矩阵的机制。
用 LSH（局部敏感哈希）注意力取代传统注意力（更多细节见下文）。这是一种避免在注意力层中计算完整的 query-key 乘积的技术。
通过使用可逆 Transformer 层来避免存储每一层的中间结果，从而在反向传播过程中获取这些结果（从下一层的输入中减去残差即可恢复它们），或者重新计算给定层内部的结果（效率低于直接存储，但能节省内存）。
按块计算前馈操作，而不是对整个批次进行计算。

轴向位置编码

Axial Positional Encodings 最初由 Google 的 trax 库实现，并由本模型论文的作者在后续工作中进一步发展。在处理超长输入序列的模型中，传统的 position id 编码会为每个位置 $i, \ldots, n_s$ 存储一个大小为 $d$ 的嵌入向量，其中 config.hidden_size 为 $n_s$ ， config.max_embedding_size. 这意味着当序列长度为 $n_s = 2^{19} \approx 0.5M$ 且 config.hidden_size 为 $d = 2^{10} \approx 1000$ 时，位置编码矩阵为： $X_{i,j}, \text{ with } i \in \left[1,\ldots, d\right] \text{ and } j \in \left[1,\ldots, n_s\right]$

仅这一项就需要存储超过 5 亿个参数。轴向位置编码将 $X_{i,j}$ 分解为两个矩阵： $X^{1}_{i,j}, \text{ with } i \in \left[1,\ldots, d^1\right] \text{ and } j \in \left[1,\ldots, n_s^1\right]$

以及 $X^{2}_{i,j}, \text{ with } i \in \left[1,\ldots, d^2\right] \text{ and } j \in \left[1,\ldots, n_s^2\right]$

其中： $d = d^1 + d^2 \text{ and } n_s = n_s^1 \times n_s^2 .$

因此，以下等式成立： $X_{i,j} = \begin{cases} X^{1}_{i, k}, & \text{if }\ i < d^1 \text{ with } k = j \mod n_s^1 \\ X^{2}_{i - d^1, l}, & \text{if } i \ge d^1 \text{ with } l = \lfloor\frac{j}{n_s^1}\rfloor \end{cases}$

直观地说，这意味着位置嵌入向量 $x_j \in \mathbb{R}^{d}$ 现在是两个分解后的嵌入向量的组合：$x^1{k, l} + x^2{l, k}$，其中 config.max_embedding_size 维度 $j$ 被分解为 $k \text{ 和 } l$。这种设计确保了每个位置嵌入向量 $x_j$ 都是唯一的。

再次使用上述示例，采用轴向位置编码，其中 $d^1 = 2^9, d^2 = 2^9, n_s^1 = 2^9, n_s^2 = 2^{10}$ 可将参数量从 500,000,000 大幅减少至 $2^{18} + 2^{19} \approx 780 000$ 个参数，这意味着内存占用减少了 85%。

在实际使用中，参数 config.axial_pos_embds_dim 被设置为一个元组 $(d^1, d^2)$ ，其元素之和必须等于 config.hidden_size ，而 config.axial_pos_shape 被设置为一个元组 $(n_s^1, n_s^2)$ ，其元素之积必须等于 config.max_embedding_size, ，在训练期间，该值必须等于 序列长度 ，即 input_ids.

LSH 自注意力

在 Locality sensitive hashing (LSH) self attention 中，key 和 query 的投影权重是共享的。因此，key 和 query 的嵌入向量也是共享的。LSH self attention 使用 Practical and Optimal LSH for Angular Distance 中提出的 locality sensitive hashing 机制，将每个共享的 key 和 query 嵌入向量分配到 config.num_buckets 个可能的桶之一。其基本假设是：key 和 query 嵌入向量之间越“相似”（以 cosine similarity 衡量），它们被分配到同一个桶的可能性就越大。

可以通过增大 config.num_hashes 或直接增大 forward 函数的参数 num_hashes 来提高 LSH 机制的精度，从而使 LSH 自注意力的输出更好地逼近“常规”完整自注意力的输出。随后，这些桶会被排序并切分为查询键嵌入向量块，每个块的长度为 config.lsh_chunk_length。对于每个块，查询嵌入向量会关注其自身的键向量（即与自身绑定），以及前 config.lsh_num_chunks_before 个相邻块的键嵌入向量和后 config.lsh_num_chunks_after 个相邻块的键嵌入向量。

有关更多信息，请参阅原始论文或这篇出色的博客文章。

请注意，config.num_buckets 也可以分解为列表 $(n{\text{buckets}}^1, n{\text{buckets}}^2)$。这样，查询键嵌入向量不再被分配到 $(1,\ldots, n{\text{buckets}})$ 中的一个，而是被分配到 $(1-1,\ldots, n{\text{buckets}}^1-1, \ldots, 1-n{\text{buckets}}^2, \ldots, n{\text{buckets}}^1-n_{\text{buckets}}^2)$ 中的一个。这对于非常长的序列至关重要，可以节省内存。

从头开始训练模型时，建议将 config.num_buckets=None 保留为默认值，这样系统会根据序列长度动态计算 num_buckets 的合适值。该值随后会自动保存到配置中，并应在推理阶段复用。

使用 LSH 自注意力机制，可以将 query-key 矩阵乘法的内存和时间复杂度从 $\mathcal{O}(n_s \times n_s)$ 降低至 $\mathcal{O}(n_s \times \log(n_s))$，这通常是 Transformer 模型中的内存和时间瓶颈，其中 $n_s$ 表示序列长度。

局部自注意力

局部自注意力本质上是一个带有 key、query 和 value 投影的“常规”自注意力层，但被分块处理，使得在长度为 config.local_chunk_length 的每个块中，query 嵌入向量仅关注其所在块内的 key 嵌入向量，以及前 config.local_num_chunks_before 个相邻块和后 config.local_num_chunks_after 个相邻块中的 key 嵌入向量。

使用 Local self attention，可以将 query-key 矩阵乘法操作的内存和时间复杂度从 $\mathcal{O}(n_s \times n_s)$ 降低至 $\mathcal{O}(n_s \times \log(n_s))$，这通常是 Transformer 模型中的内存和时间瓶颈，其中 $n_s$ 表示序列长度。

训练

在训练期间，我们必须确保序列长度设置为可被 config.lsh_chunk_length 和 config.local_chunk_length 的最小公倍数整除的值，并且 Axial Positional Encodings 的参数已按照上述说明正确设置。Reformer 的内存效率非常高，因此可以轻松地在长达 64000 个 token 的序列上训练模型。

在训练时，ReformerModelWithLMHead 应如下使用：

input_ids = tokenizer.encode("This is a sentence from the training data", return_tensors="pt").to(model.device)
loss = model(input_ids, labels=input_ids)[0]

资源

ReformerConfig

class transformers.ReformerConfig

< source >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Noneattention_head_size: int = 64attn_layers: list[str] | tuple[str, ...] = ('local', 'lsh', 'local', 'lsh', 'local', 'lsh')axial_norm_std: float = 1.0axial_pos_embds: bool = Trueaxial_pos_shape: list[int] | tuple[int, ...] = (64, 64)axial_pos_embds_dim: list[int] | tuple[int, ...] = (64, 192)chunk_size_lm_head: int = 0eos_token_id: int | list[int] | None = 2feed_forward_size: int = 512hash_seed: int | None = Nonehidden_act: str = 'relu'hidden_dropout_prob: float | int = 0.05hidden_size: int = 256initializer_range: float = 0.02is_decoder: bool = Falselayer_norm_eps: float = 1e-12local_num_chunks_before: int = 1local_num_chunks_after: int = 0local_attention_probs_dropout_prob: float | int = 0.05local_attn_chunk_length: int = 64lsh_attn_chunk_length: int | None = 64lsh_attention_probs_dropout_prob: float | None = 0.0lsh_num_chunks_before: int | None = 1lsh_num_chunks_after: int | None = 0max_position_embeddings: int = 4096num_attention_heads: int = 12num_buckets: int | list[int] | None = Nonenum_hashes: int = 1vocab_size: int = 320tie_word_embeddings: bool = Falseuse_cache: bool = Trueclassifier_dropout: float | int | None = Nonebos_token_id: int | None = Nonepad_token_id: int | None = 0 )

参数

attention_head_size (int，可选，默认为 64) — 投影后的 key、query 和 value 向量的维度
attn_layers (list[str]，可选，默认为 ["local", "lsh", "local", "lsh", "local", "lsh"]) — 按升序排列的注意力层类型列表。可在 LSHSelfAttention 层（"lsh"）和 LocalSelfAttention 层（"local"）之间进行选择。有关 LSHSelfAttention 层的更多信息，请参阅 LSH Self Attention。有关 LocalSelfAttention 层的更多信息，请参阅 Local Self Attention。
axial_norm_std (float，可选，默认为 1.0) — 用于初始化轴向位置编码权重矩阵的 normal_initializer 的标准差。
axial_pos_embds (bool，可选，默认为 True) — 是否使用轴向位置嵌入。有关轴向位置嵌入工作原理的更多信息，请参阅 Axial Position Encodings。
axial_pos_shape (list[int]，可选，默认为 [64, 64]) — 轴向位置编码的位置维度。在训练期间，位置维度的乘积必须等于序列长度。有关轴向位置嵌入工作原理的更多信息，请参阅轴向位置编码。
axial_pos_embds_dim (list[int]，可选，默认为 [64, 192]) — 轴向位置编码的嵌入维度。嵌入维度之和必须等于隐藏层大小。有关轴向位置编码工作原理的更多信息，请参阅轴向位置编码。
chunk_size_lm_head (int, 可选，默认为 0) — 最终语言模型前馈头层的分块大小。分块大小为 0 表示前馈层不进行分块。分块大小为 n 表示前馈层每次处理 n < sequence_length 个嵌入。有关前馈分块的更多信息，请参阅前馈分块是如何工作的？。
eos_token_id (Union[int, list[int]]，可选，默认为 2) — 词汇表中用于表示流结束（end-of-stream）的 token id。
feed_forward_size (int，可选，默认为 512) — 残差注意力块中前馈层的维度。
hash_seed (int, 可选) — 可用于使 LSHSelfAttention 中的局部敏感哈希具有确定性的种子。此参数仅应在测试目的下设置。对于评估和训练目的，应将 hash_seed 保留为 None，以确保局部敏感哈希方案中的旋转完全随机。
hidden_act (str，可选，默认为 relu) — 解码器中的非线性激活函数（函数或字符串）。例如，"gelu"、 "relu"、"silu" 等。
hidden_dropout_prob (Union[float, int]，可选，默认为 0.05) — 嵌入层、编码器和池化层中所有全连接层的 dropout 概率。
hidden_size (int，可选，默认为 256) — 隐藏表示的维度。
initializer_range (float，可选，默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
is_decoder (bool，可选，默认为 False) — 指示该模型是否用作解码器。如果为 False，则模型用作编码器。
layer_norm_eps (float，可选，默认为 1e-12) — 层归一化层所使用的 epsilon 值。
local_num_chunks_before (int，可选，默认为 1) — 在 LocalSelfAttention 层中，自身需要关注的前一个相邻块的数量。
local_num_chunks_after (int，可选，默认为 0) — 在 LocalSelfAttention 层中，除自身外，还需关注其后的相邻块数量。
local_attention_probs_dropout_prob (float，可选，默认值为 0.1) — LocalSelfAttention 中注意力概率的 dropout 比率。
local_attn_chunk_length (int，可选，默认为 64) — 局部注意力层中每个块（chunk）的长度。
lsh_attn_chunk_length (int，可选，默认值为 64) — LSHSelfAttention 中关注自身的块长度。分块将内存复杂度从序列长度 x 序列长度（自注意力）降低为块长度 x 块长度 x 序列长度 / 块长度（分块自注意力）。
lsh_attention_probs_dropout_prob (float, 可选，默认为 0.1) — LSHSelfAttention 中注意力概率的 dropout 比率。
lsh_num_chunks_before (int, 可选，默认为 1) — 在 LSHSelfAttention 层中，自身需要关注的前一个相邻块的数量。
lsh_num_chunks_after (int，可选，默认为 0) — 在 LSHSelfAttention 层中，自身需要关注的后续相邻块的数量。
max_position_embeddings (int，可选，默认为 4096) — 该模型可能使用的最大序列长度。
num_attention_heads (int，可选，默认为 12) — Transformer 解码器中每个注意力层的注意力头数量。
num_buckets (int 或 list[int]，可选) — 桶的数量，即使用局部敏感哈希方案将关键查询向量“哈希到”的桶数。每个查询关键向量会被哈希到 1, ..., num_buckets 中的一个哈希值。桶的数量也可以分解为一个列表，以改善内存复杂度。在这种情况下，如果 num_buckets 被分解为两个因子，则每个查询关键向量会被哈希到 1-1, 1-2, ..., num_buckets[0]-1, ..., num_buckets[0]-num_buckets[1] 中的一个哈希值。桶的数量（或各因子的乘积）应大致等于序列长度 / lsh_chunk_length。如果未设置 num_buckets，则会在运行时计算一个合适的值。
num_hashes (int, 可选，默认为 1) — 局部敏感哈希方案中的哈希轮数（例如随机旋转的次数）。num_hashes 越高，LSHSelfAttention 的精度越高，但哈希过程所需的内存和时间也越多。
vocab_size (int，可选，默认为 320) — 模型的词汇表大小。定义了 input_ids 可以表示的不同 token 数量。
tie_word_embeddings (bool，可选，默认为 False) — 是否根据模型的 tied_weights_keys 映射来绑定权重嵌入。
use_cache (bool，可选，默认为 True) — 模型是否应返回最后一次键/值注意力（并非所有模型都使用）。仅当 config.is_decoder=True 或模型为仅解码器生成式模型时相关。
classifier_dropout (Union[float, int], 可选) — 分类器的 dropout 比率。
bos_token_id (int, 可选) — 词汇表中用于表示流起始的 Token id。
pad_token_id (int, 可选，默认为 0) — 词表中用于填充的 token id。

这是用于存储 ReformerModel 配置的类。它根据指定的参数实例化 Reformer 模型，并定义模型架构。使用默认值实例化配置，将得到与 google/reformer-crime-and-punishment 类似的配置。

配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。

示例：

>>> from transformers import ReformerConfig, ReformerModel

>>> # Initializing a Reformer configuration
>>> configuration = ReformerConfig()

>>> # Initializing a Reformer model (with random weights)
>>> model = ReformerModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

ReformerTokenizer

class transformers.ReformerTokenizer

< source >

参数

vocab_file (str) — 包含实例化分词器所需词汇表的 SentencePiece 文件（通常具有 .spm 扩展名）。
eos_token (str, 可选, 默认为 "</s>") — 序列结束标记。在构建包含特殊标记的序列时，序列结束使用的并非此标记，而是 sep_token。

在构建包含特殊标记的序列时，序列末尾使用的并不是这个标记。实际使用的是 sep_token。
unk_token (str，可选，默认为 "<unk>") — 未知标记。不在词汇表中的标记无法转换为 ID，因此会被设置为该标记。
pad_token (str，可选，默认为 "<pad>") — 用于填充的标记，例如在对长度不同的序列进行批处理时。
additional_special_tokens (list[str]，可选) — 分词器使用的额外特殊标记。
vocab (str 或 dict[str, int]，可选) — 自定义词汇表。如果未提供，则从 vocab_file 加载词汇表。
merges (str 或 list[str]，可选) — 自定义 merges 列表。如果未提供，则从 vocab_file 中加载 merges。

构建一个 Reformer 分词器（基于 HuggingFace 的 tokenizers 库）。基于 BPE。

该分词器继承自 TokenizersBackend，其中包含了大部分主要方法。用户应参考该父类以获取有关这些方法的更多信息。

save_vocabulary

< source >

( save_directory: strfilename_prefix: str | None = None )

ReformerModel

class transformers.ReformerModel

< source >

( config: ReformerConfig )

参数

config (ReformerConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

未附加任何特定头部的裸 Reformer 模型，直接输出原始隐藏状态。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< source >

( input_ids: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonenum_hashes: int | None = Nonepast_buckets_states: transformers.models.reformer.modeling_reformer.ReformerDynamicCache | None = Noneuse_cache: bool | None = Noneoutput_hidden_states: bool | None = Noneoutput_attentions: bool | None = Nonereturn_dict: bool | None = None**kwargs ) → ReformerModelOutput or tuple(torch.Tensor)

参数

input_ids (torch.LongTensor 形状为 (batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。训练时，input_ids 的 sequence_length 必须是相关模型分块长度（lsh、local 或两者）的整数倍。评估时，索引会自动填充至分块长度的整数倍。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
什么是注意力掩码？
position_ids (torch.Tensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].
什么是位置 ID？
inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
num_hashes (int, 可选) — 分桶过程中应执行的哈希轮数。设置此参数会覆盖 config.num_hashes.
有关更多信息，请参阅 ReformerConfig 中的 num_hashes。
past_buckets_states (ReformerDynamicCache, 可选) — 的列表 tuple(torch.LongTensor, torch.Tensor 长度为 config.n_layers, ，其中第一个元素是之前的 buckets ，形状为 (batch_size, num_heads, num_hashes, sequence_length)) ，第二个元素是之前的 hidden_states ，形状为 (batch_size, sequence_length, hidden_size)).
包含预计算的隐藏状态和桶（仅对 LSH Self-Attention 相关）。可用于加速顺序解码。
use_cache (bool，可选) — 如果设置为 True，则会返回 past_key_values 键值状态，可用于加速解码（参见 past_key_values）。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量中的 hidden_states。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息，请参阅返回张量中的 attentions。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通元组。

返回值

ReformerModelOutput 或 tuple(torch.Tensor)

ReformerModelOutput 或 torch.Tensor 元组（当传入 return_dict=False 或 config.return_dict=False 时），包含根据配置（ReformerConfig）和输入而定的各种元素。

ReformerModel 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

last_hidden_state（形状为 (batch_size, num_predict, hidden_size) 的 torch.Tensor）— 模型最后一层的隐藏状态序列。

num_predict 对应于 target_mapping.shape[1]。如果 target_mapping 为 None，则 num_predict 对应于 sequence_length。
past_buckets_states (list[tuple(torch.LongTensor, torch.Tensor)]，可选，在传入 use_cache=True 或 config.use_cache=True 时返回) — 一个长度为 config.n_layers 的 tuple(torch.LongTensor, torch.Tensor 列表，其中第一个元素是形状为 (batch_size, num_heads, num_hashes, sequence_length) 的先前 buckets，第二个元素是形状为 (batch_size, sequence_length, hidden_size) 的先前 hidden_states。

包含预计算的桶和隐藏状态，可用于（参见 past_buckets_states 输入）加速顺序解码。
hidden_states (tuple[torch.Tensor]，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层输出的一个张量，以及每个层输出的一个张量），形状为 (batch_size, sequence_length, hidden_size)。

模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (tuple[torch.Tensor]，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

ReformerModelWithLMHead

class transformers.ReformerModelWithLMHead

< source >

( config: ReformerConfig )

参数

config (ReformerConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

带有 language modeling 头的 Reformer 模型。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< source >

( input_ids: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonenum_hashes: int | None = Nonepast_buckets_states: list[tuple[torch.Tensor]] | None = Noneuse_cache: bool | None = Noneoutput_hidden_states: bool | None = Noneoutput_attentions: bool | None = Nonereturn_dict: bool | None = Nonelabels: typing.Optional[torch.Tensor] = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs ) → CausalLMOutput or tuple(torch.Tensor)

参数

input_ids (torch.LongTensor 形状为 (batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。训练时，input_ids 的 sequence_length 必须是相关模型分块长度（lsh、local 或两者）的整数倍。评估时，索引会自动填充至分块长度的整数倍。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
position_ids (torch.Tensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].
什么是位置 ID？
attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
什么是注意力掩码？
inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
num_hashes (int, 可选) — 分桶过程中应执行的哈希轮数。设置此参数会覆盖 config.num_hashes.
有关更多信息，请参阅 ReformerConfig 中的 num_hashes。
past_buckets_states (list[tuple(torch.LongTensor, torch.Tensor)], 可选) — 一个 tuple(torch.LongTensor, torch.Tensor 长度为 config.n_layers, 的列表，其中第一个元素是之前的 buckets ，形状为 (batch_size, num_heads, num_hashes, sequence_length)) ，第二个元素是之前的 hidden_states ，形状为 (batch_size, sequence_length, hidden_size)).
包含预计算的隐藏状态和桶（仅对 LSH Self-Attention 相关）。可用于加速顺序解码。
use_cache (bool，可选) — 如果设置为 True，则会返回 past_key_values 键值状态，可用于加速解码（参见 past_key_values）。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量中的 hidden_states。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息，请参阅返回张量中的 attentions。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通元组。
labels（形状为 (batch_size,) 的 torch.LongTensor，可选） — 用于计算序列分类/回归损失的标签。索引应在 [-100, 0, ..., config.vocab_size - 1] 范围内。所有设置为 -100 的标签都会被忽略（被掩码），损失仅针对 [0, ..., config.vocab_size] 范围内的标签进行计算。
logits_to_keep (Union[int, torch.Tensor], 可选，默认为 0) — 如果为 int，则计算最后 logits_to_keep 个 token 的 logits。如果为 0，则计算所有 input_ids 的 logits（特殊情况）。生成过程只需要最后一个 token 的 logits，仅针对该 token 进行计算可以节省内存，对于长序列或大词汇量而言，这种节省相当显著。如果为 torch.Tensor，则必须是一维张量，对应于序列长度维度中需要保留的索引。这在使用打包张量格式（将批次和序列长度合并为单一维度）时非常有用。

返回值

CausalLMOutput 或 tuple(torch.Tensor)

一个 CausalLMOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（ReformerConfig）和输入。

ReformerModelWithLMHead 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

loss (torch.Tensor，形状为 (1,)，可选，当提供 labels 时返回) — 语言建模损失（用于下一个 token 预测）。
logits（形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.Tensor）— 语言建模头的预测分数（SoftMax 之前每个词汇表 token 的分数）。
hidden_states (tuple(torch.Tensor)，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为 (batch_size, sequence_length, hidden_size)。

模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (tuple(torch.Tensor)，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

示例：

>>> import torch
>>> from transformers import AutoTokenizer, ReformerModelWithLMHead

>>> tokenizer = AutoTokenizer.from_pretrained("google/reformer-crime-and-punishment")
>>> model = ReformerModelWithLMHead.from_pretrained("google/reformer-crime-and-punishment")

>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs, labels=inputs["input_ids"])
>>> loss = outputs.loss
>>> logits = outputs.logits

ReformerForMaskedLM

class transformers.ReformerForMaskedLM

< source >

( config: ReformerConfig )

参数

config (ReformerConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

顶部带有 language modeling 头的 Reformer 模型。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< source >

( input_ids: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonenum_hashes: int | None = Nonelabels: typing.Optional[torch.Tensor] = Noneoutput_hidden_states: bool | None = Noneoutput_attentions: bool | None = Nonereturn_dict: bool | None = None**kwargs ) → MaskedLMOutput or tuple(torch.Tensor)

参数

input_ids (torch.LongTensor 形状为 (batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。训练时，input_ids 的 sequence_length 必须是相关模型分块长度（lsh、local 或两者）的整数倍。评估时，索引会自动填充至分块长度的整数倍。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
position_ids (torch.Tensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].
什么是位置 ID？
attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
什么是注意力掩码？
inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
num_hashes (int, 可选) — 分桶过程中应执行的哈希轮数。设置此参数会覆盖 config.num_hashes.
有关更多信息，请参阅 ReformerConfig 中的 num_hashes。
labels (torch.LongTensor 形状为 (batch_size, sequence_length), ，可选) — 用于计算掩码语言建模损失的标签。索引应在 [-100, 0, ..., config.vocab_size] (参见 input_ids 的文档字符串）索引设置为 -100 的 token 将被忽略（掩码），损失仅针对标签为

本示例使用了一个虚假的检查点，因为我们没有可用于 Reformer 架构的掩码语言建模任务的预训练模型。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量中的 hidden_states。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息，请参阅返回张量中的 attentions。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通元组。

返回值

MaskedLMOutput 或 tuple(torch.Tensor)

一个 MaskedLMOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（ReformerConfig）和输入。

ReformerForMaskedLM 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

loss (torch.Tensor，形状为 (1,)，可选，当提供 labels 时返回) — 掩码语言建模（MLM）损失。
logits（形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.Tensor）— 语言建模头的预测分数（SoftMax 之前每个词汇表 token 的分数）。
hidden_states (tuple(torch.Tensor)，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为 (batch_size, sequence_length, hidden_size)。

模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (tuple(torch.Tensor)，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

示例：

>>> import torch
>>> from transformers import AutoTokenizer, ReformerForMaskedLM

>>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-reformer")
>>> model = ReformerForMaskedLM.from_pretrained("hf-internal-testing/tiny-random-reformer")

>>> # add mask_token
>>> tokenizer.add_special_tokens({"mask_token": "[MASK]"})
>>> inputs = tokenizer("The capital of France is [MASK].", return_tensors="pt")

>>> # resize model's embedding matrix
>>> model.resize_token_embeddings(new_num_tokens=model.config.vocab_size + 1)
>>> with torch.no_grad():
...     logits = model(**inputs).logits

>>> # retrieve index of [MASK]
>>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]

>>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
>>> predicted_token = tokenizer.decode(predicted_token_id)

>>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"]
>>> # mask labels of non-[MASK] tokens
>>> labels = torch.where(
...     inputs.input_ids == tokenizer.mask_token_id, labels[:, : inputs["input_ids"].shape[-1]], -100
... )

>>> outputs = model(**inputs, labels=labels)
>>> loss = round(outputs.loss.item(), 2)

ReformerForSequenceClassification

class transformers.ReformerForSequenceClassification

< source >

( config: ReformerConfig )

参数

config (ReformerConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

Reformer 模型，顶部带有用于序列分类/回归任务的头部（在池化输出之上添加一个线性层），例如用于 GLUE 任务。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< source >

( input_ids: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonenum_hashes: int | None = Nonelabels: typing.Optional[torch.Tensor] = Noneoutput_hidden_states: bool | None = Noneoutput_attentions: bool | None = Nonereturn_dict: bool | None = None**kwargs ) → SequenceClassifierOutput or tuple(torch.Tensor)

参数

input_ids (torch.LongTensor 形状为 (batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。训练时，input_ids 的 sequence_length 必须是相关模型分块长度（lsh、local 或两者）的整数倍。评估时，索引会自动填充至分块长度的整数倍。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
position_ids (torch.Tensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].
什么是位置 ID？
attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
什么是注意力掩码？
inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
num_hashes (int, 可选) — 分桶过程中应执行的哈希轮数。设置此参数会覆盖 config.num_hashes.
有关更多信息，请参阅 ReformerConfig 中的 num_hashes。
labels (torch.LongTensor，形状为 (batch_size,)，可选) — 用于计算序列分类/回归损失的标签。索引应在 [0, ..., config.num_labels - 1] 范围内。如果 config.num_labels == 1，则计算回归损失（均方损失）；如果 config.num_labels > 1，则计算分类损失（交叉熵）。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量中的 hidden_states。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息，请参阅返回张量中的 attentions。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通元组。

返回值

SequenceClassifierOutput 或 tuple(torch.Tensor)

一个 SequenceClassifierOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（ReformerConfig）和输入。

ReformerForSequenceClassification 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

loss (torch.Tensor，形状为 (1,)，可选，当提供 labels 时返回) — 分类损失（如果 config.num_labels==1，则为回归损失）。
logits（形状为 (batch_size, config.num_labels) 的 torch.Tensor）— 分类（或当 config.num_labels==1 时为回归）得分（SoftMax 之前）。
hidden_states (tuple(torch.Tensor)，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为 (batch_size, sequence_length, hidden_size)。

模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (tuple(torch.Tensor)，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

单标签分类示例：

>>> import torch
>>> from transformers import AutoTokenizer, ReformerForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("google/reformer-crime-and-punishment")
>>> model = ReformerForSequenceClassification.from_pretrained("google/reformer-crime-and-punishment")

>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")

>>> with torch.no_grad():
...     logits = model(**inputs).logits

>>> predicted_class_id = logits.argmax().item()
>>> label = model.config.id2label[predicted_class_id]

>>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
>>> num_labels = len(model.config.id2label)
>>> model = ReformerForSequenceClassification.from_pretrained(
...     "google/reformer-crime-and-punishment", num_labels=num_labels
... )

>>> labels = torch.tensor(1)
>>> loss = model(**inputs, labels=labels).loss

ReformerForQuestionAnswering

class transformers.ReformerForQuestionAnswering

< source >

( config: ReformerConfig )

参数

config (ReformerConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

Reformer 变压器，其顶部带有用于抽取式问答任务（如 SQuAD）的跨度分类头（在隐藏状态输出之上添加一个线性层，用于计算 span start logits 和 span end logits）。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< source >

( input_ids: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonenum_hashes: int | None = Nonestart_positions: typing.Optional[torch.Tensor] = Noneend_positions: typing.Optional[torch.Tensor] = Noneoutput_hidden_states: bool | None = Noneoutput_attentions: bool | None = Nonereturn_dict: bool | None = None**kwargs ) → QuestionAnsweringModelOutput or tuple(torch.Tensor)

参数

input_ids (torch.LongTensor 形状为 (batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。训练时，input_ids 的 sequence_length 必须是相关模型分块长度（lsh、local 或两者）的整数倍。评估时，索引会自动填充至分块长度的整数倍。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
position_ids (torch.Tensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].
什么是位置 ID？
attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
什么是注意力掩码？
inputs_embeds（形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。
num_hashes (int, 可选) — 分桶过程中应执行的哈希轮数。设置此参数会覆盖 config.num_hashes.
有关更多信息，请参阅 ReformerConfig 中的 num_hashes。
start_positions (torch.Tensor，形状为 (batch_size,)，可选) — 用于计算 token 分类损失的标注片段起始位置（索引）的标签。位置会被截断到序列长度（sequence_length）。序列之外的位置在计算损失时不予考虑。
end_positions (torch.Tensor，形状为 (batch_size,)，可选) — 用于计算 token 分类损失的标注片段结束位置（索引）的标签。位置会被截断到序列长度（sequence_length）。序列之外的位置在计算损失时不予考虑。
output_hidden_states (bool，可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量中的 hidden_states。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息，请参阅返回张量中的 attentions。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通元组。

返回值

QuestionAnsweringModelOutput 或 tuple(torch.Tensor)

一个 QuestionAnsweringModelOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（ReformerConfig）和输入。

ReformerForQuestionAnswering 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义，但之后应调用 Module 实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。

loss (torch.Tensor，形状为 (1,)，可选，当提供 labels 时返回) — 总跨度提取损失是起始位置和结束位置的交叉熵之和。
start_logits（形状为 (batch_size, sequence_length) 的 torch.Tensor）— 跨度起始分数（SoftMax 之前）。
end_logits（形状为 (batch_size, sequence_length) 的 torch.Tensor）— 跨度结束分数（SoftMax 之前）。
hidden_states (tuple(torch.Tensor)，可选，当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为 (batch_size, sequence_length, hidden_size)。

模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (tuple(torch.Tensor)，可选，当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组（每层一个），形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。

示例：

>>> from transformers import AutoTokenizer, ReformerForQuestionAnswering
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("google/reformer-crime-and-punishment")
>>> model = ReformerForQuestionAnswering.from_pretrained("google/reformer-crime-and-punishment")

>>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"

>>> inputs = tokenizer(question, text, return_tensors="pt")
>>> with torch.no_grad():
...     outputs = model(**inputs)

>>> answer_start_index = outputs.start_logits.argmax()
>>> answer_end_index = outputs.end_logits.argmax()

>>> predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1]
>>> tokenizer.decode(predict_answer_tokens, skip_special_tokens=True)
...

>>> # target is "nice puppet"
>>> target_start_index = torch.tensor([14])
>>> target_end_index = torch.tensor([15])

>>> outputs = model(**inputs, start_positions=target_start_index, end_positions=target_end_index)
>>> loss = outputs.loss
>>> round(loss.item(), 2)
...

Update on GitHub

←RecurrentGemma RemBERT→