ゼロからはじめるPython(106) プログラム生成AIのCodeLlamaを手元のPCでも動かしてみよう

ChatGPTを筆頭にした「大規模言語モデル(LLM)」と呼ばれるAIが話題に上らない日はない。このAIが得意なタスクには「プログラムの自動生成」が挙げられるが、Metaからプログラム生成に特化したCodeLlamaが商用利用可能なオープンなライセンスでリリースされた。そこで実際に使ってみよう。

CodeLlamaで素数判定のプログラムを自動生成させたところ

プログラミングは大規模言語モデルが得意とするタスク

リリースからわずか2ヶ月で1億ユーザーを達成した「ChatGPT」の公開から本稿執筆時点で8ヶ月が過ぎた。筆者も業務でChatGPTをはじめ、Github Copilotなど、大規模言語モデル(LLM)関連サービスを使わない日はないくらいだ。

特に「プログラミング」は、間違いなく大規模言語モデルが得意とするタスクであり、GitHub Copilotを利用している多くの人がその利便性を実感しているだろう。もはやAIの補完機能なしに、プログラム開発をしたくないと思えるほどだ。

なぜ、大規模言語モデルはプログラムの生成が得意なのだろうか。まず、プログラミング言語は、自然言語よりも語彙が限られているので補完が容易であることが挙げられる。また、モデル学習時においても、プログラムとコメントが密接に結びついているので、因果関係も類推しやすいことが考えられる。

CodeLlamaについて

さて、ここからが本題だ。これまでも、革新的な大規模言語モデルを公開してきたMetaが、新たに、プログラム生成AIモデルの「CodeLlama」を発表した。このモデルは、Metaが既に公開している言語モデルの「Llama 2」をベースに開発されたもので、プログラムの自動生成に特化させたものだ。

CodeLlamaの公式ページ

なお、基本モデル「Code Llama」のほか、Pythonに特化した「Code Llama-Python」、指示に沿ったタスクを行うモデル「Code Llama - Instruct」の3種類が用意されている。それぞれ、パラメータ数が7B(70億)、13B(130億)、34B(340億)のモデルが提供される。当然、パラメータ数が大きいほど性能が良いが、その分、高性能なマシンが必要だ。こちらのリポジトリで公開されている。

なお、ライセンスに関してだが、CodeLlamaはLlama 2と同じライセンスで提供されており、月間アクティブユーザーが7億人を超えない場合、商用利用も可能な自由な形態で利用できる。

LLaMa Chatで試してみよう

ブラウザから使える「LLaMa Chat」を使うと、すぐに、codellama-34b-instructのモデルを試すことができる。例えば「クイックソートのプログラムをPythonで作ってください」と入力すると次のように、一瞬でクイックソートを行う関数を作成してくれる。

LLaMa Chatでプログラムを自動生成したところ

手元のマシンで実行してみよう

しかし、せっかくのプログラム生成モデルなので、手元のマシンにダウンロードして実行できるようにしてみよう。特別なGPUをサポートしていないマシンでも実行できるよう、llama.cppを使って、このモデルを実行してみよう。

ただし、前提条件として、llama.cppを動かすためには、OSごとのビルドツールがインストールされている必要がある。Windowsであれば、Visual Studioとビルドツールを、macOSであれば、App StoreからXcodeをインストールして、Command Line Toolsのインストールが必要となる。

続いて、ターミナル(WindowsならPowerShell、macOSならターミナル.app)を起動して次のコマンドを実行しよう。すると「llama-cpp-python」がインストールされる。

# Windows
python -m pip install -U llama-cpp-python
# macOS
python3 -m pip install -U llama-cpp-python

なお、以前「llama-cpp-python」をインストールしたことがある人は、最新版をインストールするようにしよう。古いバージョンだとうまく実行できない場合があった。

CodeLlamaのモデルをダウンロードしよう

なお、llama.cppを使う場合は、GGML形式に変換したモデルをダウンロードして使う必要がある。HuggingFaceで公開されているこちらのページを開いて、適当なファイル「codellama-7b-python.ggmlv3.Q5_K_M.bin」(4.78GB)を選んでダウンロードしよう。

GGML形式のモデルをダウンロードしよう

素数判定のプログラムを作ってもらおう

さて、GGML形式のモデル(codellama-7b-python.ggmlv3.Q5_K_M.bin)をダウンロードしたら、さっそく使ってみよう。

