[`RWKV5`] Add support for RWKV5 model
Author: ArthurZucker
Pull request: yes
Labels: name: New model

# What does this PR do?
Adds RWKV5, superseeds #26963

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29095). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Not sure if this is the appropriate place to post this, but I've been running into some issues when using the tokenizer `AutoTokenizer.from_pretrained("RWKV/v5-Eagle-7B-HF")` and I was told to forward my message here. Wondering if this is a feature or a bug? 

Original issue https://github.com/BlinkDL/RWKV-LM/issues/224


😅 could you open a seperate issue and ping me there?! 🤗 

Gotcha, starting a separate issue makes sense. Thanks! https://github.com/huggingface/transformers/issues/29111

Sorry all got caught in many things, will be off next week but the week after that will be my focus! 🤗 

BTW @BBuf don't you think it would be great to have a seperate github repo with installables for the kernels? (you can track usage, easier to maintain and propagate to here!) WDYT? 

> BTW @BBuf don't you think it would be great to have a seperate github repo with installables for the kernels? (you can track usage, easier to maintain and propagate to here!) WDYT?

I fully agree. I am planning to implement the rwkv kernel using Triton and maintain it with a standalone library. This should be done after this PR is merged. I think we can remove the CUDA kernel implementation from this PR.

Yeah and same for RWKV first version that was merged, we will have a cleaner implementation! 🤗 thanks for informing me 

I encountered some errors while using the tokenizer of rwkv6.
<img width="405" alt="微信图片_20240325151811" src="https://github.com/huggingface/transformers/assets/139205286/324f91df-6629-4b69-babf-2c72c21b8541">
The code is as follows：
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("rwkv6", trust_remote_code=True, torch_dtype=torch.float32).to(0)
tokenizer = AutoTokenizer.from_pretrained("rwkv6", trust_remote_code=True, padding_side='left')


list = ['Chinese: 他补充道：“我们现在有 4 个月大没有糖尿病的老鼠，但它们曾经得过该病。”\n\nEnglish:', 
        'Chinese: 埃胡德·乌尔博士（新斯科舍省哈利法克斯市达尔豪西大学医学教授，加拿大糖尿病协会临床与科学部门教授）提醒，这项研究仍处在早期阶段。\n\nEnglish:', 
        'Chinese: 和其他一些专家一样，他对糖尿病能否治愈持怀疑态度。他指出，这些发现与已患有 1 型糖尿病的人无关。\n\nEnglish:', 
        'Chinese: 周一，瑞典学院诺贝尔文学委员会常务秘书萨拉·丹尼尔斯在瑞典广播电台的一档节目中向公众宣布，委员会因无法直接联系到鲍勃·迪伦，通知他获得了 2016 年诺贝尔文学奖，已经放弃了与他联系的尝试。\n\nEnglish:']
inputs = tokenizer(list, return_tensors="pt",padding=True).to(0)
print(inputs["input_ids"])
output = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=1.0, top_p=0.3, top_k=0, )
for o in output:
    print(tokenizer.decode(o.tolist(), skip_special_tokens=True))
