BigQuery DataFrames（BigFrames）からtext-bisonによる日本語のテキスト生成を実行してみる

データアナリティクス事業本部 機械学習チームの鈴木です。

前回はBigQueryからML.GENERATE_TEXT関数を使ってtext-bisonによるテキスト生成を試してみましたが、今回はBigQuery DataFrames（BigFramesパッケージ）からtext-bisonを利用してみたのでご共有します。

本記事の内容

BigQuery DataFramesのサンプルノートブックのうち、large_language_models.ipynbのコードを使って、text-bisonによるテキスト生成処理を実行してみました。

テキスト生成処理のためには、bigframes.ml.llm.PaLM2TextGeneratorクラスを使用しました。

実行環境は、BigFramesパッケージが利用可能なPython環境であれば特にこだわりはなかったのですが、今回はColab Enterprise（記事執筆時点でプレビュー中）のノートブックで検証しました。この環境で一連の流れを確認できればと思います。

なお、Colab EnterpriseからBigQuery DataFramesを利用する例については、以下の記事で取り上げていました。

準備作業

準備作業としては、外部データソースへの接続の作成およびサービスアカウントへのVertex AIユーザーロールの付与までが必要になりました。この作業は以下のブログの『text-bison利用のための準備』のうち、『1. 接続の追加』と『2. サービスアカウントへの権限付与』に記載しました。

手順例にしたがい、外部接続の下にus.generate_text_connectionが作成されており、この外部接続のサービスアカウントにVertex AIユーザーロールが付与されていることとします。

Python環境は、Colab Enterpriseのノートブックおよびランタイムについては、念の為に接続のロケーションに合わせて、us-west1 (オレゴン)リージョンに作成しました。

BigFrames Pythonパッケージのバージョンは以下でした。

ノートブック上から質問文を入力する場合

以下はサンプルノートブックから引用したコードになります。

1. モデルの作成

まず、PaLM2TextGeneratorクラスのインスタンスを作成しました。特に　connection_nameでは、準備で作成しておいた外部接続を指定しました。

import bigframes.pandas as bpd
import pandas as pd
from bigframes.ml.llm import PaLM2TextGenerator

session = bpd.get_global_session()
connection = "プロジェクトID.us.generate_text_connection"

model = PaLM2TextGenerator(session=session, connection_name=connection)

インスタンス作成にBigQuery側で処理が行われました。このとき、BigQuery側で実行されたSQLは以下でした。

CREATE TEMP MODEL `9948ac16e1ed45a58cc7dbc6ed1f61fa`
REMOTE WITH CONNECTION `プロジェクトID.us.generate_text_connection`
OPTIONS(
  remote_service_type="CLOUD_AI_LARGE_LANGUAGE_MODEL_V1")

BQMLでtext-bisonを使うときと同様に、CLOUD_AI_LARGE_LANGUAGE_MODEL_V1のリモートモデルが作成されていることが分かります。

2. データフレームの作成

次に、text-bisonに入力するプロンプトを持つデータフレームを作成します。このサンプルではPandas DataFrameを作った後、BigFramesのデータフレームに変換しています。

df = pd.DataFrame(
        {
            "prompt": ["BigQueryとはなんですか？", "BQMLについて教えてください。", "BigQuery DataFrameとはなんですか？"],
        })
bf_df = bpd.read_pandas(df)

テキスト生成を実施するためのpredictメソッドでは、ガイドによると、以下のようにpromptというカラム名の列を含むデータフレームを期待しているため、bf_dfはその仕様に合ったデータフレームとなっていることが分かります。

Input DataFrame or Series, which needs to contain a column with name "prompt". Only the column will be used as input. Prompts can include preamble, questions, suggestions, instructions, or examples.

3. テキスト生成の実行

predictメソッドでテキスト生成を実行しました。

pred = model.predict(bf_df).to_pandas()

一つ目のテキストである、BigQueryとはなんですか？の生成結果について取得してみます。

pred.iloc[0, 0]

BigQueryは、Google Cloud Platformの完全管理型ペタバイトスケールの分析データウェアハウスです。SQLをサポートし、ビジネスインテリジェンス（BI）や機械学習（ML）のワークロードに適しています。\n\nBigQueryは、以下の特徴を持っています。\n\n* スケーラブル：BigQueryは、ペタバイト規模のデータを処理することができます。\n* 高速：BigQueryは、高速なクエリ性能を提供します。\n* コスト効率：BigQueryは、コスト効率の良いデータウェアハウスです。\n* 使いやすい：BigQueryは、使いやすく、SQLをサポート

内容は問題なさそうです。途中で回答が途切れていたので、max_output_tokensパラメータを調整して再実行してみました。

pred = model.predict(bf_df, max_output_tokens=256).to_pandas()
pred.iloc[0, 0]

BigQueryは、Google Cloud Platformの完全管理型ペタバイトスケールの分析データウェアハウスです。SQLをサポートし、ビジネスインテリジェンス（BI）や機械学習（ML）のワークロードに適しています。\n\nBigQueryは、以下の特徴を持っています。\n\n* スケーラブル：BigQueryは、ペタバイト規模のデータを処理することができます。\n* 高速：BigQueryは、高速なクエリ性能を提供します。\n* コスト効率：BigQueryは、コスト効率の良いデータウェアハウスです。\n* 使いやすい：BigQueryは、使いやすく、SQLをサポートしています。\n\nBigQueryは、BIやMLのワークロードに適しています。BIでは、BigQueryは、大規模なデータセットの分析やレポートの作成に使用することができます。MLでは、BigQueryは、機械学習モデルのトレーニングや評価に使用することができます。\n\nBigQueryは、Google Cloud Platformの他のサービスと統合することができます。例えば、BigQueryは、Cloud StorageやCloud Dataprocと統合することができます。\n\nBigQueryは、強力なデータウェアハウスであり、BIやMLのワークロードに適しています。BigQueryは、スケーラブル、高速、コスト効率が良く、使いやすくなっています。