以下のプログラムは、素数判定のプログラム生成するようCodeLlamaに指示するものだ。「codellama.py」という名前で保存しよう。

from llama_cpp import Llama

# ここにプロンプト(指示)を記述 --- (*1)
prompt ='def is_prime(num):\n'

# Llamaのモデルを読み込む --- (*2)
llm = Llama(
    model_path='./codellama-7b-python.ggmlv3.Q5_K_M.bin',
    seed=1234)
# プロンプトを与えてプログラムを生成してもらう --- (*3)
output = llm(prompt,
             temperature=0.1,
             max_tokens=1000,
             stop=['\n\n\n\n'],
             echo=True)
# 生成結果を表示 --- (*4)
print(output['choices'][0]['text'])

下記のプログラムとダウンロードしたモデルを同じディレクトリに配置して、下記のコマンドを実行してみよう。

# Windowsの場合
python codellama.py
# macOSの場合
python3 codellama.py

すると、次のように素数判定のプログラムが作成されて表示される。

プログラムを実行したところ

残念ながら、今回試したGGML形式の7B(70億パラメータ)のモデルは、それほど性能が良くないため、生成されたプログラムは、ちょっと手直ししないとエラーで動かないものだった。2行目と7行目の末尾にある「))」から「)」を1つ削る必要がある。

それでも、手元のマシンでも、コード生成AIを動かすことができるというのは、十分可能性を感じるものだろう。

プログラムを確認してみよう。プログラムの(＊1)では、プロンプトを指定している。プロンプト(指示書)には、どんなプログラムを作るのかを指定する。GGML形式の7Bモデルはあまり日本語が得意ではないようなので、ここでは、素数判定の関数を定義する際の関数名(is_prime)と引数(num)を与えてみた。

この部分を「# 素数判定の関数」などと変えても、それなりに動いたのでいろいろ試してみよう。どうやら、「＊＊を作って欲しい」と記述するよりも、「# 以下は＊＊の関数」のようにプログラム内に書くコメントや関数の定義をプロンプトに指定することで、その続きを作ろうと頑張ってくれるようだ。

それから、(＊2)ではLlamaのモデルファイルを読み込む処理だ。seedを適当な値に変えることで異なる挙動になるので、この部分も試してみよう。

(＊3)の部分では、実際にプロンプトを与えて、言語モデルにテキスト生成を行わせる。そして、(＊4)では生成結果を画面に表示する。なお、このPythonモジュール「llama-cpp-python」は、ChatGPTのWebAPIと互換性のある動作をするように似せてあるので、ChatGPTのAPI呼び出しを、Llamaに置き換えるのも容易になるように考えられている。

Colabで試してみよう

上記の方法で一応動かすことができるのだが、あまり精度が良くないため、GGML形式ではなく、メモリ効率の良いGPTQ形式のものを試してみようと思う。

Google Colaboratoryで改めて試してみよう。Colaboratory(以後Colabと略する)を使うと、GPUを利用した機械学習のPythonのプログラムを動かすことができるので便利だ。Colabについては、本連載の26回目で詳しく紹介している。

ブラウザでColabにアクセスして新規ノートブックを作成したら、メニューから[ランタイム > ランタイムのタイプを変更]をクリックして、GPUを使用するタイプに変更しよう。筆者は、Colab Proプランに加入しているため、下記のようにハードウェアアクセレータを選択できるようになっており、今回は執筆時点で性能が最も低い「T4 GPU」を選択した。

ランタイムのタイプをGPUに設定しよう

そして、セルに次のコードを入力して実行しよう。以下のコードは必要なライブラリをインストールするものだ。

# 必要なパッケージをインストール
!pip install transformers sentencepiece accelerate
!pip install auto-gptq

インストールができたら、下記のセルを実行して、CodeLlamaのモデルをインストールしよう。

# 読み込むモデルの指定
model_name = "TheBloke/CodeLlama-7B-Python-GPTQ"

# モデルを読み込む
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoGPTQForCausalLM.from_quantized(model_name,
    use_safetensors=True, trust_remote_code=True,
    device="cuda:0", use_triton=False, quantize_config=None)

# プロンプトを実行する関数を定義
def generate(prompt, n=1024, temperature=0.7):
  input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
  output = model.generate(
    inputs=input_ids,
    temperature=temperature,
    max_new_tokens=n)
  print(tokenizer.decode(output[0]))

以上で、実行の準備が整った。下記のプログラムを実行すると、素数判定を行うプログラムを生成してくれる。