AWS Bedrockでリランキングを行う

Contents

1 はじめに
2 事前準備
- 2.1 モデルアクセス
- 2.2 Bedrockが使えるIAMユーザの作成
3 実装
4 効果

はじめに

AWSのBedRockに待望のリランキングモデルが出ました！

LaravelでAWS-PHP-SDKを利用してみたいと思います。

モデルはamazon.rerank-v1とcohere.rerank-v3-5を試します。

事前準備

モデルアクセス

まずはBedrockのリランキングモデルを使用できるようにしましょう。

Bedrock configurations/モデルアクセスページへ移動し、試したいモデルをクリックしてモデルリクエストしましょう。

Bedrockが使えるIAMユーザの作成

下記のような感じのポリシーを作ります。
Resourceには使いたいモデルを追加してください。BedRockで用意されてる基本的なモデルの場合は、arn:aws:bedrock:{リージョン}::foundation-model/{モデル名}のような感じで記述すると良いようです。

JSON

{
    "Version":"2012-10-17", 
    "Statement":
    [
        {
            "Sid": "Statement1", 
            "Effect": "Allow", 
            "Action": [
                "bedrock:InvokeModel",
                "bedrock:InvokeModelWithResponseStream"
            ], 
            "Resource": [
                "arn:aws:bedrock:ap-northeast-1::foundation-model/cohere.rerank-v3-5:0",
                "arn:aws:bedrock:ap-northeast-1::foundation-model/amazon.rerank-v1:0"
            ]
        }
    ]
}

上記のポリシーを割り当てたIAMユーザを作成し、アクセスキーとシークレットアクセスキーを作成しましょう。

実装

.env, configにアクセスキーを登録

.envとconfigに先ほどのアクセスキーとシークレットキーを登録し、configから呼び出せるようにしておきましょう

composerでAWSのライブラリをインストール

aws用のライブラリをインストールします。(laravelに特化したものもあるようですが、ここではphpの方にしておきます)

ShellScript

composer require aws/aws-sdk-php

リランキングメソッドを作ってみた

PHP

/**
 * テキストの関連度合いをrerankingモデルを通してランキング順に並べる
 * @param string $query クエリに対して
 * @param string[] $documents どのドキュメントが似てるか
 * @param int $limit 上位X件を取得
 * @return array array({index: int, relevance_score: float})
 */
public function reranking(string $query, array $documents, int $limit = 5): array
{
    $client = new BedrockRuntimeClient([
        'version' => '2023-09-30',
        'region' => 'ap-northeast-1',
        'credentials' => [
            'key' => config('登録したアクセスキーの名前'),
            'secret' => config('登録したシークレットアクセスキーの名前'),
        ],
    ]);

    $response = $client->invokeModel([
        'contentType' => 'application/json',
        'body' => json_encode([
            'query' => $query,
            'documents' => $documents,
            'top_n' => $limit,
            'api_version' => 2
        ]),
        'modelId' => 'cohere.rerank-v3-5:0',
    ]);

    $response_body = json_decode($response['body'], true);

    return $response_body['results'];
}

bodyの部分は、モデルカタログの各モデルの下の方に「Usage API リクエスト」というところがあるので、それ通りに投げると良いようです。

ちなみにAWSのリランキングモデルだと…

PHP

$response = $client->invokeModel([
    'contentType' => 'application/json',
    'body' => json_encode([
        'query' => $query,
        'documents' => $documents,
        'top_n' => $limit,
    ]),
    'modelId' => 'amazon.rerank-v1:0',
]);
$response_body = json_decode($response['body'], true);

という感じでした。

効果

RAGの設計にもよりますが、私のシステムでは効果抜群でした！　特にいろんな方法で関連ドキュメントを取得していて、最終的にどれが一番関連があるのか評価するときに本当に大事な工程です。
AWSから呼ぶことで請求がまとまるので便利ですね！

CohereとAWS独自モデルとありますが、コストはCohereが2倍になります。
細かい・大量のデータでの評価・比較はまだしていませんが、体感的にはそんなに差があるようには感じませんでした。
まあ、最終的にABテスト(って言っていいのかな？)が簡単にできるように実装しておくのが良いのではないかと思います！