Transformers 文档

RWKV

Transformers

加入 Hugging Face 社区

并获得增强的文档体验

在模型、数据集和 Spaces 上进行协作

通过加速推理获得更快的示例

切换文档主题

开始使用

此模型于2022-08-17发布，并于2023-05-09加入Hugging Face Transformers。

RWKV

概述

RWKV 模型（版本 4）是在此仓库中提出的。

它建议对传统 Transformer 的注意力机制进行调整，使其变为线性的。这样，模型就可以用作循环网络：同时输入时间戳0和时间戳1的数据与先输入时间戳0的数据，然后输入时间戳1的数据以及时间戳0的状态是相同的（参见下面的示例）。

这可能比常规的 Transformer 更高效，并且可以处理任意长度的句子（即使模型在训练时使用了固定的上下文长度）。

此模型由 sgugger 贡献。原始代码可在此处找到。

用法示例

import torch
from transformers import AutoTokenizer, RwkvConfig, RwkvModel

model = RwkvModel.from_pretrained("sgugger/rwkv-430M-pile")
tokenizer = AutoTokenizer.from_pretrained("sgugger/rwkv-430M-pile")

inputs = tokenizer("This is an example.", return_tensors="pt")
# Feed everything to the model
outputs = model(inputs["input_ids"])
output_whole = outputs.last_hidden_state

outputs = model(inputs["input_ids"][:, :2])
output_one = outputs.last_hidden_state

# Using the state computed on the first inputs, we will get the same output
outputs = model(inputs["input_ids"][:, 2:], state=outputs.state)
output_two = outputs.last_hidden_state

torch.allclose(torch.cat([output_one, output_two], dim=1), output_whole, atol=1e-5)

如果您想确保模型在检测到 `'\n\n'` 时停止生成，我们建议使用以下停止标准

from transformers import StoppingCriteria

class RwkvStoppingCriteria(StoppingCriteria):
    def __init__(self, eos_sequence = [187,187], eos_token_id = 537):
        self.eos_sequence = eos_sequence
        self.eos_token_id = eos_token_id

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        last_2_ids = input_ids[:,-2:].tolist()
        return self.eos_sequence in last_2_ids


output = model.generate(inputs["input_ids"], max_new_tokens=64, stopping_criteria = [RwkvStoppingCriteria()])

RwkvConfig

class transformers.RwkvConfig

< source >

( vocab_size = 50277 context_length = 1024 hidden_size = 4096 num_hidden_layers = 32 attention_hidden_size = None intermediate_size = None layer_norm_epsilon = 1e-05 bos_token_id = 0 eos_token_id = 0 rescale_every = 6 tie_word_embeddings = False use_cache = True **kwargs )

参数

vocab_size (int, 可选, 默认为 50277) — RWKV 模型的词汇表大小。定义了调用 RwkvModel 时传递的 inputs_ids 可以表示的不同词元数量。
context_length (int, 可选, 默认为 1024) — 此模型在单次正向传播中可以使用的最大序列长度（在 RNN 模式下，可以使用任意序列长度）。
hidden_size (int, 可选, 默认为 4096) — 嵌入和隐藏状态的维度。
num_hidden_layers (int, 可选, 默认为 32) — 模型中的隐藏层数量。
attention_hidden_size (int, 可选) — 注意力隐藏状态的维度。如果未设置，将默认为 hidden_size。
intermediate_size (int, 可选) — 内部前馈层的维度。如果未设置，将默认为 hidden_size 的 4 倍。
layer_norm_epsilon (float, 可选, 默认为 1e-05) — 在层归一化层中使用的 epsilon 值。
bos_token_id (int, 可选, 默认为 0) — 词汇表中句子开始词元的 ID。默认为 0，因为 RWKV 使用与 GPTNeoX 相同的分词器。
eos_token_id (int, 可选, 默认为 0) — 词汇表中句子结束词元的 ID。默认为 0，因为 RWKV 使用与 GPTNeoX 相同的分词器。
rescale_every (int, 可选, 默认为 6) — 在推理时，每隔 rescale_every 层，隐藏状态（以及相应输出层的权重）都会除以 2。如果设置为 0 或负数，则不进行缩放。
tie_word_embeddings (bool, 可选, 默认为 False) — 是否将词嵌入与输入词元嵌入绑定。
use_cache (bool, 可选, 默认为 True) — 模型是否应返回最后一个状态。

