MLflow Study: Prompt Optimization in Databricks

前回のStudyはこちら：

https://qiita.com/isanakamishiro2/items/857dae30e746637c43cc


# はじめに

MLflow3がリリースされ、Databricks側でもDocumentの対応含めて様々な機能が実装されています。そのうちのひとつにMLflow Prompt Optimization(プロンプト最適化)があります。
今回は勉強として、実際に触りながら学んでみます。

https://docs.databricks.com/aws/ja/mlflow3/genai/prompt-version-mgmt/prompt-registry/automatically-optimize-prompts

なお、上記の公式Documentに記載があるように、Databricksでは**2025年6月現在ベータ版**となっています。利用にはご注意ください。

検証はDatabricks Free Editionを使って行いました。

# MLflow Prompt Optimization(プロンプト最適化)とは

以下のOSS版のドキュメントより、サイト内容を一部邦訳して抜粋します。

https://mlflow.org/docs/latest/genai/prompt-version-mgmt/prompt-registry/optimize-prompts

> # プロンプト最適化（実験的機能）
> MLflowでは、```mlflow.genai.optimize_prompt()```APIを使用したMLflowの統合インターフェースを通じて、プロンプトを高度なプロンプト最適化手法に組み込むことができます。この機能は、評価指標とラベル付きデータを活用して、プロンプトを自動的に改善するのに役立ちます。現在、このAPIはDSPyのMIPROv2アルゴリズムをサポートしています。
> 
> ## 主なメリット
> - **統合インターフェース:** 中立的なインターフェースを介して最先端のプロンプト最適化アルゴリズムにアクセスできます。
> - **プロンプト管理:** MLflow プロンプト レジストリと統合して、再利用性、バージョン管理、系統を実現します。
> - **評価:** MLflow の評価機能を使用してプロンプトのパフォーマンスを総合的に評価します。

というわけで、現時点では[DSPy](https://dspy.ai/)の機能を使ってプロンプトを最適化する機能のようです。

では、実際にどのようにプロンプトが最適化できるのか試してみましょう。

# 基本的な利用

## 準備

まずはノートブックを作成し、MLflowの最新版をインストールします。
あわせて```dspy```などのパッケージもインストールしておきます。

```python
%pip install -U "mlflow[databricks]>=3.1.0" databricks-langchain langgraph dspy databricks-agents

%restart_python
```

プロンプト最適化にはPrompt Registryを利用するため、プロンプトを保存するカタログ・スキーマの場所を定義しておきます。

```python
import mlflow
mlflow.set_registry_uri("databricks-uc")

CATALOG = "workspace"
SCHEMA = "default"
```

## 最適化対象のプロンプトを作成

次に、最適化対象となるプロンプトをPrompt Registryの保管します。
プロンプトは単純に質問の回答を日本語で得るものにします。

```python
# First prompt for summarization.
qa_prompt = mlflow.genai.register_prompt(
    name=f"{CATALOG}.{SCHEMA}.qa_prompt",
    template="次の質問に対して日本語で回答してください:{{question}}",
)
```

## 最適化に利用する訓練・評価データセットの作成

プロンプトの最適化には、最適化のための訓練・評価データを利用します。
今回は質問とその回答のペアを準備します。
データはDatabricks Assistantに作ってもらいました。

```python
import pandas as pd

# 質問と回答のペアをリストとして定義
train_data = [
    {
        "inputs": {"question": "Databricksとは何ですか？"},
        "expectations": {
            "answer": "Databricksは、データエンジニアリング、データサイエンス、機械学習のための統合データ分析プラットフォームです。"
        },
    },
    {
        "inputs": {"question": "Databricksの主な機能は何ですか？"},
        "expectations": {
            "answer": "Databricksの主な機能には、データの統合、分析、機械学習モデルのトレーニングとデプロイがあります。"
        },
    },
    {
        "inputs": {"question": "Databricksで使用できるプログラミング言語は何ですか？"},
        "expectations": {
            "answer": "Databricksでは、Python、SQL、R、Scalaなどのプログラミング言語を使用できます。"
        },
    },
    {
        "inputs": {"question": "Databricksのノートブックとは何ですか？"},
        "expectations": {
            "answer": "Databricksのノートブックは、データ分析や機械学習のコードを記述、実行、共有するためのインタラクティブな環境です。"
        },
    },
    {
        "inputs": {"question": "Databricksのクラスターとは何ですか？"},
        "expectations": {
            "answer": "Databricksのクラスターは、データ処理や分析のために使用されるコンピューティングリソースの集合です。"
        },
    },
]

eval_data = [
    {
        "inputs": {"question": "DatabricksのDelta Lakeとは何ですか？"},
        "expectations": {
            "answer": "Delta Lakeは、Databricks上で提供される信頼性の高いデータレイクソリューションで、ACIDトランザクションやスキーマエンフォースメントをサポートします。"
        },
    },
    {
        "inputs": {"question": "DatabricksのMLflowとは何ですか？"},
        "expectations": {"answer": "MLflowは、機械学習モデルのライフサイクル管理を支援するオープンソースプラットフォームです。"},
    },
    {
        "inputs": {"question": "Databricksのジョブとは何ですか？"},
        "expectations": {
            "answer": "Databricksのジョブは、スケジュールされたデータ処理タスクやワークフローを自動化するための機能です。"
        },
    },
    {
        "inputs": {"question": "Databricksのワークスペースとは何ですか？"},
        "expectations": {
            "answer": "Databricksのワークスペースは、データ分析や機械学習プロジェクトを管理するためのコラボレーション環境です。"
        },
    },
    {
        "inputs": {"question": "DatabricksのUnity Catalogとは何ですか？"},
        "expectations": {
            "answer": "Unity Catalogは、Databricks上でデータガバナンスとセキュリティを提供するための統合データカタログです。"
        },
    },
]

# pandasデータフレームに変換
pdf = pd.DataFrame(train_data)

# データフレームを表示
display(pdf)
```

以下のようなデータセットが作られます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/b7662449-b181-4b78-902d-dde82eccab33.png)

