mlflow3のプロンプト最適化機能を利用してみた

今回は前回に引き続き、mlflow3で提供されているプロンプトの最適化機能を利用してみましたので紹介します。前回の記事は以下になりますので、合わせてみてもらえると嬉しいです。

https://zenn.dev/akasan/articles/b37450993c3dce

mlflow3のプロンプト最適化とは？

本機能については以下のページにまとめられています。記事執筆時点ではまだExperimental機能として提供されています。説明によると、以下のような機能とのことです。

• mlflow.genai.optimize_prompt()を使用して、MLflowの統合インターフェースからプロンプトを高度なプロンプト最適化手法に組み込むことが可能
• 評価指標とラベル付きデータを活用してプロンプトを自動的に改善するのに役立つ
• DSPyのMIPROv2アルゴリズムをサポート

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

DSPyとは？

公式サイトの説明をまとめると、DSPyとは以下のような機能を有するとのことです。

• モジュール型AIソフトウェアを構築するための宣言型フレームワークであり、脆弱な文字列ではなく構造化されたコードを高速に反復処理でき、AIプログラムを言語モデルに適した効果的なプロンプトと重みにコンパイルするアルゴリズムを提供する
  • シンプルな分類器や洗練されたRAGパイプライン、エージェントループなど、あらゆる対象に対応
• DSPy(宣言型自己改善型Python)を使用すると、プロンプトを整理したりジョブをトレーニングしたりする代わりに、自然言語モジュールからAIソフトウェアを構築し、さまざまなモデル、推論戦略、学習アルゴリズムを組み合わせて汎用的に構成可能
  • AI ソフトウェアの信頼性、保守性、そしてモデルや戦略間の移植性が向上

また機能の要約として、DSPy は、アセンブリ言語から C 言語への移行やポインタ演算から SQL への移行のように、AI プログラミング（講義）のための高水準言語と考えてください。コミュニティに参加したり、サポートを求めたり、GitHub や Discord で貢献を始めたりしてくださいとのことです。要は自然言語を用いてAIシステムを構築するために開発されているものと見受けられます。今回は解説範囲からは外しますが、後ほど解説記事が出せたらと思います。

https://dspy.ai/#__tabbed_1_1

実際に使ってみる

環境構築

uvを利用して閑居を立ち上げました。

uv init mlflow3_prompt_optimization -p 3.12
cd mlflow3_prompt_optimization
uv add dspy mlflow

プログラムの実装

まずは公式が提供しているサンプルコードを実装してみます。

math_equation.py

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

os.environ["OPENAI_API_KEY"] = "<YOUR_OPENAI_API_KEY>"


# Define a custom scorer function to evaluate prompt performance with the @scorer decorator.
# The scorer function for optimization can take inputs, outputs, and expectations.
@scorer
def exact_match(expectations: dict[str, Any], outputs: dict[str, Any]) -> bool:
    return expectations["answer"] == outputs["answer"]


# Register the initial prompt
initial_template = """
Answer to this math question: {{question}}.
Return the result in a JSON string in the format of {"answer": "xxx"}.
"""

prompt = mlflow.genai.register_prompt(
    name="math",
    template=initial_template,
)

# The data can be a list of dictionaries, a pandas DataFrame, or an mlflow.genai.EvaluationDataset
# It needs to contain inputs and expectations where each row is a dictionary.
train_data = [
    {
        "inputs": {"question": "Given that $y=3$, evaluate $(1+y)^y$."},
        "expectations": {"answer": "64"},
    },
    {
        "inputs": {
            "question": "The midpoint of the line segment between $(x,y)$ and $(-9,1)$ is $(3,-5)$. Find $(x,y)$."
        },
        "expectations": {"answer": "(15,-11)"},
    },
    {
        "inputs": {
            "question": "What is the value of $b$ if $5^b + 5^b + 5^b + 5^b + 5^b = 625^{(b-1)}$? Express your answer as a common fraction."
        },
        "expectations": {"answer": "\\frac{5}{3}"},
    },
    {
        "inputs": {"question": "Evaluate the expression $a^3\\cdot a^2$ if $a= 5$."},
        "expectations": {"answer": "3125"},
    },
    {
        "inputs": {"question": "Evaluate $\\lceil 8.8 \\rceil+\\lceil -8.8 \\rceil$."},
        "expectations": {"answer": "17"},
    },
]

