Databricksのai_query関数を用いたLLMのバッチ推論

こちらのアップデートです。

https://docs.databricks.com/en/release-notes/product/2024/october.html#mosaic-ai-model-serving-now-supports-batch-inference

> **Mosaic AIでバッチ推論をサポートしました**
>
> Mosaic AIモデルサービングで、`ai_query`を用いたバッチLLM推論をサポートしました。[`ai_query`を使用したバッチ推論の実行](https://docs.databricks.com/ja/large-language-models/ai-query-batch-inference.html)をご覧ください。

# `ai_query`関数とは

名前の通り、AI(LLM)にクエリーを行うことができる関数です。SQLで使えるところがミソです。

https://learn.microsoft.com/ja-jp/azure/databricks/sql/language-manual/functions/ai_query

これまでは、**トークン単位の従量課金(Pay per token)** が必要な[Foundation Model API](https://docs.databricks.com/ja/machine-learning/foundation-models/index.html)のモデルしか呼び出せなかったので、日本リージョンでは使えないなーと思っていたのですが、[プロビジョン済みスループットのモデルサービングエンドポイント](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)も呼び出せるようになったので、日本リージョンでも活用できます。

# 準備

**プレビュー**で**AI Query for Custom Models and External Models**がオンになっていることを確認します。

![Screenshot 2024-10-12 at 9.55.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ed92bf04-3f9f-30e6-1df0-6b3f0e3db590.png)

また、LLMに入力するデータがある場合、テーブルとして作成しておきます。

![Screenshot 2024-10-12 at 9.57.05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0b38c52c-1adc-5a0c-e0ba-8d95d7f002d0.png)

# ウォークスルー

`system.ai`配下のモデルを用いて、プロビジョン済みスループットモデルサービングエンドポイントを作成します。

![Screenshot 2024-10-12 at 9.19.58.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/f39a5653-e022-1a20-a0ad-31d6c05c6654.png)

![Screenshot 2024-10-12 at 9.52.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/d021208d-2baa-dbe5-4b6a-3dc2396a2c3a.png)

これで、`ai_function`を用いてSQLからLLMを呼び出せるようになりました。シンプルなものから試します。

```sql
SELECT
    ai_query(
      'taka_llama_3_2_8b_instruct',
      "こんにちは"
    ) AS response
```

返ってきました！

![Screenshot 2024-10-12 at 9.54.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7ff15daa-10b7-1422-e571-2ac4097f7025.png)

事前に自分のブログ記事をテーブルに登録しておいたので、こちらを使います。

![Screenshot 2024-10-12 at 9.54.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3b8ddb05-8f45-dec4-81db-8dcf8f2297ac.png)

```sql
WITH data AS (
  SELECT substr(text, 1, 1000) AS text -- プロンプト長の制限を超えないように
  FROM takaakiyayoi_catalog.ai_functions.blog
  LIMIT 1
)
  SELECT
    text,
    ai_query(
      'taka_llama_3_2_8b_instruct',
      CONCAT('提供されたテキストから50字程度の日本語の要約を作成してください。テキスト: ', text)
    ) AS summary
  FROM data
```

ここでハマったのはプロンプト長超過によるエラーでした。クエリーが返ってこないなーと思っていたら、モデルサービングエンドポイントのログに`prompt token count (8800) cannot exceed 8191. Please reduce the length of your prompt.`のエラーが出ていました。

![Screenshot 2024-10-12 at 9.58.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/9dfb27b5-7c6a-00f0-c61d-1bdb916ee48d.png)

なので、`substr`を使って入力文字列を制限しています。

要約が返ってきました！

![Screenshot 2024-10-12 at 10.04.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7f3cf6ce-3c59-f6da-223f-ee3571b70609.png)

ただ、余計なテキストも含まれています。ここからはプロンプトエンジニアリングの世界ですね。`前振りは不要です。要約テキストだけを返してください。`を追加します。

```sql
WITH data AS (
  SELECT substr(text, 1, 1000) AS text -- プロンプト長の制限を超えないように
  FROM takaakiyayoi_catalog.ai_functions.blog
  LIMIT 1
)
  SELECT
    text,
    ai_query(
      'taka_llama_3_2_8b_instruct',
      CONCAT('提供されたテキストから50字程度の日本語の要約を作成してください。前振りは不要です。要約テキストだけを返してください。テキスト: ', text)
    ) AS summary
  FROM data
```

いい感じです。

![Screenshot 2024-10-12 at 10.07.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/83f128e5-ed5d-daa9-faff-eb1a25cf78c6.png)

ここまでは、1行のレコードのみを処理していましたが、SQLなので容易に複数行を処理できます。`LIMIT 10`に変更します。

```sql
WITH data AS (
  SELECT substr(text, 1, 1000) AS text -- プロンプト長の制限を超えないように
  FROM takaakiyayoi_catalog.ai_functions.blog
  LIMIT 10
)
  SELECT
    text,
    ai_query(
      'taka_llama_3_2_8b_instruct',
      CONCAT('提供されたテキストから50字程度の日本語の要約を作成してください。前振りは不要です。要約テキストだけを返してください。テキスト: ', text)
    ) AS summary
  FROM data
```

これは楽です。

![Screenshot 2024-10-12 at 10.10.05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/83858b61-8a2d-05d6-7792-9a7c29480a15.png)

そして、LLMのレスポンスを用いてテーブルの作成も簡単にできてしまうわけです。CTE(Common Table Expression)で`CREATE TABLE`使えないので書き直してます)

```sql
CREATE TABLE takaakiyayoi_catalog.ai_functions.blog_summaries AS
SELECT
  title,
  text,
  ai_query(
    'taka_llama_3_2_8b_instruct',
    CONCAT(
      '提供されたテキストから50字程度の日本語の要約を作成してください。前振りは不要です。要約テキストだけを返してください。テキスト: ',
      text
    )
  ) AS summary
FROM
  (
    SELECT
      title,
      substr(text, 1, 1000) AS text
    FROM
      takaakiyayoi_catalog.ai_functions.blog
    LIMIT
      10
  )
```

これで要約のカラムが追加されたテーブルを作成できました。

![Screenshot 2024-10-12 at 10.18.27.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/38a40d9f-940a-36c7-aeb9-9781e9622fee.png)

SQL一本で要約の生成、テーブルの作成ができるのでジョブ化も簡単です。個人的にはバッチ処理も重要なLLMのユースケースになると思っています。ぜひご活用ください！

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)