KerasNLP を使用して Gemma を使ってみる

ai.google.dev で表示 Google Colab で実行 Vertex AI で開く GitHub でソースを表示

このチュートリアルでは、KerasNLP を使用して Gemma の使用を開始する方法を説明します。Gemma は、Gemini モデルの作成に使用されたのと同じ研究とテクノロジーから構築された、軽量で最先端のオープンモデルのファミリーです。KerasNLP は、Keras で実装された自然言語処理(NLP)モデルのコレクションで、JAX、PyTorch、TensorFlow で実行可能です。

このチュートリアルでは、Gemma を使用して複数のプロンプトに対するテキスト レスポンスを生成します。Keras を初めて使用する場合は、始める前に Keras のスタートガイドをご覧ください(必須ではありません)。このチュートリアルでは、Keras について理解を深めていきます。

設定

Gemma の設定

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

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

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

環境変数を設定する

KAGGLE_USERNAMEKAGGLE_KEY の環境変数を設定します。

import os
from google.colab import userdata

# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env
# vars as appropriate for your system.
os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

依存関係をインストールする

Keras と KerasNLP をインストールします。

# Install Keras 3 last. See https://keras.io/getting_started/ for more details.
pip install -q -U keras-nlp
pip install -q -U keras>=3

バックエンドを選択

Keras は、シンプルさと使いやすさのために設計された高レベルのマルチフレームワーク ディープ ラーニング API です。Keras 3 では、バックエンドとして TensorFlow、JAX、PyTorch を選択できます。このチュートリアルでは、この 3 つすべてを使用できます。

import os

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

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

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

import keras
import keras_nlp

モデルを作成する

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

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

gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_2b_en")
Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...

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

summary を使用して、モデルに関する詳細情報を取得します。

gemma_lm.summary()

まとめからわかるように、このモデルには 25 億個のトレーニング可能なパラメータがあります。

テキストを生成

それでは、テキストを生成してみましょう。このモデルには、プロンプトに基づいてテキストを生成する generate メソッドがあります。オプションの max_length 引数には、生成されるシーケンスの最大長を指定します。

プロンプト "What is the meaning of life?" で試してみましょう。

gemma_lm.generate("What is the meaning of life?", max_length=64)
'What is the meaning of life?\n\nThe question is one of the most important questions in the world.\n\nIt’s the question that has been asked by philosophers, theologians, and scientists for centuries.\n\nAnd it’s the question that has been asked by people who are looking for answers to their own lives'

別のプロンプトでもう一度 generate を呼び出してみてください。

gemma_lm.generate("How does the brain work?", max_length=64)
'How does the brain work?\n\nThe brain is the most complex organ in the human body. It is responsible for controlling all of the body’s functions, including breathing, heart rate, digestion, and more. The brain is also responsible for thinking, feeling, and making decisions.\n\nThe brain is made up'

JAX または TensorFlow のバックエンドで実行している場合は、2 番目の generate 呼び出しがほぼ瞬時に返されることがわかります。これは、指定されたバッチサイズと max_lengthgenerate への各呼び出しが XLA でコンパイルされるためです。初回の実行は高コストですが、後続の実行ははるかに高速です。

リストを入力として使用し、バッチ プロンプトを提供することもできます。

gemma_lm.generate(
    ["What is the meaning of life?",
     "How does the brain work?"],
    max_length=64)
['What is the meaning of life?\n\nThe question is one of the most important questions in the world.\n\nIt’s the question that has been asked by philosophers, theologians, and scientists for centuries.\n\nAnd it’s the question that has been asked by people who are looking for answers to their own lives',
 'How does the brain work?\n\nThe brain is the most complex organ in the human body. It is responsible for controlling all of the body’s functions, including breathing, heart rate, digestion, and more. The brain is also responsible for thinking, feeling, and making decisions.\n\nThe brain is made up']

省略可: 別のサンプラーを試す

compile()sampler 引数を設定すると、GemmaCausalLM の生成方法を制御できます。デフォルトでは、"greedy" サンプリングが使用されます。

テストとして、"top_k" 戦略を設定してみましょう。

gemma_lm.compile(sampler="top_k")
gemma_lm.generate("What is the meaning of life?", max_length=64)
'What is the meaning of life? That was a question I asked myself as I was driving home from work one night in 2012. I was driving through the city of San Bernardino, and all I could think was, “What the heck am I doing?”\n\nMy life was completely different. I'

デフォルトの貪欲アルゴリズムは常に最も確率が高いトークンを選択しますが、トップ K アルゴリズムは上位 K 個の確率のトークンから次のトークンをランダムに選択します。

サンプラーを指定する必要はありません。最後のコード スニペットがユースケースに役立たない場合は、無視してかまいません。使用可能なサンプラーについて詳しくは、サンプラーをご覧ください。

次のステップ

このチュートリアルでは、KerasNLP と Gemma を使用してテキストを生成する方法を学習しました。次のステップとして、以下をご確認ください。