RAGの精度向上！ CohereからRerank 3 Nimbleが登場。SageMakerでリランク試してみよう

# Cohereの最新リランク用モデルが登場！

生成AIで有名なCohere社から、リランク用途の最新モデル「Rerank 3 Nimble」がリリースされました。英語版と多言語版が用意されています。

https://cohere.com/blog/rerank-3-nimble

※ちなみにNimbleは「機敏な」という意味で、Turbo的な命名と捉えておけばよさそうです。

クラウドではAWSのみ対応。SageMaker JumpStartから利用可能です。


# リランクって何だっけ？

生成AIでよく行われる「社内文書検索」などのユースケースを実現するRAG（Retrieval-Augmented Generation）アーキテクチャにおいて、検索精度を高めるために使われる手法の一つです。

検索して得られたドキュメントのチャンク（かたまり）複数を、ユーザーの質問と関連度が近い順番に並び替える（その後、関連度が低いものは適宜切り捨てる）ことによって、より正確な回答の生成に繋げることができます。

このリランク処理を行うための専用モデルも存在し、その一つがCohereのRerankシリーズです。

https://cohere.com/rerank

ちなみにCohere社はRAGに特に力を入れており、フラッグシップLLMのCommand RシリーズもRAGに特化したモデルとなっています。

有名なRAGの論文の筆頭著者であるパトリック・ルイス氏も、現在はCohereで研究員をされているようです。

https://arxiv.org/abs/2005.11401

なお、AWSでは生成AIモデルをサーバーレスで利用できるAmazon Bedrockというサービスがありますが、残念ながらCohereのRerankシリーズは長らくラインナップされていません。


# AWSで実際に使ってみよう

前述のCohere公式ブログによると、Rerank 3 Nimbleは7/30よりSageMakerで発売予定とされていますが、すでにAWS Marketplaceからサブスクライブ可能でしたので試してみます。

https://aws.amazon.com/marketplace/pp/prodview-ea3rcr6y56jp2

上記をサブスクライブして、設定に進みます。

![スクリーン ショット 2024-07-24 に 08.37.23 午前.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1633856/5cfe3d45-390e-05db-489b-dc038728db60.png)

SageMaker用の設定を利用すると、そのままマネジメントコンソールの作成ウィザードに進めます。

![スクリーン ショット 2024-07-24 に 08.45.22 午前.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1633856/debb0604-8186-0af6-25b1-1ed695d4ca2b.png)

エンドポイント（稼働用インスタンス）を設定すると、モデルをデプロイできます。

![スクリーン ショット 2024-07-24 に 08.43.01 午前.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1633856/cde5b324-f3c9-1d2d-8f94-5dac53d4064e.png)

これを実際にPythonで利用してみましょう。コードは以下を参考にします。

https://github.com/cohere-ai/cohere-aws/blob/rerank_v3_samples/notebooks/sagemaker/rerank_v3_notebooks/Deploy%20rerank%20multilingual%20v3.0%20model.ipynb

私は手元のMacでVSCodeから実施しました。

```py: nimble.py
from cohere_aws import Client

co = Client(region_name="ap-northeast-1")
co.connect_to_endpoint(endpoint_name="ここにエンドポイント名を入れる")

documents = [
    {"Title":"カブトムシ","Content":"カブトムシは昆虫です"},
    {"Title":"人間","Content":"人間は哺乳類です"},
    {"Title":"サンマ","Content":"サンマは魚です"},
    {"Title":"ロボット","Content":"ロボットは機械です"},
    {"Title":"小麦","Content":"小麦は植物です"},
]

response = co.rerank(
    documents=documents, 
    query='この中で食べられるものはどれ？', 
    rank_fields=['Title','Content'], 
    top_n=5 #ドキュメントの上位何件をピックアップするか指定
)

print(response)
```

実行前に以下ライブラリのインストールが必要です。`distutils` が無いと怒られますが、Claude先生に聞いたところPython 3.12では `setuptools` をインストールすれば良いようでした。

```zsh
pip install cohere_aws setuptools
```

実行すると、クエリーに従ってドキュメントをいい感じに並び替えてくれました！5つの文書が、食べられそうな順にちゃんとリランクされていますね。

```zsh
[
    RerankResult<document: {'Title': '小麦', 'Content': '小麦は植物です'}, index: 4, relevance_score: 0.00075835997>, 
    RerankResult<document: {'Title': 'サンマ', 'Content': 'サンマは魚です'}, index: 2, relevance_score: 0.00031999825>, 
    RerankResult<document: {'Title': 'カブトムシ', 'Content': 'カブトムシは昆虫です'}, index: 0, relevance_score: 0.00024346047>, 
    RerankResult<document: {'Title': '人間', 'Content': '人間は哺乳類です'}, index: 1, relevance_score: 6.352189e-05>, 
    RerankResult<document: {'Title': 'ロボット', 'Content': 'ロボットは機械です'}, index: 3, relevance_score: 3.9751925e-05>
]
```

※SageMakerのエンドポイントはランニングコストが高くつきがちなので、検証後すぐに削除することを忘れないでください。