DSPyで始めるプロンプト最適化

note株式会社の推薦チームで機械学習エンジニアをしている漆山です。弊チームでは、Amazon Bedrockを活用して業務に生成AIを活用したり、MLタスクの一部をLLMに置き換えたりしています。
Amazon Bedrockはさまざまなモデルを利用することができ、用途によってモデルを切り替える柔軟性があります。
しかし、モデルのAPIのパラメータなどは統一されておらず、またClaude v2からClaude v3へのアップデートする際にAPIが大きく変わったりしました。利用するモデルによって、コードのパラメータを変更するのは少し面倒だなと感じることもあります。
利用するモデルに関わらず、プロンプトエンジニアリングは大切です。その際にはプロンプトを管理したりする必要があり、もっと楽にプロンプトの最適化ができないかなと感じることもありました。

この記事では、エンジニアが楽にプロンプトエンジニアリングができ、LLM関連の処理をいい感じにモジュール化して記述できるDSPyについて紹介します。
DSPyってなんぞや、プロンプトいい感じならないのか、LangChainが複雑だからもう少しシンプルに書けないのかなどを考えている人にとって役に立つと思います。

※この記事は note株式会社 Advent Calendar 2024 の24日目の記事です。

DSPyとは

DSPyとはStanfordNLPが開発した、LLMをいい感じに扱えるようにしてくれるPythonライブラリです。
事前に、Chain of ThoughtやReActエージェント、RAGなどが作成できるようモジュールとして用意してくれており、気軽にLLMの処理を構築できます。
また、正解のデータがあれば、プロンプトを最適化してくれる機能があるので、プロンプトエンジニアリングをある程度簡略化できます。

個人的に良いなと思っているポイントは以下の2つです。

• 複雑すぎないインターフェース

• Pytorchのように小さいModuleを組み合わせて、大きな処理をするModuleを作成していく構成

「複雑すぎないインターフェース」は完全に私の主観なのですが、かつてLangChainが出てきた時に感じた抽象化しすぎてよくわからんみたいな感覚があまりなく、非常にシンプルな設計だなと感じます。
「Pytorchのように小さいModuleを組み合わせて、大きな処理をするModuleを作成していく構成」というのは、PytorchやKerasなどの深層学習系のライブラリを使ってる人であればわかると思いますが、小さい処理をModuleとしてまとめて大きなModuleを作っていく、ビルディングブロックのようにLLMを作成できることです。
例えば、以下のようにして簡易的なRAGモジュールを作ることができます。

class RAG(dspy.Module):
    def __init__(self):
        self.respond = dspy.ChainOfThought('context, question -> response')

    def forward(self, question):
        context = search(question).passages
        result = self.respond(context=context, question=question)
        return result.response

dspy.ChainOfThoughtは事前にdspy側で用意されたChainOfThoughtを実行してくれるdspy.Moduleクラスです。
RAG.fowardメソッド内で実際に実行したい処理を記述します。ここでは
1. 質問文のquestionを受け取って、データベースを検索し検索結果を返す
2. 質問文に関する検索結果をcontextとし、questionでもとの質問文をLLMに入力しています。LLMの出力はresponseという変数に格納されています。

個人的には、LLMを含んだ処理を非常にシンプルに書くことができるなという感じがしていてとても好きです。

DSPyの使い方

DSPyを使ってLLMの処理を書くのは非常に簡単です。
多くの場合には以下の3ステップでコードを書いていきます。

1. LLMの入出力を管理するSignatureを定義する

2. LLMの処理を記述したdspy.Moduleを宣言する

3. Signatureで入力として宣言した変数名をdspy.Moduleに渡す

簡単に用語の整理をしておきます。
SignatureはLLMに渡すための入出力をクラスとして表現したものです。
dspy.ModuleはLLMの処理をクラスにしたものです。DSPyでは事前に多くのモジュールが用意されています。単にプロンプトから生成する場合にはdspy.Predictを利用し、ChainOfThoughtのテクニックを利用したい場合にはdspy.ChainOfThoughtが利用できます。
RAGやAgent用のモジュールも用意されていますが、ここでは割愛します。