## プロンプトを最適化

では、準備したプロンプトやデータセットを利用してプロンプトを最適化します。

最適化には```mlflow.genai.optimize_prompt```関数を利用します。
パラメータとして再帰化に利用するLLMや訓練用のデータセット、最適化のための評価関数(と言うのが適切なのか？）を指定します。
感覚的にはMLflow Evaluationの利用と近いです。

今回は評価関数として、```mlflow.genai.scorers```モジュールに含まれる事前定義された評価関数であるCorrectnessを利用しました。名前の通り正しさを測るもののようです。

```python
from typing import Any
from mlflow.genai.scorers import Correctness
from mlflow.genai.optimize import OptimizerConfig, LLMParams
from mlflow.genai.scorers import scorer
import os

# OpenAI Clientが利用できるように、現在のCredentialをOPENAI_API_KEYに登録
mlflow_creds = mlflow.utils.databricks_utils.get_databricks_host_creds()
os.environ["OPENAI_API_KEY"] = mlflow_creds.token

# Correctnessスコアを計算するbuilt-inオブジェクトを作成
_correctness = Correctness()

# プロンプト最適化のための評価関数（確からしさのテスト)
@scorer
def correctness(inputs, outputs, expectations):
    expectations = {"expected_response": expectations.get("answer")}
    return (
        _correctness(inputs=inputs, outputs=outputs, expectations=expectations).value
        == "yes"
    )

# 最適化対象のプロンプト
prompt = mlflow.genai.load_prompt(f"prompts:/{CATALOG}.{SCHEMA}.qa_prompt/1")

# プロンプトを最適化。LLMとしてLlama-4 Maverickを使用。
result = mlflow.genai.optimize_prompt(
    target_llm_params=LLMParams(
        model_name="openai/databricks-llama-4-maverick",
        base_uri=f"{mlflow_creds.host}/serving-endpoints",
    ),
    prompt=prompt,
    train_data=train_data,
    eval_data=eval_data,
    scorers=[correctness],
    optimizer_config=OptimizerConfig(
        num_instruction_candidates=8,
        max_few_show_examples=2,
    ),
)

# 最適化結果のプロンプトレジストリのURLを表示
print(result.prompt.uri)
```

しばらく時間がかかりますが、完了すると最適化されたプロンプトがPrompt Registryに保管されます。

今回は以下のようなプロンプトが作られました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/1ef096e4-e457-4c4d-b175-529f178829ff.png)

最初のシンプルなプロンプトから、入力情報や出力情報の定義が追加され、またJSON形式で構造化出力するようなプロンプトに変わっています。

内容的にはまだシンプルなものではありますが、データに基づいてプロンプトを更新しているようですね。

