青空文庫で自作LLM体験(GoogleColab)
今回は、Google Colabで青空文庫から自作LLMを作成していきます。
パソコンで作成したい人は、16GBのメモリ環境でしたが、パラメータを小さくして何とか出来ましたので、こちらを参考にしてください。
Google Colabを利用した理由としては、Google Colabでは、TPUのハイメモリが300GB以上使えるからです。
メモリを気にすることなく使えるというのは良いですね。
今回も、こちらのコードを参考に、Google Colab用にしています。
では、Google Colabで作成していきましょう。
まずは、リソースは、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
所感としては、Google ColabのTPUだと、クレジットの消費も少なくてメモリも300GB以上使えるのはコストパフォーマンスが高いです。
また、データセットを色々と変えてみて、どのようなLLMを作れるかを試してみるのも面白いはずです。
この記事が気に入ったらサポートをしてみませんか?