次の方法で共有

ai_query を使用してバッチ LLM 推論を実行する

  • [アーティクル]

重要

この機能はパブリック プレビュー段階にあります。

この記事では、組み込みの Databricks SQL 関数 ai_query を使用して、Foundation Model API によってプロビジョニングされたスループットを使用するエンドポイントを使用してバッチ推論を実行する方法について説明します。 この記事の例とガイダンスは、大規模な言語モデル (LLM) を使用して複数の入力を処理するバッチ推論ワークロードに推奨されます。

SQL または PySpark で ai_query を使用して、バッチ推論ワークロードを実行できます。 データに対してバッチ推論を実行するには、ai_queryで次のように指定します。

  • Unity カタログの入力テーブルと出力テーブル
  • プロビジョニングされたスループット エンドポイント名
  • モデル プロンプトとモデル パラメーター

この AI 関数の詳細についてはai_query関数を参照してください。

要件

  • Foundation Model API でサポートされているリージョン内のワークスペース。
  • 次のいずれかの項目:
    • 汎用計算は、コンピューティング サイズ i3.2xlarge 以上で、Databricks Runtime 15.4 ML LTS またはそれ以上を実行し、少なくとも 2 つのワーカーを使用します。
    • 中規模以上の SQL ウェアハウス。
  • エンドポイントを提供する既存のモデル。 プロビジョニング済みスループット基盤モデル API を参照して、プロビジョニング済みスループット エンドポイントを作成してください。
  • 使用するデータを含む Unity カタログの Delta テーブルに対するクエリアクセス許可。
  • ai_query()を使用するには、テーブルのプロパティの pipelines.channel を 'preview' に設定します。 サンプル クエリの 例 を参照してください。

ai_query と SQL を使用する

ai_query と SQL を使用したバッチ推論の例を次に示します。 この例には、max_tokens と temperature を含む 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}),
        modelParameters => named_struct('max_tokens', ${num_output_tokens},'temperature', ${temperature})
      ) as response
    FROM ${input_table_name}
    LIMIT ${input_num_rows}
)

ai_query と PySpark を使用する

Python を使用する場合は、次に示すように、ai_query と PySpark を使用してバッチ推論を実行することもできます。

df_out = df.selectExpr("ai_query('{endpoint_name}', CONCAT('{prompt}', {input_column_name}), modelParameters => named_struct('max_tokens', ${num_output_tokens},'temperature', ${temperature})) as {output_column_name}")

df_out.write.mode("overwrite").saveAsTable(output_table_name)

Python を使用したバッチ推論ノートブックの例

次のノートブック例では、プロビジョニングされたスループット エンドポイントを作成し、Python と Meta Llama 3.1 70B モデルを使用してバッチ LLM 推論を実行します。 また、バッチ推論ワークロードのベンチマークと、エンドポイントを提供するプロビジョニング済みスループット モデルの作成に関するガイダンスもあります。

プロビジョニングされたスループット エンドポイント ノートブックを使用した LLM バッチ推論

ノートブック を取得する

次のサンプル ノートブックでは、プロビジョニングされたスループット エンドポイントを作成し、Python と、GTE Large (English) または BGE Large (English) 埋め込みモデルのいずれかを使用してバッチ LLM 推論を実行します。

プロビジョニングされたスループット エンドポイント ノートブックでの LLM バッチ推論の埋め込み

ノートブック を取得する