这是用于存储 RwkvModel 配置的配置类。它用于根据指定参数实例化 RWKV 模型，定义模型架构。使用默认值实例化配置将产生与 RWVK-4 RWKV/rwkv-4-169m-pile 架构类似的配置。

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

示例

>>> from transformers import RwkvConfig, RwkvModel

>>> # Initializing a Rwkv configuration
>>> configuration = RwkvConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = RwkvModel(configuration)

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

RwkvModel

class transformers.RwkvModel

< source >

( config model_args: ~utils.generic.ModelArgs | None = None adapter_args: ~utils.generic.AdapterArgs | None = None lora_args: ~utils.generic.LoRAArgs | None = None tokenizer_args: ~utils.generic.TokenizerArgs | None = None dataset_args: ~utils.generic.DatasetArgs | None = None data_args: ~utils.generic.DataArgs | None = None training_args: ~utils.generic.TrainingArgs | None = None generation_args: ~utils.generic.GenerationArgs | None = None vision_tower_args: ~utils.generic.VisionTowerArgs | None = None qlora_args: ~utils.generic.QLoRAArgs | None = None vision_tower_template_args: ~utils.generic.VisionTowerTemplateArgs | None = None video_tower_args: ~utils.generic.VideoTowerArgs | None = None vision_config: ~utils.generic.VisionConfig | None = None video_config: ~utils.generic.VideoConfig | None = None load_dataset: bool | None = None load_data_collator: bool | None = None load_processor: bool | None = None load_lora_adapter: bool | None = None load_adapter: bool | None = None load_qlora_adapter: bool | None = None **kwargs: typing_extensions.Unpack[transformers.modeling_utils.PreTrainedModelKwargs] )

参数

config (RwkvModel) — 包含模型所有参数的模型配置类。用配置文件初始化不会加载与模型相关的权重，只加载配置。请查阅 from_pretrained() 方法以加载模型权重。

裸 Rwkv 模型，输出原始隐藏状态，不带任何特定头部。

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

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它，并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< source >

( input_ids: torch.LongTensor | None = None attention_mask: torch.LongTensor | None = None inputs_embeds: torch.FloatTensor | None = None state: list[torch.FloatTensor] | None = None use_cache: bool | None = None output_attentions: bool | None = None output_hidden_states: bool | None = None return_dict: bool | None = None **kwargs ) → transformers.models.rwkv.modeling_rwkv.RwkvOutput or tuple(torch.FloatTensor)

参数

input_ids (torch.LongTensor 形状为 (batch_size, input_ids_length)) — input_ids_length = sequence_length 如果 past_key_values 为 None，否则为 past_key_values.get_seq_length() (输入过去键值状态的 sequence_length)。词汇表中输入序列词元的索引。

如果使用了 past_key_values，则只有未计算过其过去的 input_ids 应作为 input_ids 传递。

