Databricksのai_query関数をDeltaLiveTables上で使って手軽かつ強力なデータパイプラインを作る

Posted at 2024-10-19

導入

Databricksのai_query関数、Delta Live Tables(以下、DLT)上でも使えるようなことがドキュメントに書かれていました。
簡単にLLMを使うデータパイプラインをDLTで作れるとかなり便利なので、試してみた記録です。

準備

今回は、日本語文書の感情分類(Negative/Positiveの分類)処理をai_queryで行い、その結果を保管するというシンプルなパイプラインを構築します。

データとして、以下のようなCSVファイルを適当に作成しました。

sample_001.csv

id,comment,category
1,滋賀レイクスのB.LEAGUE PREMIERへの初年度参入が決定しました！,sports
2,日々の運動は体の健康を維持するために特に重要です。,health
3,日経平均株価が7月以来の４万円台となりました,economy

このファイルをUnity Catalog管理下の適当なボリュームに保管します。

今回は/Volumes/training/llm/raws/samples/dlt_query/という場所に保管しています。

DLTパイプライン処理の定義

ノートブックを作成し、DLT用のパイプライン処理を定義します。
ai_query関数を使う関係上、全てSQLで記述しました。

まず、CSVファイルからデータを取り込む処理。
ファイルが増分で増えていくという想定として、ストリーミングテーブルにしています。

CREATE OR REFRESH STREAMING TABLE sample_comment_raw
COMMENT "ai_queryによる加工用サンプルデータ取込用ストリーミングテーブル" AS
SELECT
  *
FROM
  STREAM read_files(
    '/Volumes/training/llm/raws/samples/dlt_query/',
    format => 'csv',
    header => true,
    mode => 'FAILFAST'
  )

次にai_queryを使った感情分析処理を記述します。
内容はシンプルですね。

CREATE
OR REFRESH MATERIALIZED VIEW sentiment_added
TBLPROPERTIES('pipelines.channel' = 'PREVIEW')
COMMENT "感情分析結果の追加" AS
SELECT
  id,
  category,
  comment,
  ai_query(
    "<ai_queryで使うエンドポイント名>",
    concat(
      "次のテキストをポジティブかネガティブに分類してください。返答はPositiveかNegativeかだけで返してください。\n\n",
      comment
    )
  ) as sentiment
FROM
  live.sample_comment_raw;