青空文庫で自作LLM体験(GoogleColab)

今回は、Google Colabで青空文庫から自作LLMを作成していきます。

パソコンで作成したい人は、16GBのメモリ環境でしたが、パラメータを小さくして何とか出来ましたので、こちらを参考にしてください。

Google Colabを利用した理由としては、Google Colabでは、TPUのハイメモリが300GB以上使えるからです。

メモリを気にすることなく使えるというのは良いですね。

今回も、こちらのコードを参考に、Google Colab用にしています。

では、Google Colabで作成していきましょう。

まずは、リソースは、TPUを選択します。

TPUのリソース

次に、Google Colabで次のコードを打ちます。ライブラリが足りないなどが出てきましたら、随時ライブラリをインストールしていってください。

!git clone https://github.com/speed1313/jax-llm.git
%cd jax-llm
!pip install -r requirements.lock
!pip install etils
!pip install datasets
!pip install requests
!pip install tiktoken
!pip install jax

次に、青空文庫をダウンロードして、トークナイザーを訓練します。そして、青空文庫のデータをNanoLMで訓練します。今回は、batch_sizeを32に変更しています。

!python /content/jax-llm/src/jax_llm/prepare_aozora.py --book_num 10246
!python /content/jax-llm/src/jax_llm/train_tokenizer.py --data_name "aozora_10246"
!python /content/jax-llm/src/jax_llm/train.py --data_name "aozora_10246" --batch_size 32 --n_iterations 5000 --n_freq_eval 100 --dropout_rate 0.1 --learning_rate 0.001 --num_layers 12 --embed_size 512  --head_size 64 --num_heads 8 --block_size 256

では、200トークンで、吾輩は猫に繋がる文章を生成してもらいます。

!python /content/jax-llm/src/jax_llm/generate.py --prompt "吾輩は猫である。" --data_name "aozora_10246" --max_new_tokens 200

Output: 吾輩は 猫 である 。 「 さあ 、 この 野郎 が 、 あの 男 たち 、 一 品 だけ にして 、 その 力 、 すなわち 、 この 「 御 御 冗談 、 これは 面白い 話 でした 。」 「 いや 、 その 、 この 「 神 」 や 「 いき だ 、 と 、 と 、 その 言葉は よく 、 その 「 お 、 この 玉 琴 の音 と共に 飛び 散った 。 また 、 その 風 貌 を 、 その 性質 から 、 この 「 お ー さん 、 あんまり 、 お 幾 里 、 と っぷり 暮れて しまいました 。 その 、 大きな 船 から 、 この 道を 発見 している けれども 、 それが 最も 多い ものが 、 他の ものが この 科学 的に 、 歴史 的には 、 その 「 この 宝 』 は 、 この 私 自身 、 あの 私が 考え 出す のは 、「 いき 」 に 、 「 どう ？ 、 この 島 君は 、 その 川 内 と 、 その いわゆる 「 日本 国民 社会 、 国家 への 歴史 として 、 これを 説明 する 、 すなわち 一 、 一 章 で 本 から 、 大 金の 薬 に 砂糖 を入れて くれた 。 「 なぜ 、 僕 達は 、 その 「 この 、 あの 、 あの

生成結果

所感としては、Google ColabのTPUだと、クレジットの消費も少なくてメモリも300GB以上使えるのはコストパフォーマンスが高いです。

また、データセットを色々と変えてみて、どのようなLLMを作れるかを試してみるのも面白いはずです。