Amazon Bedrock の Rerank API を活用してRAGの精度を向上させる

こんにちは、機械学習チーム YAMALEX の駿です。
YAMALEX は Acroquest 社内で発足した、会社の未来の技術を創る、機械学習がメインテーマのデータサイエンスチームです。
（詳細はリンク先をご覧ください。）

RAG（検索拡張生成：Retrieval-Augmented Generation）で正しい情報をユーザーに提供する際に特に重要になるのが検索精度です。

そして、検索精度を上げる方法のひとつが「リランク」です。

リランクを行い、検索で取得した結果を関連度で並び替えなおすことで、ユーザーの求めている情報を回答しやすくなります。

そんなリランクを行うモデルがAmazon Bedrockに追加され、 Bedrock Knowledge Baseと組み合わせて使用できるようになりました。

従来は自前でモデルをホストする、など手間がかかりましたが、Knowledge Baseへの検索リクエストに設定を追加するだけで、検索・リランクをまとめて実行し、リランク後の結果のみを取得できるようになりました。

今回は、リランクモデルを使ってみて、検索結果がどのように変わるのかを確認します。

• 1. はじめに
  • 1.1. リランクとは
  • 1.2. 従来の実現方法
  • 1.3. Bedrock で使用できる リランクモデル
• 2. リランクモデルを適用してみる
  • 2.1. Bedrock の InvokeModel API で使う
  • 2.2. Bedrock Knowledge Base の Rerank API で使う
  • 2.3. Bedrock Knowledge Base の Retrieve API で使う
• 3. Amazon Rerank モデルと Cohere Rerank モデルの比較
• 4. その他
  • 4.1. 呼び出し速度
  • 4.2. 料金
• まとめ

1. はじめに

1.1. リランクとは

Bedrock Knowledge Baseを含むRAGの検索では、ベクトル検索が使われることが多いです。

しかし、ベクトル検索だけでは検索において十分な検索精度が出ず、適切な回答ができない、といったことも多いです。

そこで、ベクトル検索で取得したドキュメントに対してリランクを行うことで、より関連度の高いドキュメントが検索結果の上位に来るようにすることができます。

リランクに関しては当ブログの下記記事に詳しい内容が書かれているため、こちらもご参照ください。

acro-engineer.hatenablog.com

1.2. 従来の実現方法

いままでは、リランク処理を RAG システムに組み込むためには SageMaker のインスタンスを立てて、リランク用のモデルをホストし、推論を行う必要がありました。

例えば、2024年8月時点で Cohere Rerank 3 を使用するためには、下記記事のように SageMaker のインスタンスを作成する必要がありました。

aws.amazon.com

この方法だと、 SageMaker のインスタンス、リランクモデルの準備など手間や運用コストがかかる状況でした。

1.3. Bedrock で使用できる リランクモデル

2024年12月 から Bedrock を経由してリランクモデルを使用できるようになりました。

このリランクモデルを使用することで、自分でモデルをホストすることなく、APIを呼び出すだけでリランクを行うことができるようになりました。

運用管理の手間がなくなる、サーバーを起動しておく必要がなく使った分だけお金を払えばよい、など、リランクを簡単に始められるようになります。

Bedrock の InvokeModel API から使用できるほか、 Bedrock Knowledge Base の Rerank API 、 Retrieve API 、RetrieveAndGenerate API 、 RetrieveAndGenerateStream API から使用可能です。

2025年1月現在、Amazon Rerank 1.0 （以下、Amazon Rerank モデル）と Cohere Rerank 3.5 （以下、Cohere Rerank モデル）のモデルが提供されています。

2. リランクモデルを適用してみる

こちらの記事でも使用している、宿のレビュー検索を想定したデータを使って検証します。

今回は「肉のグリルがおいしい宿」で検索して、「地元産の野菜とお肉を使ったグリル料理」の No.10 、「炭火焼きでいただいた和牛ステーキ」の No.7 のレビューが上位の結果に表れることを期待します。

リランクモデルは Amazon Rerank モデルを使用しました。