かなり単純な例でコードを見ていきましょう。例えば、質疑応答をするための処理は以下のように書けます。

import dspy
question = "日本で一番高い山は?"

# 1,2 signatureの宣言とModuleの宣言
classify = dspy.Prediction("question -> answer")

# 3 Signatureで宣言した入力をmoduleに渡す
response = classify(question=question)

# 出力の確認
print(response.answer)

さて、上記のコードでは何のLLMを使うかを指定していないことに気づいたでしょうか？
どのモデルを使うかを決定するのにはdspyのconfigを使う必要があります。
dspyで利用するモデルを設定するコンフィグはグローバルに宣言するか、with句を利用して特定のスコープ内で宣言する必要があります。
先ほどのコードにconfigの設定をすると以下のようになります。

import dspy
question = "日本で一番高い山は?"

lm = dspy.LM("openai/gpt-4o-mini", api_key="API_KEY")
# デフォルトで利用するグローバルな設定
dspy.configure(lm=lm)

# 以下のコンテキストではgpt-3.5-turboが利用される。
with dspy.context(lm=dspy.LM("openai/gpt-3.5-turbo"):
  # 1,2 signatureの宣言とModuleの宣言
  classify = dspy.Prediction("question -> answer")


  # 3 Signatureで宣言した入力をmoduleに渡す
  response = classify(question=question)

# 出力の確認
print(response.answer)

dspyはLiteLLMを利用することで、LLMのモデルを統一的に利用できるようにしています。そのため、LiteLLM Providersで提供されているモデルなら簡単に使うことができます。

DSPyの使い方の最初のステップとして「DSPyを使う際に入出力を管理するSignatureを定義する」というのがありました。
しかし、先ほどのコードではSignatureを明示的に定義されていないように見えます。
DSPyでは２通りの方法でSignatureを定義することができます。
1. dspy.Moduleの引数にinput -> output形式の文字列として渡す
複数の入出力がある場合には、カンマで区切る "input1, input2 -> output1, output2"
2. dspy.Signatureを継承したクラスを作成し、dspy.Moduleに渡す。

上記のコードでは、1の文字列として渡す方で宣言しています。
簡単なタスクであれば1でも十分なのですが、自分でプロンプトを書いたりしたいケースもあります。
そのようなケースでは2のクラスを作成するものを利用します。
先ほどのコードを2のケースで書き直すと以下のようになります。

import dspy

# 1. Signatureの宣言
class CustomSignature(dspy.Signature):
  """あなたは日本の国土地理院で長年働いている、地理情報のエキスパートです。地理に関する質問に答えてください"""
  question = dspy.InputField(desc="ユーザーからの問い合わせ")
  answer = dspy.OutputField(desc="ユーザーからの質問の回答")
  
  
question = "日本で一番高い山は?"

lm = dspy.LM("openai/gpt-4o-mini", api_key="API_KEY")
# デフォルトで利用するグローバルな設定
dspy.configure(lm=lm)


# 以下のコンテキストではgpt-3.5-turboが利用される。
with dspy.context(lm=dspy.LM("openai/gpt-3.5-turbo"):
  # 2. Moduleの宣言
  classify = dspy.Prediction(CustomSignature)


  # 3 Signatureで宣言した入力をmoduleに渡す
  response = classify(question=question)

# 出力の確認
print(response.answer)

自作したSignatureクラスのdocstringを見てください。Signatureはdocstringに記載されているプロンプトをLLMに渡すプロンプトとして利用するようになっています。
複雑なタスクをLLMにやらせたい場合には、Signatureを宣言して入出力を定義しましょう。
LLMのプロンプトに含めたい動的な値はdspy.InputFieldとして宣言することができます。ここでは、ユーザーからの問い合わせはquestionという変数で指定してます。
LLMの出力結果は、dspy.OutputFieldに宣言します。

これで、1, 2, 3すべてのステップで処理を書くことができるようになりました。実装されたコードも複雑ではなく簡単にLLMを利用したコードが書くことができることがわかると思います。

プロンプトの最適化を行う。

DSPyにはLLMを最適化する様々なOptimizer(古いバージョンではteleprompterと呼ばれていました)があります。
どのように最適化したいのかによって、利用するOptimizerが異なります。
Optimizerに関するドキュメントに詳細は書かれているので、そちらを参照する方が好ましいです。
ここではドキュメントを読んでまとめた結果を書いてみます。

ユースケースの分類

DSPyのOptimizerは、大きく分けて以下の3つのユースケースに分類できます。

1. Few-Shot Learning (数ショット学習): プロンプトに含める例 (デモンストレーション) を最適化することで、モデルの性能を向上させる。

2. Instruction Optimization (指示最適化): プロンプト内の指示文を最適化することで、モデルの性能を向上させる。

3. Finetuning (ファインチューニング): プロンプトベースのDSPyプログラムを基に、モデルの重みを更新 (ファインチューニング) する。

1. Few-Shot Learning (数ショット学習)

データが少ない場合 (例: 10件程度):BootstrapFewShot

データが比較的多い場合 (例: 50件以上): BootstrapFewShotWithRandomSearch

類似したデモンストレーションを活用したい場合:
推奨Optimizer: KNNFewShot

2. Instruction Optimization (指示最適化)

指示文のみを最適化したい場合 (0-shotを維持したい場合): MIPROv2 (0-shot)

指示文とデモンストレーションの両方を最適化したい場合 (データが多い場合): MIPROv2
データが豊富（例：過学習を防ぐために200件以上）で、長い最適化時間（例：40回以上の試行）を許容できる場合に有効です。

シンプルな最適化: COPRO

3. Finetuning (ファインチューニング)

プロンプトベースのプログラムを小さなLMに適用したい場合: BootstrapFinetune

それ以外には、Ensemble Optimizerなどもあります。これは複数の小さなLLMで推論させて集計することで、精度向上を目指します。

簡単なフローチャートにすると、以下のようになります。

```mermaid
graph TD
    A[DSPy Optimizer 選択フロー] --> B{データ量は？};
    B -- 少ない (約10件) --> C{Few-Shot Learning が必要？};
    C -- はい --> D[BootstrapFewShot];
    C -- いいえ --> E{指示最適化のみ？};
    E -- はい --> F[MIPROv2 0-shot];
    E -- いいえ --> H;
    B -- 多い (50件以上) --> G{指示最適化のみ？};
    G -- はい --> F;
    G -- いいえ --> H{長い最適化時間 40試行以上 と十分なデータ 200件以上 がある？};
    H -- はい --> I{シンプルな指示最適化で良い？};
    H -- いいえ --> K;
    I -- はい --> N[COPRO];
    I -- いいえ --> J[MIPROv2];
    K{類似したデモンストレーションを活用したい？};
    K -- はい --> L[KNNFewShot];
    K -- いいえ --> M[BootstrapFewShotWithRandomSearch];
    A --> O{大規模LMで良好な結果が得られた後、<br>小さなLMをファインチューニングしたい？};
    O -- はい --> P[BootstrapFinetune];
    O -- いいえ --> Q[Ensemble];
```

ここではHugging Faceから読み込めるPolyAI/banking77データセットを利用して、プロンプトの最適化をしてみます。
このデータセットはざっくりと言えば、銀行に関する問い合わせを77のクラスに分類するタスクです。

dspyの使い方の1,2,3に加えて以下の処理が必要になります。

• 評価用のmetricsを用意する

• Optimizerを用意する

評価用のmetricを用意する

dspyには評価専用のクラスのdspy.Evaluateが用意されています。自身がこのEvaluateを直接使うことはなく、評価用のCallableなオブジェクトや関数を用意して上げる必要があります。
Evaluateが受け付ける関数は以下のようになっています。

def metric(example, pred, trace=None) -> float:
  return float(example.label == pred.label)

metricは第一引数に訓練データ、第二引数にLLMの予測値、第三引数にtraceを取る関数です。
返り値としてLLMのプロンプトの出力のスコアを返すことが想定されています。
この関数をOptimizerに渡すことで、最適化が可能になります。

実際のコードは以下のようになります。

import dspy
import random
from typing import Literal
from dspy.datasets import DataLoader
from datasets import load_dataset

# データセットの読み込み.
CLASSES = load_dataset("PolyAI/banking77", split="train", trust_remote_code=True).features['label'].names
kwargs = dict(fields=("text", "label"), input_keys=("text",), split="train", trust_remote_code=True)

# 2000件だけ読み込む
trainset = [
    dspy.Example(x, hint=CLASSES[x.label], label=CLASSES[x.label]).with_inputs("text", "hint")
    for x in DataLoader().from_huggingface(dataset_name="PolyAI/banking77", **kwargs)[:2000]
]
random.Random(0).shuffle(trainset)

# LLMの処理
lm = dspy.LM("gemini/gemini-2.0-flash-exp", num_retries=1000)
with dspy.context(lm=lm):
  signature = dspy.Signature("text -> label").with_updated_fields('label', type_=Literal[tuple(CLASSES)])
  classify = dspy.ChainOfThoughtWithHint(signature)

  # BootstrapFewShotWithRandomSearchで最適化する
  optimizer = dspy.BootstrapFewShotWithRandomSearch(metric=(lambda x, y, trace=None: x.label == y.label), num_threads=10)
  # 実行にはリトライの待ち時間含めて数時間かかる
  optimized_module = optimizer.compile(classify, trainset=trainset)
  print(optimized_module(text="What does a pending cach withdrawal mean?"))

  # 最適化したものを保存する
  optimized_module.save("polyai_banking77_optimized")

# 最適化したものを読み込み実行する

lm = dspy.LM("gemini/gemini-2.0-flash-exp", num_retries=1000)
with dspy.context(lm=lm):
  signature = dspy.Signature("text -> label").with_updated_fields('label', type_=Literal[tuple(CLASSES)])
  classify = dspy.ChainOfThoughtWithHint(signature)
  classify.load(path="polyai_banking77_optimized")

  print(classify(text="What does a pending cach withdrawal mean?"))

dspy.BootstapFewShotWithRandomSearchに予測結果のスコアを返すLambda関数を渡しています。
optimzier.compileメソッドにdspy.Moduleとtrain_setを渡すことでプロンプトの最適化を実行しています。
Optimizerやcompileメソッドに渡せる引数もどのOptimizerを利用するのかによって異なります。この辺りは実装を見にいくのが一番確実です。

終わりに

DSPyは簡単にLLMの処理を記述することができ、プロンプトの最適化なども行うことができます。
最近では、公式ドキュメントも充実してきており、チートシートも用意されています。
それだけでなく公式ページにはAIにコードを質問される機能も用意されており、どうコードを書いていいかわからない場合に質問することもできるようになっているなど、開発者がとっつきやすいように情報が整備されています。

まだまだ知名度は高くないライブラリですが、使う分にはそんなに困らない程度にはドキュメントが整備され、わからない時にはAIに聞くことができるという土台も整いつつあります。

今回、クラス分類程度にしか利用していませんが、RAGモジュールやAgent用のモジュールも用意されているため、LLMを利用したAgentシステムを作りたいなどのユースケースにも対応できます。
LLMで何らかの処理を実現したい場合のライブラリの選択肢としてはかなりいいものではないかと思っており、少しずつ色々なユースケースで試していきたいと思っています。