Databricksのai_query関数の基礎から高度な使い方

ai_queryのリファレンスで紹介されていた使い方をウォークスルーします。

ai_query関数とは

Databricksで大量のテキストデータにLLMを適用する際に使える汎用型の関数です。モデルサービングエンドポイントとプロンプトを指定することで、SQLの文脈で要約や翻訳などLLMタスクを柔軟にこなすことができます。SQLの中で使えるので、CREATE TABLE AS SELECT(CTAS)を通じて、大量のテキストデータに対してバッチのLLM推論処理を高速に行うことができます。

基本的な構文はこちらです。endpointはLLMがデプロイされているモデルサービングポイント名、requestはプロンプトを含むLLMへの指示を示すテキストです。

ai_query(endpoint, request)

準備

LLMがデプロイされたモデルサービングエンドポイントが必要です。今回はsystem.ai.llama_v3_3_70b_instructをデプロイしたプロビジョン済みスループットのモデルサービングエンドポイントを使用しています。

基本的な使い方

ある意味一番シンプルな例です。

SELECT ai_query(
    'taka-llama-pt',
    '100文字でDatabricks SQLを説明して'
  ) AS summary

モデルパラメータの指定

引数modelParametersでモデルのトークン数やtemperatureのようなパラメータを指定できます。

SELECT contents, ai_query(
    "taka-llama-pt",
    "以下の文を要約して: " || contents,
    modelParameters => named_struct('max_tokens', 100, 'temperature', 0.7)
) AS summary
FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 1;

出力の構造化による出力スキーマの強制

これは強力な機能です。通常ですとLLMのレスポンスは構造化されていないテキストですが、戻り値の構造を指定することで、それに合わせたレスポンスを得ることができます。後段でデータを加工したい場合に便利です・

responseFormatで戻り値の構造を指定します。

SELECT ai_query(
    "taka-llama-pt",
    "以下の文からマニュアルの詳細情報を抽出して: " || contents,
    responseFormat => 'STRUCT<doc_extraction:STRUCT<title:STRING, products:ARRAY<STRING>, summary:STRING, keywords:ARRAY<STRING>>>'
) AS extracted_info
FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 2;

JSONとして出力されていますが文字列なので、JSONとしてパースします。

SELECT parse_json(ai_query(
    "taka-llama-pt",
    "以下の文からマニュアルの詳細情報を抽出して: " || contents,
    responseFormat => 'STRUCT<doc_extraction:STRUCT<title:STRING, products:ARRAY<STRING>, summary:STRING, keywords:ARRAY<STRING>>>'
)) AS extracted_info
FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 2;

UDFでのai_queryの使用

関数の中にai_queryを含めることができますので、処理をモジュール化できます。

CREATE FUNCTION takaakiyayoi_catalog.doc_maintenance.create_title(text STRING)
  RETURNS STRING
  RETURN ai_query(
    'taka-llama-pt',
    CONCAT('次の記事から20文字程度のタイトルを作成して: \n', text));

SELECT
    title,
    takaakiyayoi_catalog.doc_maintenance.create_title(contents) AS generated_title
  FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 5;

title	generated_title
Databricks管理の概要 | Databricks Documentation	Databricks管理の概要と管理者の役割
外部Databricksサービスへの認証を有効にする | Databricks Documentation	外部Databricksサービスへの認証を有効にする
個人用アクセス トークンの監視と取り消し | Databricks Documentation	パーソナルアクセストークンの監視と取り消し
ステップ 4: ログ配信設定を作成する | Databricks Documentation	Databricksのログ配信設定を作成する方法
ステップ 2: 監査ログ配信の認証情報を設定する | Databricks Documentation	監査ログ配信の認証情報の設定方法

CTASによる処理結果の保存

これまでの処理では、結果が保存されていませんでした。結果を保存するにはCREATE TABLE AS SELECT(CTAS)を使うのが便利です。

CREATE TABLE takaakiyayoi_catalog.doc_maintenance.generated_titles AS
SELECT
  title,
  takaakiyayoi_catalog.doc_maintenance.create_title(contents) AS generated_title
FROM
  takaakiyayoi_catalog.doc_maintenance.databricks_docs
LIMIT 50;

テーブルに結果が保存されました。

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル