🚀 RWKV-4 | 基於Pile數據集訓練的1.69億參數模型卡片
RWKV是一個由Bo Peng領導的項目。你可以通過Johan Wind的博客文章這裡和這裡瞭解更多關於該模型架構的信息。你還可以通過加入RWKV Discord服務器來深入瞭解這個項目。

🚀 快速開始
模型簡述
以下是來自原倉庫的描述:
RWKV是一種具有Transformer級大語言模型性能的循環神經網絡(RNN)。它可以像GPT一樣直接進行訓練(可並行化)。它結合了RNN和Transformer的優點——性能出色、推理速度快、節省顯存、訓練速度快、具有“無限”上下文長度,並且能免費獲得句子嵌入。
✨ 主要特性
- 數據集:使用了EleutherAI/pile數據集進行訓練。
📚 詳細文檔
模型細節
模型架構的詳細信息可以在上述博客文章以及Hugging Face關於該模型集成的博客文章中找到。
模型使用
將原始權重轉換為Hugging Face格式
你可以使用convert_rwkv_checkpoint_to_hf.py
腳本,通過指定原始權重的倉庫ID、文件名和輸出目錄來進行轉換。你還可以選擇通過傳遞--push_to_hub
標誌和--model_name
參數,將轉換後的模型直接推送到Hugging Face Hub上。
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上運行模型
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上運行模型
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上運行模型
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))
📄 許可證
如果你使用此模型,請考慮引用原倉庫此處的原始工作。