青空文庫から自作LLMを作ろう！

青空文庫からLLMを作っている面白いプロジェクトを見つけましたので紹介していきます。

シンプルなコードで実行できるというのがうれしいところです。

今回は、上記プロジェクトで、自PC環境におけるメモリの制約によりパラメータを少し変えています。

では、下記を自PCで実行します。

git clone https://github.com/speed1313/jax-llm.git
cd jax-llm
pip install -r requirements.lock

インストールできなかったライブラリが出てきましたら、

pip install ライブラリ名でバージョン指定しなければ出来ますので、個別インストールして先に進んでください。

次に、Hugging Faceにある青空文庫のデータセットをダウンロードします。

--book_numで青空文庫のダウンロードする本が何冊か指定します。

また、私の環境ではpython3ではなく、pythonコマンドを指定しています。

python src/jax_llm/prepare_aozora.py --book_num 6000

次に、tokenizerをトレーニングします。

python src/jax_llm/train_tokenizer.py --data_name "aozora_6000"

次に、青空文庫のデータセットでNanoLMモデルをトレーニングします。

自PC環境の制約により、各種パラメータを小さくしています。

python src/jax_llm/train.py --data_name "aozora_6000" --batch_size 32 --n_iterations 5 --n_freq_eval 100 --dropout_rate 0.1 --learning_rate 0.001 --num_layers 12 --embed_size 64  --head_size 8 --num_heads 8 --block_size 64

次に、トレーニングしたモデルで、「国境の長いトンネルを抜けると雪国であった。」という川端康成の雪国に続く文章を生成してもらいます。

python src/jax_llm/generate.py --prompt "国境の長いトンネルを抜けると雪国であった。" --data_name "aozora_6000" --max_new_tokens 200

生成結果は、次のようになりました。

max_new_tokens 200
Output: 国境 の長い トンネル を抜 けると 雪 国 であった 。 址 ことがあった 後見 だけれども ゆすぶ 、 照らして 橘 坐って 。 性の 後見 、 、 欠伸 、 、 ないように 育ち 坐って 、 こういうことは んですって ないように 経歴 、 坐って 殯 、 をたた φ 運命の 育ち 、 家来の 、 を加え 鉄砲 、 ゆすぶ 欠伸 、 、 今に 、 腹も 、 、 坐って 、 、 ダージリン おら 運命の 、 家来の 忝 欠伸 、 、 ダージリン こういうことは 運命の を加え 、 、 、 、 、 今に 、 、 腹も ざりし 悪意 皆の 、 、 燬 だけれども 、 、 条約 をたた 本に 、 金沢 、 、 、 、 、 坐って 、 、 本に 、 悪意 、 、 、 欠伸 ヘリコ 址 将来を 、 、 運 命の 雨の 坐って 、 un 、 経歴 、 条約 、 行ったり 、 、 、 ゆすぶ ないように 家来の 。 、 自身に 家来の 、 勇吉は 殯 橘 燬 条約 φ 、 勇吉は 忝 、 条 約 、 、 を加え 、 ゆすぶ 、 ダージリン 、 、 、 、 、 を加え 、 、 、 φ un ゆすぶ ないように だけれども 経歴 家来の 性の 、 、 条約 、 、 、 皆の 腹 も 私娼 ともすると 、 、 悪意 この間 、 ゆすぶ 市電 瀧 燬 、 olog 、 性の 、 、 、 ゆすぶ 、 だけれども 腹も 坐って 、 ないように 、 皆の 、

出力結果

所感としては、面白いです。

手軽に試すことが出来て、実行環境におけるメモリやGPU制約などが無ければ、それなりの自作LLMの作成体験をすることができるのは素晴らしいです。

私のパソコンは16GBのメモリでも、それなりに実行できるので、多くの人が自作LLMの作成体験を積むことができるはずです。

追記:2024/06/21
Google Colabでも出来る記事を書きましたので、自PCのメモリに不安のある方は、Google Colabの記事を参考にしてみてください。