Granite Code Model で Hello World するまで

## Granite Code Model とは

プログラミングコードの生成、ドキュメンテーション、説明、バグ修正タスクなどに特化した、decoder-only な LLM である。116 もの言語でトレーニングされており、他のコード生成系オープン LLM と比較して高いパフォーマンスを示している。
`Granite Code Base Models` と `Granite Code Instruct Models` の 2 タイプあり、前者がコードに関する基本的なタスクを実行するのに適している基盤モデル (foundational models) であるのに対し、後者はコードに関する指示 (instruction) と git コミットをセットでファインチューニングしたモデルである。
また、各モデルに 3B, 8B, 20B, 34B のパラメータサイズのバリエーションがある。

https://github.com/ibm-granite/granite-code-models/tree/main

## この記事でやったこと

この Granite Code Model に対して、「XXX の関数を生成して」というシンプルなタスクを実行してもらうために行なった作業内容を紹介する。

## 導入手順


### 前提

以下のツール群、および python バージョンを使用した。
これらでなくても全然よく、とにかく実行環境があれば良いとは思う。

- 仮想環境構築: `pyenv`
- パッケージ管理: `poetry`
- python バージョン: `3.12.2`  

### 基本的な環境構築

```shell
pyenv local 3.12.2

poetry init

poetry add transformers

poetry shell

poetry install
```

### ベースとなるコード準備

