この記事は BigQuery Advent Calendar 2024 20日目の記事です。

前提

今回は、BigQuery の ML.GENERATE_TEXT 関数についてお話をします。

ML.GENERATE_TEXT 関数は、BigQuery テーブルに保存されているテキストに対して自然言語生成タスクを実行することができる関数です。

ML.GENERATE_TEXT 関数を使用するには事前準備が必要ですが、本記事では割愛させていただきます。
詳細は公式ドキュメントをご確認ください。

背景

ML_GENERATE_TEXT 関数を使用していると、一部結果が null で返ってくることがあります。
例えば、次のクエリを実行してみました。

サンプルクエリ

SELECT
  *
FROM
  ML.GENERATE_TEXT(
    MODEL `project_id.dataset.model`, -- モデル名はご自身のものに変更ください
    (
      SELECT
        id,
        CONCAT('Classify the sentiment of the following text as positive or negative. Text:', input_column) AS prompt
      FROM
        UNNEST([
          STRUCT(1 AS id, "I love this product! It's amazing." AS input_column),
          STRUCT(2 AS id, "Terrible service. I'm very disappointed." AS input_column),
          STRUCT(3 AS id, "This is an average experience." AS input_column),
          STRUCT(4 AS id, "Highly recommended! Great value for money." AS input_column),
          STRUCT(5 AS id, "I wouldn't buy this again." AS input_column),
          STRUCT(6 AS id, "Fantastic quality and speedy delivery." AS input_column),
          STRUCT(7 AS id, "Not what I expected. Needs improvement." AS input_column),
          STRUCT(8 AS id, "Customer support was very helpful." AS input_column),
          STRUCT(9 AS id, "The product arrived damaged. Unacceptable." AS input_column),
          STRUCT(10 AS id, "Great features and easy to use." AS input_column),
          STRUCT(11 AS id, "It's okay, but I've seen better." AS input_column),
          STRUCT(12 AS id, "Perfect for my needs. Absolutely love it!" AS input_column),
          STRUCT(13 AS id, "Too expensive for the quality offered." AS input_column),
          STRUCT(14 AS id, "Quick delivery and excellent packaging." AS input_column),
          STRUCT(15 AS id, "The instructions were unclear and confusing." AS input_column),
          STRUCT(16 AS id, "Exceeded my expectations! Will buy again." AS input_column),
          STRUCT(17 AS id, "The colors were not as shown in the pictures." AS input_column),
          STRUCT(18 AS id, "Very satisfied with the purchase." AS input_column),
          STRUCT(19 AS id, "Horrible experience. Would not recommend." AS input_column),
          STRUCT(20 AS id, "Five stars! Everything was perfect." AS input_column)
        ])
    ),
    STRUCT(
      0.1 AS temperature,
      1000 AS max_output_tokens,
      0.1 AS top_p,
      10 AS top_k,
      TRUE AS flatten_json_output
    )
  );

すると、一部のレコードでは以下のような結果が返されます。

ml_generate_text_status カラムを確認すると、モデルのリソースが枯渇していることが原因のようです。
これは、リクエスト量が多すぎたことが影響していると思われます。

出力に null が含まれると、その後の処理で手間が増えるなど、なにかと厄介です。
そこで、本記事の内容となります。

※本問題の解決方法として、例えば、max_output_tokens などの引数を調整することで null が返されないようにすることも可能かもしれませんが、今回は考慮しません。

解決方法

方法として2つのアプローチが考えられます。

レコード単位で関数を実行する方法
null が返されたレコードのみを再処理する方法

前者は、処理を最小単位に分割する方法です。
モデルのリソースが枯渇してしまうなら、枯渇しないようにリクエスト量を調整しようという考えです。
一度に送信するリクエスト量を抑えるため、エラーの発生確率を低くすることが期待されます。

後者は、まず全てのレコードを一括で処理し、結果として null が返されたレコードに対してのみ再処理を行う方法です。
リソースの枯渇を前提とし、リクエスト量を事前に調整するのではなく、null が返されなくなるまで繰り返しリトライを行うという考えです。

1. レコード単位で関数を実行する方法

サンプルクエリは以下の通りです。

サンプルクエリ

DECLARE max_retries INT64 DEFAULT 5; -- 最大実行回数を定義
DECLARE temp_result STRING; -- ML.GENERATE_TEXT 関数の実行結果を格納する

