CCCMKホールディングス TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

DSPy入門!RAG Pipelineの最適化を試してみました。

こんにちは、CCCMKホールディングス TECH LAB三浦です。

ここのところ本当に暑い日が続いています。暑いと自分が思っている以上に体に負担がかかっているんだな、と感じます。外に出る時はなるべく日差しを避けて歩くようにしないと、と意識するようになりました。

さて、Large Language Model(LLM)を使っていると結構頭を悩ませるのがプロンプトを作成する過程です。ちょっと言い回しを変えるだけでガラッと出力結果が変わってしまい、さっきまで上手くいっていたのに一気におかしくなってしまった!といった具合に、安定しない点に苦労することが多いです。

そんな折にDSPyという興味深いフレームワークを知りました。

dspy-docs.vercel.app

DSPyは機械学習モデルを学習する時の様に、LLMに与えるプロンプトをパラメータとして最適化することが出来ます。ドキュメントでも言及されていますがPyTorchに影響を受けているそうで、PyTorchでモデルを構築する感覚でLLMの処理を組むことが出来ます。

最初論文を読んだりしながら理解をしていこうと思ったのですが、とりあえず触ってみてどんな動きになるのかを自分の目で確かめた方が理解しやすいだろう、ということで、今回Retrieval-Augmented Generation (RAG)をDSPyで組み、最適化するところまでを試してみました。

DSPyのインストール

DSPyのインストールはpipコマンドで可能です。

pip install dspy-ai

ModuleとSignature

最初にDSPyの基本的なところから試してみます。DSPyでは最適化処理などにLLMを使用しますが、デフォルトで使用するLLMをDSPyのLanguage Models(LM)に設定しておく必要があります。ここではAzure OpenAI Serviceの"gpt-35-turbo-16k"を使用することにして、次のようなコードで設定します。

import os
import dspy

lm =dspy.AzureOpenAI(
    api_base=os.environ.get("AZURE_OPENAI_API_BASE"),
    api_version="2024-02-01",
    model="gpt-35-turbo-16k",
    api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
    model_type="chat",
    max_tokens=512
)

dspy.settings.configure(lm=lm) #これでデフォルトで使用するLMを設定

LLMに対するプロンプトエンジニアリングには、Chain of Thought(CoT)やReActなど、よく使用されるものがあります。それらはDSPyではModuleとして予め組み込まれています。Moduleにどのような入力を与え、どのような出力を得るのかはSignatureで指定します。Signatureは一番簡単な方法だと"question -> answer"のように入力と出力を矢印で結んだ文字列で指定することが可能です。

たとえば入力文字列をそのままLLMに与える場合とCoTでLLMに与える場合には次のようなコードで実現することが出来ます。

signature = "question->answer"
question ="上底が7cm下底が13cm高さ8cmの台形の面積は?"
predict = dspy.Predict(signature)
cot = dspy.ChainOfThought(signature)

print("Normal Prediction: " + predict(question=question).answer)
print("Chain of Thought: " + cot(question=question).answer)

結果は次のようになります。

Normal Prediction: Question: 上底が7cm下底が13cm高さ8cmの台形の面積は?
Answer: 60 cm²
Chain of Thought: The area of the trapezoid with a top base of 7cm, a bottom base of 13cm, and a height of 8cm is (7 + 13) * 8 / 2 = 80 square cm.

そのままの入力(Normal Predictionのパターン)だと台形の面積は間違ってますが、CoT(Chain of Thoughtのパターン)では正確に計算が出来ていることが分かります。

DSPyではこのようにあらかじめ定義されているプリミティブなModuleを組み合わせてPyTorchでモデルを組んでいくようにパイプラインを組んでいくことも出来ますし、自分でModuleそのものを作ることも出来ます。

RAG Module

ここからはRetrieval-Augmented Generation (RAG) PipelineをDSPyのModuleを使って組んでいく手順をご紹介します。RAG PipelineはRAGを実現するための一連の処理で、RAGはLLMが事前学習で学習していない事柄について、外部データを参照させることで回答可能にする手法です。

Retriever

まずドキュメントを格納し、クエリに対して関連情報を検索するRetrieverのコアの処理を作ります。ここは様々な方法がありますが、今回はLangChainとChromaを用いて構築しました。なお参照させるデータはこのブログの下書き用に私が作成したテキストファイルです。

# Create Retrieval Model with Chroma and OpenAIEmbedding

from langchain.document_loaders import DirectoryLoader
from langchain_chroma import Chroma
from langchain_openai import AzureOpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

embeddings = AzureOpenAIEmbeddings(
    model="text-embedding-ada-002",
    api_version="2024-02-01",
    azure_endpoint=os.environ.get("AZURE_OPENAI_API_BASE"),
)

data_dir = "/data/path/"
loader = DirectoryLoader(data_dir)
documents = loader.load()

text_splitter = CharacterTextSplitter(
    separator="\n\n",
    chunk_size=512,
    chunk_overlap=50,
    length_function=len,
    is_separator_regex=False,
)

splitted_documents = text_splitter.split_documents(documents)
db = Chroma.from_documents(splitted_documents, embeddings)
retriever = db.as_retriever()

