langchainとDatabricksで(私が)学ぶRAG : RAGatouilleを使ったRAG

導入

以下の記事でRAGatouilleを使ったRerankを試しましたが、今回はRAGatouilleを通常のIndex作成に利用し、RAGを実装してみます。

シリーズ一覧はこちら。

検証はDatabricks on AWS上で実施しました。
DatabricksのDBRは14.2 ML、GPUクラスタ(g4dn.xlarge)上で動作を確認しています。

RAGatouilleとは？

ColBERTのような先端モデルなどを簡単に利用できるように設計されたRAG用のモジュールです。
ColBERTモデルのファインチューニングや、ドキュメントデータのIndex作成、検索など、RAGにおける主要機能を提供しています。

READMEを一部邦訳すると、

RAGatouilleの主な動機は単純で、最先端の研究と錬金術的なRAGパイプラインの実践の間のギャップを埋めることです。RAGは複雑で、多くの可動部分があります。最高のパフォーマンスを得るには、多くのコンポーネントを最適化する必要がありますが、その中でも非常に重要なのが、検索に使用するモデルです。

高密度検索、つまり OpenAI(text-ada-002)などの埋め込みを使用することは、適切なベースラインですが、高密度埋め込みがユースケースに最適ではない可能性があることを示す多くの研究があります。

情報検索の研究分野は最近活況を呈しており、ColBERTのようなモデルは、高密度の埋め込みよりも新しいドメインや複雑なドメインに一般化することに適しており、途方もなくデータ効率が高く、データ量の少ない英語以外の言語で効率的にトレーニングするのにさらに適していることが示されています!残念ながら、これらの新しいアプローチのほとんどはあまり知られておらず、密な埋め込みよりもはるかに使いにくいものです。

ここでRAGatouilleの出番です:RAGatouilleの目的は、このギャップを埋めることです:詳細や文献の年数を気にすることなく、RAGパイプラインで最先端の方法を簡単に使用できるようにすることです!現在、RAGatouilleはColBERTを使いやすくすることに重点を置いています。今後の予定をご覧になりたい方は、幅広いロードマップをご覧ください。

というわけで、私のようなド素人であっても、ColBERTを使ったRAG構築を簡単に行うことができます。

今回はLangchainを使ってRAGatouilleと連携したRAGのChainを構築します。
基本的には以下内容のウォークスルーです。

Step0. パッケージインストール

使うパッケージをインストールします。
今回はRAGatouilleとColBERT(今回はJaColBERT)を利用するために必要なパッケージを追加で入れています。
なお、AutoAWQはCUDA 11.8版をWheelを指定してインストールしています。
(DatabricksのDBR ML 14台はCUDAのバージョンが11.8のため)

%pip install -U -qq transformers accelerate langchain faiss-gpu ragatouille fugashi unidic_lite 
# CUDA 11.8用
%pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.8/autoawq-0.1.8+cu118-cp310-cp310-linux_x86_64.whl

dbutils.library.restartPython()

Step.1 RAGatouilleでColBERTモデルをロード

ドキュメントのインデックスを作成するためのColBERTモデルをロードします。

前回同様、日本語用のColBERTモデルJaColBERTが以下に公開されていますので、こちらを利用します。

from ragatouille import RAGPretrainedModel

RAG = RAGPretrainedModel.from_pretrained("bclavie/JaColBERT")

Step.2 RAG対象のドキュメントを読み込み

日本語Wikipediaから検索対象のドキュメントを取得します。

まずは、MediaWiki APIを使ってWikipediaのデータを取得する関数を定義。

import requests

def get_wikipedia_page(title: str):
    """
    Retrieve the full text content of a Wikipedia page.

    :param title: str - Title of the Wikipedia page.
    :return: str - Full text content of the page as raw string.
    """
    # Wikipedia API endpoint
    URL = "https://ja.wikipedia.org/w/api.php"

    # Parameters for the API request
    params = {
        "action": "query",
        "format": "json",
        "titles": title,
        "prop": "extracts",
        "explaintext": True,
    }

    # Custom User-Agent header to comply with Wikipedia's best practices
    headers = {"User-Agent": "RAGatouille_tutorial/0.0.1 ([e-mail])"}

    response = requests.get(URL, params=params, headers=headers)
    data = response.json()

    # Extracting page content
    page = next(iter(data["query"]["pages"].values()))
    return page["extract"] if "extract" in page else None

定義した関数を使って、Wikipediaデータを取得。
今回はこちらの内容を利用させてもらいます。

full_document = get_wikipedia_page("葬送のフリーレン")

Step3. インデックス作成

読み込んだモデルを利用してインデックスを作成します。
感覚的にはドキュメントのチャンキングとVectorstoreの作成・格納をまとめて実行してくれます。

RAG.index(
    collection=[full_document],
    index_name="Frieren-1",
    max_document_length=256,
    split_documents=True,
)

では、作成したインデックスを利用して検索してみましょう。

results = RAG.search(query="フリーレンの声優は誰？", k=3)

results

出力