[こちらのコードブロック](https://github.com/ibm-granite/granite-code-models/tree/main?tab=readme-ov-file#inference) の内容を main.py にコピペし、3 行目 の `cuda` を `cpu` に変更する。

```py
from transformers import AutoModelForCausalLM, AutoTokenizer

device = "cpu"
model_path = "ibm-granite/granite-3b-code-base" # pick anyone from above list

tokenizer = AutoTokenizer.from_pretrained(model_path)

# drop device_map if running on CPU
model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
model.eval()

# change input text as desired
input_text = "def generate():"
# tokenize the text
input_tokens = tokenizer(input_text, return_tensors="pt")

# transfer tokenized inputs to the device
for i in input_tokens:
    input_tokens[i] = input_tokens[i].to(device)

# generate output tokens
output = model.generate(**input_tokens)
# decode output tokens into text
output = tokenizer.batch_decode(output)

# loop over the batch to print, in this example the batch size is 1
for i in output:
    print(i)
```

### 初回実行 ~ エラーがなくなるまで対処

```sh
python main.py
```

以下のエラーが表示された。

```
None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
tokenizer_config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4.13k/4.13k [00:00<00:00, 6.45MB/s]
tokenizer.json: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.06M/2.06M [00:00<00:00, 2.76MB/s]
special_tokens_map.json: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.02k/1.02k [00:00<00:00, 3.97MB/s]
Traceback (most recent call last):
  File "/<home_dir>/dev/granite-test/./main.py", line 9, in <module>
    model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<home_dir>/Library/Caches/pypoetry/virtualenvs/granite-test-cUYgsk4r-py3.12/lib/python3.12/site-packages/transformers/utils/import_utils.py", line 1500, in __getattribute__
    requires_backends(cls, cls._backends)
  File "/<home_dir>/Library/Caches/pypoetry/virtualenvs/granite-test-cUYgsk4r-py3.12/lib/python3.12/site-packages/transformers/utils/import_utils.py", line 1488, in requires_backends
    raise ImportError("".join(failed))
ImportError: 
AutoModelForCausalLM requires the PyTorch library but it was not found in your environment. Checkout the instructions on the
installation page: https://pytorch.org/get-started/locally/ and follow the ones that match your environment.
Please note that you may need to restart your runtime after installation.
```

エラー内容に従い、`https://pytorch.org/get-started/locally/` で必要なパッケージ情報を確認する。

![スクリーンショット 2024-07-11 21.38.21.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/671630/7ffb81c0-4b9c-7777-d407-2a526d53d8ee.png)


以下を実行

```shell
pip3 install torch torchvision torchaudio
```


再度実行すると、以下のエラーが表示された。

```shell
python ./main.py                         
config.json: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 680/680 [00:00<00:00, 1.20MB/s]
Traceback (most recent call last):
  File "/<home_dir>/dev/granite-test/./main.py", line 9, in <module>
    model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<home_dir>/Library/Caches/pypoetry/virtualenvs/granite-test-cUYgsk4r-py3.12/lib/python3.12/site-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained
    return model_class.from_pretrained(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<home_dir>/Library/Caches/pypoetry/virtualenvs/granite-test-cUYgsk4r-py3.12/lib/python3.12/site-packages/transformers/modeling_utils.py", line 3199, in from_pretrained
    raise ImportError(
ImportError: Using `low_cpu_mem_usage=True` or a `device_map` requires Accelerate: `pip install accelerate`
```

エラー内容に従い、以下のパッケージをインストールする。

```sh
pip install accelerate
```

再度実行する

```sh
python ./main.py         
model.safetensors.index.json: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 41.6k/41.6k [00:00<00:00, 324kB/s]
model-00001-of-00002.safetensors: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4.97G/4.97G [01:34<00:00, 52.7MB/s]
model-00002-of-00002.safetensors: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.99G/1.99G [00:31<00:00, 62.9MB/s]
Downloading shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [02:07<00:00, 63.54s/it]
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:14<00:00,  7.29s/it]
generation_config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 137/137 [00:00<00:00, 691kB/s]
/<home_dir>/Library/Caches/pypoetry/virtualenvs/granite-test-cUYgsk4r-py3.12/lib/python3.12/site-packages/transformers/generation/utils.py:1249: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.
  warnings.warn(
def generate():
    """
    Generate a new set of random numbers.
    """
    global random_
```

一応エラーなく通ったぽい。

### ここから調整

#### プロンプトを変えてみる

コピーしてきたコード内のプロンプトを今回の目的に沿うように、**「a + b を計算する関数を生成する」** というものに変更した。

```python
# before
input_text = "def generate():"

# after
input_text = "generate a function to calculate a plus b"
```

結果:

```sh
UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.
  warnings.warn(
generate a function to calculate a plus b

function add(a, b) {
  return
```

おかしい。何か最後まで出力されていない（**=トークン数が不足している**）ようだ。

#### トークン数増やす

```python
# before
output = model.generate(**input_tokens)

# after
output = model.generate(**input_tokens, max_new_tokens=500)
```

結果:

```sh
generate a function to calculate a plus b

function add(a, b) {
  return a + b;
}

//generate a function to calculate a minus b

function subtract(a, b) {
  return a - b;
}

//generate a function to calculate a times b

function multiply(a, b) {
  return a * b;
}

//generate a function to calculate a divided by b

function divide(a, b) {
  return a / b;
}

//generate a function to calculate a modulo b

function modulo(a, b) {
  return a % b;
}

//generate a function to calculate a power b

function power(a, b) {
  return Math.pow(a, b);
}

//generate a function to calculate a square root

function squareRoot(a) {
  return Math.sqrt(a);
}

//generate a function to calculate a factorial

function factorial(a) {
  if (a === 0) {
    return 1;
  } else {
    return a * factorial(a - 1);
  }
}

//generate a function to calculate a factorial

function factorial(a) {
  if (a === 0) {
    return 1;
  } else {
    return a * factorial(a - 1);
  }
}

//generate a function to calculate a factorial

function factorial(a) {
  if (a === 0) {
    return 1;
  } else {
    return a * factorial(a - 1);
  }
}

//generate a function to calculate a factorial

function factorial(a) {
  if (a === 0) {
    return 1;
  } else {
    return a * factorial(a - 1);
  }
}

//generate a function to calculate a factorial

function factorial(a) {
  if (a === 0) {
    return 1;
  } else {
    return a * factorial(a - 1);
  }
}

//generate a function to calculate a factorial

function factorial(a) {
  if (a === 0) {
    return
```

ふむ、いっぱい出てきたが、欲しいのは `function add` に当たるものだ。

#### プロンプトをさらに変えてみる

```python
# before
input_text = "generate a function to calculate a plus b"

# after
input_text = f'Question:\n{input_text}\n\nAnswer:\n\n'
```

結果:

```shell

Question:
generate a function to calculate a plus b

Answer:


``
def plus(a, b):
    return a + b
``

Question:
generate a function to calculate a minus b

Answer:


``
def minus(a, b):
    return a - b
``

Question:
generate a function to calculate a times b

Answer:


``
def times(a, b):
    return a * b
``

Question:
generate a function to calculate a divided by b

Answer:


``
def divided_by(a, b):
    return a / b
``

Question:
generate a function to calculate a modulo b

Answer:


``
def modulo(a, b):
    return a % b
``

Question:
generate a function to calculate a power b

Answer:


``
def power(a, b):
    return a ** b
``

Question:
generate a function to calculate a square root

Answer:


``
def square_root(a):
    return a ** 0.5
``

Question:
generate a function to calculate a factorial

Answer:


``
def factorial(a):
    return a * factorial(a - 1) if a > 1 else 1
``

Question:
generate a function to calculate a factorial

Answer:


``
def factorial(a):
    return a * factorial(a - 1) if a > 1 else 1
``

Question:
generate a function to calculate a factorial

Answer:


``
def factorial(a):
    return a * factorial(a - 1) if a > 1 else 1
``

Question:
generate a function to calculate a factorial

Answer:


``
def factorial(a):
    return a * factorial(a - 1) if a > 1 else 1
``

Question:
generate a function to calculate a factorial

Answer:


``
def factorial(a):
    return a * factorial(a - 1) if a > 1 else 1
``

Question:
generate a function to calculate
```

まだお願いしていないものまで出力してくれている。

#### モデルを変えてみる

```python
# before
model_path = "ibm-granite/granite-3b-code-base"

# after
model_path = "ibm-granite/granite-3b-code-instruct"
```

結果:

```
Question:
generate a function to calculate a plus b

Answer:

1. Define a function that takes two parameters, a and b.
2. Add the two parameters together.
3. Return the result.

Here's an example implementation in Python:

\```python
def add(a, b):
    return a + b
\```

This function takes two parameters, a and b, and returns their sum.

```

それっぽい回答が得られた！

#### さらに改良

```python
# prompt
prompt = "generate a function to calculate a plus b"

## before
input = f'Question:\n{prompt}\n\nAnswer:\n\n'
input = tokenizer.apply_chat_template(input, tokenize=False, add_generation_prompt=True)

## after
chat = [
  { "role": "user", "content": prompt },
]
input = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
```

結果:

```
Question:
generate a function to calculate a plus b

Answer:
Sure! Here's a Python function that calculates the sum of two numbers:

\```python
def add(a, b):
    return a + b
\```

You can call this function with two numbers as arguments, and it will return their sum. For example:

\```python
result = add(3, 5)
print(result)  # Output: 8
\```

This function uses the `+` operator to add the two numbers together. The `return` statement returns the sum of `a` and `b`.

```

使用例までつくようになった！


## 終わりに

今回の試行だけでは　Granite モデルのパワーを感じられるところまではまだ到達しなかったが、先の紹介ページにあるように、パフォーマンスは保証されていることはさらなり、**トレーニングとファインチューニングに用いられているデータやそれらの手法がオープンであることも非常に重要なことである。**
そういうわけで、今後も要チェックのモデルである。