ONNX 形式のトランスフォーマー モデルを使用してエンベディングを生成する

このチュートリアルでは、変換モデルを Open Neural Network Exchange（ONNX）形式でエクスポートし、ONNX モデルを BigQuery データセットにインポートして、SQL クエリからエンベディングを生成する方法について説明します。

このチュートリアルでは、sentence-transformers/all-MiniLM-L6-v2 モデルを使用します。この文変換モデルは、文エンベディングの生成において高速かつ効果的なパフォーマンスを発揮することで知られています。文エンベディングは、テキストの根本的な意味を捉えることで、セマンティック検索、クラスタリング、文の類似性などのタスクを可能にします。

ONNX は、機械学習（ML）フレームワークを統一された形式で表現するためのフォーマットです。ONNX に対する BigQuery ML サポートにより、次のことが可能になります。

  • 使い慣れたフレームワークを使用してモデルをトレーニングします。
  • モデルを ONNX 形式に変換します。
  • ONNX モデルを BigQuery にインポートし、BigQuery ML を使用して予測を行う。

トランスフォーマー モデルファイルを ONNX に変換する

必要に応じて、このセクションの手順に沿って、sentence-transformers/all-MiniLM-L6-v2 モデルとトークナイザーを ONNX に手動で変換できます。それ以外の場合は、すでに変換されている公開 gs://cloud-samples-data Cloud Storage バケットのサンプルファイルを使用できます。

ファイルを手動で変換する場合は、Python がインストールされたローカル コマンドライン環境が必要です。Python のインストールについて詳しくは、Python のダウンロードをご覧ください。

Transformer モデルを ONNX にエクスポートする

Hugging Face Optimum CLI を使用して、sentence-transformers/all-MiniLM-L6-v2 モデルを ONNX にエクスポートします。Optimum CLI を使用したモデルのエクスポートの詳細については、optimum.exporters.onnx を使用してモデルを ONNX にエクスポートするをご覧ください。

モデルをエクスポートするには、コマンドライン環境を開いて次の操作を行います。

  1. Optimum CLI をインストールします。

    pip install optimum[onnx]

  2. モデルをエクスポートします。--model 引数で Hugging Face モデル ID を指定します。--opset 引数は ONNXRuntime ライブラリのバージョンを指定します。BigQuery でサポートされている ONNXRuntime ライブラリとの互換性を維持するために、17 に設定されています。

    optimum-cli export onnx \
  --model sentence-transformers/all-MiniLM-L6-v2 \
  --task sentence-similarity \
  --opset 17 all-MiniLM-L6-v2/

モデルファイルは all-MiniLM-L6-v2 ディレクトリに model.onnx としてエクスポートされます。

Transformer モデルに量子化を適用する

Optimum CLI を使用して、エクスポートされた Transformer モデルに量子化を適用し、モデルサイズを縮小して推論を高速化します。詳細については、量子化をご覧ください。

モデルに量子化を適用するには、コマンドラインで次のコマンドを実行します。

optimum-cli onnxruntime quantize \
  --onnx_model all-MiniLM-L6-v2/ \
  --avx512_vnni -o all-MiniLM-L6-v2_quantized

量子化されたモデルファイルは、all-MiniLM-L6-v2_quantized ディレクトリに model_quantized.onnx としてエクスポートされます。

トークナイザーを ONNX に変換する

ONNX 形式の Transformer モデルを使用してエンベディングを生成するには、通常、トークナイザーを使用して、モデルへの 2 つの入力（input_ids と attention_mask）を生成します。

これらの入力を生成するには、onnxruntime-extensions ライブラリを使用して、sentence-transformers/all-MiniLM-L6-v2 モデルのトークナイザーを ONNX 形式に変換します。トークナイザーを変換すると、生のテキスト入力に対して直接トークン化を実行して、ONNX 予測を生成できます。

トークナイザーを変換するには、コマンドラインで次の操作を行います。

  1. Optimum CLI をインストールします。

    pip install optimum[onnx]

  2. 任意のテキスト エディタを使用して、convert-tokenizer.py という名前のファイルを作成します。次の例では、nano テキスト エディタを使用しています。

    nano convert-tokenizer.py

  3. 次の Python スクリプトをコピーして convert-tokenizer.py ファイルに貼り付けます。

    from onnxruntime_extensions import gen_processing_models

# Load the Huggingface tokenizer
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# Export the tokenizer to ONNX using gen_processing_models
onnx_tokenizer_path = "tokenizer.onnx"

# Generate the tokenizer ONNX model, and set the maximum token length.
# Ensure 'max_length' is set to a value less than the model's maximum sequence length, failing to do so will result in error during inference.
tokenizer_onnx_model = gen_processing_models(tokenizer, pre_kwargs={'max_length': 256})[0]