-- 初回処理用の一時テーブルを作成
CREATE TEMP TABLE ResultTable AS
SELECT
  id,
  CONCAT('Classify the sentiment of the following text as positive or negative. Text:', input_column) AS prompt,
  CAST(NULL AS STRING) AS ml_generate_text_llm_result, -- STRING型
  CAST(NULL AS TIMESTAMP) AS created_at -- 処理時刻を記録する列を追加
FROM
  UNNEST([
    STRUCT(1 AS id, "I love this product! It's amazing." AS input_column),
    STRUCT(2 AS id, "Terrible service. I'm very disappointed." AS input_column),
    STRUCT(3 AS id, "This is an average experience." AS input_column),
    STRUCT(4 AS id, "Highly recommended! Great value for money." AS input_column),
    STRUCT(5 AS id, "I wouldn't buy this again." AS input_column),
    STRUCT(6 AS id, "Fantastic quality and speedy delivery." AS input_column),
    STRUCT(7 AS id, "Not what I expected. Needs improvement." AS input_column),
    STRUCT(8 AS id, "Customer support was very helpful." AS input_column),
    STRUCT(9 AS id, "The product arrived damaged. Unacceptable." AS input_column),
    STRUCT(10 AS id, "Great features and easy to use." AS input_column),
    STRUCT(11 AS id, "It's okay, but I've seen better." AS input_column),
    STRUCT(12 AS id, "Perfect for my needs. Absolutely love it!" AS input_column),
    STRUCT(13 AS id, "Too expensive for the quality offered." AS input_column),
    STRUCT(14 AS id, "Quick delivery and excellent packaging." AS input_column),
    STRUCT(15 AS id, "The instructions were unclear and confusing." AS input_column),
    STRUCT(16 AS id, "Exceeded my expectations! Will buy again." AS input_column),
    STRUCT(17 AS id, "The colors were not as shown in the pictures." AS input_column),
    STRUCT(18 AS id, "Very satisfied with the purchase." AS input_column),
    STRUCT(19 AS id, "Horrible experience. Would not recommend." AS input_column),
    STRUCT(20 AS id, "Five stars! Everything was perfect." AS input_column)
]);


LOOP
  -- NULLレコードを取得して再処理
  FOR record IN (
    SELECT id, prompt
    FROM ResultTable
    WHERE ml_generate_text_llm_result IS NULL
  )
  DO
    -- ML.GENERATE_TEXTの結果を取得
    SET temp_result = (
      SELECT ml_generate_text_llm_result
      FROM
        ML.GENERATE_TEXT(
          MODEL `project_id.dataset.model`, -- モデル名はご自身のものに変更ください
          (SELECT record.id AS id, record.prompt AS prompt),
          STRUCT(0.1 AS temperature, 1000 AS max_output_tokens, 0.1 AS top_p, 10 AS top_k, TRUE AS flatten_json_output)
        )
    );

    -- 結果とタイムスタンプを更新
    UPDATE ResultTable
    SET
      ml_generate_text_llm_result = temp_result,
      created_at = CURRENT_TIMESTAMP() -- 更新時刻を設定
    WHERE id = record.id;

  END FOR;

  -- 処理終了条件: 再試行回数を超過、または入力テーブルが空
  IF max_retries = 0 OR NOT EXISTS (SELECT 1 FROM ResultTable WHERE ml_generate_text_llm_result IS NULL) THEN
    LEAVE;
  END IF;

  -- 再試行カウントを増加
  SET max_retries = max_retries - 1;
END LOOP;


-- 最終結果を表示
SELECT * FROM ResultTable;

これにより、無事に null が生じることなく結果を取得することができました。
処理時間は3分ほどでした。

但し、この方法では、1レコードごとに ML.GENERATE_TEXT 関数を実行しており、子ジョブが行数に依存して増えてしまいます。
大規模なデータセットでは処理時間が長くなることも考えられます。

そこで、2の方法です。

2. null が返されたレコードのみを再処理する方法

サンプルクエリは以下の通りです。

サンプルクエリ

DECLARE max_retries INT64 DEFAULT 5;


-- ML.GENERATE_TEXT 関数実行時に用いるプロンプト一覧を格納
CREATE TEMP TABLE requests AS
SELECT
  id,
  CONCAT('Classify the sentiment of the following text as positive or negative. Text:', input_column) AS prompt
