Google Colab で RWKV-4-World を試す - RWKVパッケージ版

npaka

2023年6月28日 07:42

「Google Colab」で「RWKV-4-World」を試したので、まとめました。

【注意】Google Colab Pro/Pro+ の A100で動作確認しています。

前回

1. RWKV-4-World

「RWKV-4-World」は、「Pile」と「RedPajama」で学習した多言語用のベースモデルです。多言語用の新トークナイザーが使われています。

「RWKV」は、TransformerレベルのLLM性能を備えたRNNで、高性能、高速推論、VRAMの節約、高速学習、長い文脈長、自由な埋め込みを実現しています

2. RWKV-4-World のモデル

「RWKV-4-World」では、次の9種類のモデルが公開されています。

・RWKV-4-World-0.1B
・RWKV-4-World-0.4B
・RWKV-4-World-1.5B
・RWKV-4-World-3B
・RWKV-4-World-7B
・RWKV-4-World-CHNtuned-0.1B
・RWKV-4-World-CHNtuned-0.4B
・RWKV-4-World-CHNtuned-1.5B
・RWKV-4-World-CHNtuned-3B

・RWKV-4-World-CHNtuned-7B (学習中)
・RWKV-4-World-JPNtuned-7B (学習中)

3. RWKV のプロンプトガイドライン

「RWKV」は、Transformerモデルよりも、プロンプト書式に影響を受けます。これは、「RWKV」には "振り返り "の機能がないためです。

そのため、次のように記述する代わりに、

{{CONTEXT}}

{{INSTRUCTION}}

{{ANSWER}}

次のように記述して、実行してください。

{{INSTRUCTION}}

{{CONTEXT}}

{{ANSWER}}

人間で例えるなら、指示を与えることなく、文脈を聞かされるようなものです。もし、指示の前に文脈を先に聞かされたら、文脈で何をすべきかわからず、指示にとって重要な部分を覚えていないかもしれません。

最初に指示を与え、次に文脈を聞かせると、最初に指示を理解し、その知識を使用して文脈を処理します。

文脈を伴う質問応答タスクの場合、最も最適なのは、次のように文脈の前後に質問を繰り返すことになります。

{{QUESTION}}

{{CONTEXT}}

{{QUESTION}}

{{ANSWER}}

4. Colabでの実行

Google Colabでの実行手順は、次のとおりです。

(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」で「A100」を選択。

(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) に配置。
今回は、「RWKV-4-World-7B-v1-20230626-ctx4096.pth」を使いました。

(5) パッケージのインストール。

# パッケージのインストール
!pip install rwkv

(6) 環境変数の準備。

# 環境変数の準備
import os
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0'

(7) モデルとパイプラインの準備。

from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS

# モデルとパイプラインの準備
model = RWKV(
    model="RWKV-4-World-7B-v1-20230626-ctx4096", 
    strategy="cuda fp32"
)
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

PIPELINEには、「20B_tokenizer.json」の代わりに「rwkv_vocab_v20230424」(新しい多言語トークナイザー)を指定します。これはrwkvパッケージに含まれています。

(9) Instructプロンプトを生成する関数の準備。
チャットにはQuestion/Answer、User/AI、Human/Botを使用します。Bob/AliceやQ/Aは使用しないでください。

# Instructプロンプトの生成
def generate_prompt(instruction, input=None):
    if input:
        return f"""Instruction: {instruction}

Input: {input}

Response: """
    else:
        return f"""Question: {instruction}

Answer: """

(10) 推論の実行。

# プロンプトの準備
prompt = generate_prompt(
    "後藤ひとりが加入するバンドの名前は何ですか？",
    "後藤ひとりはギターを愛する孤独な少女。 家では孤独でただ遊んでばかりの毎日だったが、ひょんなことから伊地知虹夏率いる「結束バンド」に加入することに。 人前で演奏することに不慣れな後藤は、立派なバンドマンになれるのか？")
print("--[prompt]--\n" + prompt + "----")

# 推論の実行
args = PIPELINE_ARGS(
    temperature = 1.0,
    top_p = 0.3, 
    top_k = 100, 
    alpha_frequency = 0.25, 
    alpha_presence = 0.25, 
    token_ban = [],
    token_stop = [0],
    chunk_len = 256) 
result = pipeline.generate(prompt, token_count=200, args=args)
print(result)

--[prompt]--
Instruction: 後藤ひとりが加入するバンドの名前は何ですか？

Input: 後藤ひとりはギターを愛する孤独な少女。家では孤独でただ遊んでばかりの毎日だったが、ひょんなことから伊地知虹夏率いる「結束バンド」に加入することに。人前で演奏することに不慣れな後藤は、立派なバンドマンになれるのか？

Response: ----

後藤ひとりが加入するバンドの名前は「結束バンド」です。