Commit
•
f82bb49
1
Parent(s):
c7b4236
WIP reposetup (from rwkv-6-world-1b6)
Browse files- .gitattributes +1 -0
- NOTES.md +171 -0
- README.md +10 -3
- added_tokens.json +3 -0
- config.json +25 -0
- configuration_rwkv6.py +118 -0
- generation_config.json +12 -0
- hf_rwkv_tokenizer.py +279 -0
- imgs/crimson-finch-unsplash-david-clode.jpg +0 -0
- modeling_rwkv6.py +746 -0
- rwkv_vocab_v20230424.txt +0 -0
- special_tokens_map.json +6 -0
- tokenizer_config.json +12 -0
.gitattributes
CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
36 |
+
pytorch_model.bin filter=lfs diff=lfs merge=lfs -text
|
NOTES.md
ADDED
@@ -0,0 +1,171 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
#### GPU
|
3 |
+
|
4 |
+
```python
|
5 |
+
import torch
|
6 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
7 |
+
|
8 |
+
def generate_prompt(instruction, input=""):
|
9 |
+
instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
|
10 |
+
input = input.strip().replace('\r\n','\n').replace('\n\n','\n')
|
11 |
+
if input:
|
12 |
+
return f"""Instruction: {instruction}
|
13 |
+
|
14 |
+
Input: {input}
|
15 |
+
|
16 |
+
Response:"""
|
17 |
+
else:
|
18 |
+
return f"""User: hi
|
19 |
+
|
20 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
21 |
+
|
22 |
+
User: {instruction}
|
23 |
+
|
24 |
+
Assistant:"""
|
25 |
+
|
26 |
+
|
27 |
+
model = AutoModelForCausalLM.from_pretrained("RWKV/v6-Finch-14B-HF", trust_remote_code=True, torch_dtype=torch.float16).to('cuda')
|
28 |
+
tokenizer = AutoTokenizer.from_pretrained("RWKV/v6-Finch-14B-HF", trust_remote_code=True, padding_side='left', pad_token="<s>")
|
29 |
+
|
30 |
+
text = "介绍一下大熊猫"
|
31 |
+
prompt = generate_prompt(text)
|
32 |
+
|
33 |
+
inputs = tokenizer(prompt, return_tensors="pt").to(0)
|
34 |
+
output = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=1.0, top_p=0.3, top_k=0, )
|
35 |
+
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))
|
36 |
+
```
|
37 |
+
|
38 |
+
output:
|
39 |
+
|
40 |
+
```shell
|
41 |
+
User: hi
|
42 |
+
|
43 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
44 |
+
|
45 |
+
User: 介绍一下大熊猫
|
46 |
+
|
47 |
+
Assistant: 大熊猫是一种中国特有的哺乳动物,也是中国的国宝之一。它们的外貌特征是圆形的黑白相间的身体,有着黑色的毛发和白色的耳朵。大熊猫的食物主要是竹子,它们会在竹林中寻找竹子,并且会将竹子放在竹笼中进行储存。大熊猫的寿命约为20至30年,但由于栖息地的丧失和人类活动的
|
48 |
+
```
|
49 |
+
|
50 |
+
#### Batch Inference
|
51 |
+
|
52 |
+
```python
|
53 |
+
import torch
|
54 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
55 |
+
|
56 |
+
def generate_prompt(instruction, input=""):
|
57 |
+
instruction = instruction.strip().replace('\r\n', '\n').replace('\n\n', '\n')
|
58 |
+
input = input.strip().replace('\r\n', '\n').replace('\n\n', '\n')
|
59 |
+
if input:
|
60 |
+
return f"""Instruction: {instruction}
|
61 |
+
|
62 |
+
Input: {input}
|
63 |
+
|
64 |
+
Response:"""
|
65 |
+
else:
|
66 |
+
return f"""User: hi
|
67 |
+
|
68 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
69 |
+
|
70 |
+
User: {instruction}
|
71 |
+
|
72 |
+
Assistant:"""
|
73 |
+
|
74 |
+
model = AutoModelForCausalLM.from_pretrained("RWKV/v6-Finch-14B-HF", trust_remote_code=True).to(torch.float32)
|
75 |
+
tokenizer = AutoTokenizer.from_pretrained("RWKV/v6-Finch-14B-HF", trust_remote_code=True, padding_side='left', pad_token="<s>")
|
76 |
+
|
77 |
+
texts = ["请介绍北京的旅游景点", "介绍一下大熊猫", "乌兰察布"]
|
78 |
+
prompts = [generate_prompt(text) for text in texts]
|
79 |
+
|
80 |
+
inputs = tokenizer(prompts, return_tensors="pt", padding=True)
|
81 |
+
outputs = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=1.0, top_p=0.3, top_k=0, )
|
82 |
+
|
83 |
+
for output in outputs:
|
84 |
+
print(tokenizer.decode(output.tolist(), skip_special_tokens=True))
|
85 |
+
|
86 |
+
```
|
87 |
+
|
88 |
+
output:
|
89 |
+
|
90 |
+
```shell
|
91 |
+
User: hi
|
92 |
+
|
93 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
94 |
+
|
95 |
+
User: 请介绍北京的旅游景点
|
96 |
+
|
97 |
+
Assistant: 北京是中国的首都,拥有丰富的旅游资源和历史文化遗产。以下是一些北京的旅游景点:
|
98 |
+
1. 故宫:位于北京市中心,是明清两代的皇宫,是中国最大的古代宫殿建筑群之一。
|
99 |
+
2. 天安门广场:位于北京市中心,是中国最著名的城市广场之一,也是中国最大的城市广场。
|
100 |
+
3. 颐和
|
101 |
+
User: hi
|
102 |
+
|
103 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
104 |
+
|
105 |
+
User: 介绍一下大熊猫
|
106 |
+
|
107 |
+
Assistant: 大熊猫是一种生活在中国中部地区的哺乳动物,也是中国的国宝之一。它们的外貌特征是圆形的黑白相间的身体,有着黑色的毛发和圆圆的眼睛。大熊猫是一种濒危物种,目前只有在野外的几个保护区才能看到它们的身影。大熊猫的食物主要是竹子,它们会在竹子上寻找食物,并且可以通
|
108 |
+
User: hi
|
109 |
+
|
110 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
111 |
+
|
112 |
+
User: 乌兰察布
|
113 |
+
|
114 |
+
Assistant: 乌兰察布是中国新疆维吾尔自治区的一个县级市,位于新疆维吾尔自治区中部,是新疆的第二大城市。乌兰察布市是新疆的第一大城市,也是新疆的重要城市之一。乌兰察布市是新疆的经济中心,也是新疆的重要交通枢纽之一。乌兰察布市的人口约为2.5万人,其中汉族占绝大多数。乌
|
115 |
+
```
|
116 |
+
|
117 |
+
#### CPU
|
118 |
+
|
119 |
+
```python
|
120 |
+
import torch
|
121 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
122 |
+
|
123 |
+
def generate_prompt(instruction, input=""):
|
124 |
+
instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
|
125 |
+
input = input.strip().replace('\r\n','\n').replace('\n\n','\n')
|
126 |
+
if input:
|
127 |
+
return f"""Instruction: {instruction}
|
128 |
+
|
129 |
+
Input: {input}
|
130 |
+
|
131 |
+
Response:"""
|
132 |
+
else:
|
133 |
+
return f"""User: hi
|
134 |
+
|
135 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
136 |
+
|
137 |
+
User: {instruction}
|
138 |
+
|
139 |
+
Assistant:"""
|
140 |
+
|
141 |
+
|
142 |
+
model = AutoModelForCausalLM.from_pretrained("RWKV/v6-Finch-14B-HF", trust_remote_code=True).to(torch.float32)
|
143 |
+
tokenizer = AutoTokenizer.from_pretrained("RWKV/v6-Finch-14B-HF", trust_remote_code=True, padding_side='left', pad_token="<s>")
|
144 |
+
|
145 |
+
text = "请介绍北京的旅游景点"
|
146 |
+
prompt = generate_prompt(text)
|
147 |
+
|
148 |
+
inputs = tokenizer(prompt, return_tensors="pt")
|
149 |
+
output = model.generate(inputs["input_ids"], max_new_tokens=333, do_sample=True, temperature=1.0, top_p=0.3, top_k=0, )
|
150 |
+
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))
|
151 |
+
```
|
152 |
+
|
153 |
+
output:
|
154 |
+
|
155 |
+
```shell
|
156 |
+
User: hi
|
157 |
+
|
158 |
+
Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
|
159 |
+
|
160 |
+
User: 请介绍北京的旅游景点
|
161 |
+
|
162 |
+
Assistant: 北京是中国的首都,拥有众多的旅游景点,以下是其中一些著名的景点:
|
163 |
+
1. 故宫:位于北京市中心,是明清两代的皇宫,内有大量的文物和艺术品。
|
164 |
+
2. 天安门广场:是中国最著名的广场之一,是中国人民政治协商会议的旧址,也是中国人民政治协商会议的中心。
|
165 |
+
3. 颐和园:是中国古代皇家园林之一,有着悠久的历史和丰富的文化内涵。
|
166 |
+
4. 长城:是中国古代的一道长城,全长约万里,是中国最著名的旅游景点之一。
|
167 |
+
5. 北京大学:是中国著名的高等教育机构之一,有着悠久的历史和丰富的文化内涵。
|
168 |
+
6. 北京动物园:是中国最大的动物园之一,有着丰富的动物资源和丰富的文化内涵。
|
169 |
+
7. 故宫博物院:是中国最著名的博物馆之一,收藏了大量的文物和艺术品,是中国最重要的文化遗产之一。
|
170 |
+
8. 天坛:是中国古代皇家
|
171 |
+
```
|
README.md
CHANGED
@@ -1,3 +1,10 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
### v6-Finch-14B-HF
|
2 |
+
|
3 |
+
> HF compatible model for Finch-14B.
|
4 |
+
> This is an early preview for benchmarking
|
5 |
+
|
6 |
+
![Crimson Finch Bird](./imgs/crimson-finch-unsplash-david-clode.jpg)
|
7 |
+
|
8 |
+
> origin pth weight at https://huggingface.co/BlinkDL/rwkv-6-world/blob/main/ .
|
9 |
+
|
10 |
+
More details to be done.
|
added_tokens.json
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"<s>": 0
|
3 |
+
}
|
config.json
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"Rwkv6ForCausalLM"
|
4 |
+
],
|
5 |
+
"auto_map": {
|
6 |
+
"AutoConfig": "configuration_rwkv6.Rwkv6Config",
|
7 |
+
"AutoModelForCausalLM": "modeling_rwkv6.Rwkv6ForCausalLM"
|
8 |
+
},
|
9 |
+
"attention_hidden_size": 4096,
|
10 |
+
"bos_token_id": 0,
|
11 |
+
"eos_token_id": 0,
|
12 |
+
"head_size": 64,
|
13 |
+
"head_size_divisor": 8,
|
14 |
+
"hidden_size": 4096,
|
15 |
+
"intermediate_size": null,
|
16 |
+
"layer_norm_epsilon": 1e-05,
|
17 |
+
"model_type": "rwkv6",
|
18 |
+
"num_attention_heads": 64,
|
19 |
+
"num_hidden_layers": 61,
|
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_rwkv6.py
ADDED
@@ -0,0 +1,118 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
RWKV6_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
|
25 |
+
|
26 |
+
|
27 |
+
class Rwkv6Config(PretrainedConfig):
|
28 |
+
"""
|
29 |
+
This is the configuration class to store the configuration of a [`Rwkv6Model`]. It is used to instantiate a RWKV6
|
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 RWKV6 model. Defines the number of different tokens that can be represented by the
|
41 |
+
`inputs_ids` passed when calling [`Rwkv6Model`].
|
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 rwkv6 self_attention module.
|
50 |
+
head_size (`int`, *optional*, defaults to 64): head_size of rwkv6 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.
|
57 |
+
eos_token_id (`int`, *optional*, defaults to 0):
|
58 |
+
The id of the end of sentence token in the vocabulary. Defaults to 0.
|
59 |
+
rescale_every (`int`, *optional*, defaults to 6):
|
60 |
+
At inference, the hidden states (and weights of the correponding output layers) are divided by 2 every
|
61 |
+
`rescale_every` layer. If set to 0 or a negative number, no rescale is done.
|
62 |
+
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
63 |
+
Whether or not to tie the word embeddings with the input token embeddings.
|
64 |
+
use_cache (`bool`, *optional*, defaults to `True`):
|
65 |
+
Whether or not the model should return the last state.
|
66 |
+
|
67 |
+
|
68 |
+
Example:
|
69 |
+
|
70 |
+
```python
|
71 |
+
>>> from transformers import Rwkv6Config, Rwkv6Model
|
72 |
+
|
73 |
+
>>> # Initializing a Rwkv6 configuration
|
74 |
+
>>> configuration = Rwkv6Config()
|
75 |
+
|
76 |
+
>>> # Initializing a model (with random weights) from the configuration
|
77 |
+
>>> model = Rwkv6Model(configuration)
|
78 |
+
|
79 |
+
>>> # Accessing the model configuration
|
80 |
+
>>> configuration = model.config
|
81 |
+
```"""
|
82 |
+
|
83 |
+
model_type = "rwkv6"
|
84 |
+
|
85 |
+
def __init__(
|
86 |
+
self,
|
87 |
+
vocab_size=65536,
|
88 |
+
hidden_size=768,
|
89 |
+
num_hidden_layers=24,
|
90 |
+
attention_hidden_size=None,
|
91 |
+
head_size=64,
|
92 |
+
head_size_divisor=8,
|
93 |
+
intermediate_size=None,
|
94 |
+
layer_norm_epsilon=1e-5,
|
95 |
+
bos_token_id=0,
|
96 |
+
eos_token_id=0,
|
97 |
+
rescale_every=6,
|
98 |
+
tie_word_embeddings=False,
|
99 |
+
use_cache=True,
|
100 |
+
**kwargs,
|
101 |
+
):
|
102 |
+
self.vocab_size = vocab_size
|
103 |
+
self.hidden_size = hidden_size
|
104 |
+
self.num_hidden_layers = num_hidden_layers
|
105 |
+
self.attention_hidden_size = attention_hidden_size if attention_hidden_size is not None else hidden_size
|
106 |
+
self.head_size = head_size
|
107 |
+
self.head_size_divisor = head_size_divisor
|
108 |
+
self.intermediate_size = None
|
109 |
+
self.layer_norm_epsilon = layer_norm_epsilon
|
110 |
+
self.rescale_every = rescale_every
|
111 |
+
self.use_cache = use_cache
|
112 |
+
|
113 |
+
self.bos_token_id = bos_token_id
|
114 |
+
self.eos_token_id = eos_token_id
|
115 |
+
|
116 |
+
super().__init__(
|
117 |
+
tie_word_embeddings=tie_word_embeddings, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs
|
118 |
+
)
|
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": 2048,
|
6 |
+
"max_new_tokens": 2048,
|
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 |
+
}
|
hf_rwkv_tokenizer.py
ADDED
@@ -0,0 +1,279 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2024 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 RWKV6."""
|
16 |
+
|
17 |
+
import os
|
18 |
+
import re
|
19 |
+
from typing import TYPE_CHECKING, List, Optional, Tuple
|
20 |
+
|
21 |
+
from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
|
22 |
+
from transformers.utils import logging
|
23 |
+
|
24 |
+
|
25 |
+
if TYPE_CHECKING:
|
26 |
+
pass
|
27 |
+
|
28 |
+
logger = logging.get_logger(__name__)
|
29 |
+
|
30 |
+
|
31 |
+
VOCAB_FILES_NAMES = {
|
32 |
+
"vocab_file": "rwkv_vocab_v20230424.txt",
|
33 |
+
}
|
34 |
+
|
35 |
+
class TRIE:
|
36 |
+
__slots__ = tuple("ch,to,values,front".split(","))
|
37 |
+
to: list
|
38 |
+
values: set
|
39 |
+
|
40 |
+
def __init__(self, front=None, ch=None):
|
41 |
+
self.ch = ch
|
42 |
+
self.to = [None for ch in range(256)]
|
43 |
+
self.values = set()
|
44 |
+
self.front = front
|
45 |
+
|
46 |
+
def __repr__(self):
|
47 |
+
fr = self
|
48 |
+
ret = []
|
49 |
+
while fr != None:
|
50 |
+
if fr.ch != None:
|
51 |
+
ret.append(fr.ch)
|
52 |
+
fr = fr.front
|
53 |
+
return "<TRIE %s %s>" % (ret[::-1], self.values)
|
54 |
+
|
55 |
+
def add(self, key: bytes, idx: int = 0, val=None):
|
56 |
+
if idx == len(key):
|
57 |
+
if val is None:
|
58 |
+
val = key
|
59 |
+
self.values.add(val)
|
60 |
+
return self
|
61 |
+
ch = key[idx]
|
62 |
+
if self.to[ch] is None:
|
63 |
+
self.to[ch] = TRIE(front=self, ch=ch)
|
64 |
+
return self.to[ch].add(key, idx=idx + 1, val=val)
|
65 |
+
|
66 |
+
def find_longest(self, key: bytes, idx: int = 0):
|
67 |
+
u: TRIE = self
|
68 |
+
ch: int = key[idx]
|
69 |
+
|
70 |
+
while u.to[ch] is not None:
|
71 |
+
u = u.to[ch]
|
72 |
+
idx += 1
|
73 |
+
if u.values:
|
74 |
+
ret = idx, u, u.values
|
75 |
+
if idx == len(key):
|
76 |
+
break
|
77 |
+
ch = key[idx]
|
78 |
+
return ret
|
79 |
+
|
80 |
+
|
81 |
+
class RWKV_TOKENIZER:
|
82 |
+
def __init__(self, file_name):
|
83 |
+
self.idx2token = {}
|
84 |
+
sorted = [] # must be already sorted
|
85 |
+
with open(file_name, "r", encoding="utf-8") as f:
|
86 |
+
lines = f.readlines()
|
87 |
+
for l in lines:
|
88 |
+
idx = int(l[: l.index(" ")])
|
89 |
+
x = eval(l[l.index(" ") : l.rindex(" ")])
|
90 |
+
x = x.encode("utf-8") if isinstance(x, str) else x
|
91 |
+
assert isinstance(x, bytes)
|
92 |
+
|
93 |
+
assert len(x) == int(l[l.rindex(" ") :])
|
94 |
+
sorted += [x]
|
95 |
+
self.idx2token[idx] = x
|
96 |
+
|
97 |
+
self.token2idx = {}
|
98 |
+
for k, v in self.idx2token.items():
|
99 |
+
self.token2idx[v] = int(k)
|
100 |
+
|
101 |
+
self.root = TRIE()
|
102 |
+
for t, i in self.token2idx.items():
|
103 |
+
_ = self.root.add(t, val=(t, i))
|
104 |
+
|
105 |
+
def encodeBytes(self, src: bytes):
|
106 |
+
idx: int = 0
|
107 |
+
tokens = []
|
108 |
+
while idx < len(src):
|
109 |
+
_idx: int = idx
|
110 |
+
idx, _, values = self.root.find_longest(src, idx)
|
111 |
+
assert idx != _idx
|
112 |
+
_, token = next(iter(values))
|
113 |
+
tokens.append(token)
|
114 |
+
return tokens
|
115 |
+
|
116 |
+
def decodeBytes(self, tokens):
|
117 |
+
return b"".join(map(lambda i: self.idx2token[i], tokens))
|
118 |
+
|
119 |
+
def encode(self, src):
|
120 |
+
if isinstance(src, str):
|
121 |
+
return [self.encodeBytes(src.encode("utf-8"))]
|
122 |
+
elif isinstance(src, list):
|
123 |
+
return [self.encodeBytes(s.encode("utf-8")) for s in src]
|
124 |
+
|
125 |
+
def decode(self, tokens):
|
126 |
+
return [self.decodeBytes(batch).decode("utf-8") for batch in tokens]
|
127 |
+
# try:
|
128 |
+
# return self.decodeBytes(tokens).decode('utf-8')
|
129 |
+
# except:
|
130 |
+
# return '\ufffd' # bad utf-8
|
131 |
+
|
132 |
+
def printTokens(self, tokens):
|
133 |
+
for i in tokens:
|
134 |
+
s = self.idx2token[i]
|
135 |
+
try:
|
136 |
+
s = s.decode("utf-8")
|
137 |
+
except:
|
138 |
+
pass
|
139 |
+
print(f"{repr(s)}{i}", end=" ")
|
140 |
+
print()
|
141 |
+
|
142 |
+
|
143 |
+
class Rwkv6Tokenizer(PreTrainedTokenizer):
|
144 |
+
vocab_files_names = VOCAB_FILES_NAMES
|
145 |
+
model_input_names = ["input_ids", "attention_mask"]
|
146 |
+
|
147 |
+
def __init__(
|
148 |
+
self, vocab_file, bos_token="<s>", eos_token="<s>", unk_token="<s>", **kwargs
|
149 |
+
):
|
150 |
+
if not os.path.isfile(vocab_file):
|
151 |
+
raise ValueError(
|
152 |
+
f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
|
153 |
+
" model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
|
154 |
+
)
|
155 |
+
|
156 |
+
with open(vocab_file, "r", encoding="utf-8") as reader:
|
157 |
+
tokens = reader.readlines()
|
158 |
+
|
159 |
+
if "add_bos_token" in kwargs:
|
160 |
+
self.add_bos_token = kwargs["add_bos_token"]
|
161 |
+
else:
|
162 |
+
self.add_bos_token = False
|
163 |
+
self.trie_tokenizer = RWKV_TOKENIZER(vocab_file)
|
164 |
+
vocab = self.trie_tokenizer.token2idx
|
165 |
+
self.encoder = vocab
|
166 |
+
self.decoder = {v: k for k, v in vocab.items()}
|
167 |
+
self._added_tokens_decoder = {0: AddedToken(str(bos_token))}
|
168 |
+
super().__init__(
|
169 |
+
bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, **kwargs
|
170 |
+
)
|
171 |
+
|
172 |
+
@property
|
173 |
+
def vocab_size(self):
|
174 |
+
return len(self.encoder)
|
175 |
+
|
176 |
+
def get_vocab(self):
|
177 |
+
vocab = {str(self.convert_ids_to_tokens(i)): i for i in range(self.vocab_size)}
|
178 |
+
vocab.update(self.added_tokens_encoder)
|
179 |
+
return vocab
|
180 |
+
|
181 |
+
def _tokenize(self, text, split_special_tokens=False):
|
182 |
+
# return self.wordpiece_tokenizer.tokenize(text.encode("utf-8"))
|
183 |
+
return self.trie_tokenizer.encode(text)[0]
|
184 |
+
|
185 |
+
def _convert_token_to_id(self, token):
|
186 |
+
return token
|
187 |
+
|
188 |
+
def _convert_id_to_token(self, index):
|
189 |
+
"""Converts an index (integer) in a token (byte) using the vocab."""
|
190 |
+
token = self.decoder.get(index, self.unk_token)
|
191 |
+
if isinstance(token, (bytes)):
|
192 |
+
token = token.decode("utf-8", errors="replace")
|
193 |
+
return token
|
194 |
+
|
195 |
+
def convert_tokens_to_string(self, tokens):
|
196 |
+
"""Converts a sequence of tokens (bytes) in a single string. Additional tokens are encoded to bytes"""
|
197 |
+
out_string = b"".join(
|
198 |
+
[k.encode(errors="replace") if isinstance(k, str) else k for k in tokens]
|
199 |
+
).decode("utf-8")
|
200 |
+
return out_string
|
201 |
+
|
202 |
+
def save_vocabulary(
|
203 |
+
self, save_directory: str, filename_prefix: Optional[str] = None
|
204 |
+
) -> Tuple[str]:
|
205 |
+
index = 0
|
206 |
+
if os.path.isdir(save_directory):
|
207 |
+
vocab_file = os.path.join(
|
208 |
+
save_directory,
|
209 |
+
(filename_prefix + "-" if filename_prefix else "") + "vocab.txt",
|
210 |
+
)
|
211 |
+
else:
|
212 |
+
vocab_file = (
|
213 |
+
filename_prefix + "-" if filename_prefix else ""
|
214 |
+
) + save_directory
|
215 |
+
with open(vocab_file, "w", encoding="utf-8") as writer:
|
216 |
+
for token, token_index in sorted(
|
217 |
+
self.encoder.items(), key=lambda kv: kv[1]
|
218 |
+
):
|
219 |
+
if index != token_index:
|
220 |
+
logger.warning(
|
221 |
+
f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
|
222 |
+
" Please check that the vocabulary is not corrupted!"
|
223 |
+
)
|
224 |
+
index = token_index
|
225 |
+
writer.write(str(token) + "\n")
|
226 |
+
index += 1
|
227 |
+
return (vocab_file,)
|
228 |
+
|
229 |
+
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
|
230 |
+
if self.add_bos_token:
|
231 |
+
bos_token_ids = [self.bos_token_id]
|
232 |
+
else:
|
233 |
+
bos_token_ids = []
|
234 |
+
|
235 |
+
output = bos_token_ids + token_ids_0
|
236 |
+
|
237 |
+
if token_ids_1 is None:
|
238 |
+
return output
|
239 |
+
|
240 |
+
return output + bos_token_ids + token_ids_1
|
241 |
+
|
242 |
+
def get_special_tokens_mask(
|
243 |
+
self,
|
244 |
+
token_ids_0: List[int],
|
245 |
+
token_ids_1: Optional[List[int]] = None,
|
246 |
+
already_has_special_tokens: bool = False,
|
247 |
+
) -> List[int]:
|
248 |
+
"""
|
249 |
+
Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
|
250 |
+
special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
|
251 |
+
|
252 |
+
Args:
|
253 |
+
token_ids_0 (`List[int]`):
|
254 |
+
List of IDs.
|
255 |
+
token_ids_1 (`List[int]`, *optional*):
|
256 |
+
Optional second list of IDs for sequence pairs.
|
257 |
+
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
|
258 |
+
Whether or not the token list is already formatted with special tokens for the model.
|
259 |
+
|
260 |
+
Returns:
|
261 |
+
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
|
262 |
+
"""
|
263 |
+
if already_has_special_tokens:
|
264 |
+
return super().get_special_tokens_mask(
|
265 |
+
token_ids_0=token_ids_0,
|
266 |
+
token_ids_1=token_ids_1,
|
267 |
+
already_has_special_tokens=True,
|
268 |
+
)
|
269 |
+
|
270 |
+
if not self.add_bos_token:
|
271 |
+
return super().get_special_tokens_mask(
|
272 |
+
token_ids_0=token_ids_0,
|
273 |
+
token_ids_1=token_ids_1,
|
274 |
+
already_has_special_tokens=False,
|
275 |
+
)
|
276 |
+
|
277 |
+
if token_ids_1 is None:
|
278 |
+
return [1] + ([0] * len(token_ids_0))
|
279 |
+
return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
|
imgs/crimson-finch-unsplash-david-clode.jpg
ADDED
![]() |
modeling_rwkv6.py
ADDED
@@ -0,0 +1,746 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2024 The RWKV team and 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 |
+
"""PyTorch RWKV6 World model."""
|
16 |
+
|
17 |
+
from dataclasses import dataclass
|
18 |
+
from typing import List, Optional, Tuple, Union
|
19 |
+
|
20 |
+
from pathlib import Path
|
21 |
+
|
22 |
+
import torch
|
23 |
+
import torch.nn.functional as F
|
24 |
+
import torch.utils.checkpoint
|
25 |
+
from torch import nn
|
26 |
+
from torch.nn import CrossEntropyLoss
|
27 |
+
|
28 |
+
from transformers.modeling_utils import PreTrainedModel
|
29 |
+
from transformers.utils import (
|
30 |
+
ModelOutput,
|
31 |
+
add_code_sample_docstrings,
|
32 |
+
add_start_docstrings,
|
33 |
+
add_start_docstrings_to_model_forward,
|
34 |
+
is_ninja_available,
|
35 |
+
is_torch_cuda_available,
|
36 |
+
logging,
|
37 |
+
)
|
38 |
+
|
39 |
+
from .configuration_rwkv6 import Rwkv6Config
|
40 |
+
try:
|
41 |
+
from fla.ops.rwkv6.recurrent_fuse import fused_recurrent_rwkv6
|
42 |
+
except ImportError:
|
43 |
+
print("Required module is not installed. Please install it using the following commands:")
|
44 |
+
print("pip install -U git+https://github.com/sustcsonglin/flash-linear-attention")
|
45 |
+
print("Additionally, ensure you have the correct version of Triton installed:")
|
46 |
+
print("pip install triton==2.2.0")
|
47 |
+
|
48 |
+
|
49 |
+
logger = logging.get_logger(__name__)
|
50 |
+
|
51 |
+
_CHECKPOINT_FOR_DOC = "RWKV/rwkv-6-world-1b6"
|
52 |
+
_CONFIG_FOR_DOC = "Rwkv6Config"
|
53 |
+
|
54 |
+
def rwkv6_linear_attention_cpu(receptance, key, value, time_decay, time_first, state):
|
55 |
+
# For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed
|
56 |
+
# within a torch.no_grad.
|
57 |
+
batch, seq_length, _ = receptance.shape
|
58 |
+
num_heads, head_size = time_first.shape
|
59 |
+
key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2).transpose(-2, -1)
|
60 |
+
value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
|
61 |
+
receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
|
62 |
+
time_decay = torch.exp(-torch.exp(time_decay.float())).view(batch, seq_length, num_heads, head_size).permute(0, 2, 3, 1)
|
63 |
+
time_first = time_first.float().reshape(-1, 1, 1).reshape(num_heads, -1, 1)
|
64 |
+
out = torch.zeros_like(key).reshape(batch, seq_length, num_heads, head_size)
|
65 |
+
|
66 |
+
for current_index in range(seq_length):
|
67 |
+
current_receptance = receptance[:, :, current_index:current_index+1, :]
|
68 |
+
current_key = key[:, :, :, current_index:current_index+1]
|
69 |
+
current_value = value[:, :, current_index:current_index+1, :]
|
70 |
+
current_time_decay = time_decay[:, :, :, current_index:current_index+1]
|
71 |
+
attention_output = current_key @ current_value
|
72 |
+
out[:, current_index] = (current_receptance @ (time_first * attention_output + state)).squeeze(2)
|
73 |
+
with torch.no_grad():
|
74 |
+
state = attention_output + current_time_decay * state
|
75 |
+
|
76 |
+
return out, state
|
77 |
+
|
78 |
+
def rwkv6_linear_attention(
|
79 |
+
training,
|
80 |
+
receptance,
|
81 |
+
key,
|
82 |
+
value,
|
83 |
+
time_decay,
|
84 |
+
time_first,
|
85 |
+
state,
|
86 |
+
):
|
87 |
+
no_cuda = any(t.device.type != "cuda" for t in [time_decay, time_first, receptance, key, value])
|
88 |
+
# Launching the CUDA kernel for just one token will actually be slower (there is no for loop in the CPU version
|
89 |
+
# in this case).
|
90 |
+
one_token = key.size(1) == 1
|
91 |
+
if not training or no_cuda or one_token:
|
92 |
+
return rwkv6_linear_attention_cpu(
|
93 |
+
receptance, key, value, time_decay, time_first, state
|
94 |
+
)
|
95 |
+
else:
|
96 |
+
batch, seq_length, _ = receptance.shape
|
97 |
+
num_heads, head_size = time_first.shape
|
98 |
+
key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, T, H, K -> B, H, T, K
|
99 |
+
value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, T, H, K - > B, H, T, V
|
100 |
+
receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, H, T, K
|
101 |
+
time_decay = -torch.exp(time_decay.float()).view(batch, seq_length, num_heads, head_size).permute(0, 2, 1, 3) # B, T, H, K -> B, H, T, K
|
102 |
+
time_first = time_first.float().reshape(num_heads, head_size) # H, K
|
103 |
+
out, state = fused_recurrent_rwkv6(receptance, key, value, time_decay, time_first, scale=1.0, initial_state=state, output_final_state=True)
|
104 |
+
return out.transpose(1, 2), state
|
105 |
+
|
106 |
+
|
107 |
+
class Rwkv6SelfAttention(nn.Module):
|
108 |
+
def __init__(self, config, layer_id=0):
|
109 |
+
super().__init__()
|
110 |
+
self.config = config
|
111 |
+
self.layer_id = layer_id
|
112 |
+
hidden_size = config.hidden_size
|
113 |
+
attention_hidden_size = config.attention_hidden_size
|
114 |
+
self.attention_hidden_size = attention_hidden_size
|
115 |
+
head_size = config.head_size
|
116 |
+
num_heads = attention_hidden_size // head_size
|
117 |
+
|
118 |
+
self.time_maa_x = nn.Parameter(torch.empty(1, 1, hidden_size))
|
119 |
+
self.time_maa_w = nn.Parameter(torch.empty(1, 1, hidden_size))
|
120 |
+
self.time_maa_k = nn.Parameter(torch.empty(1, 1, hidden_size))
|
121 |
+
self.time_maa_v = nn.Parameter(torch.empty(1, 1, hidden_size))
|
122 |
+
self.time_maa_r = nn.Parameter(torch.empty(1, 1, hidden_size))
|
123 |
+
self.time_maa_g = nn.Parameter(torch.empty(1, 1, hidden_size))
|
124 |
+
|
125 |
+
TIME_MIX_EXTRA_DIM = 32 # generate TIME_MIX for w,k,v,r,g
|
126 |
+
self.time_maa_w1 = nn.Parameter(torch.empty(hidden_size, TIME_MIX_EXTRA_DIM*5))
|
127 |
+
self.time_maa_w2 = nn.Parameter(torch.empty(5, TIME_MIX_EXTRA_DIM, hidden_size))
|
128 |
+
|
129 |
+
self.time_decay = nn.Parameter(torch.empty(1, 1, attention_hidden_size))
|
130 |
+
|
131 |
+
TIME_DECAY_EXTRA_DIM = 64
|
132 |
+
self.time_decay_w1 = nn.Parameter(torch.empty(hidden_size, TIME_DECAY_EXTRA_DIM))
|
133 |
+
self.time_decay_w2 = nn.Parameter(torch.empty(TIME_DECAY_EXTRA_DIM, attention_hidden_size))
|
134 |
+
|
135 |
+
self.time_faaaa = nn.Parameter(torch.empty(num_heads, config.head_size))
|
136 |
+
|
137 |
+
|
138 |
+
self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
|
139 |
+
self.receptance = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
140 |
+
self.key = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
141 |
+
self.value = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
142 |
+
self.gate = nn.Linear(hidden_size, attention_hidden_size, bias=False)
|
143 |
+
self.output = nn.Linear(attention_hidden_size, hidden_size, bias=False)
|
144 |
+
self.ln_x = nn.GroupNorm(num_heads, hidden_size, eps=(1e-5)*(config.head_size_divisor**2))
|
145 |
+
|
146 |
+
def extract_key_value(self, hidden, state=None):
|
147 |
+
# Mix hidden with the previous timestep to produce key, value, receptance
|
148 |
+
if hidden.size(1) == 1 and state is not None:
|
149 |
+
shifted = state[0][:, :, self.layer_id]
|
150 |
+
else:
|
151 |
+
shifted = self.time_shift(hidden)
|
152 |
+
if state is not None:
|
153 |
+
shifted[:, 0] = state[0][:, :, self.layer_id]
|
154 |
+
if len(shifted.size()) == 2:
|
155 |
+
shifted = shifted.unsqueeze(1)
|
156 |
+
|
157 |
+
x = hidden
|
158 |
+
|
159 |
+
B, T, C = hidden.shape
|
160 |
+
|
161 |
+
xx = shifted - x
|
162 |
+
|
163 |
+
xxx = x + xx * self.time_maa_x
|
164 |
+
xxx = torch.tanh(xxx @ self.time_maa_w1).view(B*T, 5, -1).transpose(0, 1)
|
165 |
+
xxx = torch.bmm(xxx, self.time_maa_w2).view(5, B, T, -1)
|
166 |
+
mw, mk, mv, mr, mg = xxx.unbind(dim=0)
|
167 |
+
|
168 |
+
time_decay = x + xx * (self.time_maa_w + mw)
|
169 |
+
key = x + xx * (self.time_maa_k + mk)
|
170 |
+
value = x + xx * (self.time_maa_v + mv)
|
171 |
+
receptance = x + xx * (self.time_maa_r + mr)
|
172 |
+
gate = x + xx * (self.time_maa_g + mg)
|
173 |
+
|
174 |
+
receptance = self.receptance(receptance)
|
175 |
+
key = self.key(key)
|
176 |
+
value = self.value(value)
|
177 |
+
gate = F.silu(self.gate(gate))
|
178 |
+
|
179 |
+
time_decay = torch.tanh(time_decay @ self.time_decay_w1) @ self.time_decay_w2
|
180 |
+
time_decay = self.time_decay + time_decay
|
181 |
+
|
182 |
+
if state is not None:
|
183 |
+
state[0][:, :, self.layer_id] = hidden[:, -1]
|
184 |
+
|
185 |
+
return receptance, key, value, gate, time_decay, state
|
186 |
+
|
187 |
+
def forward(self, hidden, state=None, use_cache=False, seq_mode=True):
|
188 |
+
receptance, key, value, gate, time_decay, state = self.extract_key_value(hidden, state=state)
|
189 |
+
|
190 |
+
B,T,C = receptance.shape
|
191 |
+
H, S = self.time_faaaa.shape
|
192 |
+
|
193 |
+
layer_state = state[1][:, :, :, :, self.layer_id] if state is not None else None
|
194 |
+
out, layer_state = rwkv6_linear_attention(
|
195 |
+
self.training, receptance, key, value, time_decay, self.time_faaaa, layer_state,
|
196 |
+
)
|
197 |
+
|
198 |
+
if layer_state is not None:
|
199 |
+
state[1][:, :, :, :, self.layer_id] = layer_state
|
200 |
+
|
201 |
+
out = out.reshape(B * T, H * S)
|
202 |
+
out = F.group_norm(out, num_groups=H, weight=self.ln_x.weight.to(out.dtype), bias=self.ln_x.bias.to(out.dtype), eps=self.ln_x.eps).reshape(B, T, H * S)
|
203 |
+
out = out.to(dtype=hidden.dtype) * gate
|
204 |
+
out = self.output(out)
|
205 |
+
return out, state
|
206 |
+
|
207 |
+
|
208 |
+
class Rwkv6FeedForward(nn.Module):
|
209 |
+
def __init__(self, config, layer_id=0):
|
210 |
+
super().__init__()
|
211 |
+
self.config = config
|
212 |
+
self.layer_id = layer_id
|
213 |
+
hidden_size = config.hidden_size
|
214 |
+
# https://github.com/BlinkDL/RWKV-LM/blob/3db37a72356b736966ddd377268f02b80963af3f/RWKV-v4neo/train.py#L168
|
215 |
+
intermediate_size = (
|
216 |
+
config.intermediate_size
|
217 |
+
if config.intermediate_size is not None
|
218 |
+
else int((config.hidden_size * 3.5) // 32 * 32)
|
219 |
+
)
|
220 |
+
|
221 |
+
self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
|
222 |
+
self.time_maa_k = nn.Parameter(torch.empty(1, 1, hidden_size))
|
223 |
+
self.time_maa_r = nn.Parameter(torch.empty(1, 1, hidden_size))
|
224 |
+
|
225 |
+
self.key = nn.Linear(hidden_size, intermediate_size, bias=False)
|
226 |
+
self.receptance = nn.Linear(hidden_size, hidden_size, bias=False)
|
227 |
+
self.value = nn.Linear(intermediate_size, hidden_size, bias=False)
|
228 |
+
|
229 |
+
def forward(self, hidden, state=None):
|
230 |
+
if hidden.size(1) == 1 and state is not None:
|
231 |
+
shifted = state[2][:, :, self.layer_id]
|
232 |
+
else:
|
233 |
+
shifted = self.time_shift(hidden)
|
234 |
+
if state is not None:
|
235 |
+
shifted[:, 0] = state[2][:, :, self.layer_id]
|
236 |
+
if len(shifted.size()) == 2:
|
237 |
+
shifted = shifted.unsqueeze(1)
|
238 |
+
|
239 |
+
delta_hidden_to_shifted = shifted - hidden
|
240 |
+
key = hidden + delta_hidden_to_shifted * self.time_maa_k
|
241 |
+
receptance = hidden + delta_hidden_to_shifted * self.time_maa_r
|
242 |
+
|
243 |
+
key = torch.square(torch.relu(self.key(key)))
|
244 |
+
value = self.value(key)
|
245 |
+
receptance = torch.sigmoid(self.receptance(receptance))
|
246 |
+
|
247 |
+
if state is not None:
|
248 |
+
state[2][:, :, self.layer_id] = hidden[:, -1]
|
249 |
+
|
250 |
+
return receptance * value, state
|
251 |
+
|
252 |
+
|
253 |
+
class Rwkv6Block(nn.Module):
|
254 |
+
def __init__(self, config, layer_id):
|
255 |
+
super().__init__()
|
256 |
+
self.config = config
|
257 |
+
self.layer_id = layer_id
|
258 |
+
|
259 |
+
if layer_id == 0:
|
260 |
+
self.pre_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
|
261 |
+
|
262 |
+
self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
|
263 |
+
self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
|
264 |
+
|
265 |
+
self.attention = Rwkv6SelfAttention(config, layer_id)
|
266 |
+
self.feed_forward = Rwkv6FeedForward(config, layer_id)
|
267 |
+
|
268 |
+
def forward(self, hidden, state=None, use_cache=False, output_attentions=False, seq_mode=True):
|
269 |
+
if self.layer_id == 0:
|
270 |
+
hidden = self.pre_ln(hidden)
|
271 |
+
attention, state = self.attention(self.ln1(hidden), state=state, use_cache=use_cache, seq_mode=seq_mode)
|
272 |
+
hidden = hidden + attention
|
273 |
+
|
274 |
+
feed_forward, state = self.feed_forward(self.ln2(hidden), state=state)
|
275 |
+
hidden = hidden + feed_forward
|
276 |
+
|
277 |
+
outputs = (hidden, state)
|
278 |
+
if output_attentions:
|
279 |
+
outputs += (attention,)
|
280 |
+
else:
|
281 |
+
outputs += (None,)
|
282 |
+
|
283 |
+
return outputs
|
284 |
+
|
285 |
+
|
286 |
+
class Rwkv6PreTrainedModel(PreTrainedModel):
|
287 |
+
"""
|
288 |
+
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
289 |
+
models.
|
290 |
+
"""
|
291 |
+
|
292 |
+
config_class = Rwkv6Config
|
293 |
+
base_model_prefix = "rwkv6"
|
294 |
+
_no_split_modules = ["Rwkv6Block"]
|
295 |
+
_keep_in_fp32_modules = ["time_decay", "time_first"]
|
296 |
+
supports_gradient_checkpointing = True
|
297 |
+
|
298 |
+
def _init_weights(self, module):
|
299 |
+
"""Initialize the weights."""
|
300 |
+
if isinstance(module, Rwkv6SelfAttention):
|
301 |
+
layer_id = module.layer_id
|
302 |
+
num_hidden_layers = module.config.num_hidden_layers
|
303 |
+
hidden_size = module.config.hidden_size
|
304 |
+
attention_hidden_size = module.attention_hidden_size
|
305 |
+
head_size = module.config.head_size
|
306 |
+
num_heads = attention_hidden_size // head_size
|
307 |
+
|
308 |
+
ratio_0_to_1 = layer_id / (num_hidden_layers - 1) # 0 to 1
|
309 |
+
ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers) # 1 to ~0
|
310 |
+
|
311 |
+
time_weight = torch.tensor(
|
312 |
+
[i / hidden_size for i in range(hidden_size)],
|
313 |
+
dtype=module.time_maa_k.dtype,
|
314 |
+
device=module.time_maa_k.device,
|
315 |
+
)
|
316 |
+
time_weight = time_weight[None, None, :]
|
317 |
+
|
318 |
+
decay_speed = [
|
319 |
+
-6.0 + 5.0 * (h / (attention_hidden_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
|
320 |
+
for h in range(attention_hidden_size)
|
321 |
+
]
|
322 |
+
decay_speed = torch.tensor(decay_speed, dtype=module.time_decay.dtype, device=module.time_decay.device)
|
323 |
+
tmp = torch.tensor(
|
324 |
+
[
|
325 |
+
(1.0 - (i / (attention_hidden_size - 1.0))) * ratio_0_to_1 + 0.1 * ((i + 1) % 3 - 1)
|
326 |
+
for i in range(attention_hidden_size)
|
327 |
+
],
|
328 |
+
dtype=module.time_faaaa.dtype,
|
329 |
+
device=module.time_faaaa.device,
|
330 |
+
)
|
331 |
+
|
332 |
+
with torch.no_grad():
|
333 |
+
module.time_maa_x.data = 1.0 - torch.pow(time_weight, ratio_1_to_almost0)
|
334 |
+
module.time_maa_w.data = 1.0 - torch.pow(time_weight, ratio_1_to_almost0)
|
335 |
+
module.time_maa_k.data = 1.0 - torch.pow(time_weight, ratio_1_to_almost0)
|
336 |
+
module.time_maa_v.data = 1.0 - (torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1)
|
337 |
+
module.time_maa_r.data = 1.0 - torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
|
338 |
+
module.time_maa_g.data = 1.0 - torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
|
339 |
+
|
340 |
+
TIME_MIX_EXTRA_DIM = 32 # generate TIME_MIX for w,k,v,r,g
|
341 |
+
module.time_maa_w1.data = torch.zeros(hidden_size, TIME_MIX_EXTRA_DIM*5, dtype=module.time_maa_w1.dtype, device=module.time_maa_w1.device).uniform_(-1e-4, 1e-4)
|
342 |
+
module.time_maa_w2.data = torch.zeros(5, TIME_MIX_EXTRA_DIM, hidden_size, dtype=module.time_maa_w2.dtype, device=module.time_maa_w2.device).uniform_(-1e-4, 1e-4)
|
343 |
+
|
344 |
+
TIME_DECAY_EXTRA_DIM = 64
|
345 |
+
module.time_decay_w1.data = torch.zeros(hidden_size, TIME_DECAY_EXTRA_DIM, dtype=module.time_decay_w1.dtype, device=module.time_decay_w1.device).uniform_(-1e-4, 1e-4)
|
346 |
+
module.time_decay_w2.data = torch.zeros(TIME_DECAY_EXTRA_DIM, attention_hidden_size, dtype=module.time_decay_w2.dtype, device=module.time_decay_w2.device).uniform_(-1e-4, 1e-4)
|
347 |
+
|
348 |
+
module.time_decay.data = decay_speed.reshape(num_heads, head_size)
|
349 |
+
module.time_faaaa.data = tmp.reshape(num_heads, head_size)
|
350 |
+
|
351 |
+
elif isinstance(module, Rwkv6FeedForward):
|
352 |
+
layer_id = module.layer_id
|
353 |
+
num_hidden_layers = module.config.num_hidden_layers
|
354 |
+
hidden_size = module.config.hidden_size
|
355 |
+
|
356 |
+
ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers) # 1 to ~0
|
357 |
+
|
358 |
+
time_weight = torch.tensor(
|
359 |
+
[i / hidden_size for i in range(hidden_size)],
|
360 |
+
dtype=module.time_maa_k.dtype,
|
361 |
+
device=module.time_maa_k.device,
|
362 |
+
)
|
363 |
+
time_weight = time_weight[None, None, :]
|
364 |
+
|
365 |
+
with torch.no_grad():
|
366 |
+
module.time_maa_k.data = 1.0 - torch.pow(time_weight, ratio_1_to_almost0)
|
367 |
+
module.time_maa_r.data = 1.0 - torch.pow(time_weight, ratio_1_to_almost0)
|
368 |
+
|
369 |
+
|
370 |
+
@dataclass
|
371 |
+
class Rwkv6Output(ModelOutput):
|
372 |
+
"""
|
373 |
+
Class for the RWKV model outputs.
|
374 |
+
|
375 |
+
Args:
|
376 |
+
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
|
377 |
+
Sequence of hidden-states at the output of the last layer of the model.
|
378 |
+
state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`):
|
379 |
+
The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
|
380 |
+
avoid providing the old `input_ids`.
|
381 |
+
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
382 |
+
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
383 |
+
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of
|
384 |
+
the model at the output of each layer plus the optional initial embedding outputs.
|
385 |
+
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
386 |
+
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
387 |
+
sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
|
388 |
+
the self-attention heads.
|
389 |
+
"""
|
390 |
+
|
391 |
+
last_hidden_state: torch.FloatTensor = None
|
392 |
+
state: Optional[List[torch.FloatTensor]] = None
|
393 |
+
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
394 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
395 |
+
|
396 |
+
|
397 |
+
@dataclass
|
398 |
+
class Rwkv6CausalLMOutput(ModelOutput):
|
399 |
+
"""
|
400 |
+
Base class for causal language model (or autoregressive) outputs.
|
401 |
+
|
402 |
+
Args:
|
403 |
+
loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
|
404 |
+
Language modeling loss (for next-token prediction).
|
405 |
+
logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
|
406 |
+
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
|
407 |
+
state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`):
|
408 |
+
The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
|
409 |
+
avoid providing the old `input_ids`.
|
410 |
+
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
411 |
+
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
412 |
+
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of
|
413 |
+
the model at the output of each layer plus the optional initial embedding outputs.
|
414 |
+
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
415 |
+
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
416 |
+
sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
|
417 |
+
the self-attention heads.
|
418 |
+
"""
|
419 |
+
|
420 |
+
loss: Optional[torch.FloatTensor] = None
|
421 |
+
logits: torch.FloatTensor = None
|
422 |
+
state: Optional[List[torch.FloatTensor]] = None
|
423 |
+
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
424 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
425 |
+
|
426 |
+
|
427 |
+
RWKV6_START_DOCSTRING = r"""
|
428 |
+
This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
|
429 |
+
library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
|
430 |
+
etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)
|
431 |
+
subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
|
432 |
+
general usage and behavior.
|
433 |
+
|
434 |
+
Parameters:
|
435 |
+
config ([`Rwkv6Config`]): Model configuration class with all the parameters of the model.
|
436 |
+
Initializing with a config file does not load the weights associated with the model, only the
|
437 |
+
configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
|
438 |
+
"""
|
439 |
+
|
440 |
+
RWKV6_INPUTS_DOCSTRING = r"""
|
441 |
+
Args:
|
442 |
+
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
|
443 |
+
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else
|
444 |
+
`past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
|
445 |
+
sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their
|
446 |
+
past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See
|
447 |
+
[`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
|
448 |
+
IDs?](../glossary#input-ids)
|
449 |
+
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
450 |
+
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
|
451 |
+
is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
|
452 |
+
model's internal embedding lookup matrix.
|
453 |
+
state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*):
|
454 |
+
If passed along, the model uses the previous state in all the blocks (which will give the output for the
|
455 |
+
`input_ids` provided as if the model add `state_input_ids + input_ids` as context).
|
456 |
+
use_cache (`bool`, *optional*):
|
457 |
+
If set to `True`, the last state is returned and can be used to quickly generate the next logits.
|
458 |
+
output_attentions (`bool`, *optional*):
|
459 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
|
460 |
+
tensors for more detail.
|
461 |
+
output_hidden_states (`bool`, *optional*):
|
462 |
+
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
|
463 |
+
more detail.
|
464 |
+
return_dict (`bool`, *optional*):
|
465 |
+
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
466 |
+
"""
|
467 |
+
|
468 |
+
|
469 |
+
@add_start_docstrings(
|
470 |
+
"The bare RWKV6 Model transformer outputting raw hidden-states without any specific head on top.",
|
471 |
+
RWKV6_START_DOCSTRING,
|
472 |
+
)
|
473 |
+
class Rwkv6Model(Rwkv6PreTrainedModel):
|
474 |
+
def __init__(self, config):
|
475 |
+
super().__init__(config)
|
476 |
+
|
477 |
+
self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
|
478 |
+
self.blocks = nn.ModuleList([Rwkv6Block(config, layer_id=idx) for idx in range(config.num_hidden_layers)])
|
479 |
+
self.ln_out = nn.LayerNorm(config.hidden_size)
|
480 |
+
|
481 |
+
self.layers_are_rescaled = False
|
482 |
+
self.gradient_checkpointing = False
|
483 |
+
|
484 |
+
# Initialize weights and apply final processing
|
485 |
+
self.post_init()
|
486 |
+
|
487 |
+
def get_input_embeddings(self):
|
488 |
+
return self.embeddings
|
489 |
+
|
490 |
+
def set_input_embeddings(self, new_embeddings):
|
491 |
+
self.embeddings = new_embeddings
|
492 |
+
|
493 |
+
@add_start_docstrings_to_model_forward(RWKV6_INPUTS_DOCSTRING)
|
494 |
+
@add_code_sample_docstrings(
|
495 |
+
checkpoint=_CHECKPOINT_FOR_DOC,
|
496 |
+
output_type=Rwkv6Output,
|
497 |
+
config_class=_CONFIG_FOR_DOC,
|
498 |
+
)
|
499 |
+
def forward(
|
500 |
+
self,
|
501 |
+
input_ids: Optional[torch.LongTensor] = None,
|
502 |
+
attention_mask: Optional[torch.LongTensor] = None, # noqa
|
503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
504 |
+
state: Optional[List[torch.FloatTensor]] = None,
|
505 |
+
use_cache: Optional[bool] = None,
|
506 |
+
output_attentions: Optional[bool] = None,
|
507 |
+
output_hidden_states: Optional[bool] = None,
|
508 |
+
return_dict: Optional[bool] = None,
|
509 |
+
) -> Union[Tuple, Rwkv6Output]:
|
510 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
511 |
+
output_hidden_states = (
|
512 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
513 |
+
)
|
514 |
+
# FIXME - training is supportable with the CUDA code
|
515 |
+
# rwkv6 only support inference in huggingface.
|
516 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
517 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
518 |
+
|
519 |
+
if self.training == self.layers_are_rescaled and (
|
520 |
+
self.embeddings.weight.dtype == torch.float16 or self.embeddings.weight.dtype == torch.bfloat16
|
521 |
+
):
|
522 |
+
self._rescale_layers()
|
523 |
+
|
524 |
+
if input_ids is not None and inputs_embeds is not None:
|
525 |
+
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
|
526 |
+
elif input_ids is None and inputs_embeds is None:
|
527 |
+
raise ValueError("You have to specify either input_ids or inputs_embeds")
|
528 |
+
|
529 |
+
if inputs_embeds is None:
|
530 |
+
inputs_embeds = self.embeddings(input_ids)
|
531 |
+
|
532 |
+
if state is None:
|
533 |
+
state = []
|
534 |
+
head_size = self.config.head_size
|
535 |
+
num_heads = self.config.attention_hidden_size // head_size
|
536 |
+
state_attn_x = torch.zeros(
|
537 |
+
(inputs_embeds.size(0), self.config.hidden_size, self.config.num_hidden_layers),
|
538 |
+
dtype=inputs_embeds.dtype,
|
539 |
+
requires_grad=False,
|
540 |
+
device=inputs_embeds.device,
|
541 |
+
).contiguous()
|
542 |
+
state_attn_kv = torch.zeros(
|
543 |
+
(
|
544 |
+
inputs_embeds.size(0),
|
545 |
+
num_heads,
|
546 |
+
head_size,
|
547 |
+
head_size,
|
548 |
+
self.config.num_hidden_layers,
|
549 |
+
),
|
550 |
+
dtype=torch.float32,
|
551 |
+
requires_grad=False,
|
552 |
+
device=inputs_embeds.device,
|
553 |
+
).contiguous()
|
554 |
+
state_ffn_x = torch.zeros(
|
555 |
+
(inputs_embeds.size(0), self.config.hidden_size, self.config.num_hidden_layers),
|
556 |
+
dtype=inputs_embeds.dtype,
|
557 |
+
requires_grad=False,
|
558 |
+
device=inputs_embeds.device,
|
559 |
+
).contiguous()
|
560 |
+
state.append(state_attn_x)
|
561 |
+
state.append(state_attn_kv)
|
562 |
+
state.append(state_ffn_x)
|
563 |
+
|
564 |
+
seq_mode = inputs_embeds.shape[1] > 1
|
565 |
+
hidden_states = inputs_embeds
|
566 |
+
|
567 |
+
all_self_attentions = () if output_attentions else None
|
568 |
+
all_hidden_states = () if output_hidden_states else None
|
569 |
+
for idx, block in enumerate(self.blocks):
|
570 |
+
hidden_states, state, attentions = block(
|
571 |
+
hidden_states, state=state, use_cache=use_cache, output_attentions=output_attentions, seq_mode=seq_mode
|
572 |
+
)
|
573 |
+
if (
|
574 |
+
self.layers_are_rescaled
|
575 |
+
and self.config.rescale_every > 0
|
576 |
+
and (idx + 1) % self.config.rescale_every == 0
|
577 |
+
):
|
578 |
+
hidden_states = hidden_states / 2
|
579 |
+
|
580 |
+
if output_hidden_states:
|
581 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
582 |
+
|
583 |
+
if output_attentions:
|
584 |
+
all_self_attentions = all_self_attentions + (attentions,)
|
585 |
+
|
586 |
+
hidden_states = self.ln_out(hidden_states)
|
587 |
+
|
588 |
+
if output_hidden_states:
|
589 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
590 |
+
|
591 |
+
if not return_dict:
|
592 |
+
return (hidden_states, state, all_hidden_states, all_self_attentions)
|
593 |
+
|
594 |
+
return Rwkv6Output(
|
595 |
+
last_hidden_state=hidden_states,
|
596 |
+
state=state,
|
597 |
+
hidden_states=all_hidden_states, # None
|
598 |
+
attentions=all_self_attentions, # None
|
599 |
+
)
|
600 |
+
|
601 |
+
def _rescale_layers(self):
|
602 |
+
# Layers should be rescaled for inference only.
|
603 |
+
if self.layers_are_rescaled == (not self.training):
|
604 |
+
return
|
605 |
+
if self.config.rescale_every > 0:
|
606 |
+
with torch.no_grad():
|
607 |
+
for block_id, block in enumerate(self.blocks):
|
608 |
+
if self.training:
|
609 |
+
block.attention.output.weight.mul_(2 ** int(block_id // self.config.rescale_every))
|
610 |
+
block.feed_forward.value.weight.mul_(2 ** int(block_id // self.config.rescale_every))
|
611 |
+
else:
|
612 |
+
# Deal with quantization statistics
|
613 |
+
if hasattr(block.attention.output.weight, "SCB"):
|
614 |
+
block.attention.output.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
|
615 |
+
block.feed_forward.value.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
|
616 |
+
elif hasattr(block.attention.output.weight, "quant_state"):
|
617 |
+
self._bnb_4bit_dequantize_and_rescale(block.attention.output, block_id)
|
618 |
+
self._bnb_4bit_dequantize_and_rescale(block.feed_forward.value, block_id)
|
619 |
+
else:
|
620 |
+
block.attention.output.weight.div_(2 ** int(block_id // self.config.rescale_every))
|
621 |
+
block.feed_forward.value.weight.div_(2 ** int(block_id // self.config.rescale_every))
|
622 |
+
|
623 |
+
self.layers_are_rescaled = not self.training
|
624 |
+
|
625 |
+
def _bnb_4bit_dequantize_and_rescale(self, target_layer, block_id):
|
626 |
+
r"""
|
627 |
+
Perform the dequantization and rescaling of the weights of a given layer. After that operation the layer will
|
628 |
+
be quantized again.
|
629 |
+
"""
|
630 |
+
if not is_bitsandbytes_available():
|
631 |
+
raise ImportError("Please install bitsandbytes to use this method.")
|
632 |
+
import bitsandbytes as bnb
|
633 |
+
|
634 |
+
dequant_weights = bnb.functional.dequantize_4bit(target_layer.weight.data, target_layer.weight.quant_state)
|
635 |
+
|
636 |
+
dequant_weights.div_(2 ** int(block_id // self.config.rescale_every))
|
637 |
+
|
638 |
+
# re-quantize the model:
|
639 |
+
# we need to put it first on CPU then back to the device
|
640 |
+
# this will create an overhead :/
|
641 |
+
# We set requires_grad=False as we cannot compute gradients on top of 4bit parameters anyway and to avoid
|
642 |
+
# bugs with bnb
|
643 |
+
quant_weight = bnb.nn.Params4bit(dequant_weights.to("cpu"), requires_grad=False).to(dequant_weights.device)
|
644 |
+
setattr(target_layer, "weight", quant_weight)
|
645 |
+
|
646 |
+
|
647 |
+
# copied from HuggingFace https://github.com/huggingface/transformers/blob/main/src/transformers/models/rwkv/modeling_rwkv.py
|
648 |
+
@add_start_docstrings(
|
649 |
+
"""
|
650 |
+
The RWKV6 Model transformer with a language modeling head on top (linear layer with weights tied to the input
|
651 |
+
embeddings).
|
652 |
+
""",
|
653 |
+
RWKV6_START_DOCSTRING,
|
654 |
+
)
|
655 |
+
class Rwkv6ForCausalLM(Rwkv6PreTrainedModel):
|
656 |
+
_tied_weights_keys = ["head.weight"]
|
657 |
+
|
658 |
+
def __init__(self, config):
|
659 |
+
super().__init__(config)
|
660 |
+
self.rwkv = Rwkv6Model(config)
|
661 |
+
self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
662 |
+
|
663 |
+
# Initialize weights and apply final processing
|
664 |
+
self.post_init()
|
665 |
+
|
666 |
+
def get_output_embeddings(self):
|
667 |
+
return self.head
|
668 |
+
|
669 |
+
def set_output_embeddings(self, new_embeddings):
|
670 |
+
self.head = new_embeddings
|
671 |
+
|
672 |
+
def prepare_inputs_for_generation(self, input_ids, state=None, inputs_embeds=None, **kwargs):
|
673 |
+
# only last token for inputs_ids if the state is passed along.
|
674 |
+
if state is not None:
|
675 |
+
input_ids = input_ids[:, -1].unsqueeze(-1)
|
676 |
+
|
677 |
+
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
678 |
+
if inputs_embeds is not None and state is None:
|
679 |
+
model_inputs = {"inputs_embeds": inputs_embeds}
|
680 |
+
else:
|
681 |
+
model_inputs = {"input_ids": input_ids}
|
682 |
+
|
683 |
+
model_inputs["state"] = state
|
684 |
+
return model_inputs
|
685 |
+
|
686 |
+
@add_start_docstrings_to_model_forward(RWKV6_INPUTS_DOCSTRING)
|
687 |
+
@add_code_sample_docstrings(
|
688 |
+
checkpoint=_CHECKPOINT_FOR_DOC,
|
689 |
+
output_type=Rwkv6CausalLMOutput,
|
690 |
+
config_class=_CONFIG_FOR_DOC,
|
691 |
+
)
|
692 |
+
def forward(
|
693 |
+
self,
|
694 |
+
input_ids: Optional[torch.LongTensor] = None,
|
695 |
+
attention_mask: Optional[torch.LongTensor] = None,
|
696 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
697 |
+
state: Optional[List[torch.FloatTensor]] = None,
|
698 |
+
labels: Optional[torch.LongTensor] = None,
|
699 |
+
use_cache: Optional[bool] = None,
|
700 |
+
output_attentions: Optional[bool] = None,
|
701 |
+
output_hidden_states: Optional[bool] = None,
|
702 |
+
return_dict: Optional[bool] = None,
|
703 |
+
) -> Union[Tuple, Rwkv6CausalLMOutput]:
|
704 |
+
r"""
|
705 |
+
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
706 |
+
Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
|
707 |
+
`labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
|
708 |
+
are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
|
709 |
+
"""
|
710 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
711 |
+
|
712 |
+
outputs = self.rwkv(
|
713 |
+
input_ids,
|
714 |
+
inputs_embeds=inputs_embeds,
|
715 |
+
state=state,
|
716 |
+
use_cache=use_cache,
|
717 |
+
output_attentions=output_attentions,
|
718 |
+
output_hidden_states=output_hidden_states,
|
719 |
+
return_dict=return_dict,
|
720 |
+
)
|
721 |
+
hidden_states = outputs[0]
|
722 |
+
|
723 |
+
logits = self.head(hidden_states)
|
724 |
+
|
725 |
+
loss = None
|
726 |
+
if labels is not None:
|
727 |
+
# move labels to correct device to enable model parallelism
|
728 |
+
labels = labels.to(logits.device)
|
729 |
+
# Shift so that tokens < n predict n
|
730 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
731 |
+
shift_labels = labels[..., 1:].contiguous()
|
732 |
+
# Flatten the tokens
|
733 |
+
loss_fct = CrossEntropyLoss()
|
734 |
+
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
|
735 |
+
|
736 |
+
if not return_dict:
|
737 |
+
output = (logits,) + outputs[1:]
|
738 |
+
return ((loss,) + output) if loss is not None else output
|
739 |
+
|
740 |
+
return Rwkv6CausalLMOutput(
|
741 |
+
loss=loss,
|
742 |
+
logits=logits,
|
743 |
+
state=outputs.state,
|
744 |
+
hidden_states=outputs.hidden_states,
|
745 |
+
attentions=outputs.attentions,
|
746 |
+
)
|
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,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"bos_token": "<s>",
|
3 |
+
"eos_token": "<s>",
|
4 |
+
"pad_token": "<s>",
|
5 |
+
"unk_token": "<s>"
|
6 |
+
}
|
tokenizer_config.json
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"name_or_path": "rwkv-6-tokenizer",
|
3 |
+
"add_prefix_space": false,
|
4 |
+
"tokenizer_class": "Rwkv6Tokenizer",
|
5 |
+
"use_fast": false,
|
6 |
+
"auto_map": {
|
7 |
+
"AutoTokenizer": [
|
8 |
+
"hf_rwkv_tokenizer.Rwkv6Tokenizer",
|
9 |
+
null
|
10 |
+
]
|
11 |
+
}
|
12 |
+
}
|