🧠

Amazon Bedrock Knowledge BasesでCohereのRerank試してみた

そのだ

2024/12/02に公開

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

事前準備

boto3のバージョンを1.35.72に更新する
CohereのRerankモデルを有効化する

実際に使ってみた

今回はretrieve APIを使用してRerankしてみます。ドキュメント通りではありますが、retrieveの項目にrerankingConfigurationを設定するだけでRerank機能を追加できます。

import boto3

region = boto3.Session().region_name
client = boto3.client('bedrock-agent-runtime',region_name=region)

modelId = "cohere.rerank-v3-5:0"
model_package_arn = f"arn:aws:bedrock:{region}::foundation-model/{modelId}"

query = 'AWSの事例について教えてください。' # 質問

response = client.retrieve(
  knowledgeBaseId='KBのID',
  retrievalConfiguration={
    'vectorSearchConfiguration': {
      'numberOfResults': 5,
      'overrideSearchType': 'SEMANTIC',
      'rerankingConfiguration': {
        'bedrockRerankingConfiguration': {
          "modelConfiguration": {
            "modelArn": model_package_arn
          },
        },
        'type': 'BEDROCK_RERANKING_MODEL',
      },
    }
  },
  retrievalQuery={
    'text': query
  },
)

for r in response['retrievalResults']:
    print(r['score'], r['location']['webLocation']['url'])

結果は以下のとおりです。弊社の公開事例について質問し、Rerankなしとありで結果が違うことがわかると思います。

# Rerankあり
0.5588332414627075 https://fusic.co.jp/works/58
0.4263460636138916 https://fusic.co.jp/stories/4
0.3762110471725464 https://fusic.co.jp/works/48
0.3417148292064667 https://fusic.co.jp/works/23

# Rerankなし
0.51292914 https://fusic.co.jp/works/48
0.506749 https://fusic.co.jp/stories/4
0.50632834 https://fusic.co.jp/works/58
0.5046154 https://fusic.co.jp/works/23

rerankingConfigurationではRerankに使用するメタデータをカスタマイズできます。selectionModeをSELECTIVEに変更し、selectiveModeConfigurationで使用するメタデータを選択したり、除外することで、より適切なRerankの結果を得ることができます。試しにWebクローリングしてるKBでurlのメタデータを削除してみました。

response = client.retrieve(
  knowledgeBaseId='KBのID',
  retrievalConfiguration={
    'vectorSearchConfiguration': {
      'numberOfResults': 5,
      'overrideSearchType': 'SEMANTIC',
      'rerankingConfiguration': {
        'bedrockRerankingConfiguration': {
          "modelConfiguration": {
            "modelArn": model_package_arn
          },
          "metadataConfiguration": {
            "selectionMode": "SELECTIVE",
            "selectiveModeConfiguration": {
              'fieldsToExclude': [
                {'fieldName': 'location'} # locationを除外
              ]
            }
          }          
        },
        'type': 'BEDROCK_RERANKING_MODEL',
      },
    }
  },
  retrievalQuery={
    'text': query
  },
)

Rerankの結果は変わりませんが、メタデータ（url）を除外したときの方がスコアは上がりました。

# Rerankあり メタデータ（url）を除外
0.6044228076934814 https://fusic.co.jp/works/58
0.4921643137931824 https://fusic.co.jp/stories/4
0.46410053968429565 https://fusic.co.jp/works/48
0.4379923343658447 https://fusic.co.jp/works/23

# Rerankあり
0.5588332414627075 https://fusic.co.jp/works/58
0.4263460636138916 https://fusic.co.jp/stories/4
0.3762110471725464 https://fusic.co.jp/works/48
0.3417148292064667 https://fusic.co.jp/works/23

推測にはなってしまいますが、おそらくデフォルトでは全てのメタデータをRerankに使用しており、Textに関係のない不要なものを削除したため、スコアが上がったと思います。

まとめ

retrieve APIの結果をわざわざ整形してRerankしなくとも、rerankingConfigurationを追加するだけでCohereのRerankモデルを使用することができました。KBを使用している人は多いと思うので、ぜひ使ってみてください。

Fusic 技術ブログPublication

さまざまな個性を受け入れて有機的につなぐ社内環境を整える。あらゆる事業機会の創出と実現を繰り返し、世の中に対する視点を絶えず増やして成長していく。あっと驚くような角度から発展できるポイントを見つけ、そこにいい感じにフィットする形でテクノロジーを組み込んで、世の中をちょっとずつ、時には大胆にアップデートしつづけていく。

事前準備

実際に使ってみた

まとめ

Discussion