Google Colab で RWKV を試す - ChatRWKV版
「Google Colab」で「RWKV」を試したので、まとめました。
1. RWKV
「RWKV」は、TransformerレベルのLLM性能を備えたRNNです。高性能、高速推論、VRAMの節約、高速学習、長い文脈長、自由な埋め込みを実現しています。
2. Colabでの実行
Colabでの実行手順は、次のとおりです。
(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」の「プレミアム」を選択。
(2) Googleドライブのマウント。
# Googleドライブのマウント
from google.colab import drive
drive.mount('/content/drive')(3) 作業フォルダへの移動。
# 作業フォルダへの移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd '/content/drive/My Drive/work'(4) モデルをダウンロードして、作業フォルダ(work)に配置。
14Bモデルの現時点での最新版「RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth」を使いました。
(5) ChatRWKVのインストール。
# ChatRWKVのインストール
!git clone https://github.com/BlinkDL/ChatRWKV
%cd ChatRWKV
!pip install -r requirements.txt(6) チャットのコードの修正。
・v2/chat.py のモデルパスの変更
親フォルダ (work) に配置した「RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth」を次のように指定します。
if CHAT_LANG == 'English':
args.MODEL_NAME = '/fsx/BlinkDL/HF-MODEL/rwkv-4-pile-14b/RWKV-4-Pile-14B-20230228-ctx4096-test663'↓
if CHAT_LANG == 'English':
args.MODEL_NAME = '../RWKV-4-Pile-14B-20230313-ctx8192-test1050'・v2/chat.py のpromptをinputに変更
msg = prompt(f'{user}{interface} ')↓
msg = input(f'{user}{interface} ')・v2/prompt/default/English-2.py の初期プロンプトを日本語に変更
interface = ":"
user = "Human"
bot = "AI"
init_prompt = f'''
あなたは人間と会話するチャットボットです。
{user}{interface} こんにちは！
{bot}{interface} こんにちは！お元気ですか？
'''(7) chat.py の実行。
初回は5分少々でプロンプトが表示されます。
!python v2/chat.py
3. RWKVのプロンプトのコマンド
RWKVのプロンプトのコマンドは、次のとおりです。
・<プロンプト> : チャット
・+qa <プロンプト> : チャット
・+gen <プロンプト> : テキスト生成
・+i <インストラクト> : インストラクト
・+reset : 会話リセット (初期プロンプトのみの状態に戻す)
・+ : チャットの代替返信
・++ : +qa/+gen/+i を再度実行
・+++ : テキスト生成の続きの生成
各コマンドは、内部で次のように変換されてモデルを呼び出してます。
・<プロンプト>、+qa <プロンプト>
「Human:」と「AI:」は初期プロンプトの設定によって変わります。
Human: <プロンプト>\n\nAI:・+gen <プロンプト>
\n<プロンプト>・+i <インストラクト>
Below is an instruction that describes a task. Write a response that appropriately completes the request.
# Instruction:
<インストラクト>
# Response:4. Raven
「Raven」は、「RWKV」を「Alpaca」や「CodeAlpaca」などのデータセットでファインチューニングしたモデルです。
以下から、ダウンロードできます。