No	文章
1	この宿の温泉は、まさに極上の癒しでした。源泉かけ流しの湯はとても柔らかく、肌がすべすべになる感じがしました。露天風呂からは美しい山々の景色が広がり、夜には満天の星空を眺めながらゆっくりと浸かることができました。何度でも訪れたくなる温泉です。
2	宿の温泉はとても気持ち良く、リラックスできました。内湯と露天風呂があり、それぞれ異なる趣があります。特に露天風呂から見える庭園の景色が素晴らしく、四季折々の美しさを楽しむことができます。お湯も適温で、長時間浸かっていても疲れませんでした。
3	温泉が自慢の宿とのことで期待していましたが、期待以上でした。天然温泉の源泉をそのまま使用しているので、お湯の質がとても良く、温まった後もポカポカが続きました。貸切風呂も予約して利用しましたが、プライベートな空間で贅沢なひとときを過ごすことができました。
4	温泉は広々としていて、開放感が抜群でした。露天風呂からは海が一望でき、波の音を聞きながらゆっくりと過ごすことができました。お湯も熱すぎず、じっくりと温まることができてとても満足です。また、日帰り利用も可能なので、気軽に訪れることができるのも嬉しいポイントです。
5	温泉は硫黄の香りが心地よく、いかにも温泉に来たという実感がありました。お湯の効能も高く、肌がツルツルになるのを感じました。複数の湯船があり、時間帯によっては貸切状態になることもあり、贅沢な気分を味わえました。また、湯上がりに提供される冷たいドリンクも嬉しいサービスでした。
6	宿の食事はまさに芸術品のようでした。地元の新鮮な食材をふんだんに使った会席料理は、見た目も美しく、一品一品が丁寧に作られているのが伝わってきました。特に旬の魚介を使った刺身は絶品で、これだけでもまた訪れたいと思いました。
7	夕食は地元の名物料理が盛りだくさんで、大満足でした。特に炭火焼きでいただいた和牛ステーキは口の中でとろける美味しさで、何度もおかわりしたくなるほどでした。朝食も種類が豊富で、地元の野菜を使ったサラダや、手作りの豆腐が美味しかったです。
8	夕食はコース料理で、どれも美味しかったのですが、特に印象に残ったのは地元で採れた野菜を使った前菜と、自家製のデザートです。食材の味を活かしたシンプルな調理法で、素材の良さが際立っていました。朝食もバランスが良く、特に焼きたてのパンが絶品でした。
9	宿の食事は期待以上でした。海の近くということもあり、新鮮な魚介類が豊富に使われていて、お刺身や煮魚がとても美味しかったです。夕食は量もたっぷりで、どの料理も心のこもった味付けでした。朝食の和食もとても美味しく、特に温泉卵が絶品でした。
10	夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。

2.1. Bedrock の InvokeModel API で使う

InvokeModel API は Bedrock で提供されているモデルを呼び出すための API です。

body にリランクを行いたいドキュメントのリストとユーザーのクエリを入力すると、レスポンスとして、ユーザーのクエリとの関連度が高い順に並び替えられたドキュメントとそれぞれの関連度（スコア）を取得できます。

コード

query = "肉のグリルがおいしい宿"
documents = [
    "この宿の温泉は、まさに極上の癒しでした。源泉かけ流しの湯はとても柔らかく、肌がすべすべになる感じがしました。露天風呂からは美しい山々の景色が広がり、夜には満天の星空を眺めながらゆっくりと浸かることができました。何度でも訪れたくなる温泉です。",
    # (略)
]

response = bedrock.invoke_model(
    modelId="amazon.rerank-v1:0",
    body=json.dumps({
        "query": query,
        "documents": documents,
        "top_n": 3,
    }),
)

body = json.loads(response["body"].read())
pprint.pprint(body["results"])

出力

[{'index': 9, 'relevance_score': 0.001466458403084568},
 {'index': 6, 'relevance_score': 0.0005013742398679934},
 {'index': 8, 'relevance_score': 0.0003640086870995012}]

※リランク結果に含まれる index は0オリジンのため、上の表と揃えるために、 index に +1 します。

結果