eval_data = [
    {
        "inputs": {
            "question": "The sum of 27 consecutive positive integers is $3^7$. What is their median?"
        },
        "expectations": {"answer": "81"},
    },
    {
        "inputs": {"question": "What is the value of $x$ if $x^2 - 10x + 25 = 0$?"},
        "expectations": {"answer": "5"},
    },
    {
        "inputs": {
            "question": "If $a\\ast b = 2a+5b-ab$, what is the value of $3\\ast10$?"
        },
        "expectations": {"answer": "26"},
    },
    {
        "inputs": {
            "question": "Given that $-4$ is a solution to $x^2 + bx -36 = 0$, what is the value of $b$?"
        },
        "expectations": {"answer": "-5"},
    },
]

# Optimize the prompt
result = mlflow.genai.optimize_prompt(
    target_llm_params=LLMParams(model_name="openai/gpt-4.1-mini"),
    prompt=prompt,
    train_data=train_data,
    eval_data=eval_data,
    scorers=[exact_match],
    optimizer_config=OptimizerConfig(
        num_instruction_candidates=8,
        max_few_show_examples=2,
    ),
)

# The optimized prompt is automatically registered as a new version
print(result.prompt.uri)

こちらを実行すると以下のような結果となりました。

uv run math_equation.py

# 結果
2025/07/29 20:21:48 INFO mlflow.genai.optimize.optimizers.dspy_mipro_optimizer: 🎯 Starting prompt optimization for: prompts:/math/1
⏱️ This may take several minutes or longer depending on dataset size...
📊 Training with 5 examples.
2025/07/29 20:21:48 INFO mlflow.genai.optimize.optimizers.dspy_mipro_optimizer: 🎯 Starting prompt optimization for: prompts:/math/1
⏱️ This may take several minutes or longer depending on dataset size...
📊 Training with 5 examples.
2025/07/29 20:21:48 INFO mlflow.genai.optimize.optimizers.dspy_mipro_optimizer: 🎉 Optimization complete! Initial score: 0.0. Final score: 50.0.
2025/07/29 20:21:48 INFO mlflow.genai.optimize.optimizers.dspy_mipro_optimizer: 🎉 Optimization complete! Initial score: 0.0. Final score: 50.0.
prompts:/math/2

何か最適化を初めているような気配はありますが何をしているかいまいちわからないと思うので、次のステップでmlflowのUIでチェックしてみます。

mlflow UI上での結果の確認

それでは先ほどの結果をmlflowのUIで確認してみましょう。以下のコマンドを実行するとmlflow UIを起動できます。

uv run mlflow ui --port 8000

実行後にlocalhost:8000にアクセスすると以下のような画面が表示されるかと思います。

トップ画面

ヘッダーにあるPromptsを選択すると以下の画面に移動します。

プロンプト画面

コード上で今回作成しているプロンプトは以下のようにしてmathという名前で作成しており、その名前でプロンプトが登録されているのが確認できます。

prompt = mlflow.genai.register_prompt(
    name="math",
    template=initial_template,
)

次にmathを選択すると、Version 1とVersion 2が表示されて、それぞれ以下のようになっているかと思います。Version 1がinitial_templateとして作成したもの、Version 2が最適化されたプロンプトになります。

初期プロンプト

最適化後プロンプト

最適化されたプロンプトは画面に映り切ってないので以下に貼っておきます。

<system>
Your input fields are:
1. `question` (str):
Your output fields are:
1. `answer` (str):
All interactions will be structured in the following way, with the appropriate values filled in.

Inputs will have the following structure:

