🤑

Amazon Bedrockで実践してみるRAG-Fusion

ken

2024/02/16に公開

みなさん、確定申告の季節ですね。（時候の挨拶）
納税、してますか？

先日この記事を見かけて、RAG-Fusion発想は面白いけど遅くなりそうだしどうなんだろう？と思ったので、実際に動かしてみることにしました。

せっかくなので僕が勝手にプッシュしているAmazon Bedrockを使ってやってみるか〜というのが本記事の趣旨です。

RAG-Fusionとは

どうも、AWS大好き芸人の@ken11です。AWS確定申告というサービスが生まれてCloudFormationでe-taxしてくれたらいいのになって思います。嘘です。

今回はRAG-Fusionの実装を試してみます。RAG-Fusionそのものの詳しいことは先述のリンク先を読んでいただくのがよいと思いますが、要は検索クエリを拡張してより広範な関連文書を取ってこよう、そのときの再評価におけるランク付けはRRF使おう、みたいな話です。

元の検索クエリを拡張する際にもLLMで生成しようみたいな話なので、LLMに対する推論回数が+1され、RAGの検索回数も+nされ、遅くなるのは宿命と考えられます。

実装

準備

まずRAGを使うにあたって、今回は題材に確定申告を使ってみましょう。(伏線回収)
今年も確定申告の季節がやって参りました、皆さん残り1ヶ月頑張ってください。僕はあとはもうお金を振り込むだけの悲しい状態に遷移しています()

というわけでRAGに入れたい元データはこちらを使ってみます。

所得税及び復興特別所得税の確定申告の手引き

Amazon Bedrockでナレッジベースを構築する

Bedrockではナレッジベースという機能を使うことでRAGを簡単に実装できてしまいます。

1. ナレッジベース作成開始

「ナレッジベースを作成」を押すとこのような画面になるので、適当に設定して次へ進みます。

2. データソースの選択

続いてデータソースを選択します。
先ほどの題材となるPDFを適当なS3に置いて、この画面で選択します。
BedrockのナレッジベースはPDF置いとくと勝手にRAGができあがるのが最高の体験だと思います。

3. 埋め込みモデルの選択

RAGのための埋め込みモデルを選択します。
ここではTitanを使いましょう。

これで作成ボタンを押すと、作成が始まります。
しばらく時間がかかるので、源泉徴収票でも探しながら気長に待ちましょう。

4. データソースを同期する

最後に、このままでは使えないので、データソースの同期を実行します。

以上でAmazon Bedrockを使ったナレッジベース(RAG)構築完了です。

RAG-Fusionをやってみる

それでは、RAG-Fusionを実際に動かしてみましょう。
ここから先のコードはGitHubにもノートブックごとアップロードしているので、好きに使ってください。

1. 準備

必要なライブラリのインストールと変数の用意をします。

import boto3
import json


# ナレッジベースのID
knowledge_base_id = ""
# データソースのID
data_source_id = ""
kb_client_runtime = boto3.client('bedrock-agent-runtime')
bedrock_runtime = boto3.client(service_name='bedrock-runtime')

# 推論用モデル(今回はClaudeのv2を使います)
model_id = "anthropic.claude-v2"
# クエリ拡張用のモデル(Claude Instantを使います。これは簡単なタスクにおける速度向上が狙いです)
q_model_id = "anthropic.claude-instant-v1"

region = "us-east-1"

ここで、ナレッジベースのIDとデータソースのIDはBedrockの以下の画面で確認できます。

また、今回モデルにはClaudeを使っています。Claudeが使えるのはAmazon Bedrockの大きな魅力ですね。ClaudeはXMLタグをいい感じに認識してくれたり、プロンプトの書きやすさも素晴らしいモデルだと思います。今回の主題から離れてしまうのでこれ以上は書きませんが、Claudeぜひ使ってみてください。

2. 推論用コード

