rwkv-4-169m-pile - RNN与Transformer的高性能结合：高效文本生成

项目介绍：rwkv-4-169m-pile

rwkv-4-169m-pile 是一个由 Bo Peng 领导的项目，目标是开发一种高效的语言模型。这个项目建立在 EleutherAI 的“Pile”数据集之上，该数据集是一个大规模、多用途的文本数据集。

项目背景

RWKV 是一种结合了循环神经网络（RNN）和 Transformer 模型性能的架构设计。它能够像 GPT 模型那样进行并行训练，结合了 RNN 和 Transformer 最优的特性，从而在快速推理、节省显存、快速训练等方面表现出色。此外，它拥有“无限”上下文长度以及免费的句子嵌入能力，使其在大文本处理任务上具有明显优势。

模型细节

关于 RWKV 模型架构的细节分析，读者可以通过 Johan Wind 的博客获取更多信息。此外，Hugging Face 也提供了关于 RWKV 和 Transformer 模型整合的博客文章。

如何使用

转换模型权重

可以使用 convert_rwkv_checkpoint_to_hf.py 脚本将 RWKV 的原始权重转换为 Hugging Face 模型格式。这需要指定原始权重的库标识、文件名和输出目录。转换完成后，还可以直接将转换后的模型推送到 Hugging Face 的模型中心。示例命令如下：

python convert_rwkv_checkpoint_to_hf.py --repo_id RAW_HUB_REPO --checkpoint_file RAW_FILE --output_dir OUTPUT_DIR --push_to_hub --model_name dummy_user/converted-rwkv

文本生成

使用 AutoModelForCausalLM 和 AutoTokenizer 类可以从模型中生成文本。以下是一些使用指南：

在 CPU 上运行模型：

如果要在 CPU 上运行，可以参考以下代码片段：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile")
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")

prompt = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."

inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(inputs["input_ids"], max_new_tokens=40)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

在单个 GPU 上运行模型：

如果希望利用 GPU 优化计算性能，可以使用下述代码：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile").to(0)
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")

prompt = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."

inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=40)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

在半精度模式下运行模型（利用GPU）：

通过启用半精度模式，可以显著减少模型所需的计算资源：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile", torch_dtype=torch.float16).to(0)
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")

prompt = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."

inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=40)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

在多个 GPU 上运行模型：

利用多 GPU 进行分布式计算，可以进一步提升大规模文本生成的效率：

# pip install accelerate
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")

prompt = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."

inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=40)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

引用

如果您在工作中使用了此模型，请考虑引用原始工作。完整的项目和详细信息可以在原始 GitHub 仓库中找到。