BigQuery MLでテキスト生成を試してみた

TL;DR

• BigQuery MLを使うと、BigQuery上のデータに対する推論処理がBigQuery内で完結できます。
• Dataformからも当然呼び出せるので、推論対象データをBigQueryで準備できるのなら、推論バッチ全体がBigQueryで完結します。
• AIモデルは、Googleマネージドなモデルや、各自で用意したローカルモデルも利用できます。

きっかけ

AI戦略室では、BigQueryで処理したデータに対して、AIモデルを用いたバッチ推論を多数実施しております。
その際、Vertex AI Pipelinesを利用して、BigQueryによるオペレーションの後にVertex AI Batch Predictionによる推論を実行する実装を行っています。しかし、このVertex AI Batch Predictionには、実装が大変だったり、オブザバビリティが低かったりと問題がいくつかありました。(IssueTrackerで要望は上がってるみたいだが。。)そこで、Batch Predictionに代わるプロダクトが無いか調べていたところ、Dataflow MLとBigQuery MLがいい感じだったので、今回はBigQuery MLの方を調査してみます。

BigQueryからモデルを呼び出すクエリのイメージ

SELECT *
FROM ML.GENERATE_TEXT( -- [1]
  MODEL `<PROJECT_ID>.sano_bqml_test.gemini_flash`, -- [2]
  TABLE `<PROJECT_ID>.sano_bqml_test.questions_for_gemini`, -- [3]
  STRUCT(  -- [4]
      1024 AS max_output_tokens,
      0.1 AS temperature,
      5 AS top_k,
      0.3 AS top_p,
      TRUE AS flatten_json_output  
   )
);

[1]: テキスト生成メソッドを生成する関数は、 ML.GENERATE_TEXT関数です。
[2]: データセット内にあらかじめ登録したモデルのみ呼び出すことができます。モデルは以下の3種類が使用できます。

• リモートモデル(Vertex AI Endpoint)
  • Geminiを含む、Vertex AI上にデプロイされたモデル
  • Cloud VisionやCloud Natural LanguageなどのマネージドAI API
• ローカルモデル
  • BigQuery内でホストされ、BigQuery内で使用されるモデル
  • 線形回帰や、ARIMAモデルなどが該当
• インポートモデル
  • 自前のモデルをGCSに配置し、インポートして使う

今回はGemini Flash 1.5を使って検証してみようと思いますので、「リモートモデル」に該当します。

[3]: プロンプトが含まれているテーブルを指定します。ML.GENERATE_TEXT関数は、promptという固定のカラム名からプロンプトを読み込み、決められたカラム名で生成結果を出力します。

[4]: いわゆる「温度」や「TopK」といったパラメータを、構造体を経由して渡します。

BigQueryからモデルを呼び出すインフラのイメージ（たぶんこう）

リモートモデルはあくまでVertex AIでデプロイされているもので、BigQueryの外の世界になります。
こういったBigQueryの外の世界にBigQueryがアクセスするときは、「外部接続」というものを経由します。

外部接続は、外部リソースへのアクセスのためにサービスアカウントを利用しているので、サービスアカウントには外部リソースに適したロールをアタッチする必要があります。

また、外部接続を作った後は、CREATE MODEL関数を使って、データセット内にモデルオブジェクトを作成する必要があります。BQからリモートモデルにアクセスするための「アクセスポイント」みたいなものです。
このアクセスポイント経由で、Vertex AIのリモートモデルに推論をリクエストし、その結果をBigQueryに返すことで、(ユーザーから見た際の)BQ内での推論処理の完結が成り立っています。

BQのクエリからGeminiを呼び出してテキスト生成する

ここからはBQ上でのデモ画像が続きます。

外部接続を作成する

まずは外部のモデルにBigQueryから接続できるように、「外部接続」を作成します。
BigQuery Studioの[+ADD]ボタンから、[Connections to external data sources]を作成。

外部接続とは、BigQueryからBigQuery外のデータソース(GCSやCloudSQLなど)に接続するためのリソースです。今回はVertex AI Remote modelsを選択します。
適当なConnection IDを設定します。

するとプロジェクトの下に、External Connectionという項目ができます。
このとき、外部接続するためのサービスアカウントが払い出されるので、このサービスアカウントに対して、Vertex AI Userロールを割り当てておきます。

リモートモデルを登録する

CREATE OR REPLACE MODEL
  `<PROJECT_ID>.sano_bqml_test.gemini_flash` --[1]
REMOTE WITH CONNECTION `<PROJECT_ID>.asia-northeast1.sano-bqml-test` -- [2]
OPTIONS (ENDPOINT = 'gemini-1.5-flash-002'); --[3]

[1]：モデルはデータセット内部に作成する
[2]：先ほど作成した外部接続のパス
[3]：エンドポイント名

入力データ

ML.GENERATE_TEXT関数ではpromptというカラムをAIに入力するように設定されているので、promptというカラムを必ず含めます。
今回は、下記のようなテーブルを作成しました。

推論を実行する

前述のクエリを少し改良しました。
ML.GENERATE_TEXT関数の出力をSELECT *で受けるとカラムがたくさん増えてしまうので、
ちゃんとカラムを選択しました。

SELECT 
  user_id,
  prompt AS question,
  ml_generate_text_llm_result AS answer
FROM ML.GENERATE_TEXT(
  MODEL `<PROJECT_ID>.sano_bqml_test.gemini_flash`,
  TABLE `<PROJECT_ID>.sano_bqml_test.questions_for_gemini`,
  STRUCT(1024 AS max_output_tokens, 0.1 AS temperature,
  5 AS top_k, 0.3 AS top_p, TRUE AS flatten_json_output)
);

実行結果

ユーザID123456さんの質問への回答は、、

BigQuery MLから直接Geminiを呼び出す方法は、現時点ではありません。

いや、今呼び出したやないかーい！！
って感じですが、おそらくグラウンディングがうまくできていないからでしょう。
外部接続を通して呼び出しているので、確かに「直接」は呼び出してないですが、外部接続によって呼び出す方法を提示して欲しいですね。

Vertex AIではグランディング系の機能も充実してきているので、いずれBQMLでもグラウンディング付きの推論を生成して、ハルシネーションも考慮した生成もできるようになると、勝手に期待しております。

料金

リモートモデルの場合、推論料金＝Vertex AIのマネージドモデルの料金。つまりBigQuery MLからGemini Flashに推論を行った場合、Vertex AIでGemini Flashで推論を行った場合の料金＋スキャン量が課金されます。

このあたりのインフラの比較等、あらかじめ知見としてためておくと、いざシステム構築するときの判断材料となるでしょう。
私はバリバリのエンジニア人間でこういったインフラはとても興味があるので、どんどん調べていきたいです。

考察

BigQueryにはDataformというサブプロダクトがあり、これはdbtのように、SQLジョブに依存関係をつけて実行できるものです。BQMLはDataformからも当然呼び出せるので、日々のデータマートの更新にも簡単に活用できます。また、従来のバッチ処理でVertex Pipelinesすら使う必要もなくなるでしょう。
Dataform自体には課金が発生しないため、これは有用ではないかと思います。

また、ONNX オニキス 形式に変換できるモデルであれば、今回紹介しませんでしたがインポートモデルでGCSからモデルをインポートして利用できるので、MLOps等で学習モデルを定期的にGCSに出力している場合と愛称は抜群によいと思います。