索引可以使用 AutoTokenizer 获取。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
attention_mask (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 用于避免对填充词元索引执行注意力的掩码。掩码值选择范围为 [0, 1]：
- 1 表示未被掩盖的词元，
- 0 表示被掩盖的词元。
什么是注意力掩码？
inputs_embeds (torch.FloatTensor 形状为 (batch_size, sequence_length, hidden_size), 可选) — 可选地，您可以选择直接传递嵌入表示，而不是传递 input_ids。如果您想对如何将 input_ids 索引转换为相关向量有更多控制，而不是模型内部的嵌入查找矩阵，这会很有用。
state (tuple，包含五个 torch.FloatTensor，形状为 (batch_size, hidden_size, num_hidden_layers), 可选) — 如果传入，模型将在所有块中使用先前的状态（这将产生提供的 input_ids 的输出，如同模型将 state_input_ids + input_ids 作为上下文）。
use_cache (bool, 可选) — 如果设置为 True，将返回最后一个状态，可用于快速生成下一个 logits。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息，请参阅返回张量下的 attentions。
output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量下的 hidden_states。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通的元组。

transformers.models.rwkv.modeling_rwkv.RwkvOutput 或 tuple(torch.FloatTensor)

一个 transformers.models.rwkv.modeling_rwkv.RwkvOutput 或一个 torch.FloatTensor 元组（如果传递了 return_dict=False 或当 config.return_dict=False 时），其中包含根据配置 (RwkvConfig) 和输入而定的各种元素。

last_hidden_state (torch.FloatTensor | None.last_hidden_state, shape (batch_size, sequence_length, hidden_size), 默认为 None) — 模型最后一层的输出的隐藏状态序列。
state (list，包含五个 torch.FloatTensor，形状为 (batch_size, hidden_size, num_hidden_layers)) — 模型在最后一个时间步的状态。可在后续的前向方法中使用，与下一个 input_ids 一起，以避免提供旧的 input_ids。
hidden_states (tuple[torch.FloatTensor, ...] | None.hidden_states, 当传入 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — 形状为 (batch_size, sequence_length, hidden_size) 的 torch.FloatTensor 元组（一个用于嵌入层的输出，如果模型有嵌入层，+ 每个层的输出）。

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

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

RwkvModel 的前向方法，覆盖了 __call__ 特殊方法。

虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用 Module 实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。

RwkvLMHeadModel

class transformers.RwkvForCausalLM

< source >

参数

config (RwkvForCausalLM) — 包含模型所有参数的模型配置类。用配置文件初始化不会加载与模型相关的权重，只加载配置。请查阅 from_pretrained() 方法以加载模型权重。

带有语言模型头部（权重与输入嵌入绑定的线性层）的 RWKV 模型 Transformer。

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

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它，并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< source >

( input_ids: torch.LongTensor | None = None attention_mask: torch.LongTensor | None = None inputs_embeds: torch.FloatTensor | None = None state: list[torch.FloatTensor] | None = None labels: torch.LongTensor | None = None use_cache: bool | None = None output_attentions: bool | None = None output_hidden_states: bool | None = None return_dict: bool | None = None logits_to_keep: int | torch.Tensor = 0 **kwargs ) → transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutput 或 tuple(torch.FloatTensor)

参数

input_ids (形状为 (batch_size, input_ids_length) 的 torch.LongTensor) — input_ids_length = sequence_length 如果 past_key_values 为 None，否则为 past_key_values.get_seq_length()（输入过去键值状态的 sequence_length）。词汇表中输入序列 token 的索引。

如果使用了 past_key_values，则只有未计算过过去值的 input_ids 应作为 input_ids 传递。

索引可以使用 AutoTokenizer 获取。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是 input ID？
attention_mask (形状为 (batch_size, sequence_length) 的 torch.LongTensor, 可选) — 掩码，用于避免对填充 token 索引执行注意力计算。掩码值选择范围为 [0, 1]:
- 1 表示未被掩码的 token，
- 0 表示被掩码的 token。
什么是注意力掩码？
inputs_embeds (形状为 (batch_size, sequence_length, hidden_size) 的 torch.FloatTensor, 可选) — 可选地，您可以选择直接传递嵌入表示，而不是传递 input_ids。如果您想对如何将 input_ids 索引转换为关联向量有比模型内部嵌入查找矩阵更强的控制，这将很有用。
state (形状为 (batch_size, hidden_size, num_hidden_layers) 的五个 torch.FloatTensor 的 tuple, 可选) — 如果传入，模型将在所有块中使用先前的状态（这将根据提供的 input_ids，如同模型将 state_input_ids + input_ids 作为上下文一样，给出输出）。
labels (形状为 (batch_size, sequence_length) 的 torch.LongTensor, 可选) — 用于语言建模的标签。请注意，标签在模型内部会进行偏移，即您可以设置 labels = input_ids。索引选择范围为 [-100, 0, ..., config.vocab_size]。所有设置为 -100 的标签都将被忽略（掩码），损失仅针对 [0, ..., config.vocab_size] 范围内的标签进行计算。
use_cache (bool, 可选) — 如果设置为 True，则返回最后一个状态，可用于快速生成下一个 logits。
output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息，请参阅返回张量中的 attentions。
output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。有关详细信息，请参阅返回张量中的 hidden_states。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通的元组。
logits_to_keep (Union[int, torch.Tensor], 可选, 默认为 0) — 如果是 int，则计算最后 logits_to_keep 个 token 的 logits。如果是 0，则计算所有 input_ids 的 logits（特殊情况）。生成时只需要最后一个 token 的 logits，只计算该 token 可以节省内存，这对于长序列或大词汇量来说非常重要。如果是 torch.Tensor，则必须是 1D，对应于序列长度维度中要保留的索引。这在使用打包张量格式（批次和序列长度为单维度）时很有用。

transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutput 或 tuple(torch.FloatTensor)

一个 transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutput 或一个 torch.FloatTensor 的元组（如果传递 return_dict=False 或 config.return_dict=False），包含根据配置 (RwkvConfig) 和输入而定的各种元素。

loss (torch.FloatTensor 形状为 (1,)，可选，当提供 labels 时返回) — 语言建模损失（用于下一个 token 预测）。
logits (形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.FloatTensor) — 语言建模头部的预测分数（SoftMax 之前的每个词汇标记的分数）。
state (list，包含五个 torch.FloatTensor，形状为 (batch_size, hidden_size, num_hidden_layers)) — 模型在最后一个时间步的状态。可在后续的前向方法中使用，与下一个 input_ids 一起，以避免提供旧的 input_ids。
hidden_states (tuple[torch.FloatTensor, ...] | None.hidden_states, 当传入 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — 形状为 (batch_size, sequence_length, hidden_size) 的 torch.FloatTensor 元组（一个用于嵌入层的输出，如果模型有嵌入层，+ 每个层的输出）。

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

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

RwkvForCausalLM 的前向方法，覆盖了 __call__ 特殊方法。

虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用 Module 实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。

示例

RWKV 注意力机制与循环公式

在传统的自回归 Transformer 中，注意力机制表示为 $O = \hbox{softmax}(QK^{T} / \sqrt{d}) V$

其中 $Q$, $K$ 和 $V$ 是形状为 seq_len x hidden_size 的矩阵，分别称为查询 (query)、键 (key) 和值 (value)（它们实际上是更大的矩阵，包含批次维度和注意力头维度，但我们只关心最后两个维度，这是进行矩阵乘积的地方，因此为简单起见，我们只考虑这两个维度）。乘积 $QK^{T}$ 的形状为 seq_len x seq_len，我们可以与 $V$ 进行矩阵乘积以获得输出 $O$ ，其形状与其他的相同。

将 softmax 替换为它的值得到 $O_{i} = \frac{\sum_{j=1}^{i} e^{Q_{i} K_{j}^{T} / \sqrt{d}} V_{j}}{\sum_{j=1}^{i} e^{Q_{i} K_{j}^{T} / \sqrt{d}}}$

注意， $QK^{T}$ 中对应 $j > i$ 的项被掩码（求和在 j 处停止），因为注意力机制不允许查看未来的 token（只能查看过去的 token）。

相比之下，RWKV 注意力机制由以下公式给出 $O_{i} = \sigma(R_{i}) \frac{\sum_{j=1}^{i} e^{W_{i-j} + K_{j}} V_{j}}{\sum_{j=1}^{i} e^{W_{i-j} + K_{j}}}$

其中 $R$ 是作者称为“感受度”的新矩阵， $K$ 和 $V$ 仍然是键和值（这里的 $\sigma$ 是 sigmoid 函数）。 $W$ 是一个表示 token 位置的新向量，由以下公式给出 $W_{0} = u \hbox{ and } W_{k} = (k-1)w \hbox{ for } k \geq 1$

，其中 $u$ 和 $w$ 是代码中分别称为 time_first 和 time_decay 的可学习参数。分子和分母都可以递归地表示。将它们命名为 $N_{i}$ 和 $D_{i}$ 我们有 $N_{i} = e^{u + K_{i}} V_{i} + \hat{N}_{i} \hbox{ where } \hat{N}_{i} = e^{K_{i-1}} V_{i-1} + e^{w + K_{i-2}} V_{i-2} \cdots + e^{(i-2)w + K_{1}} V_{1}$

因此 $\hat{N}_{i}$ （在代码中称为 numerator_state）满足 $\hat{N}_{0} = 0 \hbox{ and } \hat{N}_{j+1} = e^{K_{j}} V_{j} + e^{w} \hat{N}_{j}$

和 $D_{i} = e^{u + K_{i}} + \hat{D}_{i} \hbox{ where } \hat{D}_{i} = e^{K_{i-1}} + e^{w + K_{i-2}} \cdots + e^{(i-2)w + K_{1}}$

因此 $\hat{D}_{i}$ （在代码中称为 denominator_state）满足 $\hat{D}_{0} = 0 \hbox{ and } \hat{D}_{j+1} = e^{K_{j}} + e^{w} \hat{D}_{j}$

实际使用的递推公式要复杂一些，因为为了数值稳定性，我们不想计算大数的指数。通常softmax不是直接计算的，而是用最大项的指数去除分子和分母 $\frac{e^{x_{i}}}{\sum_{j=1}^{n} e^{x_{j}}} = \frac{e^{x_{i} - M}}{\sum_{j=1}^{n} e^{x_{j} - M}}$

，其中 $M$ 是所有$x_{j}$中的最大值。因此，除了保存分子状态 ($\hat{N}$) 和分母状态 ($\hat{D}$)，我们还跟踪指数中遇到的所有项的最大值。所以我们实际使用的是 $\tilde{N}_{i} = e^{-M_{i}} \hat{N}_{i} \hbox{ and } \tilde{D}_{i} = e^{-M_{i}} \hat{D}_{i}$

由以下递推公式定义 $\tilde{N}_{0} = 0 \hbox{ and } \tilde{N}_{j+1} = e^{K_{j} - q} V_{j} + e^{w + M_{j} - q} \tilde{N}_{j} \hbox{ where } q = \max(K_{j}, w + M_{j})$

和 $\tilde{D}_{0} = 0 \hbox{ and } \tilde{D}_{j+1} = e^{K_{j} - q} + e^{w + M_{j} - q} \tilde{D}_{j} \hbox{ where } q = \max(K_{j}, w + M_{j})$

且 $M_{j+1} = q$。有了这些，我们就可以计算 $N_{i} = e^{u + K_{i} - q} V_{i} + e^{M_{i}} \tilde{N}_{i} \hbox{ where } q = \max(u + K_{i}, M_{i})$

和 $D_{i} = e^{u + K_{i} - q} + e^{M_{i}} \tilde{D}_{i} \hbox{ where } q = \max(u + K_{i}, M_{i})$

最终得到 $O_{i} = \sigma(R_{i}) \frac{N_{i}}{D_{i}}$

在 GitHub 上更新

←RoFormer Seed-Oss→