Bedrockのrerankモデル動かしてみる

Amazon Bedrock Rerankモデルの動作検証

はじめに

AWS re:Invent 2024で発表されたAmazon Bedrock Rerankモデルについて動作検証を行いました。RAGアプリケーションの精度向上に重要なリランキングが、このアップデートで大幅に簡単になるんじゃないかなとワクワクしています。

新しく追加された2つのモデル：

• Amazon Rerank 1.0
• Cohere Rerank 3.5

リランキングが必要な背景

RAGパイプラインでは、以下の課題に直面することがあります：

1. ベクトル検索の限界

• クエリと文書のベクトル類似度が高くても意味的関連性が低いケース
• 複数の観点を含むクエリへの部分的な一致

2. 検索結果の最適化ニーズ

• コンテキストを考慮した順位付けの必要性
• ユーザー意図に沿った優先順位付け

Rerankモデルは、これらの課題に対して検索結果を再評価し、より適切な順序に並べ替えます。

Rerankモデルの利点

1. 精度向上

従来のベクトル検索に比べ、クロスエンコーディングを用いることで文書間の関連性をより詳細に評価できます。単純な単語の一致だけでなく、文脈やニュアンスを考慮した検索が可能になります。

2. 開発効率

これまでリランキング処理は独自実装が必要でしたが、APIベースで簡単に組み込めるようになりました。複雑なロジックの実装や保守が不要になり、開発リソースを他の機能に振り向けることができます。

3. スケーラビリティ

マネージドサービスとして提供されるため、インフラ管理の手間がかかりません。また、Bedrockの他の機能（ベクトル検索やナレッジベース）と組み合わせることで、より高度なRAGパイプラインを構築できます。
これらの利点により、RAGアプリケーションの品質向上とリリースまでの時間短縮が期待できます。

動作検証

データセット

• データ：fujiki/japanese_hh-rlhf-49k（Hugging Face）
• サンプル：検索対象200件、クエリ10件
• 上位3位に正解が入っているかを確認

実装コード(GoogleColab)

import boto3
import json
import pandas as pd
from google.colab import userdata
os.environ['AWS_ACCESS_KEY_ID'] = userdata.get('AWS_ACCESS_KEY_ID')
os.environ['AWS_SECRET_ACCESS_KEY'] = userdata.get('AWS_SECRET_ACCESS_KEY')
os.environ['AWS_DEFAULT_REGION'] = userdata.get('AWS_DEFAULT_REGION')
client = boto3.Session(region_name=os.environ['AWS_DEFAULT_REGION']).client("bedrock-runtime")

df = pd.read_parquet("hf://datasets/fujiki/japanese_hh-rlhf-49k/data/train-00000-of-00001-157934b4864eb8e0.parquet")
df['index'] = df['index'].astype(int)

# データ準備(データは200件。クエリ用に10件)
df_sample = df.sample(n=200).reset_index(drop=True)
df_sample2nd = df_sample.sample(n=10).reset_index(drop=True)

queries = df_sample2nd["instruction"].tolist()
sources = df_sample["output"].tolist()

# Amazon Rerank モデルによる検証
correct = 0
total = 0

for i, query in enumerate(queries):
    query_index = df_questions["index"][i]
    
    response = client.invoke_model(
        modelId="amazon.rerank-v1:0",
        body=json.dumps({
            "query": query,
            "documents": sources,
            "top_n": 3,
        })
    )
    
    response_body = json.loads(response.get("body").read())
    top_indices = df_sources.loc[[result["index"] for result in response_body["results"]]]["index"].tolist()
    is_correct = query_index in top_indices
    if is_correct:
        correct += 1
    total += 1

print(f"Accuracy: {correct/total:.2%}")

#出力の一部
Query 1:
Q: タイソンフーズカンパニーのことを聞いたことがありますか？
A: はい、私は実際にタイソンフーズカンパニーのことを聞いたことがあります。大規模な食品製造および加工会社であり、スミスフィールドフーズに次いで2番目に大きい肉会社です。鶏肉、牛肉、豚肉などの食料品を生産する米国には、多くのタイソンフーズ植物があります。タイソンフーズは、世界中で約10万人の従業員を雇用しています。
A: 確かに、私は彼のことを聞いたことがあります。
A: うーん、これは非常に興味深い作業のように聞こえます！心に留めておくべき主なことは、これがどんな秘密の社会であり、メンバーの動機と目標が何であるかだと思います。役立つかもしれないいくつかのアイデアがあります：

 - 社会はスパイのグループであり、秘密の活動に従事し、秘密の知性を互いに伝えます
 - 社会は、政治的/経済的影響を調整し、前進させるために召集する裕福で強力な人々のエリートクラブになる可能性があります
 - 社会は、既存の政治的/経済的秩序の転覆を計画する不吉な計画者の陰謀である可能性があります
 - 社会は、凶悪な犯罪を犯すことを陰謀している悪役の乗組員かもしれません
 - 社会は魔法使いや魔術師のカバルであり、いくつかのオカルトや魔法の計画を計画している可能性があります
 - 社会は、最先端のハイテクプロジェクトを開発している科学者/宇宙飛行士/エンジニアのグループになる可能性があります

