Cohere Rerank 3.5がBedrockで使用できるようになったので試してみた

Cohere Rerank 3.5がAmazon Bedrockで使用できるようになったので、試してみました。

本ブログではモデルの性能比較はしません。

https://aws.amazon.com/jp/blogs/machine-learning/cohere-rerank-3-5-is-now-available-in-amazon-bedrock-through-rerank-api/

何が嬉しいのか

AWS上でCohereのrerankを使用するには、SageMakerでエンドポイントを作成して使用する必要がありました。しかし、これらの構築なしでBedrockから使用できるようになったので、モデルのアクセス許可するだけで、rerankモデルを使用することができます。また、エンドポイントのインスタンスの料金がかからないので、費用も安くなるのがメリットです。

事前準備

Cohereのモデルアクセスを許可

Bedrockのモデルアクセス画面でCohereのRerank 3.5を許可します。
東京リージョンでもリクエストが可能でした。

Rerankとは

ざっくり解説すると、RAGにおける検索精度を高めるために使用される手法の一つです。クエリの文脈やニュアンスを考慮しながら、各ドキュメントとクエリの関連性を評価します。これにより、ベクトルを使った事前計算された単純な類似度だけでなく、クエリの意図や文脈に応じたスコアを付けることが可能になり、精度が向上すると言われています。

rerankについて具体的に知りたい方は、以下のドキュメントを参照ください。CohereではなくPineconeの記事ですが、個人的に一番わかりやすかったです。
https://www.pinecone.io/learn/series/rag/rerankers/#References

料金について

料金に関してですが、1000クエリあたり2ドルとのことです。
１クエリに対して最大100個のドキュメントチャンクを含めることができますが、100個を超えると複数クエリに分割されます。例えば、ドキュメントが350個の場合、４クエリとして計算されるそうです。

**You are charged for number of queries where a query can contain up to 100 document chunks. If the query contains more than 100 document chunks, it is counted as multiple queries. For example, if a request contains 350 documents, it will be treated as 4 queries. Please note that each document can only contain upto 500 tokens (inclusive of the query and document’s total tokens), and if the token length is higher than 500 tokens, it is broken down into multiple documents.

https://aws.amazon.com/bedrock/pricing/?nc1=h_ls

実際に使ってみた

ブログ通りにやろうとすると、「bedrock-agent-runtimeはrerankがない！」と怒られたので、マネコンに記載されている方法で使ってみました。

import boto3
import json

modelId = "cohere.rerank-v3-5:0"
client = boto3.client("bedrock-runtime")

query = "アップル製品について教えてください"
documents = [
    "アップルは、その甘さとシャキシャキした食感で知られる人気の果物である",
    "アップルはiPhoneのような革新的な製品で知られている",
    "多くの人がリンゴをヘルシーなスナックとして楽しんでいる",
    "アップル社は、その洗練されたデザインとユーザーフレンドリーなインターフェイスでハイテク業界に革命を起こした",
    "1日1個のリンゴは医者を助ける"
]

response = client.invoke_model(
    body=json.dumps({
        "query": query,
        "documents": documents,
        "top_n": 3, # 上位何個取得するか
        "api_version": 2,
    }),
    modelId=modelId,
)
results = json.loads(response['body'].read())['results']

print(query)
for r in results:
  print(r['relevance_score'], documents[r['index']])

実行結果は以下のとおりです。

0.5634506 アップルはiPhoneのような革新的な製品で知られている
0.528461 アップル社は、その洗練されたデザインとユーザーフレンドリーなインターフェイスでハイテク業界に革命を起こした
0.41478088 アップルは、その甘さとシャキシャキした食感で知られる人気の果物である

いい感じにRerankできましたね。記事によるとKnowledge Base for Amazonでも使用できるそうですが、これはまた別の記事で検証してみます。

まとめ

SageMakerを使わずにBedrockでrerankができるので、簡単にアーキテクチャに組み込める観点からとても便利だなと思いました。実際、待ち望んでた方も多いと思います。せっかくなので、実際のプロジェクトにも適用していきたいです。