プログラミング特化LLMであるStable Code 3BをMacBook(M2)で動かしてコードを書いてもらった

[Supership](https://supership.jp/)の名畑です。[勇気爆発バーンブレイバーン](https://bangbravern.com/)が平成初期のアニメに触れまくった私のツボにはまりまくりなのですが、若い方々にも刺さる内容なのかが興味あります。

## はじめに

画像生成モデル[Stable Diffusion](https://ja.stability.ai/stable-diffusion)の提供元である[stability.ai](https://www.stability.ai/)が「[Stable Code 3B - エッジでのコーディング](https://ja.stability.ai/blog/stable-code-3b)」という記事を公開していました。

> - Stable Code 3Bは、30億パラメータを持つ大規模言語モデル（LLM）であり、CodeLLaMA 7bのような2.5倍の大きさを持つモデルと同等のレベルで、正確で応答性の高いコード補完を可能にします。
> - MacBook Airのような一般的なノートパソコンでGPUがなくてもオフラインで動作します。
> - モデルを商用利用するためには Stability AI メンバーシップ へのご登録をお願いします。
> 
> 参考：[Stable Code 3B - エッジでのコーディング](https://ja.stability.ai/blog/stable-code-3b)

ということで今回は[Stable Code 3B](https://huggingface.co/stabilityai/stable-code-3b)を使ってみました。

## 環境

最初に私の環境を共有しておきます。

[MacBook Pro(Apple M2 Proチップ)](https://support.apple.com/kb/SP889?locale=ja_JP)です。
OSは[macOS 13 Ventura](https://www.apple.com/jp/macos/ventura/)です。

記事を読む限りでは**MacBook Air**でも動くとのことですが、私は試せていません。

私の環境でのPythonのバージョンは**3.10.12**です。

```shell-session
$ python --version
Python 3.10.12
```

[torch](https://pytorch.org/)と[transformers](https://huggingface.co/docs/transformers/index)を用いるので、インストールしていない方はしておきましょう。

```shell-session
$ pip install torch transformers            
```

私の環境でのバージョンはそれぞれ**2.1.2**と**4.37.1**でした。

```shell-session
$ pip list | grep -e torch -e transformers
torch              2.1.2
transformers       4.37.1
```


## サンプルコードを動かしてみる

まずは[サンプルコード](https://huggingface.co/stabilityai/stable-code-3b#usage)で試してみます。

私の環境(というか今時のMac)の**GPU**は[NVIDIA](https://ja.wikipedia.org/wiki/NVIDIA)製ではありませんので、 **model.cuda()** だけコメントアウトしました。**model.cpu()** と明示的に書いても良いです。こうすることで**CPU**が使用されます。設定している**device**は**model.device**で取得できます。

**GPU**を使う方法は後述します。せっかく「**GPUがなくてもオフラインで動作します**」という売りなので、まずは**CPU**でやってみます。

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stable-code-3b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
  "stabilityai/stable-code-3b",
  trust_remote_code=True,
  torch_dtype="auto",
)
# model.cuda()
inputs = tokenizer("import torch\nimport torch.nn as nn", return_tensors="pt").to(model.device)
tokens = model.generate(
  **inputs,
  max_new_tokens=48,
  temperature=0.2,
  do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))
```



こちらのコードがなにをするかというと、**tokenizer**に渡している 

```python
"import torch\nimport torch.nn as nn"
```

この続きを出力するというものです。

**max\_new\_tokens**は生成トークンの最大長であり、**temperature**は出力の一貫性に影響しますが、ここでは詳細は省きます。詳しくは[Text generation strategies](https://huggingface.co/docs/transformers/v4.30.0/en/generation_strategies)でも読んでみてください。

では実行してみます。
初回実行時はモデルのダウンロードで時間がかかりますのでご認識ください。

しばらく待つと以下の結果が出力されました(私の環境だと2回目以降の実行で40秒ぐらいかかりました)。

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from torch.nn.utils.weight_norm import weight_norm
```

後続のコードが生成されているので、無事に動いていそうです。

## サンプルとは違う例を渡してみる

**tokenizer**に[私の過去記事](https://qiita.com/nabata/items/00ae91dbb4a429b092ec)のコードの前半をセットしてみました。**Stable Video Diffusion**の**Web API呼び出し**の途中までです。

```python
inputs = tokenizer("""
import requests
import os
import base64
import time

api_key = os.getenv("STABILITY_API_KEY")
api_host = 'https://api.stability.ai'
original_image = "original_image.png"  # 画像ファイル名
seed = 0  # シード値(0〜2147483648)
cfg_scale = 2.5  # 画像にどれだけ忠実とするか(0〜10)
motion_bucket_id = 40  # 動画の動きの大きさ(1〜255)

# 生成API呼び出し
params = {
    "seed": seed,
    "cfg_scale": cfg_scale,
    "motion_bucket_id": motion_bucket_id,
}
file = {
    "image": (original_image, open(original_image, 'rb'), 'image/png'),
}
headers = {
        "Authorization": f"Bearer {api_key}"
    }
""", return_tensors="pt").to(model.device)
```

すると結果としては、**tokenizer**に渡したコードに続いて下記が出力されました。

```python
response = requests.post(f"{api_host}/api/v1/image", params=params, files=file, headers=headers)

# レスポンスの確認
```

準備したパラメータを用いて**post**を呼び出した上で、日本語のコメント文も記載してくれています。**max\_new\_tokens**の値に従い、コメント文で終わってしまってはいますが。
賢いですね。

実際のエンドポイントとは**url**が異なりますが、それはLLMが知らない情報でしょうからしょうがないですね。

## コメント文だけ渡してフルスクラッチで書いてもらう

これを**コード補完**と呼ぶのかとかはともかくとして、**コメント文**を渡して処理をまるまる書いてもらいました。

具体的には**tokenizer**の呼び出しに以下のように書きました。

```python
inputs = tokenizer("""
# Display positive integers from 1 to 100
# Only prime numbers are allowed
# Code is written in Python
""", return_tensors="pt").to(model.device)
```

以下を英訳したコメント文のみを渡したということです。

```shell-session
1から100の正の整数を表示する
素数のみが許容される
コードはPythonで書く
```

また、長めのコードを出力してほしいため**max\_new\_tokens**を**150**に変更しました。

```python
max_new_tokens=150,
```

このコードを実行すると、渡したコメント文に続けて以下が出力されました。

```python
# Importing the required modules
import math

# Function to check if a number is prime or not
def is_prime(num):
    if num == 1:
        return False
    for i in range(2, int(math.sqrt(num)) + 1):
        if num % i == 0:
            return False
    return True

# Function to print the prime numbers from 1 to 100
def print_prime():
    for i in range(1, 101):
        if is_prime(i):
            print(i)

# Calling the function to print the prime numbers
print_prime()
```

コードの品質は置いておくとして、実行すると無事に結果が出力されました。

```shell-session
2
3
5
7
11
13
17
19
23
29
31
37
41
43
47
53
59
61
67
71
73
79
83
89
97
```

ちゃんと1〜100の範囲での素数が出力されています。

## GPUで動かしてみた

**MacBook**の**GPU**を使用するために**device**として[MPS(Metal Performance Shaders)](https://developer.apple.com/documentation/metalperformanceshaders) を指定した場合も残しておきます。

**MPS**を指定可能かどうかは下記で判定できます。私の環境では**True**でした。

```python
torch.backends.mps.is_available()
```

**MPS**を指定するだけだと「**TypeError: BFloat16 is not supported on MPS**」のエラーが出るため、まずは **torch\_dtype** を **"auto"** から **torch.float16** に変更しました。

```python
torch_dtype=torch.float16
```

**bfloat16**と**float16**とはなにかという話は「[bfloat16 の数値形式 &nbsp;|&nbsp; Cloud TPU &nbsp;|&nbsp; Google Cloud](https://cloud.google.com/tpu/docs/bfloat16?hl=ja)」が参考になります。精度に関わります。

ちなみに「[Optimize machine learning for Metal apps - WWDC23 - Videos - Apple Developer](https://developer.apple.com/videos/play/wwdc2023/10050/)」で「**Starting with macOS Sonoma, MPSGraph adds support for a new data type, bfloat16**」と話していたので、**MacOS**を **Sonoma(バージョン14)** に上げればこのエラーは出ないのかもしれません(未確認)。

次に、**model**に対して**mps**へのデバイス切り替えを指定します。場所は**model**の取得後です。

```python
# model.cuda()
model.to(torch.device("mps"))
```

私の環境だと**GPU**を用いないときと比べて時間が**1/2〜1/3**に短縮されました。
試行回数は少ないですが、平均して**約1分**が**約20秒**になったという感じです。
          
## 最後に

「どんどん便利になっていきますね」という陳腐な感想しか出てこない。

## 宣伝

[SupershipのQiita Organization](https://qiita.com/organizations/supership-inc)を合わせてご覧いただけますと嬉しいです。他のメンバーの記事も多数あります。

Supershipではプロダクト開発やサービス開発に関わる方を絶賛募集しております。
興味がある方は[Supership株式会社 採用サイト](https://supership.jp/recruit/)よりご確認ください。