Transformers 文档

RWKV

RWKV

概述

RWKV模型是在这个仓库中提出的

它建议对传统的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

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, optional, defaults to 50277) — RWKV模型的词汇表大小。定义了调用RwkvModel时传递的inputs_ids可以表示的不同标记的数量。
  • context_length (int, optional, 默认为 1024) — 该模型在单次前向传播中可以使用的最大序列长度(在 RNN 模式下使用它允许使用任何序列长度)。
  • hidden_size (int, optional, 默认为 4096) — 嵌入和隐藏状态的维度。
  • num_hidden_layers (int, optional, 默认为 32) — 模型中的隐藏层数量。
  • attention_hidden_size (int, 可选) — 注意力隐藏状态的维度。如果未设置,将默认为 hidden_size.
  • intermediate_size (int, 可选) — 内部前馈层的维度。如果未设置,将默认为 hidden_size 的 4 倍。
  • layer_norm_epsilon (float, optional, defaults to 1e-05) — 用于层归一化层的epsilon值。
  • bos_token_id (int, optional, 默认为 0) — 词汇表中句子开始标记的ID。默认为0,因为RWKV使用与GPTNeoX相同的分词器。
  • eos_token_id (int, optional, 默认为 0) — 词汇表中句子结束标记的ID。默认为0,因为RWKV使用与GPTNeoX相同的分词器。
  • rescale_every (int, 可选, 默认为 6) — 在推理时,每经过 rescale_every 层,隐藏状态(以及相应输出层的权重)会被除以 2。如果设置为 0 或负数,则不进行缩放。
  • tie_word_embeddings (bool, optional, defaults to 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

transformers.RwkvModel

< >

( config )

参数

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

裸的RWKV模型变压器输出原始隐藏状态,没有任何特定的头部。

该模型继承自PreTrainedModel。请查看超类文档以了解库为其所有模型实现的通用方法(如下载或保存、调整输入嵌入的大小、修剪头部等)。

该模型也是一个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.RwkvOutputtuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor of shape (batch_size, input_ids_length)) — input_ids_length = sequence_length if past_key_values is None else past_key_values[0][0].shape[-2] (sequence_length of input past key value states). Indices of input sequence tokens in the vocabulary.

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

    可以使用AutoTokenizer获取索引。详情请参见PreTrainedTokenizer.encode()PreTrainedTokenizer.call()

    什么是输入ID?

  • attention_mask (torch.LongTensor of shape (batch_size, input_ids_length), optional) — Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:
    • 1 for tokens that are not masked,
    • 0 for tokens that are masked.

    目前RwkvModel尚未使用此功能,但将来会得到支持。

    什么是注意力掩码?

  • inputs_embeds (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size), optional) — 可选地,您可以选择直接传递嵌入表示,而不是传递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, optional) — 是否返回所有层的隐藏状态。有关更多详细信息,请参见返回的张量下的hidden_states
  • return_dict (bool, 可选) — 是否返回一个 ModelOutput 而不是一个普通的元组。

返回

transformers.models.rwkv.modeling_rwkv.RwkvOutputtuple(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 一起使用,以避免提供旧的 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 的前向方法,重写了 __call__ 特殊方法。

尽管前向传递的配方需要在此函数内定义,但之后应该调用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

transformers.RwkvForCausalLM

< >

( config )

参数

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

RWKV模型变压器,顶部带有语言建模头(线性层,权重与输入嵌入绑定)。

该模型继承自PreTrainedModel。请查看超类文档以了解库为其所有模型实现的通用方法(如下载或保存、调整输入嵌入的大小、修剪头部等)。

该模型也是一个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 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 ) transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutputtuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor of shape (batch_size, input_ids_length)) — input_ids_length = sequence_length if past_key_values is None else past_key_values[0][0].shape[-2] (sequence_length of input past key value states). Indices of input sequence tokens in the vocabulary.

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

    可以使用AutoTokenizer获取索引。详情请参见PreTrainedTokenizer.encode()PreTrainedTokenizer.call()

    什么是输入ID?

  • attention_mask (torch.LongTensor of shape (batch_size, input_ids_length), optional) — Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:
    • 1 for tokens that are not masked,
    • 0 for tokens that are masked.

    目前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 of shape (batch_size, sequence_length), optional) — 语言建模的标签。请注意,标签在模型内部被移位,即你可以设置 labels = input_ids 索引在 [-100, 0, ..., config.vocab_size] 中选择。所有设置为 -100 的标签 将被忽略(屏蔽),损失仅针对 [0, ..., config.vocab_size] 中的标签计算

返回

transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutputtuple(torch.FloatTensor)