# 最適化されたプロンプトを利用

では、最適化されたプロンプトを利用してみます。
比較のために、最適化前のプロンプトと最適化後のプロンプトを両方実行してみます。

```python
import mlflow
from openai import OpenAI
import codecs

# MLflowの自動ロギングを有効にして、アプリケーションにトレースを追加
mlflow.openai.autolog()

# 実行ノートブックと同じ資格情報を使用してOpenAIクライアント経由でDatabricks LLMに接続
mlflow_creds = mlflow.utils.databricks_utils.get_databricks_host_creds()
client = OpenAI(
    api_key=mlflow_creds.token, base_url=f"{mlflow_creds.host}/serving-endpoints"
)

# レジストリからプロンプトをロード
first_prompt = mlflow.genai.load_prompt(f"prompts:/{CATALOG}.{SCHEMA}.qa_prompt/1")
optimized_prompt = mlflow.genai.load_prompt(result.prompt.uri)
endpoint = "databricks-llama-4-maverick"

def predict(prompt):
    formatted_prompt = prompt.format(question="Databricksとは何ですか?")
    
    # LLMを呼び出す
    response = client.chat.completions.create(
        model="databricks-llama-4-maverick",
        messages=[
            {
                "role": "user",
                "content": formatted_prompt,
            },
        ],
    )
    return response.choices[0].message.content

# 比較のために初期プロンプトと最適化されたプロンプトの両方を実行
for p in [first_prompt, optimized_prompt]:
    print(predict(p))
```

結果は以下の通り。

```:最適化前プロンプトの結果
Databricksは、Apache Sparkを利用したビッグデータ処理と分析のためのプラットフォームおよびサービスを提供する企業、およびそのプラットフォームの名称です。Databricksは、Sparkの開発に携わった人々によって設立されました。Sparkは、大規模なデータセットを高速に処理できるオープンソースの分散処理フレームワークです。

Databricksが提供するプラットフォームは、データエンジニアリング、データサイエンス、およびビジネスユーザーを対象に設計されており、Apache Sparkを用いたデータの取り込み、変換、分析を容易に行えるようにします。主な特徴としては、以下のようなものがあります：

1. **インタラクティブなワークスペース**: データの探索、変換、分析をインタラクティブに行うための環境を提供します。ノートブック形式での作業が可能で、Python、R、Scala、SQLなどの言語に対応しています。

2. **データの取り込みと統合**: 多様なデータソースからのデータの取り込みをサポートし、データレイクやデータウェアハウスへのデータ統合を容易にします。

3. **高度な分析**: Apache Sparkの力を利用して、大規模なデータセットに対する高度な分析や機械学習を高速に実行できます。

4. **コラボレーション機能**: データサイエンティスト、エンジニア、ビジネスユーザー間のコラボレーションを促進するための機能が備わっています。

5. **セキュリティとガバナンス**: 企業がデータを安全に扱えるように、アクセス制御や監査ログなどのセキュリティ機能を提供しています。

Databricksは、クラウド（AWS、Azure、GCP）上でサービスとして提供されるほか、オンプレミス環境やプライベートクラウドへのデプロイメントもサポートしています。これにより、企業は自社のニーズや既存のITインフラに合わせた形で、ビッグデータと分析技術を活用することができます。
```

```:最適化後プロンプトの結果
{
  "answer": "Databricksは、Apache Sparkの創始者たちが開発した、ビッグデータ処理および分析プラットフォームです。データエンジニアリング、データサイエンス、ビジネスアナリティクスのための統合環境を提供し、データの処理、分析、可視化を容易にします。クラウドベースのサービスとして、AWS、Azure、Google Cloudなどの主要なクラウドプロバイダー上で動作します。"
}
```

最適化後の方は、JSON構造化出力を指定しているため、出力量としては控えめになりました。

どちらがいいかはケースバイケースですが、構造化した結果を得る際のプロンプトを最適化する際にはプロンプト最適化を利用する方が良さそうです。

# まとめ

MLflowのプロンプト最適化機能を利用してみました。
最適化時の評価関数は事前定義されたものも多くありようなので、それらを使うだけでも様々な最適化を試せそうです。
また、Prompt Registryと統合された機能になっているので、プロンプト管理の一環として最適化処理を実施、という運用もできそうですね。

まだベータ版ではあるので、まだ様々な変更が入りそうですが、安定版の登場が楽しみです。