def invoke_claude(text, model_id, max_tokens_to_sample=1000):
    body = json.dumps({
        "prompt": f"\n\nHuman:{text}\n\nAssistant: ",
        "max_tokens_to_sample": max_tokens_to_sample,
        "temperature": 0.1,
        "top_p": 0.9,
    })
    accept = 'application/json'
    content_type = 'application/json'

    response = bedrock_runtime.invoke_model(
        body=body,
        modelId=model_id,
        accept=accept,
        contentType=content_type
    )

    response_body = json.loads(response.get('body').read())
    return response_body.get('completion')[1:]

Bedrockを使ったClaudeの推論はこんな感じで簡単に実行できます。
今回は本推論時と検索クエリ拡張時でモデルを使い分けたかったので、引数で model_id をとっています。

3. クエリ拡張用コード

def generate_queries(original_query, n=4):
    f = ''
    for i in range(n):
        f += f'{n}: \n'

    prompt = f'''
以下に示すQueryはユーザの入力した検索クエリです。
このクエリに関連するクエリを多角的な視点から{n}つ生成してください
Formatに従って結果を出力してください

<Query>
{original_query}
</Query>

<Format>
{f}
</Format>
'''
    result = invoke_claude(prompt, q_model_id)
    result = result.split('<Format>')[1].split('</Format>')[0]
    # print(result)
    generated_queries = []
    for q in result.split('\n'):
        if q == '':
            continue
        generated_queries.append(q.split(' ')[1])
    return generated_queries[1:-1]

続いて、検索クエリを拡張するためのコードを用意します。
RAG-Fusionでよりよい結果を生み出すためには、ただ元のクエリの類語などにするのではなく、多角的な視点から検索クエリを用意することがよいと考えられます。
そのため、ただ似たような検索クエリにするだけならコスト等々を考えてもっとトラディショナルな手法で拡張できると思いますが、あえてLLMにリクエストして柔軟に幅広いクエリの生成を期待しています。

4. RAG検索用コード

def kb_search(query, n=5):
    res = kb_client_runtime.retrieve(
        retrievalQuery= {
            'text': query
        },
        knowledgeBaseId=knowledge_base_id,
        retrievalConfiguration= {
            'vectorSearchConfiguration': {
                'numberOfResults': n
            }
        }
    )

    # {doc: score}の形に整形します
    return_dict = {}
    for r in res['retrievalResults']:
        return_dict[r['content']['text']] = r['score']

    return return_dict

続いてRAGの検索に使うコードです。
Amazon Bedrockを使うとこのように簡単に実装できてしまいます。(n回目)
いやあ、実際Bedrock便利だから……

5. RRFのコード

def reciprocal_rank_fusion(all_results, k=50):
    fused_scores = {}
    for query, doc_scores in all_results.items():
        for rank, (doc, score) in enumerate(sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)):
            if doc not in fused_scores:
                fused_scores[doc] = 0
            previous_score = fused_scores[doc]
            fused_scores[doc] += 1 / (rank + k)

    reranked_results = {doc: score for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)}
    return reranked_results

そしてRRFのコードです。
これは複数クエリでRAG検索した結果をまとめて、再評価するためのものです。
k によってどの程度ランキングを重視するかを調整できます。

6. 実行

以上で準備はできたので、あとは実行するだけです。

original_query = '確定申告について教えて'

queries = generate_queries(original_query)
queries.insert(0, original_query)

all_results = {}
for q in queries:
    all_results[q] = kb_search(q)

# ここでは上位5件を採用している
reranked_results = list(reciprocal_rank_fusion(all_results).keys())[:5]

information = ''
for i, r in enumerate(reranked_results):
    information += f'情報{i+1}. {r}\n'