Retriever Model

DSPyではクエリを受け取って関連情報を検索して返すModuleはRetrieval Models(RM)として取り扱います。RMは自分でカスタマイズしたものを利用することが可能で、その場合はdspy.Retrieveを継承したクラスを実装します。以下の様になります。

from typing import List, Union, Optional

class ChromaRMClient(dspy.Retrieve):
    def __init__(self, k: int=3):
        super().__init__(k=k)
    
    def forward(self, query_or_queries: Union[str, List[str]], k:Optional[int])->List[dspy.Prediction]:
        retriever = db.as_retriever(search_kwargs={'k':k if k else self.k})
        response = retriever.invoke(query_or_queries)
        response = [dspy.Prediction(long_text=d.page_content) for d in response]
        return response

この実装はDSPyのドキュメントの中の"Creating Custom RM Client"を参考にしています。

dspy-docs.vercel.app

こちらの参考にしているのですが、1部そのままでは動かない箇所がありました。具体的には実行するとAttributeError: 'str' object has no attribute 'long_text'というエラーが表示されてしまう問題で、この問題を解消するため、

  • forwardの戻り値の型をList[dspy.Prediction]にする
  • forwardが返す値を[dspy.Prediction(long_text=d.page_content) for d in response]のようにdspy.Predictionの属性としてlong_textを設定し、リストにする

といった対応を取りました。

RMもLMと同様、デフォルトで何を使用するかを指定する必要があります。

 rm = ChromaRMClient(k=3)
 dspy.settings.configure(lm=lm,rm=rm)

RAG Pipeline

次にRAGのパイプラインを構築します。DSPyのこちらのチュートリアルを参考にしました。

dspy-docs.vercel.app

class RAG(dspy.Module):
    def __init__(self, num_passages=3):
        super().__init__()
        self.retrieve = dspy.Retrieve(k=num_passages)
        self.generate_answer = dspy.ChainOfThought("context, question -> answer")

    def forward(self, question):
        context = self.retrieve(question).passages
        prediction = self.generate_answer(context=context, question=question)
        return dspy.Prediction(context=context, answer=prediction.answer)

これで次のようにRAG Pipelineを実行することが出来ます。

rg = RAG()
print(rg("AutoGenとは?").answer)

結果は次のようになりました。

"AutoGen is a framework for implementing Multi-Agent Conversation. It is a Python library that allows agents to engage in conversations. The framework provides a ConversableAgent class, which can be subclassed to implement agents. The ConversableAgent class has methods such as send, receive, and generate_reply, which can be used to define the behavior of the agent. AutoGen's documentation is available online, but reading the paper on AutoGen can provide a deeper understanding of the framework."

内容は良いのですが、日本語で質問したのに英語で返ってきました。これを解消するため、学習用のデータを準備し、パラメータのチューニングをかけてみました。

チューニング

まず10件の質問と回答のデータを用意しました。作成にはAzure OpenAI Serviceの"gpt-35-turbo-16k"を利用しました。

質問,回答
1年は何週間ですか?,1年は52週間です。
地球の公転周期は何日ですか?,地球の公転周期は365日です。
...

このデータをCSV形式で保存しておき、次のコードで学習用データ(trainset)として読み込みます。

import csv
trainset = []
with open("train_samples.csv","r") as f:
    reader = csv.reader(f)
    for i, row in enumerate(reader):
        if i == 0:
            continue
        trainset.append(dspy.Example(question=row[0],answer=row[1]))

# questionがinputで、残りのカラムをtargetにする
trainset = [x.with_inputs('question') for x in trainset]

あとはDSPyのチュートリアルに従い、以下のコードでチューニングを実行しました。

from dspy.teleprompt import BootstrapFewShot

# Validation logic: check that the predicted answer is correct.
# Also check that the retrieved context does actually contain that answer.
def validate_context_and_answer(example, pred, trace=None):
    answer_EM = dspy.evaluate.answer_exact_match(example, pred)
    answer_PM = dspy.evaluate.answer_passage_match(example, pred)
    return answer_EM and answer_PM

# Set up a basic teleprompter, which will compile our RAG program.
teleprompter = BootstrapFewShot(metric=validate_context_and_answer)

# Compile!
compiled_rag = teleprompter.compile(RAG(), trainset=trainset)

チューニング後のPipelinecompiled_ragを呼び出してみると、次のようにちゃんと日本語で回答が得られるようになりました。

compiled_rag("AutoGenってなに?").answer

結果

'AutoGenはMulti-Agent Conversationを実装するFrameworkです。Pythonのライブラリとして公開されています。'

DSpyの使い方のイメージが、なんとなくついてきました!

まとめ

ということで、今回はLLMを組み込んだ処理全体に含まれるプロンプトなどを機械学習モデルのパラメータのようにチューニングすることが出来るDSPyというフレームワークについて基本的な使い方を試してみました。自分の中では"こんなやり方があるのか"とかなりインパクトを受けました。DSPyについてはまだ詳しく理解できていない点も多いので、引き続き理解を深めていきたいと思います。