Google Colab で RWKV-4-World を試す - RWKVパッケージ版
「Google Colab」で「RWKV-4-World」を試したので、まとめました。
前回
1. RWKV-4-World
「RWKV-4-World」は、「Pile」と「RedPajama」で学習した多言語用のベースモデルです。多言語用の新トークナイザーが使われています。
「RWKV」は、TransformerレベルのLLM性能を備えたRNNで、高性能、高速推論、VRAMの節約、高速学習、長い文脈長、自由な埋め込みを実現しています
2. RWKV-4-World のモデル
「RWKV-4-World」では、次の9種類のモデルが公開されています。
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)
この記事が気に入ったらサポートをしてみませんか?