prompt = """
日本では本日(2月16日)から確定申告の相談及び申告書の受付がスタートしましたね。
この国の風物詩であるこの確定申告という儀式では、毎年数多の迷える子羊が出現することで有名です。
あなたは全知全能の税神として、迷える子羊の抱いている確定申告に対する悩みを聞き、彼らを正しき納税へと導いてください。
全知全能たるあなたのために、Informationセクションには迷える子羊の悩みに関連しそうな情報を提供します。有効に活用してください。

<Information>
{information}
</Information>

<UserQuery>
{query}
</UserQuery>
"""
prompt = prompt.format(information=information, query=original_query)

print(invoke_claude(prompt, model_id))

迷える子羊を導くDeus ex Machinaの誕生です(いいえ)

結果比較

では実際に通常のRAGとRAG-Fusionでどの程度結果に差が出るでしょうか？
今回はこのDeus ex Machina(いいえ)に対して、
「確定申告について教えて」
と話しかけてみます。

RAGの結果

確定申告について、以下の点をご説明します。

確定申告の概要

所得税や復興特別所得税を納める義務がある人は、1年間の所得について確定申告を行う必要があります。

確定申告は、申告書に必要事項を記入して税務署に提出することで行います。

確定申告が必要な人

給与所得者でも、一定の条件を満たす場合は確定申告が必要です。

自営業者や農業者などの申告が必要な所得がある人も確定申告が必要です。

確定申告のメリット

税額控除などを受けることができ、納める税金を減らすことができます。

医療費控除などにより、所得税や住民税の還付を受けることができます。

確定申告の手続き

確定申告書に必要事項を記入し、添付書類を添えて税務署に提出します。

e-Taxなどの電子申告も利用できます。

申告期限は原則3月15日です。

納税

確定申告後、計算した税額を納付期限までに納める必要があります。

納付には納税通知書や振替納税などの方法があります。

確定申告の概要は以上の通りです。不明な点があればご質問ください。

RAG-Fusionの結果

確定申告に関しては、以下の点に注意が必要です。

確定申告の手続き

申告期間は2月16日から3月15日までです。郵送またはe-Taxで提出できます。

申告書の控えに収受印をもらう場合は、郵送時に返信用封筒を同封するか、受付で控えを持参する必要があります。

納税の方法

納付期限は3月15日です。キャッシュレス納付が便利です。

振替納税を利用する場合は3月15日までに手続きが必要です。

控除の記入

医療費控除、社会保険料控除、生命保険料控除などの控除を受ける場合は、第一表と第二表にそれぞれ記入が必要です。

住民税・事業税

別居家族の氏名・住所を記入する必要があります。

還付金の受取

公金受取口座を登録・利用する場合は手続きが必要です。

以上の点に注意し、確定申告の手続きを進めましょう。不明な点があれば税務署に問い合わせることをおすすめします。

ちょっと選んだ題材が悪かったかな()
しかし還付金のくだり等、RAG-Fusionの方がより広範な情報をくれている気がします。

まとめ

というわけで、Amazon Bedrockを使ってRAG-Fusionを実践してみました。
個人的にはRAGよりもコストが(時間等々)かかることが気になりましたが、より広範に情報を提供したいケースなんかでは結構有用な手法では、と思いました。
使いどころによって、選択肢の一つに含められるといいね、的な。

そしてAmazon Bedrockの便利さを実感していただけたでしょうか。RAG構築がとても容易なので、結果的にRAG-Fusionもすぐに試してみることができました。
ClaudeはChatGPTとは違った使い勝手の良さがありますし、ぜひBedrockも多くの人に使ってみてほしいです。
そして東京リージョンがどんどんいろいろ潤沢になるといいなって……

Spiral AIではそんなことを日々考えたり挑戦したりしているので、興味を持った方がいらっしゃいましたらいつでもお声がけください。

ではでは、楽しい確定申告シーズンをお過ごしください。

SpiralAIテックブログPublication

実在する芸能人との会話ができる日本初のAIサービス「NaomiAI」やカスタムChatGPTを作れる「Spiralbot」を提供するSpiralAI株式会社のテックブログです。