[Salesforce Research論文]自然言語からコードを生み出す”CodeGen”を試してみた

皆さん、こんにちは。hotoke_neko改めてクルトンです。 クルトンを大量に食べてたら社内で呼ばれ始めましたので、改名しました。 今後ともよろしくお願いいたします。

さて、今回はタイトルにもある通りSalesforce Researchが出している論文"A Conversational Paradigm for Program Synthesis"のモデルCodeGenを動かしてみました。

CodeGenは自然言語からPythonのコードが生み出せるモデルです。

今回の実行環境はGoogleColabでランタイムにはGPUを使用しています。

サンプル実行

こちらにある公式のコードを動かしてみます。

pip install

まずは必要なライブラリを使えるようにするためGitHubからリポジトリから取得します。

!pip install git+https://github.com/huggingface/transformers.git

ライブラリをインポート

必要なライブラリのインポートを行います。この時、set_default_tensor_typeメソッドで処理対象のデータがcuda(GPU上）のもので処理するよう先に設定しておきます。

なお、他の型も設定できますのでご興味おありなら、こちらの公式ドキュメントご参照ください。

import torch
torch.set_default_tensor_type(torch.cuda.FloatTensor)
from transformers import AutoTokenizer, AutoModelForCausalLM

トークナイザ読み込み

トークナイザを読み込みます。トークナイザとは一言で言うと、自然言語をモデルへ入力できるようにデータを加工するためのものです。

tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-mono")

モデルを読み込み

モデルを読み込みます。 from_pretrainedメソッドは、Hugging Face Hubに置かれているモデルへのパスを入れます。

model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-2B-mono")

トークナイザで入力データ取得

トークナイザを使用して、モデルへ入力するためのデータを取得します。 トークナイザの引数に自然言語（英語）を入れます。 もう一つの引数return_tensorsですが、PyTorchやTensorFlow、Numpy型が選べます。今回はPyTorchで処理をしているので"pt"を渡します。

return_tensorsなどの引数の設定について、詳しくは公式ドキュメントをご覧ください。

inputs = tokenizer("# this function prints hello world", return_tensors="pt").to(0)

コード生成

この一文で実際にコードをの生成を行なっています。

sample = model.generate(**inputs, max_length=128)

生成内容を確認

次のコードで生成したものを確認します。

print(tokenizer.decode(sample[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"]))

これで自然言語からコードが生成できる事を確認できました！

色々と試してみた

せっかくサンプルコードが一通り実行できましたので、少し内容を変えてみます。

"this function return x+y"

return文に引数で渡された'x'と'y'を足し算しています。

"check prime number or not"

「流石にコード生成できるとはいえ、素数を確認する処理は出来るのだろうか」と気になったので生成しました。 結果は予想以上にしっかりしたコードが生成されてました。

"this function solve fizz buzz"

有名問題である「fizz buzz」はどのように処理するのか確認しました。

なるほど、最初に15の倍数を処理したのちにそれぞれの倍数で処理していますね。有名問題とはいえ、問題の内容を知らないといけないのですが、しっかりとコードを生成してくれました！（おまけで処理を確認するprint文付きです。）

終わりに

今回は自然言語からコードを生成するCodeGenについて紹介しました。 このようにコードの生成ができるものが簡単に実行できて確認できるので、是非興味を持たれた方はご確認ください！

事例として3つの英文を実行しましたが、他の文章でも面白そうなコードが作成できそうです。その際にはよければ、こちらの記事を参考にしてくだされば幸いです。

今回はここまで。

それでは、また！