💫

StarCoder でプログラミング(コード)生成の味見メモ

2023/05/05に公開

https://huggingface.co/bigcode/starcoder

ほー, すごぽよっぽそうね...

試します!

環境

  • x299 CPU 256 GB
  • 3090 x 2
  • Ubuntu 20.04

セットアップ

執筆時点では login しないと weight 落とせません.
64 GB くらいあります.

fp16 で 3090 x 2, 8bit にすれば 3090 x 1 で動くと思われますが, とりま CPU で試します.

transformers(ややこしい名前であるが, Transformer 系の ML モデルをペロッと動かすためのライブラリ)と pytorch 入っていれば動きます.

$ python -m pip install transformers torch
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "bigcode/starcoder"
device = "cpu" # for GPU use "cuda"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True).to(device)

inputs = tokenizer.encode("bool is_odd(int a) { ", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=128)
print(tokenizer.decode(outputs[0]))

デフォの設定だと 20 tokens までで物足りないので, max_new_tokens=128 にしてみました.

The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.

bool is_odd(int a) {  // 1
    return a % 2 == 1;
}

int is_even(int a) {  // 2
    return a % 2 == 0;
}

int main() {
    int a = 1;
    int b = 2;
    int c = 3;
    int d = 4;
    int e = 5;
    int f = 6;
    int g = 7;
    int h = 8;
    int i = 9;
    int j = 10;
    int k = 11;

CPU だと生成に 3~4 分という感じでした. GPU 必須な感じですね.
一応簡単なのは生成してくれましたが, コメント変だったりでちょい期待はずれかしら.
あと, max_new_tokens=128 であまり長すぎるのもよくないかもです.

まだちょこっとしか試してませんが, 良しなにぺろっといい感じのコード吐いてくれるのは難しそうで, 現状はメソッド補完とかの延長で, 難しくはないけどコーディングするのがだるい感じなのを対応してくれると考えてたほうがいいかもですね.

あとは vscode などと連携させたときのレスポンスでしょうか.
A6000 ADA(48 GB)が欲しくなりますね!

Discussion