No	文章
10	夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。
7	夕食は地元の名物料理が盛りだくさんで、大満足でした。特に炭火焼きでいただいた和牛ステーキは口の中でとろける美味しさで、何度もおかわりしたくなるほどでした。朝食も種類が豊富で、地元の野菜を使ったサラダや、手作りの豆腐が美味しかったです。
9	宿の食事は期待以上でした。海の近くということもあり、新鮮な魚介類が豊富に使われていて、お刺身や煮魚がとても美味しかったです。夕食は量もたっぷりで、どの料理も心のこもった味付けでした。朝食の和食もとても美味しく、特に温泉卵が絶品でした。

期待通り、 No.10 、 No.7 のレビュー文章が上位に来ることを確認できました。

2.2. Bedrock Knowledge Base の Rerank API で使う

Rerank API は Knowledge Base の機能として用意されていますが、実態は上のInvokeModelと同じもので、ドキュメントのリストとユーザーのクエリを入力すると、リランク後のドキュメントのリストが得られます。

コード

region = boto3.Session().region_name
amazon_rerank_arn = f"arn:aws:bedrock:{region}::foundation-model/amazon.rerank-v1:0"

response = bedrock_agent.rerank(
    queries=[
        {
            "type": "TEXT",
            "textQuery": {
                "text": query,
            },
        },
    ],
    sources=[
        {
            "inlineDocumentSource": {
                "textDocument": {
                    "text": document,
                },
                "type": "TEXT",
            },
            "type": "INLINE",
        } for document in documents
    ],
    rerankingConfiguration={
        "type": "BEDROCK_RERANKING_MODEL",
        "bedrockRerankingConfiguration": {
            "numberOfResults": 3,
            "modelConfiguration": {
                "modelArn": amazon_rerank_arn,
            },
        },
    },
)

pprint.pprint(response["results"])

出力

[{'index': 9, 'relevanceScore': 0.0014664584305137396},
 {'index': 6, 'relevanceScore': 0.0005013742484152317},
 {'index': 8, 'relevanceScore': 0.0003640086797531694}]

InvokeModel を使用した際と全く同じ結果が得られることを確認できました。

2.3. Bedrock Knowledge Base の Retrieve API で使う

Retrieve API では、InvokeModel、Rerank APIと異なり、ドキュメントのリストは入力としては渡しません。

ユーザーのクエリを入力とし、ベクトルDBをユーザーのクエリで検索した結果をドキュメントのリストとして、リランクを行います。

Retrieve APIを使用するための準備としてナレッジベースを作成し、上記の文章をひとつずつ、１チャンクとして保存しました。

まずはリランクを 行わない 場合にどのような結果になるのかを確認します。

コード

response = bedrock_agent.retrieve(
    knowledgeBaseId=knowledgebase_id,
    retrievalConfiguration={
        "vectorSearchConfiguration": {
            "numberOfResults": 3,
            "overrideSearchType": "SEMANTIC",
        },
    },
    retrievalQuery={
        "text": query,
    },
)

pprint.pprint(response["retrievalResults"])

出力

[{'content': {'text': '宿の食事はまさに芸術品のようでした。地元の新鮮な食材をふんだんに使った会席料理は、見た目も美しく、一品一品が丁寧に作られているのが伝わってきました。特に旬の魚介を使った刺身は絶品で、これだけでもまた訪れたいと思いました。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/006.txt'},
               'type': 'S3'},
  'score': 0.43565163},
 {'content': {'text': '宿の食事は期待以上でした。海の近くということもあり、新鮮な魚介類が豊富に使われていて、お刺身や煮魚がとても美味しかったです。夕食は量もたっぷりで、どの料理も心のこもった味付けでした。朝食の和食もとても美味しく、特に温泉卵が絶品でした。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/009.txt'},
               'type': 'S3'},
  'score': 0.435101},
 {'content': {'text': '夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/010.txt'},
               'type': 'S3'},
  'score': 0.4281698}]

結果