# Modify the tokenizer ONNX model signature.
# This is because certain tokenizers don't support batch inference.
tokenizer_onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1

# Save the tokenizer ONNX model
with open(onnx_tokenizer_path, "wb") as f:
  f.write(tokenizer_onnx_model.SerializeToString())

  4. convert-tokenizer.py ファイルを保存します。

  5. Python スクリプトを実行して、トークナイザーを変換します。

    python convert-tokenizer.py

変換されたトークナイザーは、all-MiniLM-L6-v2_quantized ディレクトリに tokenizer.onnx としてエクスポートされます。

変換したモデルファイルを Cloud Storage にアップロードする

変換モデルとトークナイザーを変換したら、次の操作を行います。

データセットを作成する

ML モデルを保存する BigQuery データセットを作成します。

コンソール

  1. Google Cloud コンソールで、[BigQuery] ページに移動します。

    [BigQuery] ページに移動

  2. [エクスプローラ] ペインで、プロジェクト名をクリックします。

  3. [アクションを表示] > [データセットを作成] をクリックします。

  4. [データセットを作成する] ページで、次の操作を行います。

    • [データセット ID] に「bqml_tutorial」と入力します。

    • [ロケーション タイプ] で [マルチリージョン] を選択してから、[US（米国の複数のリージョン）] を選択します。

    • 残りのデフォルトの設定は変更せず、[データセットを作成] をクリックします。

bq

新しいデータセットを作成するには、--location フラグを指定した bq mk コマンドを使用します。使用可能なパラメータの一覧については、bq mk --dataset コマンドのリファレンスをご覧ください。

  1. データの場所が US に設定され、BigQuery ML tutorial dataset という説明の付いた、bqml_tutorial という名前のデータセットを作成します。

    bq --location=US mk -d \
 --description "BigQuery ML tutorial dataset." \
 bqml_tutorial

    このコマンドでは、--dataset フラグの代わりに -d ショートカットを使用しています。-d と --dataset を省略した場合、このコマンドはデフォルトでデータセットを作成します。

  2. データセットが作成されたことを確認します。

    bq ls

API

定義済みのデータセット リソースを使用して datasets.insert メソッドを呼び出します。

{
  "datasetReference": {
     "datasetId": "bqml_tutorial"
  }
}

BigQuery DataFrames

このサンプルを試す前に、BigQuery DataFrames を使用した BigQuery クイックスタートの手順に沿って BigQuery DataFrames を設定してください。詳細については、BigQuery DataFrames のリファレンス ドキュメントをご覧ください。

BigQuery に対する認証を行うには、アプリケーションのデフォルト認証情報を設定します。詳細については、ローカル開発環境の ADC の設定をご覧ください。

import google.cloud.bigquery

bqclient = google.cloud.bigquery.Client()
bqclient.create_dataset("bqml_tutorial", exists_ok=True)

ONNX モデルを BigQuery にインポートする

変換されたトークナイザーと文変換モデルを BigQuery ML モデルとしてインポートします。

次のオプションのいずれかを選択します。