一个 transformers.models.rwkv.modeling_rwkv.RwkvCausalLMOutput 或一个由 torch.FloatTensor 组成的元组(如果传递了 return_dict=False 或当 config.return_dict=False 时),包含各种 元素,具体取决于配置(RwkvConfig)和输入。

  • loss (torch.FloatTensor 形状为 (1,)可选,当提供 labels 时返回) — 语言建模损失(用于下一个标记预测)。

  • logits (torch.FloatTensor 形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数(SoftMax 之前每个词汇标记的分数)。

  • state (五个 torch.FloatTensor 的列表,形状为 (batch_size, hidden_size, num_hidden_layers)) — 模型在最后一个时间步的状态。可以在前向方法中与下一个 input_ids 一起使用,以避免提供旧的 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 的前向方法,重写了 __call__ 特殊方法。

尽管前向传递的配方需要在此函数内定义,但之后应该调用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中,注意力被写为 O=softmax(QKT/d)VO = \hbox{softmax}(QK^{T} / \sqrt{d}) V

使用QQ,KKVV 是形状为 seq_len x hidden_size 的矩阵,分别称为查询、键和值(它们实际上是更大的矩阵,包含批次维度和注意力头维度,但我们只对最后两个维度感兴趣,因为矩阵乘积是在这两个维度上进行的,所以为了简单起见,我们只考虑这两个维度)。乘积QKTQK^{T} 的形状为 seq_len x seq_len,然后我们可以与VV 进行矩阵乘积,得到输出OO,其形状与其他矩阵相同。

将softmax替换为其值得到: Oi=j=1ieQiKjT/dVjj=1ieQiKjT/dO_{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}}}

请注意,QKTQK^{T}中对应于j>ij > i的条目被屏蔽(求和在j处停止),因为注意力机制不允许查看未来的标记(只能查看过去的标记)。

相比之下,RWKV注意力由以下公式给出: Oi=σ(Ri)j=1ieWij+KjVjj=1ieWij+KjO_{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}}}

其中RR 是作者称为接受度矩阵的新矩阵,KKVV 仍然是键和值(这里的 \(\sigma\) 是 sigmoid 函数)。WW 是一个表示标记位置的新向量,由以下公式给出: W0=u 和 Wk=(k1)w 对于 k1W_{0} = u \hbox{ 和 } W_{k} = (k-1)w \hbox{ 对于 } k \geq 1

使用uuww 可学习的参数,在代码中分别称为 time_firsttime_decay。分子和分母都可以递归地表示。将它们命名为NiN_{i}DiD_{i} 我们有: Ni=eu+KiVi+N^i 其中 N^i=eKi1Vi1+ew+Ki2Vi2+e(i2)w+K1V1N_{i} = e^{u + K_{i}} V_{i} + \hat{N}_{i} \hbox{ 其中 } \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}

所以N^i\hat{N}_{i}(在代码中称为numerator_state)满足 N^0=0 并且 N^j+1=eKjVj+ewN^j\hat{N}_{0} = 0 \hbox{ 并且 } \hat{N}_{j+1} = e^{K_{j}} V_{j} + e^{w} \hat{N}_{j}ewN^j

并且 Di=eu+Ki+D^i 其中 D^i=eKi1+ew+Ki2+e(i2)w+K1D_{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}}

所以D^i\hat{D}_{i}(在代码中称为denominator_state)满足 D^0=0 并且 D^j+1=eKj+ewD^j\hat{D}_{0} = 0 \hbox{ 并且 } \hat{D}_{j+1} = e^{K_{j}} + e^{w} \hat{D}_{j}

实际使用的递归公式稍微复杂一些,因为为了数值稳定性,我们不希望计算大数的指数。通常softmax不会直接计算,而是将最大项的指数除以分子和分母: exij=1nexj=exiMj=1nexjM\frac{e^{x_{i}}}{\sum_{j=1}^{n} e^{x_{j}}} = \frac{e^{x_{i} - M}}{\sum_{j=1}^{n} e^{x_{j} - M}}

withMM 所有xjx_{j} 的最大值。因此,在这里除了保存分子状态 (\(\hat{N}\)) 和分母状态 (\(\hat{D}\)) 之外,我们还跟踪在指数中遇到的所有项的最大值。因此,我们实际上使用 N~i=eMiN^i 和 D~i=eMiD^i\tilde{N}_{i} = e^{-M_{i}} \hat{N}_{i} \hbox{ 和 } \tilde{D}_{i} = e^{-M_{i}} \hat{D}_{i}

由以下递推公式定义: N~0=0 and N~j+1=eKjqVj+ew+MjqN~j where q=max(Kj,w+Mj)\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})

并且 D~0=0 并且 D~j+1=eKjq+ew+MjqD~j 其中 q=max(Kj,w+Mj)\tilde{D}_{0} = 0 \hbox{ 并且 } \tilde{D}_{j+1} = e^{K_{j} - q} + e^{w + M_{j} - q} \tilde{D}_{j} \hbox{ 其中 } q = \max(K_{j}, w + M_{j})

并且Mj+1=qM_{j+1} = q。有了这些,我们可以计算 Ni=eu+KiqVi+eMiN~i 其中 q=max(u+Ki,Mi)N_{i} = e^{u + K_{i} - q} V_{i} + e^{M_{i}} \tilde{N}_{i} \hbox{ 其中 } q = \max(u + K_{i}, M_{i})

Di=eu+Kiq+eMiD~i 其中 q=max(u+Ki,Mi)D_{i} = e^{u + K_{i} - q} + e^{M_{i}} \tilde{D}_{i} \hbox{ 其中 } q = \max(u + K_{i}, M_{i})

最终我们得到 Oi=σ(Ri)NiDiO_{i} = \sigma(R_{i}) \frac{N_{i}}{D_{i}}

< > Update on GitHub