全テキストが取得できました。内容も問題なさそうです。

なお、predictメソッドを実行した場合、BigQuery側では以下のSQLが実行されていました。

        CREATE TEMP TABLE `_SESSION`.`e385a2afb8bf4931ac87f0630da119a5`
        CLUSTER BY `bigframes_unnamed_index`
        AS SELECT t0.*
FROM (
  SELECT * FROM ML.GENERATE_TEXT(MODEL `プロジェクトID._366d90d0aa0b1bf65124baf63b3c0cd43b11d0da._bd135e94_6c69_487b_9790_e498143fb249_5af2922f348c411bbf66edb5b120bc9a`,
    (SELECT t0.`bigframes_unnamed_index`, t0.`prompt`
  FROM (
    SELECT t1.`bigframes_unnamed_index`, t1.`prompt`
    FROM (
      SELECT * FROM `1fbac5950ff24e3c9027a252f89d030d`
    ) t1
  ) t0), STRUCT(
    0.0 AS temperature,
    128 AS max_output_tokens,
    40 AS top_k,
    0.95 AS top_p,
    True AS flatten_json_output))
) t0

max_output_tokensなど、predictメソッドからPaLM2のAPIに渡すパラメータもここに現れていますね。

既存のテーブルから質問文を取り出してプロンプトを作りたい場合

先ほどはPandas DataFrameを手作りしてBigQuery DataFramesのDataFrameに変換し、PaML2のAPIを利用しました。

BigQueryの機能を使ってtext-bisonを使う場合は、BigQueryのテーブルに既にデータが入っているケースも多いと思うので、それを意識した例も試してみました。APIの使い方は基本的に先に試したものと相違ありません。

なお、記事執筆時点でtext-bisonは入力・出力共に文字数に対して課金（1000文字あたり$0.0010）が発生します。個人的には、BigQueryのテーブルに既に入っているデータを使用する場合は、最初はLIMIT句などでレコード数を少なめにして動作検証していき、要件に適していることが確認できれば、リクエストするテキストのレコード数は見積もった上で利用されるとよいかなと思います。

1. テーブル・データ作成

まず、サンプルのsample_quiestionsテーブルを作成し、先ほども利用したテキストを格納しておきました。bqml_llmデータセットはあらかじめUSロケーションに作成しておきました。

CREATE TABLE `bqml_llm.sample_quiestions` (
  `question_id` INT64,
  `question` STRING
)

INSERT INTO `bqml_llm.sample_quiestions` (question_id, question)
VALUES 
  (1, 'BigQueryとはなんですか？'),
  (2, 'BQMLについて教えてください。'),
  (3, 'BigQuery DataFrameとはなんですか？')

2. テーブルからのデータ取得とプロンプト作成

以下のようにテーブルのデータを読み取り、promptカラムを持つデータレイクを作成しました。プロンプトは、テーブルの質問文をそのまま使ってもよいですが、以降の質問に回答してください。：　をつけて質問回答を意識したテキストが生成ができるようにしてみました。

bf_df_sample_quiestions = bpd.read_gbq("プロジェクトID.bqml_llm.sample_quiestions")
bf_df_sample_quiestions["prompt_prefix"] = "以降の質問に回答してください。：　"
bf_df_sample_quiestions["prompt"] = bf_df_sample_quiestions["prompt_prefix"] + bf_df_sample_quiestions["question"] 
bf_df_sample_quiestions = bf_df_sample_quiestions[["prompt"]]

bf_df_sample_quiestions.head()

3. テキスト生成の実行

predictメソッドで、以下のようにテキストを生成しました。

# 念の為件数を絞る操作を明記しておきます。
# テキスト生成したいプロンプト数を見積もって実施している場合はスキップして問題ありません。
bf_df_sample_quiestions = bf_df_sample_quiestions.head(5)

# テキスト生成
pred = model.predict(bf_df_sample_quiestions).to_pandas()

以降の質問に回答してください。：　BigQueryとはなんですか？というプロンプトに対しては以下のテキストが生成されました。

BigQueryとは、Google Cloud Platformが提供する、ペタバイト級のデータを高速に分析できる、完全管理型のペタバイトスケールの分析データウェアハウスです。BigQueryは、ビジネスインテリジェンス、データウェアハウス、機械学習などのユースケースに最適です。BigQueryは、標準SQLをサポートしており、他のGoogle Cloud Platformのサービスと統合することができます。

最後に

BigQuery DataFrames（BigFrames）からtext-bisonによる日本語のテキスト生成を実行してみました。

BQMLからの場合はリモートモデルの作成をする必要がありましたが、BigFramesでは一時モデルを自動で作ってくれるのでより直感的にtext-bisonを利用できる印象を受けました。

また、今回はプレビュー中のColab Enterpriseのノートブックから実行しましたが、ゆくゆくはBigQuery Studioに統合されたノートブックから画面で利用できるようになると思うので、非常に便利になることと思います。