Gemmaモデルを使ったKerasNLPクイックスタート

チュートリアル

はじめに

この記事では、GoogleのGemmaモデルをKerasNLPで使用する方法について、Docker環境を利用して初心者にもわかりやすく解説します。Gemmaは強力な言語モデルで、テキスト生成タスクに優れた性能を発揮します。

Docker環境の構築

まず、Docker環境を構築します。これにより、環境の再現性が確保され、どのマシンでも同じ設定で開発を行うことができます。

Dockerファイルの準備

プロジェクトのルートディレクトリに以下の2つのファイルを作成します。

docker-compose.yml:

version: '3'
services:
  jax-keras:
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - ./:/workspace
    working_dir: /workspace
    command: jupyter lab --ip 0.0.0.0 --port 8989 --no-browser --allow-root
    tty: true
    ports:
      - "8989:8989"

Dockerfile:

FROM nvcr.io/nvidia/jax:24.04-py3

RUN pip install jupyterlab
RUN pip install jupytext

# requirements.txt
RUN pip install keras-nlp
RUN pip install keras==3.3.3

Docker環境の起動

ターミナルでプロジェクトのルートディレクトリに移動し、以下のコマンドを実行します：

docker-compose up --build

これにより、Docker環境が構築され、Jupyter Labが起動します。

Jupyter Labへのアクセス

ブラウザで http://localhost:8989 にアクセスすると、Jupyter Labの画面が表示されます。

Jupyter Labで環境設定

まず、環境を設定します。この設定は、Google ColabとローカルPython環境の両方に対応しています。Kaggleの認証情報を設定し、KerasのバックエンドをJAXに設定します。

!pip install python-dotenv
import os
from dotenv import load_dotenv

# .envファイルが存在する場合は読み込む
load_dotenv()

try:
    # Google Colabの場合
    from google.colab import userdata

    # KaggleのユーザーネームとAPIキーを設定
    os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
    os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')
except ImportError:
    # ローカル環境の場合、環境変数から読み込む
    kaggle_username = os.getenv("KAGGLE_USERNAME")
    kaggle_key = os.getenv("KAGGLE_KEY")

    if kaggle_username and kaggle_key:
        os.environ["KAGGLE_USERNAME"] = kaggle_username
        os.environ["KAGGLE_KEY"] = kaggle_key
    else:
        print("警告: KaggleのユーザーネームとAPIキーが設定されていません。")
        print(".envファイルまたは環境変数にKAGGLE_USERNAMEとKAGGLE_KEYを設定してください。")

# KerasのバックエンドをJAXに設定
os.environ["KERAS_BACKEND"] = "jax"

このコードは以下の動作をします：

  1. .envファイルがある場合、その内容を読み込みます。
  2. Google Colab環境ではuserdataを使用してKaggleの認証情報を取得します。
  3. ローカル環境では環境変数またはロードされた.envファイルからKaggleの認証情報を取得します。
  4. 認証情報が見つからない場合は警告メッセージを表示します。
  5. KerasのバックエンドをJAXに設定します。

注意点：

  • ローカル環境で実行する場合は、python-dotenvライブラリをインストールしてください： 
    pip install python-dotenv
  • .envファイルを使用する場合、プロジェクトのルートディレクトリに以下の内容で作成してください： 
    KAGGLE_USERNAME=your_username
KAGGLE_KEY=your_api_key

この設定により、Google ColabとローカルPython環境の両方で、KaggleのAPIを使用するコードがスムーズに動作するようになります。環境に応じて適切な方法でKaggleの認証情報を取得し、Kerasの設定を行います。

必要なライブラリのインストール

次に、必要なライブラリをインストールします。

# KerasNLPとKerasの最新バージョンをインストール
!pip install -U keras-nlp
#!pip install -U jax jaxlib
#!pip install -qU keras==3.3.3
#!pip install -U keras
!pip install -q -U "keras>=3"
!pip list

このコマンドでKerasNLPとKerasの最新バージョンがインストールされます。バージョンの互換性を確保するために、Kerasのバージョンを3.3.3に指定しています。

Gemmaモデルのロード

必要なライブラリをインポートし、Gemmaモデルをロードします。