これらのアイデアが役立つことを願っています！あなたの物語を書くことで頑張ってください！
Top 1 Score: 0.9967013487421672
Correct_index: 34190
Result_index: [34190, 5771, 21974]
Correct: True

[40619, 38219, 37121]
Query 2:
Q: ロケットとジェットエンジンの違いは何ですか？
A: ジェットエンジンとロケットは、両方ともスペースを通してオブジェクトを推進します。しかし、ジェットエンジンは燃焼の反応に依存していますが、ロケットは爆発物の反応に依存しています。
A: 良い質問！エンジンは、さまざまな可動部品をスムーズに動作させるために潤滑剤を必要とし、特定のエンジンの機能と部分に応じて、適切なタイプの潤滑剤が必要です。異なるエンジンシステムは、さまざまな理由で異なるオイルを使用しています。たとえば、自動車エンジンのオイルは、エンジンの可動部品と厳しい動作環境向けに特別に設計された添加剤を配合する必要があります。ジェット機やバックホーマシンなどの一部のエンジンでは、極端な性能が必要なため、使用する潤滑剤は、効率を最大限にし、パフォーマンスを向上させるために非常に薄くする必要があります。航空機では、エンジンの問題が発生した場合に火災を避けるために、薄い油も非常に良好な耐熱性を持つ必要があります。
A: 理解できない。
Top 1 Score: 0.9858276338086983
Correct_index: 40619
Result_index: [40619, 38219, 37121]
Correct: True

correct = 0
total = 0

for i, query in enumerate(queries):
    query_index = df_questions["index"][i]
    
    response = client.invoke_model(
        modelId="cohere.rerank-v3-5:0",
        body=json.dumps({
            "query": query,
            "documents": sources,
            "top_n": 3,
             "api_version": 2
        })
    )
    
    response_body = json.loads(response.get("body").read())
    top_indices = df_sources.loc[[result["index"] for result in response_body["results"]]]["index"].tolist()
    print(top_indices)
    
    is_correct = query_index in top_indices
    if is_correct:
        correct += 1
    total += 1
    
    # 結果表示
    print(f"Query {i+1}:")
    print(f"Q: {query}")
    print(f"A: {df_sources.iloc[response_body['results'][0]['index']]['output']}")
    print(f"A: {df_sources.iloc[response_body['results'][1]['index']]['output']}")
    print(f"A: {df_sources.iloc[response_body['results'][2]['index']]['output']}")
    print(f"Top 1 Score: {response_body['results'][0]['relevance_score']}")
    print(f"Correct_index: {query_index}")
    print(f"Result_index: {top_indices}")
    print(f"Correct: {is_correct}\n")

accuracy = correct / total
print(f"Overall Accuracy: {accuracy:.2%}")

Query 1:
Q: タイソンフーズカンパニーのことを聞いたことがありますか？
A: はい、私は実際にタイソンフーズカンパニーのことを聞いたことがあります。大規模な食品製造および加工会社であり、スミスフィールドフーズに次いで2番目に大きい肉会社です。鶏肉、牛肉、豚肉などの食料品を生産する米国には、多くのタイソンフーズ植物があります。タイソンフーズは、世界中で約10万人の従業員を雇用しています。
A: 確かに、私は彼のことを聞いたことがあります。
A: 確かに、ここに行きます：

材料：

•2カップのライスヌードル、半分に壊れます

•大さじ2杯の植物油

•ピーナッツバター大さじ2

•大さじ1杯のオイスターソース

•小さじ1杯のフィッシュソース

•小さじ1杯のライムジュース

•大さじ2杯新鮮または缶詰の竹芽を薄いストリップにカットします

•1⁄2カップの刻んだネギ

•卵1個、軽くbeatられます

•コリアンダー1⁄2カップ、刻んだ

•2カップで調理されたクラブミート

•2カップで調理した鶏肉を調理しました

手順：

1.柔らかくなるまで沸騰したお湯で麺を調理します。排水して脇に置きます。

2.中サイズのボウルに、ピーナッツバター、オイスターソース、魚のソース、ライムジュース、竹の芽、ネギを混ぜます。beatられた卵を加えて、よくかき混ぜて混ぜます。麺を加え、均等にコーティングされるまで静かに投げます。

3.サービングプラッターの上に麺を置き、クラブミートと鶏肉をトップにします。コリアンダーを飾ります。
Top 1 Score: 0.9253988
Correct_index: 34190
Result_index: [34190, 5771, 34632]
Correct: True

1問目の第3位は両モデルで違う結果でした。

結果

• このケースでは両モデルとも100%の正解率？？（どこかコード間違ってるかも）
• cohere.rerank-v3-5:0は特に応答速度が速い印象

考察

高い正解率の要因：

1. 限定的なサンプルサイズだったこと
2. クエリが検索対象から抽出され、必ず正解が含まれていること
3. データセットの関連性が明確すぎて簡単に正解できる

実務データではより複雑な結果が予想されますが、少なくとも基本的な機能は期待通りに動作することが確認できました。今後、より実践的な条件での検証を行っていきたいと思います。

参考にさせていただきました