CLOVER🍀

That was when it all began.

Transformersでコード生成を行ってみる

これは、なにをしたくて書いたもの?

llama.cppやLocalAIを使って、コード生成を試せないのかな?と思ったドキュメントを見ると、llama-cpp-pythonのドキュメントに
コード生成(コード補完)について書かれているのを見つけまして。

OpenAI Compatible Server / Guides / Code Completion

そういえば、llama-cpp-pythonにはGitHub Copilot互換のエンドポイントらしきものがあったなというのを思い出しました。

llama-cpp-pythonで、OpenAI API互換のサーバーを試す - CLOVER🍀

いきなりllama-cpp-pythonで試してみてもいいのですが、OpenAI Python APIライブラリーではGitHub Copilotのエンドポイントはカバー
していませんし、curlでアクセスして確認するのもなんなのでまずはHugging FaceのTransformersで試してみることにしました。

Replit Code V1.5 3B

llama-cpp-pythonのコード生成のドキュメントで紹介されているのは、Replit Code V1.5 3Bというモデルです。

OpenAI Compatible Server / Guides / Code Completion

こちらですね。正確には、Replit Code V1.5 3BというモデルをGGUFフォーマットに変換したものです。

abetlen/replit-code-v1_5-3b-GGUF · Hugging Face

これは、Replit社がApache License 2.0で公開しているコード生成用のLLMのようです。

なのですが、手元の環境ではメモリ(16GB)が足りずに動かせなかったので、別のモデルを探すことにしました…。
モデルのサイズは約7Gほどだったんですけどね。

SalesforceのCodeGenを使う

Replit Code V1.5 3Bが動かせなかったので、もうちょっとパラメーター数の小さいモデルを探すことにしました。

SalesforceのCodeGenというモデルに、小さめのものがありました。

GitHub - salesforce/CodeGen: CodeGen is a family of open-source model for program synthesis. Trained on TPU-v4. Competitive with OpenAI Codex.

CodeGenには1.0、2.0、2.5のバージョンがあるようですが、パラメーター数が少ないモデルはCodeGen 1.0にしかなさそうだったので
こちらを使うことにしました。

Salesforce/codegen-350M-multi · Hugging Face

Salesforce/codegen-350M-mono · Hugging Face

multiとmonoの差は対応プログラミング言語とトレーニングで、multiはC、C++、Go、Java、JavaScript、Python、monoはmultiをベースに
Python向けにトレーニングされたものです。

今回はmultiを使うことにしましょう。

余談:Code Llama

LLMで有名なモデルといえばはMeta社のLlamaですが、コード生成用途ではCode Llamaというものがあるようです。

codellama (Code Llama)

GitHub - facebookresearch/codellama: Inference code for CodeLlama models

これはLlama 2をベースにプログラミング向けにトレーニングされたモデルのようです。

Code Llamaをベースに日本語向けにトレーニングしたものとしては、「ELYZA-japanese-CodeLlama-7b」があります。

「Code Llama」をベースとした商用利用可能な日本語LLM「ELYZA-japanese-CodeLlama-7b」を公開しました

このあたりも、そのうち試してみたいところです。

環境

今回の環境は、こちら。

$ python3 --version
Python 3.10.12


$ pip3 --version
pip 22.0.2 from /usr/lib/python3/dist-packages/pip (python 3.10)

CPU環境です。

Transformersでコード生成を試してみる

まずはライブラリーのインストール。

$ pip3 install transformers[torch]

依存関係。

$ pip3 list
Package                  Version
------------------------ ----------
accelerate               0.26.1
certifi                  2023.11.17
charset-normalizer       3.3.2
filelock                 3.13.1
fsspec                   2023.12.2
huggingface-hub          0.20.2
idna                     3.6
Jinja2                   3.1.3
MarkupSafe               2.1.3
mpmath                   1.3.0
networkx                 3.2.1
numpy                    1.26.3
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        8.9.2.26
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.18.1
nvidia-nvjitlink-cu12    12.3.101
nvidia-nvtx-cu12         12.1.105
packaging                23.2
pip                      22.0.2
psutil                   5.9.7
PyYAML                   6.0.1
regex                    2023.12.25
requests                 2.31.0
safetensors              0.4.1
setuptools               59.6.0
sympy                    1.12
tokenizers               0.15.0
torch                    2.1.2
tqdm                     4.66.1
transformers             4.36.2
triton                   2.1.0
typing_extensions        4.9.0
urllib3                  2.1.0

サンプルに習って、こんな感じでソースコードを作成。

code_completion_autoclasses.py

import time
from transformers import AutoModelForCausalLM, AutoTokenizer

start_time = time.perf_counter()

tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-multi")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-multi")

source_code = """// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {"""

inputs = tokenizer(source_code, return_tensors="pt")

outputs = model.generate(**inputs, max_length=128)
generated_source_codes = tokenizer.batch_decode(outputs, skip_special_tokens=True)

for generated_source_code in generated_source_codes:
    print("input:")
    print(source_code)
    print()
    print("generated:")
    print(generated_source_code)

print()

elapsed_time = time.perf_counter() - start_time
print(f"elapsed time = {elapsed_time:.3f} sec")

Javaのソースコードを続きを作成してもらうことにしました。

source_code = """// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {"""

確認。

$ python3 code_completion_autoclasses.py
Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
input:
// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {

generated:
// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {
        System.out.println("Hello World");
    }
}


elapsed time = 3.737 sec

それっぽい感じになりました。

パイプラインでも試してみましょう。

code_completion_pipeline.py

import time
from transformers import pipeline

start_time = time.perf_counter()

code_generator = pipeline("text-generation", model="Salesforce/codegen-350M-multi")

source_codes = [
    """// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {"""
]


for source_code in source_codes:
    generated_source_codes = code_generator(source_code, max_length=128)
    print("input:")
    print(source_code)
    print()

    for generated_source_code in generated_source_codes:
        print("generated:")
        print(generated_source_code["generated_text"])

    print()

elapsed_time = time.perf_counter() - start_time
print(f"elapsed time = {elapsed_time:.3f} sec")

結果。

$ python3 code_completion_pipeline.py`
Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
input:
// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {

generated:
// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {
        System.out.println("Hello World");
    }
}

elapsed time = 3.146 sec

良さそうですね。

パイプラインの方は毎回けっこう違う結果になって、謎の処理を書くことも多かったです。

// App.java
public class App {
    // print "Hello World"
    public static void main(String[] args) {
        // print "Hello World"
        /*
        for (int i = 0; i < 100_000 * 10 * 10; i++) {
            System.out.println("i:" + i);
            System.out.println("i:" + Thread.currentThread().getStackTrace()[13].getFileName() + ", " + i);
            System.out.println("i:" + Thread.currentThread().getStackTr

とりあえず、コード生成としては試せたのでよしとしましょう。

おわりに

Transformersでコード生成を試してみました。

コード生成用のモデルを見つけ、動かすことはできたのでそこは良かったのですが。

手元の環境ではReplit Code V1.5 3Bを全然動かせなかったのでちょっと驚きました。
やっぱり、こういうことをやろうとするとそれなりのリソースがある環境が必要だな、と思いました…。