No	文章
6	宿の食事はまさに芸術品のようでした。地元の新鮮な食材をふんだんに使った会席料理は、見た目も美しく、一品一品が丁寧に作られているのが伝わってきました。特に旬の魚介を使った刺身は絶品で、これだけでもまた訪れたいと思いました。
9	宿の食事は期待以上でした。海の近くということもあり、新鮮な魚介類が豊富に使われていて、お刺身や煮魚がとても美味しかったです。夕食は量もたっぷりで、どの料理も心のこもった味付けでした。朝食の和食もとても美味しく、特に温泉卵が絶品でした。
10	夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。

上位３件まで取得した場合、No.10 のレビューが３位、 No.7 のレビューは検索結果に含まれない結果になりました。

この検索結果を使って RAG を行っても、精度の高い回答は得られなそうです。

---

次に Retrieve API でリランクモデルを指定して、どのように検索結果が変化するのかを確認します。

コード

response = bedrock_agent.retrieve(
    knowledgeBaseId=knowledgebase_id,
    retrievalConfiguration={
        "vectorSearchConfiguration": {
            # (1) 最初の検索では10件取得する
            "numberOfResults": 10,
            "overrideSearchType": "SEMANTIC",
            "rerankingConfiguration": {
                "bedrockRerankingConfiguration": {
                    "modelConfiguration": {
                        "modelArn": amazon_rerank_arn,
                    },
                    # (2) 検索で取得した10件でリランクを行い、上位3件を返却する
                    "numberOfRerankedResults": 3,
                },
                "type": "BEDROCK_RERANKING_MODEL",
            },
        },
    },
    retrievalQuery={
        "text": query,
    },
)

pprint.pprint(response)

出力

[{'content': {'text': '夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/010.txt'},
               'type': 'S3'},
  'score': 0.0014721895568072796},
 {'content': {'text': '夕食は地元の名物料理が盛りだくさんで、大満足でした。特に炭火焼きでいただいた和牛ステーキは口の中でとろける美味しさで、何度もおかわりしたくなるほどでした。朝食も種類が豊富で、地元の野菜を使ったサラダや、手作りの豆腐が美味しかったです。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/007.txt'},
               'type': 'S3'},
  'score': 0.0004994205664843321},
 {'content': {'text': '宿の食事は期待以上でした。海の近くということもあり、新鮮な魚介類が豊富に使われていて、お刺身や煮魚がとても美味しかったです。夕食は量もたっぷりで、どの料理も心のこもった味付けでした。朝食の和食もとても美味しく、特に温泉卵が絶品でした。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/009.txt'},
               'type': 'S3'},
  'score': 0.0003640086797531694}]

結果

No	文章
10	夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。
7	夕食は地元の名物料理が盛りだくさんで、大満足でした。特に炭火焼きでいただいた和牛ステーキは口の中でとろける美味しさで、何度もおかわりしたくなるほどでした。朝食も種類が豊富で、地元の野菜を使ったサラダや、手作りの豆腐が美味しかったです。
9	宿の食事は期待以上でした。海の近くということもあり、新鮮な魚介類が豊富に使われていて、お刺身や煮魚がとても美味しかったです。夕食は量もたっぷりで、どの料理も心のこもった味付けでした。朝食の和食もとても美味しく、特に温泉卵が絶品でした。

リランクを行うことで、 No.10 と No.7 の文章が上位２件を占める結果となりました。

ユーザーが求める情報を、より多く回答することができるようになります。

3. Amazon Rerank モデルと Cohere Rerank モデルの比較

次に、同じ内容を、同じく Bedrock で使用可能な Cohere Rerank モデルで試してみます。

modelArn を Cohere Rerank モデルのものに差し替えるだけで、使用するリランクモデルを変更できます。

簡単でいいですね。

コード

cohere_rerank_arn = f"arn:aws:bedrock:{region}::foundation-model/cohere.rerank-v3-5:0"
# （略）

出力

