Amazon Bedrockでリランクモデルを有効化する方法とログ確認手順

内容

Amazon Bedrockのナレッジベースで構築したRAGシステムの回答精度を向上させるため、リランク機能の検証を行います。以下記事に記載のリランクモデルAmazon Rerank 1.0を使用してみます。

イメージ

イメージは下記のようになります。ユーザからの入力をエンベディングモデルでベクトル変換後、リランクモデルを使用して、ベクトル検索でドキュメントのリスト取得・ランク付けを実施します。ランク付けされた上位コンテンツをもとに基盤モデルで回答生成後、最終的な回答をユーザへ返す流れです。

設定

モデル呼出ログの有効化

リランク処理が実際に行われているかを確認するため、モデル呼び出し時のログを有効化します。

事前にIAMロールとCloudWatch Logsのロググループを作成後、Bedrockの設定画面からログ記録を有効にします。

ロググループ内にログストリームが作成されました。

リランクの有効化

RetrieveAndGenerateAPIでリランクの設定を行います。以下はテスト用プログラムの全文です。

import boto3
import json

REGION = "ap-northeast-1"
MODEL_ARN = "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0"
MODEL_ARN_RERANK = "arn:aws:bedrock:ap-northeast-1::foundation-model/amazon.rerank-v1:0"
KNOWLEDGEBASE_ID = '********'

bedrock_agent = boto3.client("bedrock-agent-runtime", region_name=REGION)
user_question = "質問内容を入力"

context_prompt = f"""
以下はユーザーからの質問です：

<question>
{user_question}
</question>
"""

# retrieve_and_generate_stream 呼び出し
response = bedrock_agent.retrieve_and_generate_stream(
    input={"text": context_prompt},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": KNOWLEDGEBASE_ID,
            "modelArn": MODEL_ARN,
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "rerankingConfiguration": {
                        "bedrockRerankingConfiguration": {
                            "metadataConfiguration": {
                                "selectionMode": "ALL",
                            },
                            "modelConfiguration": {
                                "modelArn": MODEL_ARN_RERANK
                            },
                        },
                        "type": "BEDROCK_RERANKING_MODEL"
                    }
                }
            }
        }
    }
)

stream = response.get("stream")
raw_events = []

if stream:
    for event in stream:
        raw_events.append(event)

print(json.dumps(raw_events, indent=2, ensure_ascii=False))

変更箇所は下記になります。まずリランクモデルを定義を追加します。

MODEL_ARN_RERANK = "arn:aws:bedrock:ap-northeast-1::foundation-model/amazon.rerank-v1:0"

retrieveAndGenerateConfigurationの中に下記を追加しました。

"retrievalConfiguration": {
    "vectorSearchConfiguration": {
        "rerankingConfiguration": {
            "bedrockRerankingConfiguration": {
                "metadataConfiguration": {
                    "selectionMode": "ALL",
                },
                "modelConfiguration": {
                    "modelArn": MODEL_ARN_RERANK
                },
            },
            "type": "BEDROCK_RERANKING_MODEL"
        }
    }
}

結果の確認

CloudWatch Logsで実行結果を確認すると、リランク処理が行われていることが確認出来ます。レスポンスの形式は以下のようになります。

 "output": {
        "outputContentType": "application/json",
        "outputBodyJson": {
            "results": [
                {
                    "index": 1,
                    "relevance_score": 0.00013765463********
                },
                {
                    "index": 3,
                    "relevance_score": 0.000016061610********
                },
                {
                    "index": 4,
                    "relevance_score": 0.0000098953323********
                },
                {
                    "index": 0,
                    "relevance_score": 0.00000297113812********
                },
                {
                    "index": 2,
                    "relevance_score": 0.0000028350801********
                }
            ]
        }