ai_queryを使用したバッチ推論の実行

プレビュー

この機能はパブリックプレビュー段階です。

この記事では、組み込みの Databricks SQL 関数 ai_queryを使用してバッチ推論を実行する方法について説明します。 この AI 関数の詳細については、ai_query 関数 を参照してください。

Databricks では、バッチ推論に ai_query とモデルサービングを使用することをお勧めします。 迅速な実験のために、 ai_query を トークン単位の従量課金エンドポイントと共に使用できます。

大規模なデータまたは本番運用データに対してバッチ推論を実行する準備ができたら、 Databricks はプロビジョニング スループット エンドポイントを使用してパフォーマンスを高速化することをお勧めします。 ai_query は、数十億トークンの範囲のデータセットを確実かつ一貫して処理することが確認されています。 プロビジョニング スループット エンドポイントを作成する方法については、 プロビジョニング スループット 基盤モデル APIs を参照してください。

Unity Catalog テーブルで LLM を使用したバッチ推論を開始するには、「基盤モデルを使用したバッチ推論 APIs プロビジョニング スループット」のノートブックの例を参照してください。

要件

  • ai_query 関数の要件を参照してください。

  • 使用するデータを含む Unity Catalog の Delta テーブルに対するクエリ アクセス許可。

バッチ推論の例のクエリ

このセクションの例では、クエリを実行する既存のエンドポイントにデプロイされたモデルがあることを前提としています。 Serving UI を使用している場合は、エンドポイントを選択し、右上の [使用] ボタンをクリックして [バッチ推論に使用] を選択できます。この選択により、SQL エディターが開き、 ai_queryを使用してバッチ推論の SQL クエリを記述して実行できます。

以下は、max_tokens と temperatureで failOnError と modelParameters を使用した一般的な例です。この例では、 concat()を使用してモデルのプロンプトと推論列を連結する方法も示しています。 連結を実行するには、 ||、 concat()、 format_string()など、複数の方法があります。

CREATE OR REPLACE TABLE ${output_table_name} AS (
  SELECT
      ${input_column_name},
      AI_QUERY(
        "${endpoint}",
        CONCAT("${prompt}", ${input_column_name}),
        failOnError => True,
        modelParameters => named_struct('max_tokens', ${num_output_tokens},'temperature', ${temperature})
      ) as response
    FROM ${input_table_name}
    LIMIT ${input_num_rows}
)

次の例では、comment_text データセットを使用して、llama_3_1_8b エンドポイントの背後にあるモデルに対してクエリを実行します。

WITH data AS (
  SELECT *
  FROM ml.sentiment.comments
  LIMIT 10000
)
  SELECT
    comment_text,
    ai_query(
      'llama_3_1_8b_batch',
      CONCAT('You are provided with text. Classify the text into one of these labels: "Positive", "Neutral", "Negative". Do not explain. Do not output any confidence score. Do not answer questions. Text: ', comment_text)
    ) AS label
  FROM data

次の例は、データ前処理ステップと後処理ステップを示しています。

WITH temp AS (
  SELECT *
  FROM ml.sentiment.comments
  LIMIT 10000
),
pre_process AS (
  SELECT comment_text
  FROM temp
  WHERE length(comment_text) > 50
),
sentiment AS (
  SELECT
    comment_text,
    ai_query(
      'llama_3_1_8b_batch',
      Concat('You are provided with text. Classify the text into one of these labels: "Positive", "Neutral", "Negative". Do not explain. Do not output any confidence score. Do not answer questions. Text: ', comment_text)
    ) AS label
  FROM pre_process
)
SELECT
  comment_text,
  label,
  CASE
    WHEN label NOT IN ("Positive", "Neutral", "Negative") THEN True
    ELSE FALSE
  END AS error
FROM sentiment

ジョブのスケジュール

SQLスクリプトの準備ができたら、必要な頻度でジョブを実行するようにスケジュールできます。 「スケジュールされたノートブック ジョブの作成と管理」を参照してください。