Transformers 文件

RWKV

Transformers

加入 Hugging Face 社群

並獲得增強的文件體驗

在模型、資料集和 Spaces 上進行協作

透過加速推理獲得更快的示例

切換文件主題

開始使用

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

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 可以表示的不同標記的數量。
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

< 源 >

( config )

引數

config (RwkvModel) — 模型的配置類，包含模型的所有引數。用配置檔案初始化不會載入與模型相關的權重，只加載配置。檢視 from_pretrained() 方法以載入模型權重。

裸露的 Rwkv 模型，輸出原始隱藏狀態，頂部沒有任何特定頭部。

此模型繼承自 PreTrainedModel。查閱超類文件以瞭解庫為其所有模型實現的通用方法（例如下載或儲存、調整輸入嵌入大小、修剪頭部等）。

此模型也是 PyTorch torch.nn.Module 子類。將其作為常規 PyTorch 模組使用，並參考 PyTorch 文件以瞭解所有與一般使用和行為相關的事項。

forward

< 源 >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None state: typing.Optional[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)) — 如果 past_key_values 為 None，則 input_ids_length = sequence_length，否則為 past_key_values[0][0].shape[-2] (輸入過去的鍵值狀態的 sequence_length)。詞彙表中輸入序列標記的索引。

如果使用了 past_key_values，則只有未計算過其過去的 input_ids 應作為 input_ids 傳遞。

索引可以使用 AutoTokenizer 獲取。有關詳細資訊，請參見 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什麼是 input 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 (五個 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 (形狀為 (batch_size, sequence_length, hidden_size) 的 torch.FloatTensor, 可選，預設為 None) — 模型最後一層輸出的隱藏狀態序列。
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 例項而不是此函式，因為前者負責執行預處理和後處理步驟，而後者則默默地忽略它們。

RwkvLMHeadModel

類 transformers.RwkvForCausalLM

< 源 >

( config )

引數

config (RwkvForCausalLM) — 帶有模型所有引數的模型配置類。使用配置檔案初始化不會載入與模型相關的權重，只加載配置。請檢視 from_pretrained() 方法載入模型權重。

RWKV 模型Transformer，頂部帶有一個語言建模頭（權重與輸入嵌入繫結在一起的線性層）。

此模型繼承自 PreTrainedModel。查閱超類文件以瞭解庫為其所有模型實現的通用方法（例如下載或儲存、調整輸入嵌入大小、修剪頭部等）。

此模型也是 PyTorch torch.nn.Module 子類。將其作為常規 PyTorch 模組使用，並參考 PyTorch 文件以瞭解所有與一般使用和行為相關的事項。

forward

< 源 >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None state: typing.Optional[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)) — 如果 past_key_values 為 None，則 input_ids_length = sequence_length；否則為 past_key_values[0][0].shape[-2]（輸入過去鍵值狀態的 sequence_length）。詞彙表中輸入序列標記的索引。

如果使用了 past_key_values，則只有未計算其過去的 input_ids 應作為 input_ids 傳入。

索引可以使用 AutoTokenizer 獲取。有關詳細資訊，請參閱 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什麼是 input 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 (形狀為 (batch_size, hidden_size, num_hidden_layers) 的五個 torch.FloatTensor 元組，可選) — 如果傳入，模型將在所有塊中使用先前的狀態（這將為提供的 input_ids 提供輸出，就像模型將 state_input_ids + input_ids 作為上下文一樣）。
labels (torch.LongTensor，形狀為 (batch_size, sequence_length)，可選) — 用於語言建模的標籤。請注意，標籤在模型內部是移位的，即您可以設定 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 而不是純元組。

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 (五個 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__ 特殊方法。

示例

RWKV 注意力與迴圈公式

在傳統的自迴歸 Transformer 中，注意力可以表示為 $O = \hbox{softmax}(QK^{T} / \sqrt{d}) V$

，其中 $Q$ , $K$ 和 $V$ 是形狀為 seq_len x hidden_size 的矩陣，分別稱為查詢、鍵和值（它們實際上是更大的矩陣，具有批次維度和注意力頭維度，但我們只對最後兩個感興趣，因為矩陣乘法發生在那裡，因此為簡化起見，我們只考慮這兩個）。乘積 $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 處停止），因為注意力不允許檢視未來的標記（只能檢視過去的標記）。

相比之下，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$ 是一個表示標記位置的新向量，其定義為 $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{N}$ ) 和分母狀態 ( $\hat{D}$ $\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 Splinter→