FROM
  UNNEST([
    STRUCT(1 AS id, "I love this product! It's amazing." AS input_column),
    STRUCT(2 AS id, "Terrible service. I'm very disappointed." AS input_column),
    STRUCT(3 AS id, "This is an average experience." AS input_column),
    STRUCT(4 AS id, "Highly recommended! Great value for money." AS input_column),
    STRUCT(5 AS id, "I wouldn't buy this again." AS input_column),
    STRUCT(6 AS id, "Fantastic quality and speedy delivery." AS input_column),
    STRUCT(7 AS id, "Not what I expected. Needs improvement." AS input_column),
    STRUCT(8 AS id, "Customer support was very helpful." AS input_column),
    STRUCT(9 AS id, "The product arrived damaged. Unacceptable." AS input_column),
    STRUCT(10 AS id, "Great features and easy to use." AS input_column),
    STRUCT(11 AS id, "It's okay, but I've seen better." AS input_column),
    STRUCT(12 AS id, "Perfect for my needs. Absolutely love it!" AS input_column),
    STRUCT(13 AS id, "Too expensive for the quality offered." AS input_column),
    STRUCT(14 AS id, "Quick delivery and excellent packaging." AS input_column),
    STRUCT(15 AS id, "The instructions were unclear and confusing." AS input_column),
    STRUCT(16 AS id, "Exceeded my expectations! Will buy again." AS input_column),
    STRUCT(17 AS id, "The colors were not as shown in the pictures." AS input_column),
    STRUCT(18 AS id, "Very satisfied with the purchase." AS input_column),
    STRUCT(19 AS id, "Horrible experience. Would not recommend." AS input_column),
    STRUCT(20 AS id, "Five stars! Everything was perfect." AS input_column)
  ]);


-- 出力結果格納用のテーブルを作成
CREATE TEMP TABLE ResultTable(
  id INT64,
  ml_generate_text_llm_result STRING,
  created_at TIMESTAMP
);


-- ループ処理
LOOP
  -- 処理終了条件: 入力テーブルが空、または再試行回数を超過
  IF max_retries = 0 OR NOT EXISTS (SELECT 1 FROM requests) THEN
    LEAVE;
  END IF;

  -- ML.GENERATE_TEXT 実行結果を ResultTable に格納
  INSERT INTO ResultTable
  SELECT
    *
  FROM(
    SELECT
      id,
      ml_generate_text_llm_result,
      CURRENT_TIMESTAMP() AS created_at
    FROM ML.GENERATE_TEXT(
      MODEL `project_id.dataset.model`, -- モデル名はご自身のものに変更ください
      (SELECT id, prompt FROM requests),
      STRUCT(0.1 AS temperature, 1000 AS max_output_tokens, 0.1 AS top_p, 10 AS top_k, TRUE AS flatten_json_output)
    )
  )
  WHERE
    ml_generate_text_llm_result IS NOT NULL;

  -- 次回処理のリクエストを準備
  DELETE FROM requests
    WHERE
      EXISTS (
        SELECT *
        FROM ResultTable
        WHERE requests.id = ResultTable.id
      );

  -- 残りリトライ回数を減算
  SET max_retries = max_retries - 1;
END LOOP;


-- 処理結果の確認
SELECT * FROM ResultTable;

非常にシンプルな構成になりました。
１の方法と比較して、処理時間も短くなっていました。

まとめ

このように、今回は ML.GENERATE_TEXT 関数の出力結果に null が含まれないような方法を考えてみました。

クエリの書き方は様々であると思います。
今回は、LOOP と FOR...IN を使用していますが、 WHILE などを使用されてみてもよいかもしれません。

より良いクエリの書き方などあれば、是非ご意見いただけますと幸いです！

最後に

ML.GENERATE_TEXT 関数はとても便利です。
クエリを実行するだけで、BigQuery テーブル内のテキストデータに対して自然言語生成タスクを簡単に実行できます。さまざまなユースケースで役立つと考えています。

この記事が少しでも皆さんのお役に立ち、ML.GENERATE_TEXT 関数の更なる活用に繋がれば幸いです。

余談

公式ドキュメントを確認すると、ML.GENERATE_TEXT 関数のパラメータに GROUND_WITH_GOOGLE_SEARCH が最近追加されていました。（2024/12/21 現在）

グラウンディングを活用できるので、より精度の高いレスポンスが期待できそうですね。