Dataflow で Gemma オープンモデルを使用する

Gemma は、Gemini モデルの研究とそこで使用されたテクノロジーに基づいて構築された、軽量で最先端のオープンモデルのファミリーです。Gemma モデルは、Apache Beam 推論パイプラインで使用できます。「オープン ウェイト」とは、モデルの事前トレーニング済みパラメータ（重み）から解放されることを意味します。元のデータセット、モデル アーキテクチャ、トレーニング コードなどの詳細は提供されません。

ユースケース

Dataflow で Gemma モデルを使用して感情分析を行うことができます。Dataflow と Gemma モデルを使用すると、カスタマー レビューなどのイベントを発生時に処理できます。モデルでレビューを実行して分析し、推奨事項を生成します。Gemma と Apache Beam を組み合わせることで、このワークフローをシームレスに完了できます。

サポートと制限事項

Gemma オープンモデルは、Apache Beam と Dataflow でサポートされていますが、次の要件があります。

  • Apache Beam Python SDK バージョン 2.46.0 以降を使用するバッチ パイプラインとストリーミング パイプラインで使用できます。
  • Dataflow ジョブで Runner v2 を使用する必要があります。
  • Dataflow ジョブで GPU を使用する必要があります。Dataflow でサポートされている GPU タイプの一覧については、可用性をご覧ください。推奨の GPU タイプは T4 と L4 です。
  • モデルは、.keras ファイル形式でダウンロードして保存する必要があります。
  • TensorFlow モデルハンドラの使用をおすすめします。ただし、必須ではありません。

前提条件

  • Kaggle から Gemma モデルにアクセスできる必要があります。
  • 同意フォームに記入して利用規約に同意してください。
  • Gemma モデルをダウンロードします。Cloud Storage バケットなど、Dataflow ジョブがアクセスできる場所に .keras ファイル形式で保存します。モデルパス変数の値を指定する場合は、この保存場所のパスを使用します。
  • Dataflow でジョブを実行するには、カスタム コンテナ イメージを作成します。この手順により、Dataflow サービスで GPU を使用してパイプラインを実行できるようになります。

パイプラインで Gemma を使用する

Apache Beam パイプラインで Gemma モデルを使用するには、次の操作を行います。

  1. Apache Beam コードで、パイプラインの依存関係をインポートした後、保存したモデルのパスを指定します。

    model_path = "MODEL_PATH"

    MODEL_PATH は、ダウンロードしたモデルの保存先のパスに置き換えます。たとえば、モデルを Cloud Storage バケットに保存した場合、パスの形式は gs://STORAGE_PATH/FILENAME.keras になります。

  2. Gemma モデルの Keras 実装には、プロンプトに基づいてテキストを生成する generate() メソッドがあります。generate() メソッドに要素を渡すには、カスタム推論関数を使用します。

    def gemma_inference_function(model, batch, inference_args, model_id):
  vectorized_batch = np.stack(batch, axis=0)
  # The only inference_arg expected here is a max_length parameter to
  # determine how many words are included in the output.
  predictions = model.generate(vectorized_batch, **inference_args)
  return utils._convert_to_result(batch, predictions, model_id)

  3. トレーニング済みモデルのパスを指定してパイプラインを実行します。この例では、TensorFlow モデルハンドラを使用します。

    class FormatOutput(beam.DoFn):
  def process(self, element, *args, **kwargs):
    yield "Input: {input}, Output: {output}".format(input=element.example, output=element.inference)

# Instantiate a NumPy array of string prompts for the model.
examples = np.array(["Tell me the sentiment of the phrase 'I like pizza': "])
# Specify the model handler, providing a path and the custom inference function.
model_handler = TFModelHandlerNumpy(model_path, inference_fn=gemma_inference_function)
with beam.Pipeline() as p:
  _ = (p | beam.Create(examples) # Create a PCollection of the prompts.
         | RunInference(model_handler, inference_args={'max_length': 32}) # Send the prompts to the model and get responses.
         | beam.ParDo(FormatOutput()) # Format the output.
         | beam.Map(print) # Print the formatted output.
  )

次のステップ