Keras CodeGemma クイックスタート

ai.google.dev で表示 Google Colab で実行 GitHub のソースを表示

CodeGemma は、Gemini モデルの作成に使用されたものと同じ研究とテクノロジーに基づいて構築された、軽量で最先端のオープンモデルのファミリーです。

CodeGemma モデルは、2,500 億トークン(約 5,000 億トークン)以上の主なコードで アーキテクチャは Gemma モデルファミリーと同じですその結果、CodeGemma モデルは、どちらの補完においても最先端のコード パフォーマンスを達成しています。 生成タスクをサポートしながら、強力な 大規模に理解し、推論する能力が必要です。

CodeGemma には次の 3 つのバリアントがあります。

  • 7B コードの事前トレーニング済みモデル
  • 7B 命令チューニング済みのコードモデル
  • コード インフィルとオープンエンド生成のために特別にトレーニングされた 2B モデル。

このガイドでは、KerasNLP で CodeGemma 2B モデルを使用してコード補完タスクを実行する手順について説明します。

セットアップ

CodeGemma にアクセスする

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

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

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

ランタイムの選択

このチュートリアルを完了するには、CodeGemma 2B モデルを実行するのに十分なリソースを備えた Colab ランタイムが必要です。この場合は、T4 GPU を使用できます。

  1. Colab ウィンドウの右上にある ▾(その他の接続オプション)を選択します。
  2. [ランタイムのタイプを変更] を選択します。
  3. [ハードウェア アクセラレータ] で [T4 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_USERNAMEKAGGLE_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 つのバックエンドのいずれかでワークフローを実行できます。

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

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

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

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

import keras_nlp
import keras

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

モデルの読み込み

KerasNLP には、多くの一般的なモデル アーキテクチャの実装が用意されています。このチュートリアルでは、因果言語モデリング用のエンドツーエンドの Gemma モデルである GemmaCausalLM を使用してモデルを作成します。因果言語モデルは、以前のトークンに基づいて次のトークンを予測します。

from_preset メソッドを使用してモデルを作成します。

gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("code_gemma_2b_en")
gemma_lm.summary()
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/config.json...
100%|██████████| 554/554 [00:00<00:00, 1.41MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/model.weights.h5...
100%|██████████| 4.67G/4.67G [05:06<00:00, 16.4MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/tokenizer.json...
100%|██████████| 401/401 [00:00<00:00, 382kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/assets/tokenizer/vocabulary.spm...
100%|██████████| 4.04M/4.04M [00:01<00:00, 2.41MB/s]

from_preset メソッドは、プリセットのアーキテクチャと重みからモデルをインスタンス化します。上記のコードでは、文字列 code_gemma_2b_en でプリセット アーキテクチャ(20 億のパラメータを持つ CodeGemma モデル)を指定しています。

中間入力のコード補完

この例では、CodeGemma の中間埋め(FIM)機能を使用して、周囲のコンテキストに基づいてコードを完成します。これは、コードエディタ アプリケーションで、テキスト カーソルが周囲のコードに基づいている場所(カーソルの前後)にコードを挿入する場合に特に便利です。

CodeGemma では、4 つのユーザー定義のトークンを使用できます。3 つは FIM 用で、もう 1 つはマルチファイル コンテキスト サポート用の <|file_separator|> トークンです。これらを使用して定数を定義します。

BEFORE_CURSOR = "<|fim_prefix|>"
AFTER_CURSOR = "<|fim_suffix|>"
AT_CURSOR = "<|fim_middle|>"
FILE_SEPARATOR = "<|file_separator|>"

モデルのストップ トークンを定義します。

END_TOKEN = gemma_lm.preprocessor.tokenizer.end_token

stop_tokens = (BEFORE_CURSOR, AFTER_CURSOR, AT_CURSOR, FILE_SEPARATOR, END_TOKEN)

stop_token_ids = tuple(gemma_lm.preprocessor.tokenizer.token_to_id(x) for x in stop_tokens)

コード補完のプロンプトをフォーマットします。次のことに注意してください。

  • FIM トークンと接頭辞 / 接尾辞の間には空白を入れないでください
  • モデルが入力を継続できるように、FIM 中央トークンを最後に配置する必要があります。
  • ファイル内の現在のカーソル位置、またはモデルに提供するコンテキストの量に応じて、接頭辞または接尾辞を空にできます。

ヘルパー関数を使用してプロンプトを書式設定します。

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

before = "import "
after = """if __name__ == "__main__":\n    sys.exit(0)"""
prompt = format_completion_prompt(before, after)
print(prompt)
<|fim_prefix|>import <|fim_suffix|>if __name__ == "__main__":
    sys.exit(0)<|fim_middle|>

プロンプトを実行します。レスポンス トークンをストリーミングすることをおすすめします。ユーザー定義のトークンまたはターン終了/セネテンス トークンのいずれかに達したらストリーミングを停止し、結果のコード補完を取得します。

gemma_lm.generate(prompt, stop_token_ids=stop_token_ids, max_length=128)
'<|fim_prefix|>import <|fim_suffix|>if __name__ == "__main__":\n    sys.exit(0)<|fim_middle|>sys\n<|file_separator|>'

モデルは、提案されたコード補完として sys を提供します。

概要

このチュートリアルでは、CodeGemma を使用して、周囲のコンテキストに基づいてコードを記述する方法について説明しました。次に、CodeGemma と KerasNLP による AI アシスト プログラミングのノートブックで、CodeGemma のその他の使用例をご確認ください。

CodeGemma モデルの技術仕様については、CodeGemma モデルカードもご覧ください。