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

プレビュー

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

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

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

  • Unity Catalog の入力テーブルと出力テーブル

  • プロビジョニングされたスループットエンドポイント名

  • モデル プロンプトと任意のモデル パラメーター

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

要件

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 バッチ推論

ノートブックを新しいタブで開く