導入
Databricksのリリースノートを見ていたら、Mosaic AI Model Servingでの構造化出力サポート(Public Preview)がアナウンスされていました。
関連ドキュメントも公開されています。
今やLLMの構造化出力は欠かせないものになっていると個人的には思っており、Mosaic AI Model Servingで正式に対応されたのは喜ばしいです。
現時点ではPay-per-Token式のエンドポイント（日本リージョンまだ未対応）と、プロビジョニングされたスループット基盤モデルのエンドポイントの両方に対応している模様。
後者は私が使っているDatrabricks環境下でも利用できるので構造化出力を試してみます。
Step1. 準備
Mosaic AI Model ServingのエンドポイントをUIを使って用意します。
今回はsystem.aiカタログ・スキーマに公開されているmeta_llama_v3_1_8b_instructを利用してModel Servingのエンドポイントを作成しました。
(Llama 3.2の3Bで最初試したのですが、ある程度の性能を持ったモデルでないとうまく構造化出力できなかったので8Bモデルにしました)
しばらく後にエンドポイントが出来上がります。
Step2. 構造化出力を試してみる。
ノートブックを作成し、サーバレスクラスタにアタッチします。
その上で、OpenAI関連パッケージをインストールします。
%pip install -U openai typing_extensions
dbutils.library.restartPython()
Databricks 公式ドキュメントのコードを若干修正して構造化出力をともなった推論を実行。
変更部分はシークレットを利用する部分とプロンプト部分です。
import os
import json
from openai import OpenAI
# SecretsからAPIキーとホストのURLを取得
DATABRICKS_TOKEN = dbutils.secrets.get("llm-serve", "api_token")
DATABRICKS_BASE_URL = dbutils.secrets.get("llm-serve", "host")
# OpenAIのクライアントを作成
client = OpenAI(
api_key=DATABRICKS_TOKEN, base_url=DATABRICKS_BASE_URL + "serving-endpoints/"
)
# LLMが返すフォーマットをJSONスキーマで指定
response_format = {
"type": "json_schema",
"json_schema": {
"name": "research_paper_extraction",
"schema": {
"type": "object",
"properties": {
"title": {"type": "string"},
"authors": {"type": "array", "items": {"type": "string"}},
"abstract": {"type": "string"},
"keywords": {"type": "array", "items": {"type": "string"}},
},
},
"strict": True,
},
}
# クエリ内容
messages = [
{
"role": "system",
"content": "あなたは構造化データ抽出の専門家です。研究論文からの非構造化テキストが与えられ、それを指定された構造に変換する必要があります。",
},
{
"role": "user",
"content": """データウェアハウス（DWH）の父と称されるビル・インモン（Bill Inmon）氏は、効率的な機械学習やビジネス分析をデータレイクで直接行うことを可能にするデータレイクハウスの誕生を支持しています。
インモン氏は本書の中で次のように述べています。「データレイクハウスのアーキテクチャは、DWH 市場の黎明期に見られたものに匹敵するインパクトをもたらすものである。オープンな環境でのデータ管理、エンタープライズの各部門から得られる多様なデータの集約、データレイクのデータサイエンスとデータウェアハウスのエンドユーザー分析を組み合わせた機能が相乗効果を発揮し、データレイクハウスを利用する企業におけるデータ活用の可能性を引き出す。」
本書は、データレイクハウスの構築を成功させるための 5 つの重要な要素について詳しく解説しています。
エンタープライズにおけるデータの大半が既に格納されているデータレイクをまず活用
データレイクのデータ品質の向上とガバナンスの強化
データの最適化によるクエリの高速化
機械学習のネイティブなサポート
オープンなデータフォーマットと API で、ロックインを回避""",
},
]
# 推論の実行
response = client.chat.completions.create(
model="llama-v3-1-8b-instruct-endpoint",
messages=messages,
response_format=response_format,
)
json.loads(response.choices[0].message.model_dump()["content"])
{'abstract': 'データレイクハウスの概念と、ビル・インモンが提唱した5つの重要な要素について説明しています。',
'authors': ['?'],
'keywords': ['データレイクハウス', 'データ活用', 'データ品質', 'ガバナンス', '機械学習', 'オープン情報', 'API'],
'title': 'データレイクハウスの都合の良い製品'}
生成された内容はともかく、response_formatで指定したJSONスキーマに従ったJSON形式のデータとして結果を取得できました。
実行してみている感じ、速度上のオーバーヘッドも体感的には感じません。
ちゃんと確認してないのですがrespoinse_formatのAPI仕様はOpenAI APIの構造化出力仕様と似せてるのではないかと。
ただ、一部対応していないJSONスキーマの書式があるようですね。
また、現状？の制限事項としてmaxLengthなど要素数の強制などはプロビジョニングされたスループット基盤モデルの場合まだ未対応の模様。
とはいえ、Function Callingに代表されるような構造化出力用途では十分な気がします。
まとめ
Databricks Mosaic AI Model Servingの構造化出力対応を簡単に試してみました。
こういう重要なところをきちんと抑えてくるところがDatabricksの良いところですね。
AIエージェント活用において重要な機能だから、ということもあるとは思いますが。
これからの様々なアップデートも楽しみです。