😺

BigQuery MLのML.GENERATE_EMBEDDING関数を活用したテキストエンベディング生成

に公開

この記事では、BigQuery MLにおける ML.GENERATE_EMBEDDING 関数に焦点を当て、BigQuery内で直接テキストエンベディングを効率的に生成する方法を解説します。


1. ML.GENERATE_EMBEDDINGの概要とメリット

ML.GENERATE_EMBEDDING 関数は、BigQuery MLにおける組み込み関数で、テキストデータからエンベディング(数値ベクトル表現)を直接生成できます。これにより、前処理から埋め込み生成までのプロセスをBigQuery上で完結することができて、データ転送のコストを省いてパイプラインを簡素化できます。

メリット

  • 高速性: 約100万件のテキストに対しても、数分で処理可能な高いスループット。(手元で試したところ6分程度で処理完了しました。)
  • BigQueryで完結: ETL処理などの実装を必要とせず、BigQueryのデータとシームレスに連携可能。

2. 実装手順

2.1 BigQuery ML接続の作成

まず、BigQueryからVertex AIモデルにアクセスするために BigQuery ML接続 を作成します。この接続は、エンベディングを実行するリージョンと同じリージョンで作成する必要があります。

  • 作成した接続のサービスアカウントには、Vertex AI ユーザー ロールを付与します。

2.2 エンベディングモデルの定義

Vertex AIの埋め込みモデルを呼び出すために、BigQuery MLモデルを定義します。

CREATE OR REPLACE MODEL `project_id.dataset.embedding_model`
REMOTE WITH CONNECTION `project_id.region.connection_name`
OPTIONS (endpoint = 'text-multilingual-embedding-002');

置き換えポイント:

  • project_id: 利用しているプロジェクトIDを指定
  • dataset: モデルを作成するBigQueryデータセット名
  • region: 作成した接続と同一のリージョン(例: us-west1
  • connection_name: 作成したBigQuery ML接続名

2.3 エンベディングの実行クエリ

作成したモデルを使って、テキストデータからエンベディングを生成します。

SELECT
  *
FROM ML.GENERATE_EMBEDDING(
  MODEL `project_id.dataset.embedding_model`, -- 2.2で作成したモデル名
  (
    SELECT
      SAFE_CAST(
        JSON_EXTRACT_SCALAR(response, '$.candidates[0].content.parts[0].text')
        AS STRING
      ) AS content
    FROM `project.dataset.input_table` -- テキストを含むテーブル
  ),
  STRUCT(
    'RETRIEVAL_DOCUMENT' AS task_type
  )
)

この例では、別の処理であらかじめエンベディング対象のデータを格納したテーブルのJSON型カラムからテキストを抽出し、エンベディングを生成しています。


まとめ

元々この辺の処理をPythonライブラリ経由で行っていたのですが、こちらの機能を知り試してみました。数件程度でしたらBigQueryを使わなくてもいいですが、ある程度データ量が多くなってきたらこちらの機能に置き換えるのも良いかもしれません。
ML.GENERATE_EMBEDDING 関数は、BigQueryユーザーにとって非常に強力なツールです。BigQuery内で完結できるため、開発・運用の効率化にも大きく貢献します。

参考文献

Discussion