[[ ## question ## ]]
{question}

Outputs will be a JSON object with the following fields.

{
  "answer": "{answer}"
}
In adhering to this structure, your objective is: 
        Read the math question carefully, perform the necessary algebraic or arithmetic calculations to solve it, then return the final answer as a JSON string in the format {"answer": "xxx"}. Ensure the answer is concise and correctly formatted, such as a number, fraction, or coordinate pair.
</system>

<user>
[[ ## question ## ]]
The midpoint of the line segment between $(x,y)$ and $(-9,1)$ is $(3,-5)$. Find $(x,y)$.
</user>

<assistant>
{
  "answer": "(15,-11)"
}
</assistant>

<user>
[[ ## question ## ]]
Given that $y=3$, evaluate $(1+y)^y$.
</user>

<assistant>
{
  "answer": "64"
}
</assistant>

<user>
[[ ## question ## ]]
Evaluate $\lceil 8.8 \rceil+\lceil -8.8 \rceil$.
</user>

<assistant>
{
  "answer": "17"
}
</assistant>

<user>
[[ ## question ## ]]
What is the value of $b$ if $5^b + 5^b + 5^b + 5^b + 5^b = 625^{(b-1)}$? Express your answer as a common fraction.
</user>

<assistant>
{
  "answer": "\\frac{5}{3}"
}
</assistant>

<user>
[[ ## question ## ]]
{{question}}

Respond with a JSON object in the following order of fields: `answer`.
</user>

最適化されたプロンプトをみると以下のような構成になっているようです。

• システムプロンプト：ユーザ入力および生成AIの出力のスキーマが提示されている

  <system>
  Your input fields are:
  1. `question` (str):
  Your output fields are:
  1. `answer` (str):
  All interactions will be structured in the following way, with the appropriate values filled in.
  
  Inputs will have the following structure:
  
  [[ ## question ## ]]
  {question}
  
  Outputs will be a JSON object with the following fields.
  
  {
    "answer": "{answer}"
  }
  In adhering to this structure, your objective is: 
          Read the math question carefully, perform the necessary algebraic or arithmetic calculations to solve it, then return the final answer as a JSON string in the format {"answer": "xxx"}. Ensure the answer is concise and correctly formatted, such as a number, fraction, or coordinate pair.
  </system>
  

• 解答例：最適化の学習時に受け渡したデータを回答のサンプルとして提示

  <user>
  [[ ## question ## ]]
  The midpoint of the line segment between $(x,y)$ and $(-9,1)$ is $(3,-5)$. Find $(x,y)$.
  </user>
  
  <assistant>
  {
    "answer": "(15,-11)"
  }
  </assistant>
  
  <user>
  [[ ## question ## ]]
  Given that $y=3$, evaluate $(1+y)^y$.
  </user>
  
  <assistant>
  {
    "answer": "64"
  }
  </assistant>
  
  <user>
  [[ ## question ## ]]
  Evaluate $\lceil 8.8 \rceil+\lceil -8.8 \rceil$.
  </user>
  
  <assistant>
  {
    "answer": "17"
  }
  </assistant>
  
  <user>
  [[ ## question ## ]]
  What is the value of $b$ if $5^b + 5^b + 5^b + 5^b + 5^b = 625^{(b-1)}$? Express your answer as a common fraction.
  </user>
  
  <assistant>
  {
    "answer": "\\frac{5}{3}"
  }
  </assistant>
  
  <user>
  [[ ## question ## ]]
  {{question}}
  
  Respond with a JSON object in the following order of fields: `answer`.
  </user>
  

確かに最初のプロンプトの内容をもとに、指定したサンプルに適合する形で調整されていることが確認できます。

最適化後のプロンプトを使って推論してみる

それでは最適化されたプロンプトを用いて実行させてみましょう。

optimized_prompt_inference.py

import openai
import json
import os
import mlflow

os.environ["OPENAI_API_KEY"] = "..."

def predict(question: str, prompt_uri: str) -> str:
    prompt = mlflow.genai.load_prompt(prompt_uri)
    content = prompt.format(question=question)
    completion = openai.chat.completions.create(
        model="gpt-4.1-mini",
        messages=[{"role": "user", "content": content}],
        temperature=0.1,
    )

    return json.loads(completion.choices[0].message.content)


questions = [
    "What is the smallest prime number greater than or equal to 100?",
    "What is the derivative of sin(x)?"
]

results = [
    predict(question, "prompts:/math/2")
    for question in questions
]

for question, result in zip(questions, results):
    print(f"{question} -> {result}")

まずは推論を実行する関数としてpredictを作成し、OpenAIのクライアントを利用して推論させています。なお、プロンプトはmlflowで生成されたものをベースにしています。質問については試しに2問作ってみました。predictにそれぞれのプロンプトを受け渡してその結果を得ると以下のようになりました。その結果、我々が望んでいたフォーマットで結果が返ってくることを確認しました。ちなみに、問題は優しめにしておいたのであっています。

uv run optimized_prompt_inference.py

# 結果
What is the smallest prime number greater than or equal to 100? -> {'answer': '101'}
What is the derivative of sin(x)? -> {'answer': 'cos(x)'}

データセットのフォーマット

最適化するためのデータセットのフォーマットは以下が対応しているようです。

• 辞書のリスト
• Pandas DataFrame
• Spark DataFrame
• mlflow.genai.EvaluationDataset

なお、このうちどのフォーマットを利用しても問題ないとのことですが、入力と期待出力のペアがふくまれている必要があるとのことです。各データに利用できるのはプリミティブ方やリスト、ネストされた辞書またはPydanticのモデルを利用することができるようです。

ちなみに、最初の例を無理やりpydanticを使って書き換えてみました。

math_equation_pydantic.py

import os
from typing import Any
import mlflow
from mlflow.genai.scorers import scorer
from mlflow.genai.optimize import OptimizerConfig, LLMParams
from pydantic import BaseModel

os.environ["OPENAI_API_KEY"] = "..."


# Define a custom scorer function to evaluate prompt performance with the @scorer decorator.
# The scorer function for optimization can take inputs, outputs, and expectations.
@scorer
def exact_match(expectations: dict[str, Any], outputs: dict[str, Any]) -> bool:
    return expectations["answer"] == outputs["answer"]


# Register the initial prompt
initial_template = """
Answer to this math question: {{question}}.
Return the result in a JSON string in the format of {"answer": "xxx"}.
"""

prompt = mlflow.genai.register_prompt(
    name="math_pydantic",
    template=initial_template,
)

class Answer(BaseModel):
    answer: str

# The data can be a list of dictionaries, a pandas DataFrame, or an mlflow.genai.EvaluationDataset
# It needs to contain inputs and expectations where each row is a dictionary.
train_data = [
    {
        "inputs": {"question": "Given that $y=3$, evaluate $(1+y)^y$."},
        "expectations": {"answer": Answer(answer="64")},
    },
    {
        "inputs": {
            "question": "The midpoint of the line segment between $(x,y)$ and $(-9,1)$ is $(3,-5)$. Find $(x,y)$."
        },
        "expectations": {"answer": Answer(answer="15, -11")},
    },
    {
        "inputs": {
            "question": "What is the value of $b$ if $5^b + 5^b + 5^b + 5^b + 5^b = 625^{(b-1)}$? Express your answer as a common fraction."
        },
        "expectations": {"answer": Answer(answer="\\frac{5}{3}")},
    },
    {
        "inputs": {"question": "Evaluate the expression $a^3\\cdot a^2$ if $a= 5$."},
        "expectations": {"answer": Answer(answer="3125")},
    },
    {
        "inputs": {"question": "Evaluate $\\lceil 8.8 \\rceil+\\lceil -8.8 \\rceil$."},
        "expectations": {"answer": Answer(answer="17")},
    },
]

eval_data = [
    {
        "inputs": {
            "question": "The sum of 27 consecutive positive integers is $3^7$. What is their median?"
        },
        "expectations": {"answer": Answer(answer="81")},
    },
    {
        "inputs": {"question": "What is the value of $x$ if $x^2 - 10x + 25 = 0$?"},
        "expectations": {"answer": Answer(answer="5")},
    },
    {
        "inputs": {
            "question": "If $a\\ast b = 2a+5b-ab$, what is the value of $3\\ast10$?"
        },
        "expectations": {"answer": Answer(answer="26")},
    },
    {
        "inputs": {
            "question": "Given that $-4$ is a solution to $x^2 + bx -36 = 0$, what is the value of $b$?"
        },
        "expectations": {"answer": Answer(answer="-5")},
    },
]

# Optimize the prompt
result = mlflow.genai.optimize_prompt(
    target_llm_params=LLMParams(model_name="openai/gpt-4.1-mini"),
    prompt=prompt,
    train_data=train_data,
    eval_data=eval_data,
    scorers=[exact_match],
    optimizer_config=OptimizerConfig(
        num_instruction_candidates=8,
        max_few_show_examples=2,
    ),
)

# The optimized prompt is automatically registered as a new version
print(result.prompt.uri)

expectationsの値を指定するときにanswerに該当する部分にpydanticクラスのインスタンスを指定することで利用できます。

{
    "inputs": {
        "question": "Given that $-4$ is a solution to $x^2 + bx -36 = 0$, what is the value of $b$?"
    },
    "expectations": {"answer": Answer(answer="-5")},
},

このコードを実行すると先ほどの例と同じようにプロンプトが登録されます。登録されたプロンプトをもとにテストコードを以下のように実装しました。

optimized_prompt_pydantic_inference.py

import openai
import json
import os
import mlflow

os.environ["OPENAI_API_KEY"] = "..."

def predict(question: str, prompt_uri: str) -> str:
    prompt = mlflow.genai.load_prompt(prompt_uri)
    content = prompt.format(question=question)
    completion = openai.chat.completions.create(
        model="gpt-4.1-mini",
        messages=[{"role": "user", "content": content}],
        temperature=0.1,
    )

    return json.loads(completion.choices[0].message.content)


questions = [
    "What is the smallest prime number greater than or equal to 100?",
    "What is the derivative of sin(x)?"
]

results = [
    predict(question, "prompts:/math_pydantic/3")
    for question in questions
]

for question, result in zip(questions, results):
    print(f"{question} -> {result}")

それでは実行してみます。結果を見るとプロンプトも問題なく生成され推論も実行できるていることを確認できました。

uv run math_equation_pydantic.py  # プロンプト最適化
uv run optimized_prompt_pydantic_inference.py

# 結果
What is the smallest prime number greater than or equal to 100? -> {'answer': '101'}
What is the derivative of sin(x)? -> {'answer': 'The derivative of sin(x) with respect to x is cos(x).'}

Pydanticを利用して生成されたプロンプト

<system>
Your input fields are:
1. `question` (str):
Your output fields are:
1. `answer` (Answer):
All interactions will be structured in the following way, with the appropriate values filled in.

Inputs will have the following structure:

[[ ## question ## ]]
{question}

Outputs will be a JSON object with the following fields.

{
  "answer": "{answer}        # note: the value you produce must adhere to the JSON schema: {\"type\": \"object\", \"properties\": {\"answer\": {\"type\": \"string\", \"title\": \"Answer\"}}, \"required\": [\"answer\"], \"title\": \"Answer\"}"
}
In adhering to this structure, your objective is: 
        Answer the math question provided in the input, then return the result as a JSON string formatted like {"answer": "xxx"}.
</system>

<user>
[[ ## question ## ]]
{{question}}

Respond with a JSON object in the following order of fields: `answer` (must be formatted as a valid Python Answer).
</user>

まとめ

今回はmlflowのプロンプト最適化手法を試してみました。プロンプトを考えるのはかなり負荷が高いタスクだと思うので、やりたいタスクとその回答サンプルを用意すればプロンプト候補を生成してくれるということで、とても有用な機能だと思います。ぜひ皆さんも試してみてください。