chatRWKV を chatux を使って Web チャットボットにする

Apache 2.0。 本プロジェクト内に内包されているコードは https://github.com/BlinkDL/ChatRWKV をベースにしています。

• Python: 3.9
• GPU CUDA Version: 12
• GPU memory: 24GB
• OS: Ubuntu Desktop 22.04 / Windows 11

本プロジェクトをクローンする

https://github.com/riversun/chatux-server-rwkv

関連パッケージをインストールします。(※)

pip install rwkv fastapi uvicorn

（既に CUDA版 pytorch が動作する環境がある前提です。もし実行時エラーが出たら、必要に応じて足りないパッケージをインストールします）

モデルデータには 3b,7b,14b と種類があり、数値が大きいほど、多くのパラメータをもっているので、表現力が高いということになります。 パラメータ数が多いほどGPUのメモリを消費しますので、実行環境にあわせてデータをダウンロードします。

以下は、モデルデータの種類と消費GPUメモリです。

fp16i8 (fp16学習済データを int8 に量子化する　という意味と思われる) をつかえば、 Accuracy はやや落ちるようですが、使用GPUメモリ容量を減らすことができます。

コード内で指定するときは cuda fp16 や cuda fp16i8 のように指定します

以下から、重みデータ(*.pth)のファイルをダウンロードします。

args.n_layer = 40
args.n_embd = 5120
args.ctx_len = 8192

args.n_layer = 32
args.n_embd = 4096
args.ctx_len = 4096

args.n_layer = 32
args.n_embd = 2560
args.ctx_len = 4096

ダウンロードしたファイルは [project]/data フォルダに置きます

chatux-server-rwkv.py を開いて

• #specify RWKV strategy,model(weight data) のあたりに、以下のように STRATEGY= と MODEL_NAME があるので、それぞれ入力します。 MODEL_NAME はファイル名だけでOK。
• # specify params for weight data のあたりに、学習済モデルデータ用のパラメータを記載します。↑に示したもののコピペでもOK

# specify RWKV strategy,model(weight data)
STRATEGY = 'cuda fp16i8'
MODEL_NAME = 'RWKV-4-Pile-7B-20230109-ctx4096.pth'

# specify params for weight data
args.n_layer = 32
args.n_embd = 4096
args.ctx_len = 4096

python chatux-server-rwkv.py

サーバーが起動したらブラウザで http://localhost:8001 を開く

デモ動画