カスタムモデル周りはハマってしまった。。。
導入
以下の記事で、Databricksのai_query関数がバッチ推論やサービングエンドポイントを利用できるようになったのを知りました。
同様のことをやろうとするとPythonでSparkのUDFを作ったりして・・・ということをしないといけなかったのでかなり便利です。SQL単体でできるし。
元記事ではプロビジョン済みスループットのモデルサービングエンドポイントを利用されていましたが、ドキュメントを読むにOpenAIなどの外部モデルやMLflowのカスタムpyfuncモデルのエンドポイントも利用できるようなのでそちらを試してみました。
結構ハマりポイントが多かったので、備忘録も兼ねて記事に残します。
Step1. 機能有効化
qi_queryからの外部モデルやカスタムモデルの利用については、「プレビュー」メニューから以下の項目を有効化しておく必要があります。
当初オフになっているのに気づかず、ai_query実行時に403エラーが出てハマってました。
上記記事でもしっかり解説されてますので、ちゃんと読みましょう。
Step2. 外部モデルを利用してみる
外部モデルのサービングエンドポイントを利用して、ai_queryを実行します。
まずは試験用のエンドポイントをサービングメニューから作成。
以下のような感じで、Azure上のOpenAI APIを利用するエンドポイントを作成しました。
（画像には出ていませんが、gpt-4o-miniを利用しています)
こちらのエンドポイントを利用してqi_queryを実行してみます。
Databricksのサンプルデータとして提供されているsamples.tpch.partテーブルを利用して商品名を日本語に翻訳してみました。
WITH tmp AS (
SELECT p_name FROM samples.tpch.part LIMIT 10
)
SELECT
p_name,
ai_query(
'test-external-openai-endpoint',
CONCAT(
'次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
p_name
)
)
AS translated
FROM
tmp
;
内容の妥当感はともかく、ちゃんと処理された結果が出力されました。
Step3. カスタムChatModelのモデルを利用してみる
次にMLflowのカスタムモデルを利用してみます。
MLflowのカスタムChatModelを使ったサービングエンドポイントならそのまま使えるかなと思い、以下の記事で作成したエンドポイントを流用します。
エンドポイント名だけを変えて実行。
WITH tmp AS (
SELECT p_name FROM samples.tpch.part LIMIT 10
)
SELECT
p_name,
ai_query(
'gemma-2-2b-jpn-it-endpoint',
CONCAT(
'次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
p_name
)
)
AS translated
FROM
tmp
;
結果、ダメでした。
エラー
[AI_FUNCTION_MODEL_SCHEMA_PARSE_ERROR] Failed to parse the schema for the serving endpoint "gemma-2-2b-jpn-it-endpoint": Failed to convert the schema object returned from the model endpoint into a valid SQL data type, response JSON was: \"{"openapi":"3.1.0","info":{"title":"gemma-2-2b-jpn-it-endpoint","version":"2"},"servers":[{"url":"https://tokyo.cloud.databricks.com/serving-endpoints/gemma-2-2b-jpn-it-endpoint"}],"paths":{"/served-models/gemma-2-2b-jpn-it-2/invocations":{"post":{"requestBody":{"content":{"application/json":{"schema":{"oneOf":[{"type":"object","properties":{"dataframe_split":{"type":"object","properties":{"index":{"type":"array","items":{"type":"integer"}},"columns":{"description":"required fields: messages","type":"array","items":{"type":"string","enum":["messages","temperature","max_tokens","stop","n","stream","top_p","top_k","frequency_penalty","presence_penalty","tools","metadata"]}},"data":{"type":"array","items":{"type":"array","prefixItems":[{"type":"array","items":{"required":["role"],"type":"object","properties":{"name":{"type":"string"},"role":{"type":"string"},"tool_calls":{"type":"array","items":{"required":["function","id","type"],"type":"object","properties":{"function":{"required":["arguments","name"],"type":"object","properties":{"arguments":{"type":"string"},"name":{"type":"string"}}},"id":{"type":"string"},"type":{"type":"string"}}}},"refusal":{"type":"string"},"content":{"type":"string"},"tool_call_id":{"type":"string"}}}},{"type":"number","format":"double"},{"type":"integer","format":"int64"},{"type":"array","items":{"type":"string"}},{"type":"integer","format":"int64"},{"type":"boolean"},{"type":"number","format":"double"},{"type":"integer","format":"int64"},{"type":"number","format":"double"},{"type":"number","format":"double"},{"type":"array","items":{"required":["type"],"type":"object","properties":{"function":{"required":["name","parameters"],"type":"object","properties":{"description":{"type":"string"},"name":{"type":"string"},"parameters":{"required":["properties"],"type":"object","properties":{"additionalProperties":{"type":"boolean"},"properties":{"type":"object"},"required":{"type":"array","items":{"type":"string"}},"type":{"type":"string"}}},"strict":{"type":"boolean"}}},"type":{"type":"string"}}}},{"type":"object"}]}}}}}},{"type":"object","properties":{"dataframe_records":{"type":"array","items":{"required":["messages"],"type":"object","properties":{"messages":{"type":"array","items":{"required":["role"],"type":"object","properties":{"name":{"type":"string"},"role":{"type":"string"},"tool_calls":{"type":"array","items":{"required":["function","id","type"],"type":"object","properties":{"function":{"required":["arguments","name"],"type":"object","properties":{"arguments":{"type":"string"},"name":{"type":"string"}}},"id":{"type":"string"},"type":{"type":"string"}}}},"refusal":{"type":"string"},"content":{"type":"string"},"tool_call_id":{"type":"string"}}}},"frequency_penalty":{"type":"number","format":"double"},"n":{"type":"integer","format":"int64"},"max_tokens":{"type":"integer","format":"int64"},"top_p":{"type":"number","format":"double"},"presence_penalty":{"type":"number","format":"double"},"temperature":{"type":"number","format":"double"},"tools":{"type":"array","items":{"required":["type"],"type":"object","properties":{"function":{"required":["name","parameters"],"type":"object","properties":{"description":{"type":"string"},"name":{"type":"string"},"parameters":{"required":["properties"],"type":"object","properties":{"additionalProperties":{"type":"boolean"},"properties":{"type":"object"},"required":{"type":"array","items":{"type":"string"}},"type":{"type":"string"}}},"strict":{"type":"boolean"}}},"type":{"type":"string"}}}},"stop":{"type":"array","items":{"type":"string"}},"stream":{"type":"boolean"},"metadata":{"type":"object"},"top_k":{"type":"integer","format":"int64"}}}}}}]}}}},"responses":{"200":{"description":"Successful operation","content":{"application/json":{"schema":{"type":"object","properties":{"predictions":{"type":"array","items":{"type":"object","properties":{"model":{"type":"string"},"choices":{"type":"array","items":{"required":["finish_reason","index","message"],"type":"object","properties":{"finish_reason":{"type":"string"},"index":{"type":"integer","format":"int64"},"message":{"required":["role"],"type":"object","properties":{"name":{"type":"string"},"role":{"type":"string"},"tool_calls":{"type":"array","items":{"required":["function","id","type"],"type":"object","properties":{"function":{"required":["arguments","name"],"type":"object","properties":{"arguments":{"type":"string"},"name":{"type":"string"}}},"id":{"type":"string"},"type":{"type":"string"}}}},"refusal":{"type":"string"},"content":{"type":"string"},"tool_call_id":{"type":"string"}}}}}},"usage":{"required":["completion_tokens","prompt_tokens","total_tokens"],"type":"object","properties":{"completion_tokens":{"type":"integer","format":"int64"},"prompt_tokens":{"type":"integer","format":"int64"},"total_tokens":{"type":"integer","format":"int64"}}},"object":{"type":"string"},"id":{"type":"string"},"metadata":{"type":"object"},"created":{"type":"integer","format":"int64"}}}}}}}}}}}}}}\".
Set the `returnType` parameter manually in the AI_QUERY function to override schema resolution. SQLSTATE: 2203G
エラーメッセージを見るに、returnTypeのスキーマ指定が必須の模様。
エラー内容に基づいてreturnTypeを指定、その後は入力スキーマの指定でもエラーが出るなど試行錯誤し、最終的には以下のSQLの実行まで行いました。
ただし、こちらでもエラー。
WITH tmp AS (
SELECT
p_name
FROM
samples.tpch.part
LIMIT
10
)
SELECT
p_name,
ai_query(
'gemma-2-2b-jpn-it-endpoint',
request => named_struct(
'messages',
array(
named_struct(
'role',
'user',
'content',
CONCAT(
'次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
p_name
)
)
)
),
returnType => 'STRUCT<id:STRING, object:STRING, created:LONG, model:STRING, choices:ARRAY<STRUCT<finish_reason:STRING, index:LONG, message:STRUCT<content:STRING, role:STRING>>>, usage:STRUCT<completion_tokens:LONG, prompt_tokens:LONG, total_tokens:LONG>>'
) AS translated
FROM
tmp;
エラー
[REMOTE_FUNCTION_HTTP_RESULT_PARSE_ERROR] Failed to evaluate the ai_query SQL function due to inability to parse the JSON result from the remote HTTP response; the error message is 'Can not parse the response from model serving endpoint. ai_query expect the model prediction output to be arrays, and the array size must be 1. However, the received predictions value is: {"model":"training.llm.gemma-2-2b-jpn-it","choices":[{"index":0,"message":{"role":"assistant","content":"ミステリーのような輝く青緑色。 \\n\\n\\n\\n<end_of_turn>"},"finish_reason":"stop"}],"usage":{"prompt_tokens":39,"completion_tokens":13,"total_tokens":52},"object":"chat.completion","id":"2","created":1728808851}.'. Check API documentation: https://docs.databricks.com/en/machine-learning/model-serving/index.html. Please fix the problem indicated in the error message and retry the query again. SQLSTATE: 22032
どうも、カスタムモデルだと出力は要素が一つのリスト(ARRAY)の結果である必要がある模様。
ChatModelだと出力が単一の辞書型になるので、ここで一旦断念。
通常のPyfuncモデルだと出力がリストになるので、現時点ではそこを想定した仕様になっているのだと思います。
Step4. カスタムChatModelのモデルを利用してみる(リトライ)
カスタムChatモデルをそのまま利用できないなら、カスタムモデルエンドポイントを外部モデルのエンドポイントとして使えば いけるんじゃないか？と思ったのでトライ。
外部モデルのプロバイダとしてDatabricks Model Servingを指定し、上記のカスタムモデルgemma-2-2b-jpn-it-endpointを利用するように外部モデルエンドポイントを作成。
そう、カスタムChatModelで作成したモデルは外部モデルのエンドポイントとして利用できるようです。（今回初めて知った）
作成したエンドポイントを使ってai_queryを実行。
WITH tmp AS (
SELECT p_name FROM samples.tpch.part LIMIT 10
)
SELECT
p_name,
ai_query(
'test-external-model-endpoint',
CONCAT(
'次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
p_name
)
)
AS translated
FROM
tmp
;
通常の外部モデルエンドポイントと同様に動作します。
こちらも結果内容はともかく、動きました！
まとめ
ai_queryを外部モデル/カスタムモデルを使って試してみました。
SQLで手軽にLLMのパワーを発揮できるのは非常に良いですね。カスタムモデルも利用できるので、自作のエージェントを用いたバッチ実行など応用範囲も広そうです。
エージェント関連はもっと勉強して、いろんなエージェントを今後試作してみたいと思っています。