[{'content': '== 登場人物 ==\n声の項はテレビアニメ版の声優。\n\n\n=== フリーレン一行（主要人物） ===\nフリーレン (Frieren)\n声 - 種﨑敦美\n本作の主人公。魔王を討伐した勇者パーティーの魔法使い。長命なエルフ族の出身で、少女のような外見に反して1000年以上の歳月を生き続けている。人間とは時間の感覚が大きく異なるため、数か月から数年単位の作業をまったく苦にせず、ヒンメルらかつての仲間たちとの再会も50年の月日が経ってからのことだった。',
  'score': 27.0625,
  'rank': 1,
  'document_id': 'f8fbd821-44a9-4545-a392-2370f9a1e7eb'},
 {'content': 'よミトく！」特別版として、作品の内容や制作現場の模様が公開された。出演は水卜麻美（日本テレビアナウンサー）。インタビュー出演はかまいたち、若月佑美、種﨑敦美（フリーレン役）、市ノ瀬加那（フェルン役）、小林千晃（シュタルク役）、YOASOBI（OPアーティスト）、milet（EDアーティスト）。葬送のフリーレン 〜トークの魔法〜\n2023年9月27日よりYouTube「TOHO animation チャンネル」にて不定期で配信中のWeb番組。',
  'score': 24.84375,
  'rank': 2,
  'document_id': 'f8fbd821-44a9-4545-a392-2370f9a1e7eb'},
 {'content': '幼いころに森で道に迷ったときにフリーレンと出会っており、花畑の魔法を見せてもらったことがある。エルフであるがゆえにはるかに長命であるフリーレンの未来をおもんばかって銅像に細かく注文をしたり、故郷を魔族に奪われて孤独に生きてきた彼女を思いやったりと、何かと気にかけていた。女神の石碑に触れて過去に戻ったフリーレンと出会った際は、彼女が未来から来たという事実を素直に受け入れている。\nハイター (Heiter)\n声 - 東地宏樹\n人間出身の僧侶。',
  'score': 24.015625,
  'rank': 3,
  'document_id': 'f8fbd821-44a9-4545-a392-2370f9a1e7eb'}]

rank1, 2には正解が含まれているのが見てわかりますね。

また、as_langchain_retrieverメソッドを利用することで、LangchainのRetrieverを取得することもできます。

retriever = RAG.as_langchain_retriever(k=1)

retriever.invoke("フリーレンの声優は誰？")

出力

[Document(page_content='== 登場人物 ==\n声の項はテレビアニメ版の声優。\n\n\n=== フリーレン一行（主要人物） ===\nフリーレン (Frieren)\n声 - 種﨑敦美\n本作の主人公。魔王を討伐した勇者パーティーの魔法使い。長命なエルフ族の出身で、少女のような外見に反して1000年以上の歳月を生き続けている。人間とは時間の感覚が大きく異なるため、数か月から数年単位の作業をまったく苦にせず、ヒンメルらかつての仲間たちとの再会も50年の月日が経ってからのことだった。')]

Step4. Chain作成

では、RAGのChainをLangchainを使って作成してみます。

まずは利用するLLMの読み込み。
以下のモデルを事前にダウンロードして利用しました。
モデルはこちらに記載しているLangchainのカスタムチャットモデルで読み込みます。

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers_chat import ChatHuggingFaceModel

model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--openchat-3.5-0106-AWQ"

generator = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)

単純なRAGのChainをLangchain LCELで定義します。

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts.chat import (
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.runnables import RunnablePassthrough, RunnableLambda

prompt_template = """Answer the question based only on the following context:
{context}

Question: {question}
"""

prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template(prompt_template),
        AIMessagePromptTemplate.from_template(""),
    ]
)

chat_model = ChatHuggingFaceModel(
    generator=generator,
    tokenizer=tokenizer,
    human_message_template="GPT4 Correct User: {}<|end_of_turn|>",
    ai_message_template="GPT4 Correct Assistant: {}",
    repetition_penalty=1.2,
    temperature=0.1,
    max_new_tokens=400,
)

# RAGatouilleを使ったRetriever
retriever = RAG.as_langchain_retriever(k=3)

chain = (
    {"context": retriever, "question": RunnablePassthrough()} 
    | prompt 
    | chat_model 
    | StrOutputParser()
)

Step5. 推論実行

準備が終わりましたので推論させてみます。

for c in chain.stream("フリーレンの声優は誰？"):
    print(c, end="", flush=True)

出力

フリーレンの声優は種﨑敦美です。

別の質問も。

for c in chain.stream("魔族が使う代表的な魔法を３つ箇条書きで教えてください。"):
    print(c, end="", flush=True)

出力

1. 魔族を殺す魔法（ゾルトラーク）
2. 血を操る魔法（バルテーリエ）
3. 模倣する魔法（エアフアーゼン）

ゾルトラークの日本語名がちょっと怪しいですが、正しく回答されているように思います。

まとめ

Langchain + RAGatouilleを使ってRAGパイプラインを構築しました。

RAGatouille自体は非常にシンプルなインターフェースを提供しているパッケージであり、LangchainのRetrieverも提供しているため、とても扱いやすいと思います。
ColBERTを使ったIndex性能がOpenAIのembeddingモデル対比でどれだけの優位性があるかは正直わからないのですが、簡単に触った限りでは軽量かつよい性能で検索できているように思いました。
また、ColBERTはLlamaIndexとの統合も既にされているようで、セマンティック検索における選択肢の一つになるといいなと思います。