[{'content': {'text': '夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/010.txt'},
               'type': 'S3'},
  'score': 0.3279808461666107},
 {'content': {'text': '宿の食事はまさに芸術品のようでした。地元の新鮮な食材をふんだんに使った会席料理は、見た目も美しく、一品一品が丁寧に作られているのが伝わってきました。特に旬の魚介を使った刺身は絶品で、これだけでもまた訪れたいと思いました。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/006.txt'},
               'type': 'S3'},
  'score': 0.1456373631954193},
 {'content': {'text': '夕食は地元の名物料理が盛りだくさんで、大満足でした。特に炭火焼きでいただいた和牛ステーキは口の中でとろける美味しさで、何度もおかわりしたくなるほどでした。朝食も種類が豊富で、地元の野菜を使ったサラダや、手作りの豆腐が美味しかったです。',
              'type': 'TEXT'},
  'location': {'s3Location': {'uri': 's3://xxx/007.txt'},
               'type': 'S3'},
  'score': 0.11919290572404861}]

結果

No	文章
10	夕食は地元の食材をふんだんに使った創作料理で、どの料理も工夫が感じられました。特に地元産の野菜とお肉を使ったグリル料理が絶品で、素材の味がしっかりと引き立っていました。朝食も手作りのジャムや焼き立てのパンなど、こだわりが感じられる内容で、大変満足しました。
6	宿の食事はまさに芸術品のようでした。地元の新鮮な食材をふんだんに使った会席料理は、見た目も美しく、一品一品が丁寧に作られているのが伝わってきました。特に旬の魚介を使った刺身は絶品で、これだけでもまた訪れたいと思いました。
7	夕食は地元の名物料理が盛りだくさんで、大満足でした。特に炭火焼きでいただいた和牛ステーキは口の中でとろける美味しさで、何度もおかわりしたくなるほどでした。朝食も種類が豊富で、地元の野菜を使ったサラダや、手作りの豆腐が美味しかったです。

Amazon Rerank モデルを使った場合と比べて、 No.7 の順位がひとつ落ちる結果となりましたが、依然上位３件に含まれることは変わりません。

No.6 の文章は肉料理ではなく魚介料理についてのレビューですが、温泉についてではなく、美味しい料理についてのレビューである、という点でスコアが高くなっているのだと思います。

これなら RAG で回答する際も、情報が不足することなく、生成を行うことができそうです。

4. その他

4.1. 呼び出し速度

Amazon Rerank モデルとCohere Rerank モデルでレスポンススピードに差分があるのか、確認しました。

それぞれ、オレゴンリージョンのモデルで同じリクエストを5回ずつ実行し、レスポンス時間の平均を比較しました。

Amazon Rerank モデル

No	レスポンス時間（秒）
1	0.895
2	0.687
3	0.734
4	0.828
5	0.775
平均	0.784

Cohere Rrerank モデル

No	レスポンス時間（秒）
1	0.454
2	0.508
3	0.533
4	0.495
5	0.453
平均	0.489

Amazon Rerank モデルと Cohere Rrerank モデル を比べるとCohere Rerank モデルの方が1.5倍ほどスピードが速い、という結果になりました。

4.2. 料金

今回用いたモデルの料金は下記の通りです。

リランクモデルではないですが、 Amazon Nova Lite が1,000出力トークンにつき$0.00024 というのと比べると、少し高く感じてしまいましたが、それだけ複雑なことをAPI呼び出しだけで使えている、ということでもあります。

No	モデル	料金
1	Amazon Rerank モデル	$1 / 1000クエリ
2	Cohere Rerank モデル	$2 / 1000クエリ

まとめ

Bedrock に追加されたリランクモデルの検証を行い、検索結果の改善に有用であることを確認できました。

リランクを行うことで、よりユーザー入力に即した文章を上位の検索結果として得られることが分かりました。

また、 Bedrock Knowledge Base の利点は、自前で実装することなく、設定を変更するだけで大きな改善をすることができることだと思います。

今回は retrieve までしか行いませんでしたが、 retrieve_and_generate を使えば、回答の生成まで Bedrock にお任せすることも可能です。

Bedrock のリランク機能を活用して、ユーザーの意図に寄り添った RAG システムを開発していきたいです。

Acroquest Technologyでは、キャリア採用を行っています。

• Azure OpenAI／Amazon Bedrock等を使った生成AIソリューションの開発
• ディープラーニング等を使った自然言語／画像／音声／動画解析の研究開発
• マイクロサービス、DevOps、最新のOSSやクラウドサービスを利用する開発プロジェクト
• 書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長

少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。 www.wantedly.com