import keras_nlp
import keras

# 計算精度をbfloat16に設定
keras.config.set_floatx("bfloat16")

# Gemma 9Bモデルをロード
gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma2_9b_en")
gemma_lm.summary()
Preprocessor: "gemma_causal_lm_preprocessor"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Tokenizer (type)                                   ┃                                             Vocab # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ gemma_tokenizer (GemmaTokenizer)                   │                                             256,000 │
└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘
Model: "gemma_causal_lm"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                  ┃ Output Shape              ┃         Param # ┃ Connected to               ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ padding_mask (InputLayer)     │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_ids (InputLayer)        │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ gemma_backbone                │ (None, None, 3584)        │   9,241,705,984 │ padding_mask[0][0],        │
│ (GemmaBackbone)               │                           │                 │ token_ids[0][0]            │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_embedding               │ (None, None, 256000)      │     917,504,000 │ gemma_backbone[0][0]       │
│ (ReversibleEmbedding)         │                           │                 │                            │
└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘
 Total params: 9,241,705,984 (17.21 GB)
 Trainable params: 9,241,705,984 (17.21 GB)
 Non-trainable params: 0 (0.00 B)

この部分について詳しく説明します：

  1. keras_nlpとkerasをインポートします。
  2. 計算精度をbfloat16に設定します。これにより、メモリ使用量を削減しつつ、高速な計算が可能になります。
  3. keras_nlp.models.GemmaCausalLM.from_preset()を使用して、Gemma 9Bモデル（英語版）をロードします。
  4. gemma_lm.summary()でモデルの構造を表示します。

テキスト生成の実行

モデルのロードが完了したら、テキスト生成を試してみましょう。

# 生成するテキストのプロンプト
prompt = "It was a dark and stormy night."

# テキスト生成を実行
result = gemma_lm.generate(prompt, max_length=256)

# 生成されたテキストを表示
print(result)
It was a dark and stormy night.

The wind was howling, the rain was pouring, and the thunder was rumbling.

I was sitting in my living room, watching the storm rage outside.

Suddenly, I heard a knock at the door.

I got up and opened it, and there stood a man in a black cloak.

He had a strange look in his eyes, and he was holding a lantern.

"Who are you?" I asked.

"I am the storm," he replied.

"And I have come to take you away."

I was terrified, but I couldn't move.

The man in the black cloak grabbed my arm and pulled me out into the storm.

We walked for what seemed like hours, until we came to a clearing in the woods.

There, the man in the black cloak stopped and turned to me.

"You are mine now," he said.

"And I will take you to my castle."

I tried to fight him off, but he was too strong.

He dragged me into the castle, and I was never seen again.

The end.

このコードの説明：

  1. prompt変数に生成の起点となるテキストを設定します。
  2. gemma_lm.generate()メソッドを使用してテキストを生成します。max_lengthパラメータで生成するテキストの最大長を指定します。
  3. 生成されたテキストを表示します。

生成されるテキストは毎回異なる可能性がありますが、与えられたプロンプトを元に、物語のような文章が続くでしょう。

まとめ

この記事では、Google ColabでGemmaモデルを使用してテキスト生成を行う方法を学びました。主なステップは以下の通りです：

  1. 環境設定
  2. 必要なライブラリのインストール
  3. Gemmaモデルのロード
  4. テキスト生成の実行

Gemmaモデルは非常に強力で、様々なタスクに応用できます。今回は基本的なテキスト生成を行いましたが、質問応答や要約など、他のタスクにも活用できます。

注意点

  • Gemmaモデルの実行には、大容量のGPU（例：A100）が必要です。
  • モデルの出力は確率的であり、同じプロンプトでも毎回異なる結果が得られる可能性があります。
  • 生成されたテキストの内容や品質は、入力プロンプトや設定パラメータによって大きく変わります。

以上で、GemmaモデルをKerasNLPで使用するクイックスタートガイドは終了です。このガイドを通じて、強力な言語モデルを簡単に使用できることがお分かりいただけたかと思います。さらに詳しく学びたい方は、KerasやKerasNLPの公式ドキュメントを参照することをおすすめします。

コメント

タイトルとURLをコピーしました