CodeGemma と KerasNLP を使用した AI 支援プログラミング

ai.google.dev で表示

Google Colab で実行

GitHub のソースを表示

概要

CodeGemma は Gemma のバリエーションであり、コーディング タスク用にファインチューニングされています。このチュートリアルは、Keras CodeGemma クイックスタートをベースにしたもので、CodeGemma がプログラミング タスクを支援するその他の方法を紹介します。

セットアップ

CodeGemma にアクセスする

このチュートリアルを完了するには、まず Gemma の設定にある設定手順を完了する必要があります。Gemma の設定手順では、次の方法について説明します。

• kaggle.com で Gemma にアクセスしてください。
• Gemma 7B モデルを実行するのに十分なリソースを備えた Colab ランタイムを選択します。
• Kaggle のユーザー名と API キーを生成して構成します。

Gemma の設定が完了したら、次のセクションに進み、Colab 環境の環境変数を設定します。

ランタイムの選択

CodeGemma 7B モデルを実行するには、A100 GPU を使用するランタイムを提供する有料の Colab Pro プランに加入する必要があります。

1. Colab ウィンドウの右上にある ▾（その他の接続オプション）を選択します。
2. [ランタイムのタイプを変更] を選択します。
3. [ハードウェア アクセラレータ] で、[A100 GPU] を選択します。

API キーを設定する

Gemma を使用するには、Kaggle ユーザー名と Kaggle API キーを指定する必要があります。

Kaggle API キーを生成するには、Kaggle ユーザー プロフィールの [Account] タブに移動し、[Create New Token] を選択します。これにより、API 認証情報を含む kaggle.json ファイルのダウンロードがトリガーされます。

Colab の左側のペインで Secrets（🅰?）を選択し、Kaggle ユーザー名と Kaggle API キーを追加します。ユーザー名は KAGGLE_USERNAME、API キーは KAGGLE_KEY という名前で保存します。

環境変数を設定する

KAGGLE_USERNAME と KAGGLE_KEY の環境変数を設定します。

import os
from google.colab import userdata

os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

依存関係のインストール

pip install -q -U keras-nlp

バックエンドを選択

Keras は、シンプルさと使いやすさを考慮して設計された高レベルのマルチフレームワーク ディープ ラーニング API です。Keras 3 を使用すると、TensorFlow、JAX、PyTorch の 3 つのバックエンドのいずれかでワークフローを実行できます。

このチュートリアルでは、JAX のバックエンドを構成します。

os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch".

パッケージをインポートする

Keras と KerasNLP をインポートする。

import keras_nlp
import keras

# Run at half precision.
keras.config.set_floatx("bfloat16")

CodeGemma 7B モデルの例

このセクションでは、事前トレーニング済みの 7B CodeGemma モデルをコーディング タスクに使用する例を紹介します。

モデルを読み込む

KerasNLP は、因果言語モデリング用のエンドツーエンドの Gemma モデルである GemmaCausalLM を使用して、3 つの CodeGemma バリアント（2B と 7B の事前トレーニング済み（PT）と 7B の命令チューニング済み（IT））の実装を提供します。因果言語モデルは、以前のトークンに基づいて次のトークンを予測します。

この例では、from_preset メソッドを使用して code_gemma_7b_en モデルを読み込みます。