```

> I encountered some errors while using the tokenizer of rwkv6. <img alt="微信图片_20240325151811" width="405" src="https://private-user-images.githubusercontent.com/139205286/316417628-324f91df-6629-4b69-babf-2c72c21b8541.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTEzNTE1MTMsIm5iZiI6MTcxMTM1MTIxMywicGF0aCI6Ii8xMzkyMDUyODYvMzE2NDE3NjI4LTMyNGY5MWRmLTY2MjktNGI2OS1iYWJmLTJjNzJjMjFiODU0MS5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQwMzI1JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MDMyNVQwNzIwMTNaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT01NGRmODhhZDU2YjM1ZjRhNGNlNzEyMTRhNWQ0MWI2YTdhZWU3MTQ4NTllYTI2OWMxOWIyYzg0OGUzMDE0Mjg5JlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCZhY3Rvcl9pZD0wJmtleV9pZD0wJnJlcG9faWQ9MCJ9.O1q1iFiITN2UwE09HORSj798PeO6wtZD8YUAwF_1E4g"> The code is as follows：
> 
> ```
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer
> 
> model = AutoModelForCausalLM.from_pretrained("rwkv6", trust_remote_code=True, torch_dtype=torch.float32).to(0)
> tokenizer = AutoTokenizer.from_pretrained("rwkv6", trust_remote_code=True, padding_side='left')
> 
> 
> list = ['Chinese: 他补充道：“我们现在有 4 个月大没有糖尿病的老鼠，但它们曾经得过该病。”\n\nEnglish:', 
>         'Chinese: 埃胡德·乌尔博士（新斯科舍省哈利法克斯市达尔豪西大学医学教授，加拿大糖尿病协会临床与科学部门教授）提醒，这项研究仍处在早期阶段。\n\nEnglish:', 
>         'Chinese: 和其他一些专家一样，他对糖尿病能否治愈持怀疑态度。他指出，这些发现与已患有 1 型糖尿病的人无关。\n\nEnglish:', 
>         'Chinese: 周一，瑞典学院诺贝尔文学委员会常务秘书萨拉·丹尼尔斯在瑞典广播电台的一档节目中向公众宣布，委员会因无法直接联系到鲍勃·迪伦，通知他获得了 2016 年诺贝尔文学奖，已经放弃了与他联系的尝试。\n\nEnglish:']
> inputs = tokenizer(list, return_tensors="pt",padding=True).to(0)
> print(inputs["input_ids"])
> output = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=1.0, top_p=0.3, top_k=0, )
> for o in output:
>     print(tokenizer.decode(o.tolist(), skip_special_tokens=True))
> ```

@ArthurZucker  It seems that the `padding_side='left'` parameter is not working when tokenizer is initialized. I am uncertain whether this input parameter needs to be taken into account during the implementation of the tokenizer.

They! Should be I'll debug this!

@BBuf for left padding here is what I have:
```python
>>> tokenizer = AutoTokenizer.from_pretrained("ArthurZ/rwkv-5-utf", padding_side='left', pad_token="
")

>>> prompt = ["Hey how are you?","男：听说你们公司要派你去南方工作"]
>>> inputs_ = tokenizer(prompt, padding = True)
{'input_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6007, 21835, 21286, 22799, 64], [0, 14601, 19151, 11065, 16735, 10464, 10402, 10678, 11029, 16503, 13818, 10464, 10985, 10934, 13036, 12137, 10460]], 'attention_mask': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
```

or 
```python 
>>> in_list = ['Chinese: 他补充道：“我们现在有 4 个月大没有糖尿病的老鼠，但它们曾经得过该病。”\n\nEnglish:', Chinese: 埃胡德·乌尔博士（新斯科舍省哈利法克斯市达尔豪西大学医学教授，加拿大糖尿病协会临床与科学部门教授）提醒，这项研究仍处在早期阶段。\n\nEnglish:', 'Chinese: 和其他一些专家一样，他对糖尿病能否治愈持怀疑态度。他指出，这些发现与已患有 1 型糖尿病的人无关。\n\nEnglish:', 'Chinese: 周一，瑞典学院诺贝尔文学委员会常务秘书萨拉·丹尼尔斯在瑞典广播电台的一档节目中向公众宣布，委员会因无法直接联系到鲍勃·迪伦，通知他获得了 2016 年诺贝尔文学奖，已经放弃了与他联系的尝试。\n\nEnglish:']

>>> inputs_ = tokenizer(in_list, padding = True)
{'input_ids': [[0, 48407, 59, 33, 10390, 16416, 10655, 17222, 43899, 12605, 10402, 14446, 11454, 13191, 287, 33, 10283, 13190, 11638, 13734, 13191, 15294, 12004, 14662, 14734, 15640, 18459, 19137, 10444, 11885, 10402, 13186, 15486, 12348, 17141, 16721, 14662, 10080, 9823, 261, 48487, 59], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48407, 59, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48407, 59, 0, 284, 33, 11496, 15294, 12004, 14662, 14734, 10370, 13051, 10684, 28329, 11, 48487, 59], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48407, 59, 0, 3493, 636, 33, 12201, 16741, 16873, 11979, 13012, 11871, 11665, 19137, 12145, 15486, 12983, 12269, 10333, 10261, 10390, 15671, 15304, 14734, 11986, 16707, 28329, 11, 48487, 59]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
```
seems to work as expected. Left padding means the zeros are on the left, "right" padding they will be on the right

> @BBuf for left padding here is what I have:
> 
> ```python
> >>> tokenizer = AutoTokenizer.from_pretrained("ArthurZ/rwkv-5-utf", padding_side='left', pad_token="
")
> 
> >>> prompt = ["Hey how are you?","男：听说你们公司要派你去南方工作"]
> >>> inputs_ = tokenizer(prompt, padding = True)
> {'input_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6007, 21835, 21286, 22799, 64], [0, 14601, 19151, 11065, 16735, 10464, 10402, 10678, 11029, 16503, 13818, 10464, 10985, 10934, 13036, 12137, 10460]], 'attention_mask': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
> ```
> 
> or
> 
> ```python
> >>> in_list = ['Chinese: 他补充道：“我们现在有 4 个月大没有糖尿病的老鼠，但它们曾经得过该病。”\n\nEnglish:', Chinese: 埃胡德·乌尔博士（新斯科舍省哈利法克斯市达尔豪西大学医学教授，加拿大糖尿病协会临床与科学部门教授）提醒，这项研究仍处在早期阶段。\n\nEnglish:', 'Chinese: 和其他一些专家一样，他对糖尿病能否治愈持怀疑态度。他指出，这些发现与已患有 1 型糖尿病的人无关。\n\nEnglish:', 'Chinese: 周一，瑞典学院诺贝尔文学委员会常务秘书萨拉·丹尼尔斯在瑞典广播电台的一档节目中向公众宣布，委员会因无法直接联系到鲍勃·迪伦，通知他获得了 2016 年诺贝尔文学奖，已经放弃了与他联系的尝试。\n\nEnglish:']
> 
> >>> inputs_ = tokenizer(in_list, padding = True)
> {'input_ids': [[0, 48407, 59, 33, 10390, 16416, 10655, 17222, 43899, 12605, 10402, 14446, 11454, 13191, 287, 33, 10283, 13190, 11638, 13734, 13191, 15294, 12004, 14662, 14734, 15640, 18459, 19137, 10444, 11885, 10402, 13186, 15486, 12348, 17141, 16721, 14662, 10080, 9823, 261, 48487, 59], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48407, 59, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48407, 59, 0, 284, 33, 11496, 15294, 12004, 14662, 14734, 10370, 13051, 10684, 28329, 11, 48487, 59], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48407, 59, 0, 3493, 636, 33, 12201, 16741, 16873, 11979, 13012, 11871, 11665, 19137, 12145, 15486, 12983, 12269, 10333, 10261, 10390, 15671, 15304, 14734, 11986, 16707, 28329, 11, 48487, 59]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
> ```
> 
> seems to work as expected. Left padding means the zeros are on the left, "right" padding they will be on the right

For the second sentence of the second example, only 3 tokens are normal, and the rest are all left padded with zeros, which is not normal?

Ah ok!  Will fix today! 

haha that was `max_input_chars_per_word` 

Fixed 😉 

> Fixed 😉

thanks.

> Fixed 😉

After changing it to 500, if there is input longer than 500, would it also be truncated again?

Yeah, we can just set to infinite, juste used 500 for a quick fix, I don't know if the original one has a limit as well

> Yeah, we can just set to infinite, juste used 500 for a quick fix, I don't know if the original one has a limit as wel

 I think it would be a good idea to set it to infinite, because rwkv don't have sequence length limit in theory.

Okay my slow tests are all green for the `tokenizer` time to focus on the model! 

The fast CUDA path works thanks to @kashif , but the cpu does not yet

![image](https://github.com/huggingface/transformers/assets/139205286/f65da6d4-a40d-4254-90b2-eb96954b8d6b)
I set trust_remote_code=True but still get an error.
![image](https://github.com/huggingface/transformers/assets/139205286/56016981-7ef3-40d0-80c4-1663b168bee7)
code：
python3 utils/prepare_dataset.py -i JeanKaddour/minipile -o /users/aigclab/copilot/data/Msample
```
import click

import datasets
from transformers import AutoTokenizer

import torch


def _rechunk_tokenize(rechunk_size: int, input_column: str, output_column: str, examples):
    tokenizer = AutoTokenizer.from_pretrained(
        "RWKV/rwkv-6-world-1b6", trust_remote_code=True
    )

    special_token = torch.tensor([0], dtype=torch.long)
    seqs = []
    for e in examples[input_column]:
        seq = tokenizer(e, padding=False, truncation=False, return_tensors="pt")
        seqs.append(seq.input_ids[0])
        seqs.append(special_token)
    seqs = torch.cat(seqs)
    rechunked = seqs[: (seqs.size(0) // rechunk_size) * rechunk_size].view(
        -1, rechunk_size
    )
    return {output_column: rechunked}


@click.command()
@click.option("--rechunk_size", default=513, help="Rechunk size for the dataset")
@click.option("--input_column", default="text", help="Column to tokenize")
@click.option(
    "--output_column", default="input_ids", help="Output column for the tokenized dataset"
)
@click.option(
    "-i",
    "--input_name",
    help="HuggingFace dataset name to tokenize, accept format:"
    + '"dataset_name" or "json:file_a,file_b,..."',
)
@click.option("-o", "--output_dir", help="Output directory for the tokenized dataset")
def main(rechunk_size, input_column, output_column, input_name, output_dir):
    print(f"Tokenizing HuggingFace dataset {input_name} to locally saved {output_dir}")
    if ":" in input_name:
        input_name, input_data_file = input_name.split(":")
        if "," in input_data_file:
            input_data_file = input_data_file.split(",")
    else:
        input_data_file = None
    dataset = datasets.load_dataset(
        input_name, data_files=input_data_file, trust_remote_code=True
    )
    dataset.shuffle().flatten_indices(num_proc=8).map(
        lambda x: _rechunk_tokenize(rechunk_size, input_column, output_column, x),
        batched=True,
        remove_columns=dataset["train"].column_names,
        num_proc=8,
    ).save_to_disk(output_dir, num_proc=8)


if __name__ == "__main__":
    main()

```


One month ago, there was no problem.

Important: maybe less problematic for v5 (or maybe not!), but I found that for v6 the following line is absolutely terrible for inference accuracy:

`out = self.ln_x(rwkv.to(hidden.dtype)).view(batch, seq_length, -1)`

verus the original @BBuf version (which I tweaked and adapted to v6):

```
out = F.group_norm(out / self.config.head_size_divisor, 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)
out = out.to(dtype=hidden.dtype) * gate
```

The issue is that the potential down-cast to bf16 **prior** to the groupnorm causes really bad inference quality. If you look closely this is written differently than in the original @BBuf version where the down-cast occurs *after* the groupnorm during non-cuda inference.

Please see these lines of Bo Peng's original ChatRWKV code for reference about this groupnorm needing float32:
https://github.com/BlinkDL/ChatRWKV/blob/28ed01a8423842c3082f668922a1b45ac182dff0/rwkv_pip_package/src/rwkv/model.py#L377
https://github.com/BlinkDL/ChatRWKV/blob/28ed01a8423842c3082f668922a1b45ac182dff0/rwkv_pip_package/src/rwkv/model.py#L669


Will update the group norm! 


> ![image](https://private-user-images.githubusercontent.com/139205286/319592664-f65da6d4-a40d-4254-90b2-eb96954b8d6b.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTIzODc5MzUsIm5iZiI6MTcxMjM4NzYzNSwicGF0aCI6Ii8xMzkyMDUyODYvMzE5NTkyNjY0LWY2NWRhNmQ0LWE0MGQtNDI1NC05MGIyLWViOTY5NTRiOGQ2Yi5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQwNDA2JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MDQwNlQwNzEzNTVaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0wMjhmYjg1YTYzYmYxMmE2MzQ5YThkY2VhY2IwYTMzYzFiYzNlNDdjMzBmMzU4Y2U4N2Y0MmUwNjJlMDk2YzQ1JlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCZhY3Rvcl9pZD0wJmtleV9pZD0wJnJlcG9faWQ9MCJ9.zUjNhyOLMbVD3AE5tYY77gfnd6yKvQv30THAQDtftLE) I set trust_remote_code=True but still get an error. ![image](https://private-user-images.githubusercontent.com/139205286/319594351-56016981-7ef3-40d0-80c4-1663b168bee7.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTIzODc5MzUsIm5iZiI6MTcxMjM4NzYzNSwicGF0aCI6Ii8xMzkyMDUyODYvMzE5NTk0MzUxLTU2MDE2OTgxLTdlZjMtNDBkMC04MGM0LTE2NjNiMTY4YmVlNy5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQwNDA2JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MDQwNlQwNzEzNTVaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT02NDA1Y2U2OTA3YjVjZmVkNjc1ZTk0ZDA2Y2IyMGY0YWE5MDY4M2EzMzZhM2JhYjI1N2E0NjllMmE0ZDIwODgyJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCZhY3Rvcl9pZD0wJmtleV9pZD0wJnJlcG9faWQ9MCJ9.aOrVMdtdUUuGDQE4quhRd0feOYr_bMBJMuJQeZJde60) code： python3 utils/prepare_dataset.py -i JeanKaddour/minipile -o /users/aigclab/copilot/data/Msample
> 
> ```
> import click
> 
> import datasets
> from transformers import AutoTokenizer
> 
> import torch
> 
> 
> def _rechunk_tokenize(rechunk_size: int, input_column: str, output_column: str, examples):
>     tokenizer = AutoTokenizer.from_pretrained(
>         "RWKV/rwkv-6-world-1b6", trust_remote_code=True
>     )
> 
>     special_token = torch.tensor([0], dtype=torch.long)
>     seqs = []
>     for e in examples[input_column]:
>         seq = tokenizer(e, padding=False, truncation=False, return_tensors="pt")
>         seqs.append(seq.input_ids[0])
>         seqs.append(special_token)
>     seqs = torch.cat(seqs)
>     rechunked = seqs[: (seqs.size(0) // rechunk_size) * rechunk_size].view(
>         -1, rechunk_size
>     )
>     return {output_column: rechunked}
> 
> 
> @click.command()
> @click.option("--rechunk_size", default=513, help="Rechunk size for the dataset")
> @click.option("--input_column", default="text", help="Column to tokenize")
> @click.option(
>     "--output_column", default="input_ids", help="Output column for the tokenized dataset"
> )
> @click.option(
>     "-i",
>     "--input_name",
>     help="HuggingFace dataset name to tokenize, accept format:"
>     + '"dataset_name" or "json:file_a,file_b,..."',
> )
> @click.option("-o", "--output_dir", help="Output directory for the tokenized dataset")
> def main(rechunk_size, input_column, output_column, input_name, output_dir):
>     print(f"Tokenizing HuggingFace dataset {input_name} to locally saved {output_dir}")
>     if ":" in input_name:
>         input_name, input_data_file = input_name.split(":")
>         if "," in input_data_file:
>             input_data_file = input_data_file.split(",")
>     else:
>         input_data_file = None
>     dataset = datasets.load_dataset(
>         input_name, data_files=input_data_file, trust_remote_code=True
>     )
>     dataset.shuffle().flatten_indices(num_proc=8).map(
>         lambda x: _rechunk_tokenize(rechunk_size, input_column, output_column, x),
>         batched=True,
>         remove_columns=dataset["train"].column_names,
>         num_proc=8,
>     ).save_to_disk(output_dir, num_proc=8)
> 
> 
> if __name__ == "__main__":
>     main()
> ```

Bug fixed by https://github.com/BBuf/RWKV-World-HF-Tokenizer/commit/6dd44c8aba3321a9b27aa988aaef331a003f6bbf , it has no relation with this pr, I will update hf repo RWKV/rwkv-6-world-1b6 later .

> BTW @BBuf don't you think it would be great to have a seperate github repo with installables for the kernels? (you can track usage, easier to maintain and propagate to here!) WDYT?

It has been solved in `https://huggingface.co/RWKV/rwkv-5-world-1b5/blob/main/modeling_rwkv5.py` , we need pip install flash-rwkv first,  `modeling_rwkv5.py` in this pr can be replaced directly.