Transformers 文档
RWKV
并获得增强的文档体验
开始使用
RWKV
概述
RWKV 模型在 此仓库 中提出
它建议调整传统的 Transformer 注意力机制,使其线性化。这样,该模型可以用作循环网络:一起传递时间戳 0 和时间戳 1 的输入与传递时间戳 0 的输入,然后传递时间戳 1 的输入以及时间戳 0 的状态相同(见下面的示例)。
这可能比常规 Transformer 更有效,并且可以处理任意长度的句子(即使模型使用固定的上下文长度进行训练)。
使用示例
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
< 源代码 >( 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
传递的不同 token 的数量。 - 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) — 词汇表中句子起始 token 的 id。 默认为 0,因为 RWKV 使用与 GPTNeoX 相同的分词器。 - eos_token_id (
int
, 可选, 默认为 0) — 词汇表中句子结束 token 的 id。 默认为 0,因为 RWKV 使用与 GPTNeoX 相同的分词器。 - rescale_every (
int
, 可选, 默认为 6) — 在推理时,每隔rescale_every
层,隐藏状态(以及相应输出层的权重)除以 2。 如果设置为 0 或负数,则不进行重新缩放。 - tie_word_embeddings (
bool
, 可选, 默认为False
) — 是否将词嵌入与输入 token 嵌入绑定。 - 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
< 源代码 >( config )
参数
- config (RwkvConfig) — 带有模型所有参数的模型配置类。 使用配置文件初始化不会加载与模型关联的权重,仅加载配置。 查看 from_pretrained() 方法以加载模型权重。
裸 RWKV 模型 Transformer 输出原始隐藏状态,顶部没有任何特定的 head。
此模型继承自 PreTrainedModel。 查看超类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝 head 等)。
此模型也是 PyTorch torch.nn.Module 子类。 将其用作常规 PyTorch 模块,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。
前向传播
< 源代码 >( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None state: typing.Optional[typing.List[torch.FloatTensor]] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None ) → transformers.models.rwkv.modeling_rwkv.RwkvOutput
或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor
,形状为(batch_size, input_ids_length)
) —input_ids_length
=sequence_length
如果past_key_values
为None
,否则为past_key_values[0][0].shape[-2]
(输入过去键值状态的sequence_length
)。 词汇表中输入序列 token 的索引。如果使用
past_key_values
,则只有未计算其过去的input_ids
应作为input_ids
传递。索引可以使用 AutoTokenizer 获得。 有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.LongTensor
,形状为(batch_size, input_ids_length)
, 可选) — 掩码,以避免在 padding token 索引上执行注意力机制。 掩码值在[0, 1]
中选择:- 1 表示 未掩码 的 token,
- 0 表示 已掩码 的 token。
RwkvModel
当前未使用此功能,但将来会支持。 - inputs_embeds (
torch.FloatTensor
,形状为(batch_size, sequence_length, hidden_size)
, 可选) — 可选地,您可以选择直接传递嵌入表示而不是传递input_ids
。 如果您希望比模型的内部嵌入查找矩阵更精细地控制如何将input_ids
索引转换为关联的向量,这将非常有用。 - state (五个
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
,形状为(batch_size, sequence_length, hidden_size)
) — 模型最后一层输出端的隐藏状态序列。 -
state (五个
torch.FloatTensor
的列表,形状为(batch_size, hidden_size, num_hidden_layers)
) — 模型在最后时间步的状态。 可以在带有下一个input_ids
的 forward 方法中使用,以避免提供旧的input_ids
。 -
hidden_states (
tuple(torch.FloatTensor)
, 可选, 当传入output_hidden_states=True
或config.output_hidden_states=True
时返回) —torch.FloatTensor
的元组 (如果模型有嵌入层,则为嵌入输出;每层输出一个),形状为(batch_size, sequence_length, hidden_size)
。模型在每一层输出端的隐藏状态,加上可选的初始嵌入输出。
-
attentions (
tuple(torch.FloatTensor)
, 可选, 当传入output_attentions=True
或config.output_attentions=True
时返回) —torch.FloatTensor
的元组 (每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。注意力 softmax 之后的权重,用于计算自注意力头中的加权平均值。
RwkvModel 的 forward 方法,覆盖了 __call__
特殊方法。
虽然 forward 传递的配方需要在该函数内定义,但应该调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。
示例
>>> from transformers import AutoTokenizer, RwkvModel
>>> import torch
>>> tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")
>>> model = RwkvModel.from_pretrained("RWKV/rwkv-4-169m-pile")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)
>>> last_hidden_states = outputs.last_hidden_state
RwkvLMHeadModel
class transformers.RwkvForCausalLM
< source >( config )
参数
- config (RwkvConfig) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重,仅加载配置。请查看 from_pretrained() 方法来加载模型权重。
带有语言建模头的 RWKV 模型 Transformer(线性层,其权重与输入嵌入层绑定)。
此模型继承自 PreTrainedModel。 查看超类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝 head 等)。
此模型也是 PyTorch torch.nn.Module 子类。 将其用作常规 PyTorch 模块,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。
前向传播
< source >( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None state: typing.Optional[typing.List[torch.FloatTensor]] = None labels: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None **kwargs ) → transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutput
或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor
,形状为(batch_size, input_ids_length)
) —input_ids_length
=sequence_length
,如果past_key_values
为None
;否则为past_key_values[0][0].shape[-2]
(输入过去键值状态的sequence_length
)。词汇表中输入序列 tokens 的索引。如果使用
past_key_values
,则仅应将未计算过去的input_ids
作为input_ids
传递。索引可以使用 AutoTokenizer 获取。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.LongTensor
,形状为(batch_size, input_ids_length)
, 可选) — 掩码,用于避免在 padding token 索引上执行注意力。掩码值在[0, 1]
中选择:- 1 表示 未被掩码 的 tokens,
- 0 表示 被掩码 的 tokens。
RwkvModel
当前未使用此项,但未来将支持。 - inputs_embeds (
torch.FloatTensor
,形状为(batch_size, sequence_length, hidden_size)
, 可选) — 可选地,您可以选择直接传递嵌入表示,而不是传递input_ids
。如果您希望比模型的内部嵌入查找矩阵更好地控制如何将input_ids
索引转换为关联的向量,这将非常有用。 - state (五个
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 而不是纯粹的元组。 - labels (
torch.LongTensor
,形状为(batch_size, sequence_length)
, 可选) — 用于语言建模的标签。请注意,标签在模型内部被移动,即您可以设置labels = input_ids
。索引在[-100, 0, ..., config.vocab_size]
中选择。所有设置为-100
的标签都将被忽略(掩码),损失仅针对[0, ..., config.vocab_size]
中的标签计算
返回值
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 (
torch.FloatTensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数(SoftMax 之前每个词汇表 token 的分数)。 -
state (五个
torch.FloatTensor
的列表,形状为(batch_size, hidden_size, num_hidden_layers)
) — 模型在最后时间步的状态。 可以在带有下一个input_ids
的 forward 方法中使用,以避免提供旧的input_ids
。 -
hidden_states (
tuple(torch.FloatTensor)
, 可选, 当传入output_hidden_states=True
或config.output_hidden_states=True
时返回) —torch.FloatTensor
的元组 (如果模型有嵌入层,则为嵌入输出;每层输出一个),形状为(batch_size, sequence_length, hidden_size)
。模型在每一层输出端的隐藏状态,加上可选的初始嵌入输出。
-
attentions (
tuple(torch.FloatTensor)
, 可选, 当传入output_attentions=True
或config.output_attentions=True
时返回) —torch.FloatTensor
的元组 (每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。注意力 softmax 之后的权重,用于计算自注意力头中的加权平均值。
RwkvForCausalLM 的 forward 方法,覆盖了 __call__
特殊方法。
虽然 forward 传递的配方需要在该函数内定义,但应该调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。
示例
>>> import torch
>>> from transformers import AutoTokenizer, RwkvForCausalLM
>>> tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")
>>> model = RwkvForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs, labels=inputs["input_ids"])
>>> loss = outputs.loss
>>> logits = outputs.logits
Rwkv 注意力和循环公式
在传统的自回归 Transformer 中,注意力被写成:
其中,和是形状为 seq_len x hidden_size
的矩阵,分别命名为 query、key 和 value(它们实际上是更大的矩阵,带有批次维度和注意力头维度,但我们只对最后两个维度感兴趣,即进行矩阵乘法的位置,因此为了简单起见,我们只考虑这两个维度)。乘积然后具有形状 seq_len x seq_len
,我们可以与进行矩阵乘法得到输出与其他矩阵形状相同。
用其值替换 softmax 得到:
请注意,在中对应的条目被掩码(总和在 j 处停止),因为注意力不允许查看未来的 token(仅限过去的 token)。
相比之下,RWKV 注意力由下式给出:
其中是由作者提出的一个名为感受野的新矩阵,和仍然是键和值(\(\sigma\) 这里是 sigmoid 函数)。是一个表示 token 位置的新向量,由下式给出
其中和分别是代码中称为 time_first
和 time_decay
的可学习参数。分子和分母都可以递归地表示。将它们命名为和我们有
因此(在代码中称为 numerator_state
)满足
和$\\~\\$ $D_{i} = e^{u + K_{i}} + \hat{D}_{i} \hbox{ 其中 } \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{ 和 } \hat{D}_{j+1} = e^{K_{j}} + e^{w} \hat{D}_{j}$
实际使用的递归公式稍微复杂一些,因为为了数值稳定性,我们不希望计算大数的指数。通常,softmax 不是直接计算的,而是将分子和分母都除以最大项的指数。
其中$\\~\\$ $M$所有 的最大值$\\~\\$ $x_{j}$. 因此,除了保存分子状态 (\(\hat{N}\)) 和分母状态 (\(\hat{D}\)) 之外,我们还跟踪指数中遇到的所有项的最大值。所以我们实际上使用$\\~\\$ $\tilde{N}_{i} = e^{-M_{i}} \hat{N}_{i} \hbox{ 和 } \tilde{D}_{i} = e^{-M_{i}} \hat{D}_{i}$
由以下递归公式定义$\\~\\$ $\tilde{N}_{0} = 0 \hbox{ 和 } \tilde{N}_{j+1} = e^{K_{j} - q} V_{j} + e^{w + M_{j} - q} \tilde{N}_{j} \hbox{ 其中 } q = \max(K_{j}, w + M_{j})$
和
和. 借助这些,我们可以计算
和
最终得到
< > 在 GitHub 上更新