gemma_lm_7b = keras_nlp.models.GemmaCausalLM.from_preset("code_gemma_7b_en")

Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/config.json...
100%|██████████| 556/556 [00:00<00:00, 790kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/model.weights.h5...
100%|██████████| 15.9G/15.9G [02:39<00:00, 107MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/tokenizer.json...
100%|██████████| 401/401 [00:00<00:00, 587kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/assets/tokenizer/vocabulary.spm...
100%|██████████| 4.04M/4.04M [00:00<00:00, 16.4MB/s]

gemma_lm_7b.summary()

from_preset メソッドは、プリセットのアーキテクチャと重みからモデルをインスタンス化します。

複数行 FIM によるコード補完

PT CodeGemma モデルは、コード埋め込みタスクでトレーニングされます。このセクションでは、CodeGemma の複数行の中央入力（FIM）機能を使用して、周囲のコンテキストに基づいて指定されたカーソル位置のコードを自動入力する例を示します。

まず、定数とプロンプト形式設定ヘルパー関数を定義します。

# Formatting control tokens to specify cursor location
BEFORE_CURSOR = "<|fim_prefix|>"
AFTER_CURSOR = "<|fim_suffix|>"
AT_CURSOR = "<|fim_middle|>"
FILE_SEPARATOR = "<|file_separator|>"

# Define model stop tokens
END_TOKEN = gemma_lm_7b.preprocessor.tokenizer.end_token
stop_tokens = (BEFORE_CURSOR, AFTER_CURSOR, AT_CURSOR, FILE_SEPARATOR, END_TOKEN)
stop_token_ids = tuple(gemma_lm_7b.preprocessor.tokenizer.token_to_id(x) for x in stop_tokens)

def format_completion_prompt(before, after):
    return f"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}"

例 1 - 不足している条件を挿入する

フィボナッチ数列を生成する以下のサンプルコードは、n=1 の場合に正しく実行されません。

def fibonacci(n: int) -> int:
  if n == 0:
    return 0
  # The cursor is right before the e in the following line
  else:
    return fibonacci(n - 1) + fibonacci(n - 2)

カーソルが 4 行目（else 句のある場所）の先頭にあると仮定すると、カーソルの前後の内容は次のようになります。

before = """def fibonacci(n: int) -> int:\n  if n == 0:\n    return 0\n""" # Mind the spaces!
after = """\n  else:\n    return fibonacci(n - 1) + fibonacci(n-2)\n"""

prompt = format_completion_prompt(before, after)
print(prompt)

<|fim_prefix|>def fibonacci(n: int) -> int:
  if n == 0:
    return 0
<|fim_suffix|>
  else:
    return fibonacci(n - 1) + fibonacci(n-2)
<|fim_middle|>

プロンプトを実行します。

print(gemma_lm_7b.generate(prompt, stop_token_ids=stop_token_ids, max_length=128))

<|fim_prefix|>def fibonacci(n: int) -> int:
  if n == 0:
    return 0
<|fim_suffix|>
  else:
    return fibonacci(n - 1) + fibonacci(n-2)
<|fim_middle|>elif n == 1:
    return 1<|file_separator|>

モデルは、n=1 の正しい elif 条件をカーソルの位置に挿入します。

例 2 - 完全な DFS 走査アルゴリズム

深さ優先検索（DFS）ツリー トラバーサル アルゴリズムの予測入力コード。

before = """void dfs(node* root) {
  if (root->left) {
    dfs(root->left);
  }"""
after = """\nprintf("%d", root->value);
}"""

prompt = format_completion_prompt(before, after)
print(prompt)

<|fim_prefix|>void dfs(node* root) {
  if (root->left) {
    dfs(root->left);
  }<|fim_suffix|>
printf("%d", root->value);
}<|fim_middle|>

プロンプトを実行します。

print(gemma_lm_7b.generate(prompt, stop_token_ids=stop_token_ids, max_length=128))

<|fim_prefix|>void dfs(node* root) {
  if (root->left) {
    dfs(root->left);
  }<|fim_suffix|>
printf("%d", root->value);
}<|fim_middle|>
  if (root->right) {
    dfs(root->right);
  }<|file_separator|>

コード生成

CodeGemma 7B PT は、コードの補完に加えて、自然言語コーパスでもモデルのトレーニングが行われます。これを使用して、コードを生成するようにモデルに指示できます。

generation_prompt= """Write a rust function to identify non-prime numbers.
Examples:
>>> is_not_prime(2)
False
>>> is_not_prime(10)
True
pub fn is_not_prime(n: i32) -> bool {"""

print(gemma_lm_7b.generate(generation_prompt, max_length=500))

Write a rust function to identify non-prime numbers.
Examples:
>>> is_not_prime(2)
False
>>> is_not_prime(10)
True
pub fn is_not_prime(n: i32) -> bool {
    if n <= 1 {
        return true;
    }
    for i in 2..n {
        if n % i == 0 {
            return true;
        }
    }
    false
}

70 億 IT モデルの例

このセクションでは、より高度なコーディング タスクに CodeGemma 7B 命令チューニング済みモデルを使用します。CodeGemma 7B IT モデルは、CodeGemma 7B PT モデルから、コードに対する教師ありファインチューニングと、人間からのフィードバックを用いた強化学習によって派生しています。このセクションでは、このモデルを自由回答形式で生成する例を紹介します。

IT モデルを読み込む

from_preset メソッドを使用して、code_gemma_instruct_7b_en モデルを読み込みます。

gemma_lm_7b_it = keras_nlp.models.GemmaCausalLM.from_preset("code_gemma_instruct_7b_en")
gemma_lm_7b_it.summary()

Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/config.json...
100%|██████████| 556/556 [00:00<00:00, 754kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/model.weights.h5...
100%|██████████| 15.9G/15.9G [03:18<00:00, 86.2MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/tokenizer.json...
100%|██████████| 401/401 [00:00<00:00, 593kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/assets/tokenizer/vocabulary.spm...
100%|██████████| 4.04M/4.04M [00:00<00:00, 16.8MB/s]

IT モデルは特定のフォーマッタでトレーニングされています。フォーマッタは、チューニングを行うすべてのサンプルに追加情報でアノテーションを付けて、会話における役割を示し、ターンを区切るものです。

最初のステップとして、定数とプロンプト形式設定ヘルパー関数を定義します。

# Formatting control tokens for instruction tuning
START_OF_TURN_USER = "<start_of_turn>user"
END_OF_TURN = "<end_of_turn>"
START_OF_TURN_MODEL = "<start_of_turn>model"

# Formatting helper function
def format_instruction_prompt(context):
    return f"{START_OF_TURN_USER}\n{context}{END_OF_TURN}\n{START_OF_TURN_MODEL}\n"

コード変換

context1 = """
You are an experienced C and Python programmer. Convert the following Python code into C.
```python
def factorial(n):
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result
```\n"""

プロンプトをフォーマットします。

prompt1 = format_instruction_prompt(context1)
print(prompt1)

<start_of_turn>user

You are an experienced C and Python programmer. Convert the following Python code into C.

```python
def factorial(n):
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result
```
<end_of_turn>
<start_of_turn>model

プロンプトを実行します。

print(gemma_lm_7b_it.generate(prompt1, max_length=500))

<start_of_turn>user

You are an experienced C and Python programmer. Convert the following Python code into C.

```python
def factorial(n):
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result
```
<end_of_turn>
<start_of_turn>model
Here is the C code equivalent of the Python code:

```c
int factorial(int n) {
  int result = 1;
  for (int i = 2; i <= n; i++) {
    result *= i;
  }
  return result;
}
```

Here is a breakdown of the changes:

* The function is declared with the `int` return type, as in Python.
* The `for` loop is converted to a `for` loop with an `int` variable `i` initialized to 2 and incremented by 1 in each iteration.
* The `range` function is replaced with a simple loop that iterates from 2 to `n` (inclusive).
* The `result *= i` statement is used to multiply `result` by `i` in each iteration.
* The `return` statement is used to return the final value of `result`.

コードの脆弱性の検出

context2 = """
You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.
```cpp
int i;
unsigned int numWidgets;
Widget **WidgetList;

numWidgets = GetUntrustedSizeValue();
if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {
    ExitError("Incorrect number of widgets requested!");
}
WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));
printf("WidgetList ptr=%p\n", WidgetList);
for (i = 0; i < numWidgets; i++) {
    WidgetList[i] = InitializeWidget();
}
WidgetList[numWidgets] = NULL;
showWidgets(WidgetList);
```\n"""

プロンプトをフォーマットします。

prompt2 = format_instruction_prompt(context2)
print(prompt2)

<start_of_turn>user

You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.

```cpp
int i;
unsigned int numWidgets;
Widget **WidgetList;

numWidgets = GetUntrustedSizeValue();
if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {
    ExitError("Incorrect number of widgets requested!");
}
WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));
printf("WidgetList ptr=%p
", WidgetList);
for (i = 0; i < numWidgets; i++) {
    WidgetList[i] = InitializeWidget();
}
WidgetList[numWidgets] = NULL;
showWidgets(WidgetList);
```
<end_of_turn>
<start_of_turn>model

print(gemma_lm_7b_it.generate(prompt2, max_length=1000))

<start_of_turn>user

You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.

```cpp
int i;
unsigned int numWidgets;
Widget **WidgetList;

numWidgets = GetUntrustedSizeValue();
if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {
    ExitError("Incorrect number of widgets requested!");
}
WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));
printf("WidgetList ptr=%p
", WidgetList);
for (i = 0; i < numWidgets; i++) {
    WidgetList[i] = InitializeWidget();
}
WidgetList[numWidgets] = NULL;
showWidgets(WidgetList);
```
<end_of_turn>
<start_of_turn>model
Yes, the code is vulnerable to a memory access error.

**Reasoning:**

* The code allocates memory for `WidgetList` using `malloc` based on the value of `numWidgets`.
* However, the loop iterates from `0` to `numWidgets`, which is one element beyond the allocated memory.
* This means that accessing `WidgetList[numWidgets]` will result in a memory access error, as it is outside the bounds of the allocated memory.

**Example of Memory Access Error:**

When `numWidgets` is 5, the code allocates memory for `WidgetList` as follows:

```
WidgetList = (Widget **) malloc(5 * sizeof(Widget *));
```

The loop iterates from 0 to 4, accessing the following elements:

* `WidgetList[0]`
* `WidgetList[1]`
* `WidgetList[2]`
* `WidgetList[3]`
* `WidgetList[4]`

However, the code then attempts to access `WidgetList[5]`, which is outside the allocated memory range. This will result in a memory access error.

**Solution:**

To resolve this vulnerability, the loop should be modified to iterate from 0 to `numWidgets - 1`:

```cpp
for (i = 0; i < numWidgets - 1; i++) {
    WidgetList[i] = InitializeWidget();
}
```

This ensures that the loop does not access elements beyond the allocated memory range.

モデルはコード内の潜在的な脆弱性を検出し、その脆弱性を軽減するためのコード変更を提供します。

概要

このチュートリアルでは、さまざまなコーディング タスクで CodeGemma を使用する方法について説明しました。CodeGemma の詳細を確認する:

• CodeGemma モデルの技術仕様については、CodeGemma モデルカードをご覧ください。
• Vertex AI で CodeGemma を使用する方法について詳しくは、こちらをご覧ください。
• Keras CodeGemma クイックスタートを確認する。