CyberAgentが公開した日本語LLMをGoogle Colabで動かす方法

概要

機械学習の開発などを行なったことがない場合、そもそもどうすればいいのかまったく分からなかったので、やり方まとめました。
分かってしまえば非常に簡単です。
内容は浅いので、感触を掴んだら別の専門記事を探してください。

準備

Google Colab Pro

Google Colab Pro のアカウントを作成する。

Pro じゃなくても大丈夫かもしれませんが、GPU を使うので Pro を契約しておくのが安心です。

モデルの確認

使用したいモデルのサイズを確認してください。

7B だと約 7GB になります

Google Colab のランタイムタイプの選択

必ずGPUを含めてください。
今回試した条件は下記になります。

• GPUタイプ：A100
• ランタイムの仕様：ハイメモリ

上記条件にしないと、ハードウェア周りでエラーが多発します。

Google Colab で動かす

以下のスクリプトを Google Colab で実行してください。

※補足：PyTorch は既に Google Colab 環境に含まれているので別途インストールしなくても大丈夫です。

!pip install transformers accelerate

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-7b", device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-7b")

inputs = tokenizer("AIによって私達の暮らしは、", return_tensors="pt").to(model.device)
with torch.no_grad():
    tokens = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.05,
        pad_token_id=tokenizer.pad_token_id,
    )
    
output = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(output)

python のコードは、各モデルの Usage に書いてあるコードです。
動作しない場合は最新のサンプルコードを確認してください。

実行するとモデルがダウンロードされ、「AIによって私たちの暮らしは、」の続きが生成されているのが確認できると思います。

GPT-NeoX互換のモデルは似た方法で別のものも実行できるみたい？なので、色々試してみてください。
大AI時代をがんばって生き抜いていきましょう。