コンソール

  1. Google Cloud コンソールで、BigQuery Studio を開きます。

    [BigQuery Studio] に移動

  2. クエリエディタで次の CREATE MODEL ステートメントを実行して、tokenizer モデルを作成します。

     CREATE OR REPLACE MODEL `bqml_tutorial.tokenizer`
  OPTIONS (MODEL_TYPE='ONNX',
   MODEL_PATH='TOKENIZER_BUCKET_PATH')

    TOKENIZER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用する場合は、TOKENIZER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/tokenizer.onnx に置き換えます。

    オペレーションが完了すると、[クエリ結果] ペインに Successfully created model named tokenizer のようなメッセージが表示されます。

  3. [モデルに移動] をクリックして、[詳細] ペインを開きます。

  4. [特徴列] セクションでモデルの入力を確認し、[ラベル列] セクションでモデルの出力を確認します。

    `tokenizer model の [**詳細**] ペイン

  5. クエリエディタで次の CREATE MODEL ステートメントを実行して、all-MiniLM-L6-v2 モデルを作成します。

     CREATE OR REPLACE MODEL `bqml_tutorial.all-MiniLM-L6-v2`
  OPTIONS (MODEL_TYPE='ONNX',
   MODEL_PATH='TRANSFORMER_BUCKET_PATH')

    TRANSFORMER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用する場合は、TRANSFORMER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/model_quantized.onnx に置き換えます。

    オペレーションが完了すると、[クエリ結果] ペインに Successfully created model named all-MiniLM-L6-v2 のようなメッセージが表示されます。

  6. [モデルに移動] をクリックして、[詳細] ペインを開きます。

  7. [特徴列] セクションでモデルの入力を確認し、[ラベル列] セクションでモデルの出力を確認します。

    `all-MiniLM-L6-v2` モデルの [**詳細**] ペイン

bq

bq コマンドライン ツールの query コマンドを使用して、CREATE MODEL ステートメントを実行します。

  1. コマンドラインで、次のコマンドを実行して tokenizer モデルを作成します。

    bq query --use_legacy_sql=false \
"CREATE OR REPLACE MODEL
`bqml_tutorial.tokenizer`
OPTIONS
(MODEL_TYPE='ONNX',
MODEL_PATH='TOKENIZER_BUCKET_PATH')"

    TOKENIZER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用する場合は、TOKENIZER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/tokenizer.onnx に置き換えます。

    処理が完了すると、「Successfully created model named tokenizer」のようなメッセージが表示されます。

  2. コマンドラインで、次のコマンドを実行して all-MiniLM-L6-v2 モデルを作成します。

    bq query --use_legacy_sql=false \
"CREATE OR REPLACE MODEL
`bqml_tutorial.all-MiniLM-L6-v2`
OPTIONS
(MODEL_TYPE='ONNX',
  MODEL_PATH='TRANSFORMER_BUCKET_PATH')"

    TRANSFORMER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用する場合は、TRANSFORMER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/model_quantized.onnx に置き換えます。

    処理が完了すると、「Successfully created model named all-MiniLM-L6-v2」のようなメッセージが表示されます。

  3. モデルをインポートしたら、モデルがデータセットに表示されていることを確認します。

    bq ls -m bqml_tutorial

    出力は次のようになります。

    tableId            Type
------------------------
tokenizer          MODEL
all-MiniLM-L6-v2   MODEL

API

jobs.insert メソッドを使用してモデルをインポートします。リクエスト本文の QueryRequest リソースの query パラメータに CREATE MODEL ステートメントを入力します。

  1. 次の query パラメータ値を使用して、tokenizer モデルを作成します。

    {
"query": "CREATE MODEL `PROJECT_ID :bqml_tutorial.tokenizer` OPTIONS(MODEL_TYPE='ONNX' MODEL_PATH='TOKENIZER_BUCKET_PATH')"
}

    次のように置き換えます。

    • PROJECT_ID: プロジェクト ID。
    • TOKENIZER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用する場合は、TOKENIZER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/tokenizer.onnx に置き換えます。

  2. 次の query パラメータ値を使用して、all-MiniLM-L6-v2 モデルを作成します。

    {
"query": "CREATE MODEL `PROJECT_ID :bqml_tutorial.all-MiniLM-L6-v2` OPTIONS(MODEL_TYPE='ONNX' MODEL_PATH='TRANSFORMER_BUCKET_PATH')"
}

    次のように置き換えます。

    • PROJECT_ID: プロジェクト ID。
    • TRANSFORMER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用している場合は、TRANSFORMER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/model_quantized.onnx に置き換えます。

BigQuery DataFrames

このサンプルを試す前に、BigQuery DataFrames を使用した BigQuery クイックスタートの手順に沿って BigQuery DataFrames を設定してください。詳細については、BigQuery DataFrames のリファレンス ドキュメントをご覧ください。

BigQuery に対する認証を行うには、アプリケーションのデフォルト認証情報を設定します。詳細については、ローカル開発環境の ADC の設定をご覧ください。

ONNXModel オブジェクトを使用して、トークナイザー モデルと文変換モデルをインポートします。

import bigframes
from bigframes.ml.imported import ONNXModel

bigframes.options.bigquery.project = PROJECT_ID

bigframes.options.bigquery.location = "US"

tokenizer = ONNXModel(
  model_path= "TOKENIZER_BUCKET_PATH"
)
imported_onnx_model = ONNXModel(
  model_path="TRANSFORMER_BUCKET_PATH"
)

次のように置き換えます。

  • PROJECT_ID: プロジェクト ID。
  • TOKENIZER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用する場合は、TOKENIZER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/tokenizer.onnx に置き換えます。
  • TRANSFORMER_BUCKET_PATH は、Cloud Storage にアップロードしたモデルのパスに置き換えます。サンプルモデルを使用している場合は、TRANSFORMER_BUCKET_PATH を値 gs://cloud-samples-data/bigquery/ml/onnx/all-MiniLM-L6-v2/model_quantized.onnx に置き換えます。

インポートした ONNX モデルを使用してエンベディングを生成する

インポートしたトークナイザーと文変換モデルを使用して、bigquery-public-data.imdb.reviews 一般公開データセットのデータに基づいてエンベディングを生成します。

次のオプションのいずれかを選択します。

コンソール

ML.PREDICT 関数を使用して、モデルでエンベディングを生成します。

このクエリでは、ネストされた ML.PREDICT 呼び出しを使用して、次のとおりトークナイザーとエンベディング モデルを介して未加工のテキストを直接処理します。

  • トークン化（内部クエリ）: 内部 ML.PREDICT 呼び出しで bqml_tutorial.tokenizer モデルが使用されます。bigquery-public-data.imdb.reviews 一般公開データセットの title 列を text 入力として使用します。tokenizer モデルは、input_ids 入力や attention_mask 入力など、メインモデルが必要とする数値トークン入力に未加工のテキスト文字列を変換します。
  • エンベディング生成（外部クエリ）: 外部 ML.PREDICT 呼び出しで bqml_tutorial.all-MiniLM-L6-v2 モデルを使用します。このクエリは、内部クエリの出力から input_ids 列と attention_mask 列を入力として取得します。

SELECT ステートメントは、テキストのセマンティック エンベディングを表す FLOAT 値の配列である sentence_embedding 列を取得します。

  1. Google Cloud コンソールで、BigQuery Studio を開きます。

    [BigQuery Studio] に移動

  2. クエリエディタで、次のクエリを実行します。

    SELECT
sentence_embedding
FROM
ML.PREDICT (MODEL `bqml_tutorial.all-MiniLM-L6-v2`,
  (
  SELECT
    input_ids, attention_mask
  FROM
    ML.PREDICT(MODEL `bqml_tutorial.tokenizer`,
      (
      SELECT
        title AS text
      FROM
        `bigquery-public-data.imdb.reviews` limit 10))))

    次のような結果になります。

    +-----------------------+
| sentence_embedding    |
+-----------------------+
| -0.02361682802438736  |
| 0.02025664784014225   |
| 0.005168713629245758  |
| -0.026361213997006416 |
| 0.0655381828546524    |
| ...                   |
+-----------------------+

bq

bq コマンドライン ツールの query コマンドを使用して、クエリを実行します。このクエリは、ML.PREDICT 関数を使用して、モデルでエンベディングを生成します。

このクエリでは、ネストされた ML.PREDICT 呼び出しを使用して、次のとおりトークナイザーとエンベディング モデルを介して未加工のテキストを直接処理します。

  • トークン化（内部クエリ）: 内部 ML.PREDICT 呼び出しで bqml_tutorial.tokenizer モデルが使用されます。bigquery-public-data.imdb.reviews 一般公開データセットの title 列を text 入力として使用します。tokenizer モデルは、input_ids 入力や attention_mask 入力など、メインモデルが必要とする数値トークン入力に未加工のテキスト文字列を変換します。
  • エンベディング生成（外部クエリ）: 外部 ML.PREDICT 呼び出しで bqml_tutorial.all-MiniLM-L6-v2 モデルを使用します。このクエリは、内部クエリの出力から input_ids 列と attention_mask 列を入力として取得します。

SELECT ステートメントは、テキストのセマンティック エンベディングを表す FLOAT 値の配列である sentence_embedding 列を取得します。

コマンドラインで、次のコマンドを実行してクエリを実行します。

bq query --use_legacy_sql=false \
'SELECT
sentence_embedding
FROM
ML.PREDICT (MODEL `bqml_tutorial.all-MiniLM-L6-v2`,
  (
  SELECT
    input_ids, attention_mask
  FROM
    ML.PREDICT(MODEL `bqml_tutorial.tokenizer`,
      (
      SELECT
        title AS text
      FROM
        `bigquery-public-data.imdb.reviews` limit 10))))'

次のような結果になります。

+-----------------------+
| sentence_embedding    |
+-----------------------+
| -0.02361682802438736  |
| 0.02025664784014225   |
| 0.005168713629245758  |
| -0.026361213997006416 |
| 0.0655381828546524    |
| ...                   |
+-----------------------+

BigQuery DataFrames

このサンプルを試す前に、BigQuery DataFrames を使用した BigQuery クイックスタートの手順に沿って BigQuery DataFrames を設定してください。詳細については、BigQuery DataFrames のリファレンス ドキュメントをご覧ください。

BigQuery に対する認証を行うには、アプリケーションのデフォルト認証情報を設定します。詳細については、ローカル開発環境の ADC の設定をご覧ください。

predict メソッドを使用して、ONNX モデルでエンベディングを生成します。

import bigframes.pandas as bpd

df = bpd.read_gbq("bigquery-public-data.imdb.reviews", max_results=10)
df_pred = df.rename(columns={"title": "text"})
tokens = tokenizer.predict(df_pred)
predictions = imported_onnx_model.predict(tokens)
predictions.peek(5)

出力は次のようになります。

Transformer モデルの出力。