Google Colab で RWKV を試す - RWKVパッケージ版
「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) に配置。
RWKVのモデルとして、Pileのみ学習させたベースモデルとInstructionチューニングしたRavenがあります。Ravenはバージョン (v1〜v9…) と言語種別 (Eng、EngAndMore、EngChnJpn、ChnEngなど) もあります。
今回は、「RWKV-4-Raven-7B-v7-EngAndMore-20230404-ctx4096.pth」を使いました。
(5) トークンファイルをダウンロードして、作業フォルダ (work) に配置。
(6) パッケージのインストール。
# パッケージのインストール
!pip install rwkv
(7) 環境変数の準備。
# 環境変数の準備
import os
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0'
「RWKV_CUDA_ON」は、C++コンパイラとCUDAが必要ですが、CUDA カーネルをコンパイルして10倍高速にします。
(8) モデルとパイプラインの準備。
from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
# モデルとパイプラインの準備
model = RWKV(
model="./RWKV-4-Raven-7B-v7-EngAndMore-20230404-ctx4096",
strategy="cuda fp16")
pipeline = PIPELINE(model, "./20B_tokenizer.json")
◎ model
モデルパスを指定します。セパレータには「\」ではなく「/」を使用します。長いctxが必要な場合は、「ctx4096」モデルを使用します。
◎ strategy
cpu or cuda と精度を指定します。
設定例は、次のとおりです。
(9) パイプライン引数の準備。
# パイプライン引数の準備
args = PIPELINE_ARGS(
temperature = 1.0,
top_p = 0.7,
top_k = 100,
alpha_frequency = 0.25,
alpha_presence = 0.25,
token_ban = [],
token_stop = [0],
chunk_len = 256)
・frequency_penalty : 出力トークンに課すペナルティ。同じテキスト出力の可能性を減らす(デフォルト:0, -2〜2)
・presence_penalty : 出力トークンに課すペナルティ。同じトークンの繰り返し可能性を減らす(デフォルト:0, -2〜2)
「alpha_frequency」「alpha_presence」について詳しくは、以下を参照。
(9) Instructプロンプトを生成する関数の準備。
「ChatRWKV」の「+i」と同じ処理になります。
# Instructプロンプトの生成
def generate_prompt(instruction, input=None):
if input:
return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
# Instruction:
{instruction}
# Input:
{input}
# Response:
"""
else:
return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
# Instruction:
{instruction}
# Response:
"""
(10) パイプラインの実行。
# Instructプロンプトの生成
prompt = generate_prompt("日本で一番人気のアニメは?")
print("--[prompt]--\n" + prompt + "----")
# パイプラインの実行
result = pipeline.generate(prompt, token_count=200, args=args)
print(result)
【おまけ】 ストリーミング出力
pipeline.generate()のcallbackでストリーミング出力も可能です。
# パイプラインの実行
def my_print(s):
print(s, end='', flush=True)
pipeline.generate(prompt, token_count=200, args=args, callback=my_print)
関連
次回
この記事が気に入ったらサポートをしてみませんか?