Commit
•
c39c52a
1
Parent(s):
c98976c
model launch
Browse files- README.md +6 -0
- config.json +25 -0
- configuration_rwkv5.py +120 -0
- generation_config.json +12 -0
- modeling_rwkv5.py +862 -0
- pytorch_model.bin +3 -0
- rwkv_vocab_v20230424.txt +0 -0
- special_tokens_map.json +1 -0
- tokenization_rwkv_world.py +549 -0
- tokenizer_config.json +12 -0
README.md
CHANGED
@@ -1,3 +1,9 @@
|
|
1 |
---
|
2 |
license: apache-2.0
|
3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
license: apache-2.0
|
3 |
---
|
4 |
+
|
5 |
+
![An eagle soaring above a transformer robot](https://substackcdn.com/image/fetch/w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bbd31a7-21b4-4ff6-b43f-8735d1decf25_2048x1652.png)
|
6 |
+
|
7 |
+
### HIGHLY EXPERIMENTAL LINE OF EAGLE MODELS
|
8 |
+
|
9 |
+
NOT MEANT FOR PUBLIC RELEASE - WE USE THIS REPO FOR TESTING PURPOSES
|
config.json
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"RwkvForCausalLM"
|
4 |
+
],
|
5 |
+
"auto_map": {
|
6 |
+
"AutoConfig": "configuration_rwkv5.Rwkv5Config",
|
7 |
+
"AutoModelForCausalLM": "modeling_rwkv5.Rwkv5ForCausalLM"
|
8 |
+
},
|
9 |
+
"attention_hidden_size": 4096,
|
10 |
+
"bos_token_id": 0,
|
11 |
+
"context_length": 4096,
|
12 |
+
"eos_token_id": 0,
|
13 |
+
"head_size": 64,
|
14 |
+
"hidden_size": 4096,
|
15 |
+
"intermediate_size": null,
|
16 |
+
"layer_norm_epsilon": 1e-05,
|
17 |
+
"model_type": "rwkv5",
|
18 |
+
"model_version": "5_2",
|
19 |
+
"num_hidden_layers": 32,
|
20 |
+
"rescale_every": 6,
|
21 |
+
"tie_word_embeddings": false,
|
22 |
+
"transformers_version": "4.34.0",
|
23 |
+
"use_cache": true,
|
24 |
+
"vocab_size": 65536
|
25 |
+
}
|
configuration_rwkv5.py
ADDED
@@ -0,0 +1,120 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The OpenAI Team Authors and HuggingFace Inc. team.
|
3 |
+
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
4 |
+
#
|
5 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
6 |
+
# you may not use this file except in compliance with the License.
|
7 |
+
# You may obtain a copy of the License at
|
8 |
+
#
|
9 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
10 |
+
#
|
11 |
+
# Unless required by applicable law or agreed to in writing, software
|
12 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
13 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
14 |
+
# See the License for the specific language governing permissions and
|
15 |
+
# limitations under the License.
|
16 |
+
""" RWKV configuration"""
|
17 |
+
|
18 |
+
from transformers.configuration_utils import PretrainedConfig
|
19 |
+
from transformers.utils import logging
|
20 |
+
|
21 |
+
|
22 |
+
logger = logging.get_logger(__name__)
|
23 |
+
|
24 |
+
RWKV5_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
|
25 |
+
|
26 |
+
|
27 |
+
class Rwkv5Config(PretrainedConfig):
|
28 |
+
"""
|
29 |
+
This is the configuration class to store the configuration of a [`Rwkv5Model`]. It is used to instantiate a RWKV5
|
30 |
+
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
|
31 |
+
defaults will yield a similar configuration to that of the RWVK-4
|
32 |
+
[RWKV/rwkv-5-world-1b5](https://huggingface.co/RWKV/rwkv-5-world-1b5) architecture.
|
33 |
+
|
34 |
+
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
35 |
+
documentation from [`PretrainedConfig`] for more information.
|
36 |
+
|
37 |
+
|
38 |
+
Args:
|
39 |
+
vocab_size (`int`, *optional*, defaults to 65536):
|
40 |
+
Vocabulary size of the RWKV5 model. Defines the number of different tokens that can be represented by the
|
41 |
+
`inputs_ids` passed when calling [`Rwkv5Model`].
|
42 |
+
hidden_size (`int`, *optional*, defaults to 768):
|
43 |
+
Dimensionality of the embeddings and hidden states.
|
44 |
+
num_hidden_layers (`int`, *optional*, defaults to 24):
|
45 |
+
Number of hidden layers in the model.
|
46 |
+
attention_hidden_size (`int`, *optional*):
|
47 |
+
Dimensionality of the attention hidden states. Will default to `hidden_size` if unset.
|
48 |
+
num_attention_heads (`int`, *optional*, defaults to 64):
|
49 |
+
The attention heads to use in rwkv5 self_attention module.
|
50 |
+
head_size (`int`, *optional*, defaults to 64): head_size of rwkv5 self_attention module.
|
51 |
+
intermediate_size (`int`, *optional*):
|
52 |
+
Dimensionality of the inner feed-forward layers. Will default to 4 times `hidden_size` if unset.
|
53 |
+
layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
|
54 |
+
The epsilon to use in the layer normalization layers.
|
55 |
+
bos_token_id (`int`, *optional*, defaults to 0):
|
56 |
+
The id of the beginning of sentence token in the vocabulary. Defaults to 0 as RWKV5 uses the same tokenizer
|
57 |
+
as GPTNeoX.
|
58 |
+
eos_token_id (`int`, *optional*, defaults to 0):
|
59 |
+
The id of the end of sentence token in the vocabulary. Defaults to 0 as RWKV5 uses the same tokenizer as
|
60 |
+
GPTNeoX.
|
61 |
+
rescale_every (`int`, *optional*, defaults to 6):
|
62 |
+
At inference, the hidden states (and weights of the correponding output layers) are divided by 2 every
|
63 |
+
`rescale_every` layer. If set to 0 or a negative number, no rescale is done.
|
64 |
+
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
65 |
+
Whether or not to tie the word embeddings with the input token embeddings.
|
66 |
+
use_cache (`bool`, *optional*, defaults to `True`):
|
67 |
+
Whether or not the model should return the last state.
|
68 |
+
|
69 |
+
|
70 |
+
Example:
|
71 |
+
|
72 |
+
```python
|
73 |
+
>>> from transformers import Rwkv5Config, Rwkv5Model
|
74 |
+
|
75 |
+
>>> # Initializing a Rwkv5 configuration
|
76 |
+
>>> configuration = Rwkv5Config()
|
77 |
+
|
78 |
+
>>> # Initializing a model (with random weights) from the configuration
|
79 |
+
>>> model = Rwkv5Model(configuration)
|
80 |
+
|
81 |
+
>>> # Accessing the model configuration
|
82 |
+
>>> configuration = model.config
|
83 |
+
```"""
|
84 |
+
|
85 |
+
model_type = "rwkv5"
|
86 |
+
|
87 |
+
def __init__(
|
88 |
+
self,
|
89 |
+
vocab_size=65536,
|
90 |
+
hidden_size=768,
|
91 |
+
num_hidden_layers=24,
|
92 |
+
attention_hidden_size=None,
|
93 |
+
num_attention_heads=64,
|
94 |
+
head_size=64,
|
95 |
+
intermediate_size=None,
|
96 |
+
layer_norm_epsilon=1e-5,
|
97 |
+
bos_token_id=0,
|
98 |
+
eos_token_id=0,
|
99 |
+
rescale_every=6,
|
100 |
+
tie_word_embeddings=False,
|
101 |
+
use_cache=True,
|
102 |
+
**kwargs,
|
103 |
+
):
|
104 |
+
self.vocab_size = vocab_size
|
105 |
+
self.hidden_size = hidden_size
|
106 |
+
self.num_hidden_layers = num_hidden_layers
|
107 |
+
self.attention_hidden_size = attention_hidden_size if attention_hidden_size is not None else hidden_size
|
108 |
+
self.num_attention_heads = num_attention_heads
|
109 |
+
self.head_size = head_size
|
110 |
+
self.intermediate_size = None
|
111 |
+
self.layer_norm_epsilon = layer_norm_epsilon
|
112 |
+
self.rescale_every = rescale_every
|
113 |
+
self.use_cache = use_cache
|
114 |
+
|
115 |
+
self.bos_token_id = bos_token_id
|
116 |
+
self.eos_token_id = eos_token_id
|
117 |
+
|
118 |
+
super().__init__(
|
119 |
+
tie_word_embeddings=tie_word_embeddings, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs
|
120 |
+
)
|
generation_config.json
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"chat_format": "chatml",
|
3 |
+
"eos_token_id": 0,
|
4 |
+
"pad_token_id": 0,
|
5 |
+
"max_window_size": 4096,
|
6 |
+
"max_new_tokens": 4096,
|
7 |
+
"do_sample": true,
|
8 |
+
"top_k": 0,
|
9 |
+
"top_p": 0.1,
|
10 |
+
"repetition_penalty": 1.0,
|
11 |
+
"transformers_version": "4.31.1"
|
12 |
+
}
|
modeling_rwkv5.py
ADDED
@@ -0,0 +1,862 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 Bo Peng and HuggingFace Inc. team.
|
3 |
+
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
4 |
+
#
|
5 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
6 |
+
# you may not use this file except in compliance with the License.
|
7 |
+
# You may obtain a copy of the License at
|
8 |
+
#
|
9 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
10 |
+
#
|
11 |
+
# Unless required by applicable law or agreed to in writing, software
|
12 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
13 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
14 |
+
# See the License for the specific language governing permissions and
|
15 |
+
# limitations under the License.
|
16 |
+
"""PyTorch RWKV5 World model."""
|
17 |
+
|
18 |
+
from dataclasses import dataclass
|
19 |
+
from typing import List, Optional, Tuple, Union
|
20 |
+
|
21 |
+
import torch
|
22 |
+
import torch.nn.functional as F
|
23 |
+
import torch.utils.checkpoint
|
24 |
+
from torch import nn
|
25 |
+
from torch.nn import CrossEntropyLoss
|
26 |
+
|
27 |
+
from transformers.modeling_utils import PreTrainedModel
|
28 |
+
from transformers.utils import (
|
29 |
+
ModelOutput,
|
30 |
+
add_code_sample_docstrings,
|
31 |
+
add_start_docstrings,
|
32 |
+
add_start_docstrings_to_model_forward,
|
33 |
+
is_ninja_available,
|
34 |
+
is_torch_cuda_available,
|
35 |
+
logging,
|
36 |
+
)
|
37 |
+
|
38 |
+
from .configuration_rwkv5 import Rwkv5Config
|
39 |
+
|
40 |
+
|
41 |
+
logger = logging.get_logger(__name__)
|
42 |
+
|
43 |
+
_CHECKPOINT_FOR_DOC = "RWKV/rwkv-5-world-1b5"
|
44 |
+
_CONFIG_FOR_DOC = "Rwkv5Config"
|
45 |
+
|
46 |
+
RWKV5_PRETRAINED_MODEL_ARCHIVE_LIST = [
|
47 |
+
"RWKV/rwkv-5-world-1b5",
|
48 |
+
"RWKV/rwkv-5-world-3b",
|
49 |
+
# See all RWKV models at https://huggingface.co/models?filter=rwkv
|
50 |
+
]
|
51 |
+
|
52 |
+
rwkv5_cuda_kernel = None
|
53 |
+
|
54 |
+
|
55 |
+
def load_wkv5_cuda_kernel(head_size):
|
56 |
+
from torch.utils.cpp_extension import load as load_kernel
|
57 |
+
|
58 |
+
global rwkv5_cuda_kernel
|
59 |
+
|
60 |
+
kernel_folder = Path(__file__).resolve().parent.parent.parent / "kernels" / "rwkv5"
|
61 |
+
cuda_kernel_files = [kernel_folder / f for f in ["wkv5_op.cpp", "wkv5_cuda.cu"]]
|
62 |
+
|
63 |
+
# Only load the kernel if it's not been loaded yet or if we changed the context length
|
64 |
+
if rwkv5_cuda_kernel is not None and rwkv5_cuda_kernel.head_size == head_size:
|
65 |
+
return
|
66 |
+
|
67 |
+
logger.info(f"Loading CUDA kernel for RWKV at head size of {head_size}.")
|
68 |
+
|
69 |
+
flags = [
|
70 |
+
"-res-usage",
|
71 |
+
"--maxrregcount 60",
|
72 |
+
"--use_fast_math",
|
73 |
+
"-O3",
|
74 |
+
"-Xptxas -O3",
|
75 |
+
"--extra-device-vectorization",
|
76 |
+
f"-D_N_={head_size}",
|
77 |
+
]
|
78 |
+
rwkv5_cuda_kernel = load_kernel(
|
79 |
+
name=f"wkv_{head_size}",
|
80 |
+
sources=cuda_kernel_files,
|
81 |
+
verbose=(logging.get_verbosity() == logging.DEBUG),
|
82 |
+
extra_cuda_cflags=flags,
|
83 |
+
)
|
84 |
+
rwkv5_cuda_kernel.head_size = head_size
|
85 |
+
|
86 |
+
|
87 |
+
class WKV_5(torch.autograd.Function):
|
88 |
+
@staticmethod
|
89 |
+
def forward(ctx, B, T, C, H, r, k, v, w, u, s):
|
90 |
+
with torch.no_grad():
|
91 |
+
assert r.dtype == torch.bfloat16
|
92 |
+
assert k.dtype == torch.bfloat16
|
93 |
+
assert v.dtype == torch.bfloat16
|
94 |
+
assert w.dtype == torch.bfloat16
|
95 |
+
assert u.dtype == torch.bfloat16
|
96 |
+
assert s.dtype == torch.float32
|
97 |
+
ctx.B = B
|
98 |
+
ctx.T = T
|
99 |
+
ctx.C = C
|
100 |
+
ctx.H = H
|
101 |
+
assert r.is_contiguous()
|
102 |
+
assert k.is_contiguous()
|
103 |
+
assert v.is_contiguous()
|
104 |
+
assert w.is_contiguous()
|
105 |
+
assert u.is_contiguous()
|
106 |
+
ew = (-torch.exp(w.float())).contiguous()
|
107 |
+
eew = (torch.exp(ew)).contiguous()
|
108 |
+
ctx.save_for_backward(r, k, v, eew, ew, u)
|
109 |
+
y = torch.empty(
|
110 |
+
(B, T, C), device=r.device, dtype=torch.bfloat16, memory_format=torch.contiguous_format
|
111 |
+
) # .uniform_(-1, 1)
|
112 |
+
rwkv5_cuda_kernel.forward(B, T, C, H, r, k, v, eew, u, y, s)
|
113 |
+
return y, s
|
114 |
+
|
115 |
+
@staticmethod
|
116 |
+
def backward(ctx, gy):
|
117 |
+
with torch.no_grad():
|
118 |
+
assert gy.dtype == torch.bfloat16
|
119 |
+
B = ctx.B
|
120 |
+
T = ctx.T
|
121 |
+
C = ctx.C
|
122 |
+
H = ctx.H
|
123 |
+
assert gy.is_contiguous()
|
124 |
+
r, k, v, eew, ew, u = ctx.saved_tensors
|
125 |
+
gr = torch.empty(
|
126 |
+
(B, T, C),
|
127 |
+
device=gy.device,
|
128 |
+
requires_grad=False,
|
129 |
+
dtype=torch.bfloat16,
|
130 |
+
memory_format=torch.contiguous_format,
|
131 |
+
) # .uniform_(-1, 1)
|
132 |
+
gk = torch.empty(
|
133 |
+
(B, T, C),
|
134 |
+
device=gy.device,
|
135 |
+
requires_grad=False,
|
136 |
+
dtype=torch.bfloat16,
|
137 |
+
memory_format=torch.contiguous_format,
|
138 |
+
) # .uniform_(-1, 1)
|
139 |
+
gv = torch.empty(
|
140 |
+
(B, T, C),
|
141 |
+
device=gy.device,
|
142 |
+
requires_grad=False,
|
143 |
+
dtype=torch.bfloat16,
|
144 |
+
memory_format=torch.contiguous_format,
|
145 |
+
) # .uniform_(-1, 1)
|
146 |
+
gw = torch.empty(
|
147 |
+
(B, C),
|
148 |
+
device=gy.device,
|
149 |
+
requires_grad=False,
|
150 |
+
dtype=torch.bfloat16,
|
151 |
+
memory_format=torch.contiguous_format,
|
152 |
+
) # .uniform_(-1, 1)
|
153 |
+
gu = torch.empty(
|
154 |
+
(B, C),
|
155 |
+
device=gy.device,
|
156 |
+
requires_grad=False,
|
157 |
+
dtype=torch.bfloat16,
|
158 |
+
memory_format=torch.contiguous_format,
|
159 |
+
) # .uniform_(-1, 1)
|
160 |
+
rwkv5_cuda_kernel.backward(B, T, C, H, r, k, v, eew, ew, u, gy, gr, gk, gv, gw, gu)
|
161 |
+
gw = torch.sum(gw, 0).view(H, C // H)
|
162 |
+
gu = torch.sum(gu, 0).view(H, C // H)
|
163 |
+
return (None, None, None, None, gr, gk, gv, gw, gu)
|
164 |
+
|
165 |
+
|
166 |
+
def rwkv_linear_attention_v5_cpu(
|
167 |
+
B,
|
168 |
+
H,
|
169 |
+
S,
|
170 |
+
T,
|
171 |
+
n_head,
|
172 |
+
hidden,
|
173 |
+
time_decay,
|
174 |
+
time_first,
|
175 |
+
receptance,
|
176 |
+
key,
|
177 |
+
value,
|
178 |
+
gate,
|
179 |
+
lxw,
|
180 |
+
lxb,
|
181 |
+
ow,
|
182 |
+
state,
|
183 |
+
):
|
184 |
+
key = key.to(torch.float32).view(B, T, H, S).transpose(1, 2).transpose(-2, -1)
|
185 |
+
value = value.to(torch.float32).view(B, T, H, S).transpose(1, 2)
|
186 |
+
receptance = receptance.to(torch.float32).view(B, T, H, S).transpose(1, 2)
|
187 |
+
time_decay = torch.exp(-torch.exp(time_decay.float())).reshape(-1, 1, 1).reshape(n_head, -1, 1)
|
188 |
+
time_first = time_first.float().reshape(-1, 1, 1).reshape(n_head, -1, 1)
|
189 |
+
lxw = lxw.float()
|
190 |
+
lxb = lxb.float()
|
191 |
+
out = torch.zeros_like(key).reshape(B, T, H, S)
|
192 |
+
for t in range(T):
|
193 |
+
rt = receptance[:, :, t : t + 1, :]
|
194 |
+
kt = key[:, :, :, t : t + 1]
|
195 |
+
vt = value[:, :, t : t + 1, :]
|
196 |
+
at = kt @ vt
|
197 |
+
out[:, t] = (rt @ (time_first * at + state)).squeeze(2)
|
198 |
+
with torch.no_grad():
|
199 |
+
state = at + time_decay * state
|
200 |
+
|
201 |
+
out = out.reshape(B * T, H * S)
|
202 |
+
out = F.group_norm(out, num_groups=H, weight=lxw, bias=lxb).reshape(B, T, H * S)
|
203 |
+
out = out.to(dtype=hidden.dtype) * gate
|
204 |
+
out = out @ ow
|
205 |
+
|
206 |
+
return out, state
|
207 |
+
|
208 |
+
|
209 |
+
def rwkv_linear_attention(
|
210 |
+
B,
|
211 |
+
H,
|
212 |
+
S,
|
213 |
+
T,
|
214 |
+
n_head,
|
215 |
+
hidden,
|
216 |
+
time_decay,
|
217 |
+
time_first,
|
218 |
+
receptance,
|
219 |
+
key,
|
220 |
+
value,
|
221 |
+
gate,
|
222 |
+
lxw,
|
223 |
+
lxb,
|
224 |
+
ow,
|
225 |
+
state,
|
226 |
+
):
|
227 |
+
no_cuda = any(t.device.type != "cuda" for t in [time_decay, time_first, receptance, key, value])
|
228 |
+
# Launching the CUDA kernel for just one token will actually be slower (there is no for loop in the CPU version
|
229 |
+
# in this case).
|
230 |
+
one_token = key.size(1) == 1
|
231 |
+
if rwkv5_cuda_kernel is None or no_cuda or one_token:
|
232 |
+
return rwkv_linear_attention_v5_cpu(
|
233 |
+
B,
|
234 |
+
H,
|
235 |
+
S,
|
236 |
+
T,
|
237 |
+
n_head,
|
238 |
+
hidden,
|
239 |
+
time_decay,
|
240 |
+
time_first,
|
241 |
+
receptance,
|
242 |
+
key,
|
243 |
+
value,
|
244 |
+
gate,
|
245 |
+
lxw,
|
246 |
+
lxb,
|
247 |
+
ow,
|
248 |
+
state,
|
249 |
+
)
|
250 |
+
else:
|
251 |
+
out, state = WKV_5.apply(B, T, H * S, H, receptance, key, value, time_decay, time_first, state)
|
252 |
+
out = out.reshape(B * T, H * S)
|
253 |
+
out = F.group_norm(out, num_groups=H, weight=lxw, bias=lxb).reshape(B, T, H * S)
|
254 |
+
out = out.to(dtype=hidden.dtype) * gate
|
255 |
+
out = out @ ow
|
256 |
+
return out, state
|
257 |
+
|
258 |
+
|
259 |
+
class RwkvSelfAttention(nn.Module):
|
260 |
+
def __init__(self, config, layer_id=0):
|
261 |
+
super().__init__()
|
262 |
+
self.config = config
|
263 |
+
kernel_loaded = rwkv5_cuda_kernel is not None and rwkv5_cuda_kernel.head_size == config.head_size
|
264 |
+
if is_ninja_available() and is_torch_cuda_available() and not kernel_loaded:
|
265 |
+
try:
|
266 |
+
load_wkv5_cuda_kernel(config.context_length)
|
267 |
+
except Exception:
|
268 |
+
logger.info("Could not load the custom CUDA kernel for RWKV5 attention.")
|
269 |
+
self.layer_id = layer_id
|
270 |
+
hidden_size = config.hidden_size
|
271 |
+
# https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py#L146
|
272 |
+
num_attention_heads = hidden_size // config.head_size
|
273 |
+
self.num_attention_heads = num_attention_heads
|
274 |
+
attention_hidden_size = (
|
275 |
+
config.attention_hidden_size if config.attention_hidden_size is not None else hidden_size
|
276 |
+
)
|
277 |
+
self.attention_hidden_size = attention_hidden_size
|
278 |
+
|
279 |
+
self.time_decay = nn.Parameter(torch.empty(num_attention_heads, config.head_size))
|
280 |
+
self.time_faaaa = nn.Parameter(torch.empty(num_attention_heads, config.head_size))
|
281 |
+
self.time_mix_gate = nn.Parameter(torch.empty(1, 1, hidden_size))
|
282 |
+
|
283 |
+
self.time_mix_key = nn.Parameter(torch.empty(1, 1, hidden_size))
|
284 |
+
self.time_mix_value = nn.Parameter(torch.empty(1, 1, hidden_size))
|
285 |
+
self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, hidden_size))
|
286 |
+
|
287 |
+
self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
|
288 |
+
self.key = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
289 |
+
self.value = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
290 |
+
self.receptance = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
291 |
+
self.gate = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
292 |
+
self.output = nn.Linear(attention_hidden_size, hidden_size, bias=False)
|
293 |
+
# https://github.com/BlinkDL/RWKV-LM/blob/3db37a72356b736966ddd377268f02b80963af3f/RWKV-v4neo/src/model.py#L190C1-L190C1
|
294 |
+
self.ln_x = nn.GroupNorm(hidden_size // config.head_size, hidden_size)
|
295 |
+
|
296 |
+
# TODO: maybe jit, otherwise move inside forward
|
297 |
+
def extract_key_value(self, B, H, S, T, hidden, state=None):
|
298 |
+
# Mix hidden with the previous timestep to produce key, value, receptance
|
299 |
+
if hidden.size(1) == 1 and state is not None:
|
300 |
+
shifted = state[0][:, :, self.layer_id]
|
301 |
+
else:
|
302 |
+
shifted = self.time_shift(hidden)
|
303 |
+
if state is not None:
|
304 |
+
shifted[:, 0] = state[0][:, :, self.layer_id]
|
305 |
+
if len(shifted.size()) == 2:
|
306 |
+
shifted = shifted.unsqueeze(1)
|
307 |
+
key = hidden * self.time_mix_key + shifted * (1 - self.time_mix_key)
|
308 |
+
value = hidden * self.time_mix_value + shifted * (1 - self.time_mix_value)
|
309 |
+
receptance = hidden * self.time_mix_receptance + shifted * (1 - self.time_mix_receptance)
|
310 |
+
gate = hidden * self.time_mix_gate + shifted * (1 - self.time_mix_gate)
|
311 |
+
|
312 |
+
# https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L693
|
313 |
+
key = self.key(key)
|
314 |
+
value = self.value(value)
|
315 |
+
receptance = self.receptance(receptance)
|
316 |
+
gate = F.silu(self.gate(gate))
|
317 |
+
|
318 |
+
if state is not None:
|
319 |
+
state[0][:, :, self.layer_id] = hidden[:, -1]
|
320 |
+
|
321 |
+
return receptance, key, value, gate, state
|
322 |
+
|
323 |
+
def forward(self, hidden, state=None, use_cache=False, seq_mode=True):
|
324 |
+
B = hidden.shape[0]
|
325 |
+
H = self.time_decay.shape[0]
|
326 |
+
S = hidden.shape[-1] // H
|
327 |
+
T = hidden.shape[1]
|
328 |
+
|
329 |
+
receptance, key, value, gate, state = self.extract_key_value(B, H, S, T, hidden, state=state)
|
330 |
+
layer_state = state[1][:, :, :, :, self.layer_id] if state is not None else None
|
331 |
+
rwkv, layer_state = rwkv_linear_attention(
|
332 |
+
B,
|
333 |
+
H,
|
334 |
+
S,
|
335 |
+
T,
|
336 |
+
self.num_attention_heads,
|
337 |
+
hidden,
|
338 |
+
self.time_decay,
|
339 |
+
self.time_faaaa,
|
340 |
+
receptance,
|
341 |
+
key,
|
342 |
+
value,
|
343 |
+
gate,
|
344 |
+
self.ln_x.weight,
|
345 |
+
self.ln_x.bias,
|
346 |
+
self.output.weight.t(),
|
347 |
+
state=layer_state,
|
348 |
+
)
|
349 |
+
|
350 |
+
if layer_state is not None:
|
351 |
+
state[1][:, :, :, :, self.layer_id] = layer_state
|
352 |
+
|
353 |
+
return rwkv, state
|
354 |
+
|
355 |
+
|
356 |
+
class RwkvFeedForward(nn.Module):
|
357 |
+
def __init__(self, config, layer_id=0):
|
358 |
+
super().__init__()
|
359 |
+
self.config = config
|
360 |
+
self.layer_id = layer_id
|
361 |
+
hidden_size = config.hidden_size
|
362 |
+
# https://github.com/BlinkDL/RWKV-LM/blob/3db37a72356b736966ddd377268f02b80963af3f/RWKV-v4neo/train.py#L168
|
363 |
+
intermediate_size = (
|
364 |
+
config.intermediate_size
|
365 |
+
if config.intermediate_size is not None
|
366 |
+
else int((config.hidden_size * 3.5) // 32 * 32)
|
367 |
+
)
|
368 |
+
|
369 |
+
self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
|
370 |
+
self.time_mix_key = nn.Parameter(torch.empty(1, 1, hidden_size))
|
371 |
+
self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, hidden_size))
|
372 |
+
|
373 |
+
self.key = nn.Linear(hidden_size, intermediate_size, bias=False)
|
374 |
+
self.receptance = nn.Linear(hidden_size, hidden_size, bias=False)
|
375 |
+
self.value = nn.Linear(intermediate_size, hidden_size, bias=False)
|
376 |
+
|
377 |
+
def forward(self, hidden, state=None):
|
378 |
+
if hidden.size(1) == 1 and state is not None:
|
379 |
+
shifted = state[2][:, :, self.layer_id]
|
380 |
+
else:
|
381 |
+
shifted = self.time_shift(hidden)
|
382 |
+
if state is not None:
|
383 |
+
shifted[:, 0] = state[2][:, :, self.layer_id]
|
384 |
+
if len(shifted.size()) == 2:
|
385 |
+
shifted = shifted.unsqueeze(1)
|
386 |
+
key = hidden * self.time_mix_key + shifted * (1 - self.time_mix_key)
|
387 |
+
receptance = hidden * self.time_mix_receptance + shifted * (1 - self.time_mix_receptance)
|
388 |
+
|
389 |
+
key = torch.square(torch.relu(self.key(key)))
|
390 |
+
value = self.value(key)
|
391 |
+
receptance = torch.sigmoid(self.receptance(receptance))
|
392 |
+
|
393 |
+
if state is not None:
|
394 |
+
state[2][:, :, self.layer_id] = hidden[:, -1]
|
395 |
+
|
396 |
+
return receptance * value, state
|
397 |
+
|
398 |
+
|
399 |
+
class RwkvBlock(nn.Module):
|
400 |
+
def __init__(self, config, layer_id):
|
401 |
+
super().__init__()
|
402 |
+
self.config = config
|
403 |
+
self.layer_id = layer_id
|
404 |
+
|
405 |
+
if layer_id == 0:
|
406 |
+
self.pre_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
|
407 |
+
|
408 |
+
self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
|
409 |
+
self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
|
410 |
+
|
411 |
+
self.attention = RwkvSelfAttention(config, layer_id)
|
412 |
+
self.feed_forward = RwkvFeedForward(config, layer_id)
|
413 |
+
|
414 |
+
def forward(self, hidden, state=None, use_cache=False, output_attentions=False, seq_mode=True):
|
415 |
+
if self.layer_id == 0:
|
416 |
+
hidden = self.pre_ln(hidden)
|
417 |
+
attention, state = self.attention(self.ln1(hidden), state=state, use_cache=use_cache, seq_mode=seq_mode)
|
418 |
+
hidden = hidden + attention
|
419 |
+
|
420 |
+
feed_forward, state = self.feed_forward(self.ln2(hidden), state=state)
|
421 |
+
hidden = hidden + feed_forward
|
422 |
+
|
423 |
+
outputs = (hidden, state)
|
424 |
+
if output_attentions:
|
425 |
+
outputs += (attention,)
|
426 |
+
else:
|
427 |
+
outputs += (None,)
|
428 |
+
|
429 |
+
return outputs
|
430 |
+
|
431 |
+
|
432 |
+
class Rwkv5PreTrainedModel(PreTrainedModel):
|
433 |
+
"""
|
434 |
+
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
435 |
+
models.
|
436 |
+
"""
|
437 |
+
|
438 |
+
config_class = Rwkv5Config
|
439 |
+
base_model_prefix = "rwkv"
|
440 |
+
_no_split_modules = ["RwkvBlock"]
|
441 |
+
_keep_in_fp32_modules = ["time_decay", "time_first"]
|
442 |
+
supports_gradient_checkpointing = True
|
443 |
+
|
444 |
+
def _init_weights(self, module):
|
445 |
+
"""Initialize the weights."""
|
446 |
+
if isinstance(module, RwkvSelfAttention):
|
447 |
+
layer_id = module.layer_id
|
448 |
+
num_hidden_layers = module.config.num_hidden_layers
|
449 |
+
hidden_size = module.config.hidden_size
|
450 |
+
attention_hidden_size = module.attention_hidden_size
|
451 |
+
num_attention_heads = hidden_size // module.config.num_attention_heads
|
452 |
+
|
453 |
+
ratio_0_to_1 = layer_id / (num_hidden_layers - 1) # 0 to 1
|
454 |
+
ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers) # 1 to ~0
|
455 |
+
|
456 |
+
time_weight = torch.tensor(
|
457 |
+
[i / hidden_size for i in range(hidden_size)],
|
458 |
+
dtype=module.time_mix_key.dtype,
|
459 |
+
device=module.time_mix_key.device,
|
460 |
+
)
|
461 |
+
time_weight = time_weight[None, None, :]
|
462 |
+
|
463 |
+
# https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py#L398
|
464 |
+
decay_speed = [
|
465 |
+
-6.0 + 5.0 * (h / (attention_hidden_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
|
466 |
+
for h in range(attention_hidden_size)
|
467 |
+
]
|
468 |
+
decay_speed = torch.tensor(decay_speed, dtype=module.time_decay.dtype, device=module.time_decay.device)
|
469 |
+
tmp = torch.tensor(
|
470 |
+
[
|
471 |
+
(1.0 - (i / (attention_hidden_size - 1.0))) * ratio_0_to_1 + 0.1 * ((i + 1) % 3 - 1)
|
472 |
+
for i in range(attention_hidden_size)
|
473 |
+
],
|
474 |
+
dtype=module.time_faaaa.dtype,
|
475 |
+
device=module.time_faaaa.device,
|
476 |
+
)
|
477 |
+
|
478 |
+
with torch.no_grad():
|
479 |
+
module.time_decay.data = decay_speed.reshape(num_attention_heads, module.config.num_attention_heads)
|
480 |
+
module.time_faaaa.data = tmp.reshape(num_attention_heads, module.config.num_attention_heads)
|
481 |
+
module.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0)
|
482 |
+
|
483 |
+
module.time_mix_value.data = torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1
|
484 |
+
module.time_mix_receptance.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
|
485 |
+
module.time_mix_gate.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
|
486 |
+
|
487 |
+
elif isinstance(module, RwkvFeedForward):
|
488 |
+
layer_id = module.layer_id
|
489 |
+
num_hidden_layers = module.config.num_hidden_layers
|
490 |
+
hidden_size = module.config.hidden_size
|
491 |
+
|
492 |
+
ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers) # 1 to ~0
|
493 |
+
|
494 |
+
time_weight = torch.tensor(
|
495 |
+
[i / hidden_size for i in range(hidden_size)],
|
496 |
+
dtype=module.time_mix_key.dtype,
|
497 |
+
device=module.time_mix_key.device,
|
498 |
+
)
|
499 |
+
time_weight = time_weight[None, None, :]
|
500 |
+
|
501 |
+
with torch.no_grad():
|
502 |
+
module.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0)
|
503 |
+
module.time_mix_receptance.data = torch.pow(time_weight, ratio_1_to_almost0)
|
504 |
+
|
505 |
+
|
506 |
+
@dataclass
|
507 |
+
class Rwkv5Output(ModelOutput):
|
508 |
+
"""
|
509 |
+
Class for the RWKV model outputs.
|
510 |
+
|
511 |
+
Args:
|
512 |
+
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
|
513 |
+
Sequence of hidden-states at the output of the last layer of the model.
|
514 |
+
state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`):
|
515 |
+
The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
|
516 |
+
avoid providing the old `input_ids`.
|
517 |
+
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
518 |
+
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
519 |
+
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of
|
520 |
+
the model at the output of each layer plus the optional initial embedding outputs.
|
521 |
+
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
522 |
+
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
523 |
+
sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
|
524 |
+
the self-attention heads.
|
525 |
+
"""
|
526 |
+
|
527 |
+
last_hidden_state: torch.FloatTensor = None
|
528 |
+
state: Optional[List[torch.FloatTensor]] = None
|
529 |
+
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
530 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
531 |
+
|
532 |
+
|
533 |
+
@dataclass
|
534 |
+
class Rwkv5CausalLMOutput(ModelOutput):
|
535 |
+
"""
|
536 |
+
Base class for causal language model (or autoregressive) outputs.
|
537 |
+
|
538 |
+
Args:
|
539 |
+
loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
|
540 |
+
Language modeling loss (for next-token prediction).
|
541 |
+
logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
|
542 |
+
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
|
543 |
+
state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`):
|
544 |
+
The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
|
545 |
+
avoid providing the old `input_ids`.
|
546 |
+
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
547 |
+
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
548 |
+
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of
|
549 |
+
the model at the output of each layer plus the optional initial embedding outputs.
|
550 |
+
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
551 |
+
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
552 |
+
sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
|
553 |
+
the self-attention heads.
|
554 |
+
"""
|
555 |
+
|
556 |
+
loss: Optional[torch.FloatTensor] = None
|
557 |
+
logits: torch.FloatTensor = None
|
558 |
+
state: Optional[List[torch.FloatTensor]] = None
|
559 |
+
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
560 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
561 |
+
|
562 |
+
|
563 |
+
RWKV_START_DOCSTRING = r"""
|
564 |
+
This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
|
565 |
+
library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
|
566 |
+
etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)
|
567 |
+
subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
|
568 |
+
general usage and behavior.
|
569 |
+
|
570 |
+
Parameters:
|
571 |
+
config ([`Rwkv5Config`]): Model configuration class with all the parameters of the model.
|
572 |
+
Initializing with a config file does not load the weights associated with the model, only the
|
573 |
+
configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
|
574 |
+
"""
|
575 |
+
|
576 |
+
RWKV_INPUTS_DOCSTRING = r"""
|
577 |
+
Args:
|
578 |
+
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
|
579 |
+
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else
|
580 |
+
`past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
|
581 |
+
sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their
|
582 |
+
past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See
|
583 |
+
[`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
|
584 |
+
IDs?](../glossary#input-ids)
|
585 |
+
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
586 |
+
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
|
587 |
+
is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
|
588 |
+
model's internal embedding lookup matrix.
|
589 |
+
state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*):
|
590 |
+
If passed along, the model uses the previous state in all the blocks (which will give the output for the
|
591 |
+
`input_ids` provided as if the model add `state_input_ids + input_ids` as context).
|
592 |
+
use_cache (`bool`, *optional*):
|
593 |
+
If set to `True`, the last state is returned and can be used to quickly generate the next logits.
|
594 |
+
output_attentions (`bool`, *optional*):
|
595 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
|
596 |
+
tensors for more detail.
|
597 |
+
output_hidden_states (`bool`, *optional*):
|
598 |
+
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
|
599 |
+
more detail.
|
600 |
+
return_dict (`bool`, *optional*):
|
601 |
+
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
602 |
+
"""
|
603 |
+
|
604 |
+
|
605 |
+
@add_start_docstrings(
|
606 |
+
"The bare RWKV Model transformer outputting raw hidden-states without any specific head on top.",
|
607 |
+
RWKV_START_DOCSTRING,
|
608 |
+
)
|
609 |
+
class Rwkv5Model(Rwkv5PreTrainedModel):
|
610 |
+
def __init__(self, config):
|
611 |
+
super().__init__(config)
|
612 |
+
|
613 |
+
self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
|
614 |
+
self.blocks = nn.ModuleList([RwkvBlock(config, layer_id=idx) for idx in range(config.num_hidden_layers)])
|
615 |
+
self.ln_out = nn.LayerNorm(config.hidden_size)
|
616 |
+
|
617 |
+
self.layers_are_rescaled = False
|
618 |
+
self.gradient_checkpointing = False
|
619 |
+
|
620 |
+
# Initialize weights and apply final processing
|
621 |
+
self.post_init()
|
622 |
+
|
623 |
+
def get_input_embeddings(self):
|
624 |
+
return self.embeddings
|
625 |
+
|
626 |
+
def set_input_embeddings(self, new_embeddings):
|
627 |
+
self.embeddings = new_embeddings
|
628 |
+
|
629 |
+
@add_start_docstrings_to_model_forward(RWKV_INPUTS_DOCSTRING)
|
630 |
+
@add_code_sample_docstrings(
|
631 |
+
checkpoint=_CHECKPOINT_FOR_DOC,
|
632 |
+
output_type=Rwkv5Output,
|
633 |
+
config_class=_CONFIG_FOR_DOC,
|
634 |
+
)
|
635 |
+
def forward(
|
636 |
+
self,
|
637 |
+
input_ids: Optional[torch.LongTensor] = None,
|
638 |
+
attention_mask: Optional[torch.LongTensor] = None, # noqa
|
639 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
640 |
+
state: Optional[List[torch.FloatTensor]] = None,
|
641 |
+
use_cache: Optional[bool] = None,
|
642 |
+
output_attentions: Optional[bool] = None,
|
643 |
+
output_hidden_states: Optional[bool] = None,
|
644 |
+
return_dict: Optional[bool] = None,
|
645 |
+
) -> Union[Tuple, Rwkv5Output]:
|
646 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
647 |
+
output_hidden_states = (
|
648 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
649 |
+
)
|
650 |
+
# rwkv5 only support inference in huggingface.
|
651 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
652 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
653 |
+
|
654 |
+
if self.training == self.layers_are_rescaled and (
|
655 |
+
self.embeddings.weight.dtype == torch.float16 or self.embeddings.weight.dtype == torch.bfloat16
|
656 |
+
):
|
657 |
+
self._rescale_layers()
|
658 |
+
|
659 |
+
if input_ids is not None and inputs_embeds is not None:
|
660 |
+
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
|
661 |
+
elif input_ids is None and inputs_embeds is None:
|
662 |
+
raise ValueError("You have to specify either input_ids or inputs_embeds")
|
663 |
+
|
664 |
+
if inputs_embeds is None:
|
665 |
+
inputs_embeds = self.embeddings(input_ids)
|
666 |
+
|
667 |
+
if use_cache and state is None:
|
668 |
+
# https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L904-L906
|
669 |
+
state = []
|
670 |
+
num_attention_heads = self.config.hidden_size // self.config.num_attention_heads
|
671 |
+
state.append(
|
672 |
+
torch.zeros(
|
673 |
+
(inputs_embeds.size(0), self.config.hidden_size, self.config.num_hidden_layers),
|
674 |
+
dtype=inputs_embeds.dtype,
|
675 |
+
requires_grad=False,
|
676 |
+
device=inputs_embeds.device,
|
677 |
+
).contiguous()
|
678 |
+
)
|
679 |
+
state.append(
|
680 |
+
torch.zeros(
|
681 |
+
(
|
682 |
+
inputs_embeds.size(0),
|
683 |
+
num_attention_heads,
|
684 |
+
self.config.hidden_size // num_attention_heads,
|
685 |
+
self.config.hidden_size // num_attention_heads,
|
686 |
+
self.config.num_hidden_layers,
|
687 |
+
),
|
688 |
+
dtype=torch.float32,
|
689 |
+
requires_grad=False,
|
690 |
+
device=inputs_embeds.device,
|
691 |
+
).contiguous()
|
692 |
+
)
|
693 |
+
state.append(
|
694 |
+
torch.zeros(
|
695 |
+
(inputs_embeds.size(0), self.config.hidden_size, self.config.num_hidden_layers),
|
696 |
+
dtype=inputs_embeds.dtype,
|
697 |
+
requires_grad=False,
|
698 |
+
device=inputs_embeds.device,
|
699 |
+
).contiguous()
|
700 |
+
)
|
701 |
+
|
702 |
+
seq_mode = inputs_embeds.shape[1] > 1
|
703 |
+
hidden_states = inputs_embeds
|
704 |
+
|
705 |
+
all_self_attentions = () if output_attentions else None
|
706 |
+
all_hidden_states = () if output_hidden_states else None
|
707 |
+
for idx, block in enumerate(self.blocks):
|
708 |
+
hidden_states, state, attentions = block(
|
709 |
+
hidden_states, state=state, use_cache=use_cache, output_attentions=output_attentions, seq_mode=seq_mode
|
710 |
+
)
|
711 |
+
if (
|
712 |
+
self.layers_are_rescaled
|
713 |
+
and self.config.rescale_every > 0
|
714 |
+
and (idx + 1) % self.config.rescale_every == 0
|
715 |
+
):
|
716 |
+
hidden_states = hidden_states / 2
|
717 |
+
|
718 |
+
if output_hidden_states:
|
719 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
720 |
+
|
721 |
+
if output_attentions:
|
722 |
+
all_self_attentions = all_self_attentions + (attentions,)
|
723 |
+
|
724 |
+
hidden_states = self.ln_out(hidden_states)
|
725 |
+
|
726 |
+
if output_hidden_states:
|
727 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
728 |
+
|
729 |
+
if not return_dict:
|
730 |
+
return (hidden_states, state, all_hidden_states, all_self_attentions)
|
731 |
+
|
732 |
+
return Rwkv5Output(
|
733 |
+
last_hidden_state=hidden_states,
|
734 |
+
state=state,
|
735 |
+
hidden_states=all_hidden_states, # None
|
736 |
+
attentions=all_self_attentions, # None
|
737 |
+
)
|
738 |
+
|
739 |
+
def _rescale_layers(self):
|
740 |
+
# Layers should be rescaled for inference only.
|
741 |
+
if self.layers_are_rescaled == (not self.training):
|
742 |
+
return
|
743 |
+
if self.config.rescale_every > 0:
|
744 |
+
with torch.no_grad():
|
745 |
+
for block_id, block in enumerate(self.blocks):
|
746 |
+
if self.training:
|
747 |
+
block.attention.output.weight.mul_(2 ** int(block_id // self.config.rescale_every))
|
748 |
+
block.feed_forward.value.weight.mul_(2 ** int(block_id // self.config.rescale_every))
|
749 |
+
else:
|
750 |
+
# Deal with quantization statistics
|
751 |
+
if hasattr(block.attention.output.weight, "SCB"):
|
752 |
+
block.attention.output.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
|
753 |
+
block.feed_forward.value.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
|
754 |
+
elif hasattr(block.attention.output.weight, "quant_state"):
|
755 |
+
self._bnb_4bit_dequantize_and_rescale(block.attention.output, block_id)
|
756 |
+
self._bnb_4bit_dequantize_and_rescale(block.feed_forward.value, block_id)
|
757 |
+
else:
|
758 |
+
block.attention.output.weight.div_(2 ** int(block_id // self.config.rescale_every))
|
759 |
+
block.feed_forward.value.weight.div_(2 ** int(block_id // self.config.rescale_every))
|
760 |
+
|
761 |
+
self.layers_are_rescaled = not self.training
|
762 |
+
|
763 |
+
|
764 |
+
@add_start_docstrings(
|
765 |
+
"""
|
766 |
+
The RWKV Model transformer with a language modeling head on top (linear layer with weights tied to the input
|
767 |
+
embeddings).
|
768 |
+
""",
|
769 |
+
RWKV_START_DOCSTRING,
|
770 |
+
)
|
771 |
+
class Rwkv5ForCausalLM(Rwkv5PreTrainedModel):
|
772 |
+
_tied_weights_keys = ["head.weight"]
|
773 |
+
|
774 |
+
def __init__(self, config):
|
775 |
+
super().__init__(config)
|
776 |
+
self.rwkv = Rwkv5Model(config)
|
777 |
+
self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
778 |
+
|
779 |
+
# Initialize weights and apply final processing
|
780 |
+
self.post_init()
|
781 |
+
|
782 |
+
def get_output_embeddings(self):
|
783 |
+
return self.head
|
784 |
+
|
785 |
+
def set_output_embeddings(self, new_embeddings):
|
786 |
+
self.head = new_embeddings
|
787 |
+
|
788 |
+
def prepare_inputs_for_generation(self, input_ids, state=None, inputs_embeds=None, **kwargs):
|
789 |
+
# only last token for inputs_ids if the state is passed along.
|
790 |
+
if state is not None:
|
791 |
+
input_ids = input_ids[:, -1].unsqueeze(-1)
|
792 |
+
|
793 |
+
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
794 |
+
if inputs_embeds is not None and state is None:
|
795 |
+
model_inputs = {"inputs_embeds": inputs_embeds}
|
796 |
+
else:
|
797 |
+
model_inputs = {"input_ids": input_ids}
|
798 |
+
|
799 |
+
model_inputs["state"] = state
|
800 |
+
return model_inputs
|
801 |
+
|
802 |
+
@add_start_docstrings_to_model_forward(RWKV_INPUTS_DOCSTRING)
|
803 |
+
@add_code_sample_docstrings(
|
804 |
+
checkpoint=_CHECKPOINT_FOR_DOC,
|
805 |
+
output_type=Rwkv5CausalLMOutput,
|
806 |
+
config_class=_CONFIG_FOR_DOC,
|
807 |
+
)
|
808 |
+
def forward(
|
809 |
+
self,
|
810 |
+
input_ids: Optional[torch.LongTensor] = None,
|
811 |
+
attention_mask: Optional[torch.LongTensor] = None,
|
812 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
813 |
+
state: Optional[List[torch.FloatTensor]] = None,
|
814 |
+
labels: Optional[torch.LongTensor] = None,
|
815 |
+
use_cache: Optional[bool] = None,
|
816 |
+
output_attentions: Optional[bool] = None,
|
817 |
+
output_hidden_states: Optional[bool] = None,
|
818 |
+
return_dict: Optional[bool] = None,
|
819 |
+
) -> Union[Tuple, Rwkv5CausalLMOutput]:
|
820 |
+
r"""
|
821 |
+
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
822 |
+
Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
|
823 |
+
`labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
|
824 |
+
are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
|
825 |
+
"""
|
826 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
827 |
+
|
828 |
+
rwkv_outputs = self.rwkv(
|
829 |
+
input_ids,
|
830 |
+
inputs_embeds=inputs_embeds,
|
831 |
+
state=state,
|
832 |
+
use_cache=use_cache,
|
833 |
+
output_attentions=output_attentions,
|
834 |
+
output_hidden_states=output_hidden_states,
|
835 |
+
return_dict=return_dict,
|
836 |
+
)
|
837 |
+
hidden_states = rwkv_outputs[0]
|
838 |
+
|
839 |
+
logits = self.head(hidden_states)
|
840 |
+
|
841 |
+
loss = None
|
842 |
+
if labels is not None:
|
843 |
+
# move labels to correct device to enable model parallelism
|
844 |
+
labels = labels.to(logits.device)
|
845 |
+
# Shift so that tokens < n predict n
|
846 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
847 |
+
shift_labels = labels[..., 1:].contiguous()
|
848 |
+
# Flatten the tokens
|
849 |
+
loss_fct = CrossEntropyLoss()
|
850 |
+
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
|
851 |
+
|
852 |
+
if not return_dict:
|
853 |
+
output = (logits,) + rwkv_outputs[1:]
|
854 |
+
return ((loss,) + output) if loss is not None else output
|
855 |
+
|
856 |
+
return Rwkv5CausalLMOutput(
|
857 |
+
loss=loss,
|
858 |
+
logits=logits,
|
859 |
+
state=rwkv_outputs.state,
|
860 |
+
hidden_states=rwkv_outputs.hidden_states,
|
861 |
+
attentions=rwkv_outputs.attentions,
|
862 |
+
)
|
pytorch_model.bin
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:e39e8d9658c636306b29625c84352af514923989cbaa4030ab02a4074f8d75cf
|
3 |
+
size 15036175548
|
rwkv_vocab_v20230424.txt
ADDED
The diff for this file is too large to render.
See raw diff
|
|
special_tokens_map.json
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
{}
|
tokenization_rwkv_world.py
ADDED
@@ -0,0 +1,549 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
"""Tokenization classes for RWKV5."""
|
16 |
+
|
17 |
+
import json
|
18 |
+
import os
|
19 |
+
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
|
20 |
+
|
21 |
+
from transformers.tokenization_utils import PreTrainedTokenizer
|
22 |
+
from transformers.tokenization_utils_base import (
|
23 |
+
BatchEncoding,
|
24 |
+
EncodedInput,
|
25 |
+
TextInput,
|
26 |
+
TruncationStrategy,
|
27 |
+
)
|
28 |
+
from transformers.utils import PaddingStrategy, TensorType, logging, to_py_obj
|
29 |
+
|
30 |
+
|
31 |
+
if TYPE_CHECKING:
|
32 |
+
from transformers.pipelines.conversational import Conversation
|
33 |
+
|
34 |
+
logger = logging.get_logger(__name__)
|
35 |
+
|
36 |
+
VOCAB_FILES_NAMES = {
|
37 |
+
"vocab_file": "rwkv_vocab_v20230424.txt",
|
38 |
+
}
|
39 |
+
PRETRAINED_VOCAB_FILES_MAP = {
|
40 |
+
"vocab_file": {
|
41 |
+
"RWKV/rwkv-5-world-169m": "https://huggingface.co/RWKV/rwkv-5-world-169m/blob/main/rwkv_vocab_v20230424.txt",
|
42 |
+
},
|
43 |
+
}
|
44 |
+
|
45 |
+
|
46 |
+
class TRIE:
|
47 |
+
__slots__ = tuple("ch,to,values,front".split(","))
|
48 |
+
to: list
|
49 |
+
values: set
|
50 |
+
|
51 |
+
def __init__(self, front=None, ch=None):
|
52 |
+
self.ch = ch
|
53 |
+
self.to = [None for ch in range(256)]
|
54 |
+
self.values = set()
|
55 |
+
self.front = front
|
56 |
+
|
57 |
+
def __repr__(self):
|
58 |
+
fr = self
|
59 |
+
ret = []
|
60 |
+
while fr is not None:
|
61 |
+
if fr.ch is not None:
|
62 |
+
ret.append(fr.ch)
|
63 |
+
fr = fr.front
|
64 |
+
return "<TRIE %s %s>" % (ret[::-1], self.values)
|
65 |
+
|
66 |
+
def add(self, key: bytes, idx: int = 0, val=None):
|
67 |
+
if idx == len(key):
|
68 |
+
if val is None:
|
69 |
+
val = key
|
70 |
+
self.values.add(val)
|
71 |
+
return self
|
72 |
+
ch = key[idx]
|
73 |
+
if self.to[ch] is None:
|
74 |
+
self.to[ch] = TRIE(front=self, ch=ch)
|
75 |
+
return self.to[ch].add(key, idx=idx + 1, val=val)
|
76 |
+
|
77 |
+
def find_longest(self, key: bytes, idx: int = 0):
|
78 |
+
u: TRIE = self
|
79 |
+
ch: int = key[idx]
|
80 |
+
|
81 |
+
while u.to[ch] is not None:
|
82 |
+
u = u.to[ch]
|
83 |
+
idx += 1
|
84 |
+
if u.values:
|
85 |
+
ret = idx, u, u.values
|
86 |
+
if idx == len(key):
|
87 |
+
break
|
88 |
+
ch = key[idx]
|
89 |
+
return ret
|
90 |
+
|
91 |
+
|
92 |
+
class RWKVWorldTokenizer(PreTrainedTokenizer):
|
93 |
+
vocab_files_names = VOCAB_FILES_NAMES
|
94 |
+
model_input_names = ["input_ids", "attention_mask"]
|
95 |
+
|
96 |
+
def __init__(self, vocab_file, errors="replace", pad_token="0", **kwargs):
|
97 |
+
self.add_bos_token = False
|
98 |
+
self.encoder = {}
|
99 |
+
sorted = [] # must be already sorted
|
100 |
+
with open(vocab_file, "r", encoding="utf-8") as f:
|
101 |
+
lines = f.readlines()
|
102 |
+
for l in lines:
|
103 |
+
idx = int(l[: l.index(" ")])
|
104 |
+
x = eval(l[l.index(" ") : l.rindex(" ")])
|
105 |
+
x = x.encode("utf-8") if isinstance(x, str) else x
|
106 |
+
assert isinstance(x, bytes)
|
107 |
+
assert len(x) == int(l[l.rindex(" ") :])
|
108 |
+
sorted += [x]
|
109 |
+
self.encoder[idx] = x
|
110 |
+
|
111 |
+
self.decoder = {}
|
112 |
+
for k, v in self.encoder.items():
|
113 |
+
self.decoder[v] = int(k)
|
114 |
+
|
115 |
+
self.trie = TRIE()
|
116 |
+
for t, i in self.decoder.items():
|
117 |
+
_ = self.trie.add(t, val=(t, i))
|
118 |
+
self.errors = errors # how to handle errors in decoding
|
119 |
+
self.cache = {}
|
120 |
+
self.first_max_length = 0
|
121 |
+
super().__init__(
|
122 |
+
errors=errors,
|
123 |
+
**kwargs,
|
124 |
+
)
|
125 |
+
|
126 |
+
@property
|
127 |
+
def eos_token_id(self) -> Optional[int]:
|
128 |
+
return 0
|
129 |
+
|
130 |
+
@property
|
131 |
+
def eot_token_id(self) -> Optional[int]:
|
132 |
+
return 0
|
133 |
+
|
134 |
+
@property
|
135 |
+
def pad_token_id(self) -> Optional[int]:
|
136 |
+
return 0
|
137 |
+
|
138 |
+
@property
|
139 |
+
def vocab_size(self):
|
140 |
+
return len(self.encoder)
|
141 |
+
|
142 |
+
def get_vocab(self):
|
143 |
+
return dict(self.encoder, **self.added_tokens_encoder)
|
144 |
+
|
145 |
+
def add_tokens(self, new_tokens, special_tokens: bool = False):
|
146 |
+
for token in new_tokens:
|
147 |
+
token_id = self.convert_tokens_to_ids(token)
|
148 |
+
self.added_tokens_decoder[token_id] = token
|
149 |
+
|
150 |
+
def convert_ids_to_tokens(self, ids, skip_special_tokens=False):
|
151 |
+
if isinstance(ids, int):
|
152 |
+
ids = [ids]
|
153 |
+
tokens = []
|
154 |
+
for id_ in ids:
|
155 |
+
if id_ in self.added_tokens_decoder:
|
156 |
+
tokens.append(self.added_tokens_decoder[id_])
|
157 |
+
else:
|
158 |
+
tokens.append(self._convert_id_to_token(id_))
|
159 |
+
return tokens
|
160 |
+
|
161 |
+
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
|
162 |
+
if self.add_bos_token:
|
163 |
+
bos_token_ids = [self.bos_token_id]
|
164 |
+
else:
|
165 |
+
bos_token_ids = []
|
166 |
+
|
167 |
+
output = bos_token_ids + token_ids_0
|
168 |
+
|
169 |
+
if token_ids_1 is None:
|
170 |
+
return output
|
171 |
+
|
172 |
+
return output + bos_token_ids + token_ids_1
|
173 |
+
|
174 |
+
def get_special_tokens_mask(
|
175 |
+
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
|
176 |
+
) -> List[int]:
|
177 |
+
"""
|
178 |
+
Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
|
179 |
+
special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
|
180 |
+
|
181 |
+
Args:
|
182 |
+
token_ids_0 (`List[int]`):
|
183 |
+
List of IDs.
|
184 |
+
token_ids_1 (`List[int]`, *optional*):
|
185 |
+
Optional second list of IDs for sequence pairs.
|
186 |
+
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
|
187 |
+
Whether or not the token list is already formatted with special tokens for the model.
|
188 |
+
|
189 |
+
Returns:
|
190 |
+
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
|
191 |
+
"""
|
192 |
+
if already_has_special_tokens:
|
193 |
+
return super().get_special_tokens_mask(
|
194 |
+
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
|
195 |
+
)
|
196 |
+
|
197 |
+
if not self.add_bos_token:
|
198 |
+
return super().get_special_tokens_mask(
|
199 |
+
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=False
|
200 |
+
)
|
201 |
+
|
202 |
+
if token_ids_1 is None:
|
203 |
+
return [1] + ([0] * len(token_ids_0))
|
204 |
+
return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
|
205 |
+
|
206 |
+
def encodeBytes(self, src: bytes):
|
207 |
+
idx: int = 0
|
208 |
+
tokens = []
|
209 |
+
while idx < len(src):
|
210 |
+
_idx: int = idx
|
211 |
+
idx, _, values = self.trie.find_longest(src, idx)
|
212 |
+
assert idx != _idx
|
213 |
+
_, token = next(iter(values))
|
214 |
+
tokens.append(token)
|
215 |
+
return tokens
|
216 |
+
|
217 |
+
def decodeBytes(self, tokens):
|
218 |
+
return b"".join(map(lambda i: self.encoder[i], tokens)) # noqa
|
219 |
+
|
220 |
+
def _tokenize(self, text, **kwargs):
|
221 |
+
"""Tokenize a string."""
|
222 |
+
return self.encodeBytes(text.encode("utf-8"))
|
223 |
+
|
224 |
+
def _decode_tokens(self, tokens):
|
225 |
+
try:
|
226 |
+
return self.decodeBytes(tokens).decode("utf-8")
|
227 |
+
except Exception:
|
228 |
+
return "\ufffd" # bad utf-8
|
229 |
+
|
230 |
+
def _decode(
|
231 |
+
self,
|
232 |
+
token_ids: Union[int, List[int]],
|
233 |
+
skip_special_tokens: bool = False,
|
234 |
+
**kwargs,
|
235 |
+
) -> str:
|
236 |
+
def remove_zeros_from_first_segment(token_ids, first_max_length):
|
237 |
+
first_segment = token_ids[:first_max_length]
|
238 |
+
first_segment_cleaned = [token for token in first_segment if token != 0]
|
239 |
+
return first_segment_cleaned + token_ids[first_max_length:]
|
240 |
+
|
241 |
+
# Convert inputs to python lists
|
242 |
+
token_ids = to_py_obj(token_ids)
|
243 |
+
token_ids = remove_zeros_from_first_segment(token_ids, self.first_max_length)
|
244 |
+
if isinstance(token_ids, int):
|
245 |
+
if token_ids in self.all_special_ids and skip_special_tokens:
|
246 |
+
return ""
|
247 |
+
return self.encoder.get(token_ids, self.unk_token)
|
248 |
+
elif isinstance(token_ids, list):
|
249 |
+
self.first_max_length
|
250 |
+
out_str = ""
|
251 |
+
out_last = 0
|
252 |
+
out_tokens = []
|
253 |
+
for i, token in enumerate(token_ids):
|
254 |
+
if token == 0:
|
255 |
+
break
|
256 |
+
out_tokens += [token]
|
257 |
+
tmp = self._decode_tokens(out_tokens[out_last:])
|
258 |
+
if "\ufffd" not in tmp:
|
259 |
+
out_str += tmp
|
260 |
+
out_last = i + 1
|
261 |
+
return out_str
|
262 |
+
else:
|
263 |
+
return token_ids
|
264 |
+
|
265 |
+
def _convert_token_to_id(self, token):
|
266 |
+
"""Converts a token (str) in an id using the vocab."""
|
267 |
+
return self.encoder.get(token, self.encoder.get(self.unk_token))
|
268 |
+
|
269 |
+
def _convert_id_to_token(self, index):
|
270 |
+
"""Converts an index (integer) in a token (str) using the vocab."""
|
271 |
+
return self.decoder.get(index)
|
272 |
+
|
273 |
+
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
274 |
+
if not os.path.exists(save_directory):
|
275 |
+
os.mkdir(save_directory)
|
276 |
+
if not os.path.isdir(save_directory):
|
277 |
+
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
|
278 |
+
return
|
279 |
+
vocab_file = os.path.join(
|
280 |
+
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
|
281 |
+
)
|
282 |
+
|
283 |
+
with open(vocab_file, "w", encoding="utf-8") as f:
|
284 |
+
for idx, x in self.encoder.items():
|
285 |
+
if isinstance(x, str):
|
286 |
+
x = x.decode("utf-8")
|
287 |
+
line = f"{idx} {repr(x)} {len(x)}\n"
|
288 |
+
f.write(line)
|
289 |
+
|
290 |
+
return (vocab_file,)
|
291 |
+
|
292 |
+
def prepare_for_tokenization(self, text, **kwargs):
|
293 |
+
return (text, kwargs)
|
294 |
+
|
295 |
+
def _get_padding_truncation_strategies(
|
296 |
+
self, padding=False, truncation=None, max_length=None, pad_to_multiple_of=None, verbose=True, **kwargs
|
297 |
+
):
|
298 |
+
return PaddingStrategy.LONGEST, TruncationStrategy.DO_NOT_TRUNCATE, -1, kwargs
|
299 |
+
|
300 |
+
def _encode_plus(
|
301 |
+
self,
|
302 |
+
text: Union[TextInput, EncodedInput],
|
303 |
+
add_special_tokens: bool = True,
|
304 |
+
padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
|
305 |
+
truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
|
306 |
+
max_length: Optional[int] = None,
|
307 |
+
stride: int = 0,
|
308 |
+
pad_to_multiple_of: Optional[int] = None,
|
309 |
+
return_tensors: Optional[Union[str, TensorType]] = None,
|
310 |
+
return_token_type_ids: Optional[bool] = None,
|
311 |
+
return_attention_mask: Optional[bool] = None,
|
312 |
+
return_overflowing_tokens: bool = False,
|
313 |
+
return_special_tokens_mask: bool = False,
|
314 |
+
return_offsets_mapping: bool = False,
|
315 |
+
return_length: bool = False,
|
316 |
+
verbose: bool = True,
|
317 |
+
**kwargs,
|
318 |
+
) -> BatchEncoding:
|
319 |
+
def get_input_ids(text, max_length=None, pad_token_id=0):
|
320 |
+
def pad_sequence(seq, max_len, pad_tok):
|
321 |
+
return [pad_tok] * (max_len - len(seq)) + seq
|
322 |
+
|
323 |
+
if isinstance(text, str):
|
324 |
+
tokens = self._tokenize(text)
|
325 |
+
if max_length is not None:
|
326 |
+
tokens = pad_sequence(tokens, max_length, pad_token_id)
|
327 |
+
return tokens
|
328 |
+
|
329 |
+
elif isinstance(text, list) and len(text) > 0 and isinstance(text[0], str):
|
330 |
+
tokenized_texts = [self._tokenize(t) for t in text]
|
331 |
+
if max_length is None:
|
332 |
+
max_length = max(len(t) for t in tokenized_texts)
|
333 |
+
return [pad_sequence(t, max_length, pad_token_id) for t in tokenized_texts]
|
334 |
+
|
335 |
+
elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
|
336 |
+
if max_length is not None and len(text) < max_length:
|
337 |
+
return pad_sequence(text, max_length, pad_token_id)
|
338 |
+
return text
|
339 |
+
|
340 |
+
else:
|
341 |
+
raise ValueError(
|
342 |
+
"Input is not valid. Should be a string, a list/tuple of strings or a list/tuple of integers."
|
343 |
+
)
|
344 |
+
|
345 |
+
if return_offsets_mapping:
|
346 |
+
raise NotImplementedError(
|
347 |
+
"return_offset_mapping is not available when using Python tokenizers. "
|
348 |
+
"To use this feature, change your tokenizer to one deriving from "
|
349 |
+
"transformers.PreTrainedTokenizerFast. "
|
350 |
+
"More information on available tokenizers at "
|
351 |
+
"https://github.com/huggingface/transformers/pull/2674"
|
352 |
+
)
|
353 |
+
|
354 |
+
first_ids = get_input_ids(text)
|
355 |
+
|
356 |
+
return self.prepare_for_model(
|
357 |
+
first_ids,
|
358 |
+
pair_ids=None,
|
359 |
+
add_special_tokens=add_special_tokens,
|
360 |
+
padding=padding_strategy.value,
|
361 |
+
truncation=truncation_strategy.value,
|
362 |
+
max_length=max_length,
|
363 |
+
stride=stride,
|
364 |
+
pad_to_multiple_of=pad_to_multiple_of,
|
365 |
+
return_tensors=return_tensors,
|
366 |
+
prepend_batch_axis=True,
|
367 |
+
return_attention_mask=return_attention_mask,
|
368 |
+
return_token_type_ids=return_token_type_ids,
|
369 |
+
return_overflowing_tokens=return_overflowing_tokens,
|
370 |
+
return_special_tokens_mask=return_special_tokens_mask,
|
371 |
+
return_length=return_length,
|
372 |
+
verbose=verbose,
|
373 |
+
)
|
374 |
+
|
375 |
+
def _batch_encode_plus(
|
376 |
+
self,
|
377 |
+
batch_text_or_text_pairs: Union[
|
378 |
+
List[TextInput],
|
379 |
+
List[EncodedInput],
|
380 |
+
],
|
381 |
+
add_special_tokens: bool = True,
|
382 |
+
padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
|
383 |
+
truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
|
384 |
+
max_length: Optional[int] = None,
|
385 |
+
stride: int = 0,
|
386 |
+
pad_to_multiple_of: Optional[int] = None,
|
387 |
+
return_tensors: Optional[Union[str, TensorType]] = None,
|
388 |
+
return_token_type_ids: Optional[bool] = None,
|
389 |
+
return_attention_mask: Optional[bool] = None,
|
390 |
+
return_overflowing_tokens: bool = False,
|
391 |
+
return_special_tokens_mask: bool = False,
|
392 |
+
return_offsets_mapping: bool = False,
|
393 |
+
return_length: bool = False,
|
394 |
+
verbose: bool = True,
|
395 |
+
**kwargs,
|
396 |
+
) -> BatchEncoding:
|
397 |
+
def get_input_ids(text, max_length=None, pad_token_id=0):
|
398 |
+
def pad_sequence(seq, max_len, pad_tok):
|
399 |
+
return [pad_tok] * (max_len - len(seq)) + seq
|
400 |
+
|
401 |
+
if isinstance(text, str):
|
402 |
+
tokens = self._tokenize(text)
|
403 |
+
if max_length is not None:
|
404 |
+
tokens = pad_sequence(tokens, max_length, pad_token_id)
|
405 |
+
return tokens
|
406 |
+
|
407 |
+
elif isinstance(text, list) and len(text) > 0 and isinstance(text[0], str):
|
408 |
+
tokenized_texts = [self._tokenize(t) for t in text]
|
409 |
+
if max_length is None:
|
410 |
+
max_length = max(len(t) for t in tokenized_texts)
|
411 |
+
return [pad_sequence(t, max_length, pad_token_id) for t in tokenized_texts]
|
412 |
+
|
413 |
+
elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
|
414 |
+
if max_length is not None and len(text) < max_length:
|
415 |
+
return pad_sequence(text, max_length, pad_token_id)
|
416 |
+
return text
|
417 |
+
|
418 |
+
else:
|
419 |
+
raise ValueError(
|
420 |
+
"Input is not valid. Should be a string, a list/tuple of strings or a list/tuple of integers."
|
421 |
+
)
|
422 |
+
|
423 |
+
if return_offsets_mapping:
|
424 |
+
raise NotImplementedError(
|
425 |
+
"return_offset_mapping is not available when using Python tokenizers. "
|
426 |
+
"To use this feature, change your tokenizer to one deriving from "
|
427 |
+
"transformers.PreTrainedTokenizerFast."
|
428 |
+
)
|
429 |
+
|
430 |
+
first_max_length = 0
|
431 |
+
second_max_length = 0
|
432 |
+
for ids_or_pair_ids in batch_text_or_text_pairs:
|
433 |
+
if not isinstance(ids_or_pair_ids, (list, tuple)):
|
434 |
+
ids, pair_ids = ids_or_pair_ids, None
|
435 |
+
else:
|
436 |
+
ids, pair_ids = ids_or_pair_ids
|
437 |
+
first_ids = get_input_ids(ids)
|
438 |
+
second_ids = get_input_ids(pair_ids) if pair_ids is not None else None
|
439 |
+
first_max_length = max(first_max_length, len(first_ids))
|
440 |
+
if second_ids is not None:
|
441 |
+
second_max_length = max(second_max_length, len(second_ids))
|
442 |
+
|
443 |
+
self.first_max_length = first_max_length
|
444 |
+
input_ids = []
|
445 |
+
for ids_or_pair_ids in batch_text_or_text_pairs:
|
446 |
+
if not isinstance(ids_or_pair_ids, (list, tuple)):
|
447 |
+
ids, pair_ids = ids_or_pair_ids, None
|
448 |
+
else:
|
449 |
+
ids, pair_ids = ids_or_pair_ids
|
450 |
+
|
451 |
+
first_ids = get_input_ids(ids, max_length=first_max_length)
|
452 |
+
second_ids = get_input_ids(pair_ids, max_length=second_max_length) if pair_ids is not None else None
|
453 |
+
input_ids.append((first_ids, second_ids))
|
454 |
+
|
455 |
+
batch_outputs = self._batch_prepare_for_model(
|
456 |
+
input_ids,
|
457 |
+
add_special_tokens=add_special_tokens,
|
458 |
+
padding_strategy=padding_strategy,
|
459 |
+
truncation_strategy=truncation_strategy,
|
460 |
+
max_length=max_length,
|
461 |
+
stride=stride,
|
462 |
+
pad_to_multiple_of=pad_to_multiple_of,
|
463 |
+
return_attention_mask=return_attention_mask,
|
464 |
+
return_token_type_ids=return_token_type_ids,
|
465 |
+
return_overflowing_tokens=return_overflowing_tokens,
|
466 |
+
return_special_tokens_mask=return_special_tokens_mask,
|
467 |
+
return_length=return_length,
|
468 |
+
return_tensors=return_tensors,
|
469 |
+
verbose=verbose,
|
470 |
+
)
|
471 |
+
|
472 |
+
return BatchEncoding(batch_outputs)
|
473 |
+
|
474 |
+
def decode(
|
475 |
+
self,
|
476 |
+
token_ids: Union[int, List[int]],
|
477 |
+
skip_special_tokens: bool = False,
|
478 |
+
clean_up_tokenization_spaces: bool = None,
|
479 |
+
**kwargs,
|
480 |
+
) -> str:
|
481 |
+
"""
|
482 |
+
Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special
|
483 |
+
tokens and clean up tokenization spaces.
|
484 |
+
|
485 |
+
Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`.
|
486 |
+
|
487 |
+
Args:
|
488 |
+
token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
|
489 |
+
List of tokenized input ids. Can be obtained using the `__call__` method.
|
490 |
+
skip_special_tokens (`bool`, *optional*, defaults to `False`):
|
491 |
+
Whether or not to remove special tokens in the decoding.
|
492 |
+
clean_up_tokenization_spaces (`bool`, *optional*):
|
493 |
+
Whether or not to clean up the tokenization spaces. If `None`, will default to
|
494 |
+
`self.clean_up_tokenization_spaces`.
|
495 |
+
kwargs (additional keyword arguments, *optional*):
|
496 |
+
Will be passed to the underlying model specific decode method.
|
497 |
+
|
498 |
+
Returns:
|
499 |
+
`str`: The decoded sentence.
|
500 |
+
"""
|
501 |
+
# Convert inputs to python lists
|
502 |
+
return self._decode(
|
503 |
+
token_ids=token_ids,
|
504 |
+
skip_special_tokens=skip_special_tokens,
|
505 |
+
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
506 |
+
**kwargs,
|
507 |
+
)
|
508 |
+
|
509 |
+
def batch_decode(
|
510 |
+
self,
|
511 |
+
sequences: Union[List[int], List[List[int]]],
|
512 |
+
skip_special_tokens: bool = False,
|
513 |
+
clean_up_tokenization_spaces: bool = None,
|
514 |
+
**kwargs,
|
515 |
+
) -> List[str]:
|
516 |
+
"""
|
517 |
+
Convert a list of lists of token ids into a list of strings by calling decode.
|
518 |
+
|
519 |
+
Args:
|
520 |
+
sequences (`Union[List[int], List[List[int]], np.ndarray, torch.Tensor, tf.Tensor]`):
|
521 |
+
List of tokenized input ids. Can be obtained using the `__call__` method.
|
522 |
+
skip_special_tokens (`bool`, *optional*, defaults to `False`):
|
523 |
+
Whether or not to remove special tokens in the decoding.
|
524 |
+
clean_up_tokenization_spaces (`bool`, *optional*):
|
525 |
+
Whether or not to clean up the tokenization spaces. If `None`, will default to
|
526 |
+
`self.clean_up_tokenization_spaces`.
|
527 |
+
kwargs (additional keyword arguments, *optional*):
|
528 |
+
Will be passed to the underlying model specific decode method.
|
529 |
+
|
530 |
+
Returns:
|
531 |
+
`List[str]`: The list of decoded sentences.
|
532 |
+
"""
|
533 |
+
return [
|
534 |
+
self.decode(
|
535 |
+
seq,
|
536 |
+
skip_special_tokens=skip_special_tokens,
|
537 |
+
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
538 |
+
**kwargs,
|
539 |
+
)
|
540 |
+
for seq in sequences
|
541 |
+
]
|
542 |
+
|
543 |
+
def _build_conversation_input_ids(self, conversation: "Conversation") -> List[int]:
|
544 |
+
input_ids = []
|
545 |
+
for is_user, text in conversation.iter_texts():
|
546 |
+
input_ids.extend(self.encode(text, add_special_tokens=False) + [self.eos_token_id])
|
547 |
+
if len(input_ids) > self.model_max_length:
|
548 |
+
input_ids = input_ids[-self.model_max_length :]
|
549 |
+
return input_ids
|
tokenizer_config.json
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"name_or_path": "rwkv-world",
|
3 |
+
"add_prefix_space": false,
|
4 |
+
"tokenizer_class": "RWKVWorldTokenizer",
|
5 |
+
"use_fast": false,
|
6 |
+
"auto_map": {
|
7 |
+
"AutoTokenizer": [
|
8 |
+
"tokenization_rwkv_world.RWKVWorldTokenizer",
|
9 |
+
null
|
10 |
+
]
|
11 |
+
}
|
12 |
+
}
|