見出し画像

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」などのデータセットでファインチューニングしたモデルです。

以下から、ダウンロードできます。


次回



この記事が気に入ったらサポートをしてみませんか?