[アップデート] Amazon Bedrockで新モデル「Cohere Command R/R+」が利用可能になったので、RAGで使ってみた

Command Rの「R」はRAGの「R」！！ (かどうかは知りません)

#Amazon Bedrock

#生成AI

#Retrieval Augmented Generation

#RAG

#AWS

青柳英明

2024.05.01

みなさん、こんにちは！
福岡オフィスの青柳です。

Amazon Bedrockの基礎モデル (FM) で、新しいモデル「Cohere Command R」「Cohere Command R+」が利用可能になりました。

Cohere Command R and Command R+ now available in Amazon Bedrock

どのようなものなのか、さっそく使ってみました。

Cohere Command R/R+ってどんなモデル？

リリースノートでは、次のように説明されています。

Command R+ は、RAG やマルチステップツールの使用など、長いコンテキストのタスク向けに最適化された Cohere の最も強力な生成言語モデルです。

Command R は、RAG やツールなどの長いコンテキストのタスクや大規模な運用ワークロード向けに最適化された生成言語モデルです。

Command R+とCommand Rの違いがよく分かりませんが、どちらも「RAG」や「ツール」(自律エージェントで使われる) での利用に向いているようですね。長いコンテキストのタスク向けということで、長文の翻訳や要約なんかも得意そうです。

サポートするトークン長は「最大128K (128,000)」と、従来のCohere Command (最大4,000) から大幅に増加しています。 (と言っても、Claude 3やGPT-4 Turboと同レベルなのですが)

Command R+はCommand Rの強化版であるようです。その分、料金はCommand R+の方がお高めです。

なんだかいろいろなことができそうな予感・・・

Command R/R+のリリースを受けて「じゃあ使ってみよう」と思いドキュメントを見て驚いたのが、、、呼び出す際に指定できるパラメーターの多様さです。

Cohere Command R and Command R+ models - Amazon Bedrock

{
    "message": string,
    "chat_history": [
        {
            "role":"USER or CHATBOT",
            "message": string
        }
  
    ],
    "documents": [
        {"title": string, "snippet": string},
    ],
    "search_queries_only" : boolean,
    "preamble" : string,
    "stream" : boolean,
    "max_tokens": int,
    "temperature": float,
    "p": float,
    "k": float,
    "prompt_truncation" : string,
    "frequency_penalty" : float,
    "presence_penalty" : float,
    "seed" : int,
    "return_prompt" : boolean,
    "tools" : [
        {
            "name": string,
            "description": string,
            "parameter_definitions": {
                "parameter name": {
                    "description": string,
                    "type": string,
                    "required": boolean
                }
            }
        }
    ],
    "tool_results" : [
        {
            "call": {
                "name": string,
                "parameters": {
                "parameter name": string
                }
            },
        "outputs": [
                {
                "text": string
                }
            ]
        }
    ],
    "stop_sequences": [string],
    "raw_prompting" : boolean

}

chat_historyなんか気になりますし、他にもいろんなことができそうな予感がします。

一度に全部を使いこなすのは無理なので、まずはリリースノートにも謳われている「RAG」を試してみることにします。

Cohere Command R/R+を使ったRAGの処理概要

Cohere社の公式ドキュメントに解説がありますので、こちらに沿って確認していきます。

https://docs.cohere.com/docs/retrieval-augmented-generation-rag

Command R/R+を使ったRAGは、以下の3ステップで行われます。

Step 1:: ユーザー入力テキストを基に検索クエリーを生成する
Step 2:: Retrieverに対してクエリーを実行して検索結果を得る (Retrieve)
Step 3:: ユーザー入力テキストと検索結果を基に応答テキストを生成する (Generate)

これらのステップのうち、Command R/R+を使用するのは「Step 1」と「Step 3」です。

Command R/R+が特徴的なのは、「Step 1」を行う機能を標準で備えているという点ではないかと思います。

Retrieverから情報を検索するためのクエリーを用意する際、通常はユーザー入力テキストをそのまま用いたり、クエリをLLMに考えさせたりする手法があると思います。「良い検索結果を得られるクエリ」を用意するのはなかなか難しいと思いますが、Command R/R+では検索クエリーを生成させるための専用のオプションが用意されているのです。

なお、「Step 2」は、Command Rの機能を使用するのではなく、一般的なRetrieverを使って行います。

例えば、マネージドなサービスである「Amazon Kendra」「Vertex AI Search」「Azure AI Search」などを使用する方法もありますし、生成AIのEmbeddingモデルとベクトルデータベースを組み合わせる方法もあります。

やってみる

それでは、実際にCommand Rを使ってRAGの処理を行ってみます。

今回、「Step 2」で使用する「Retriever」には、準備の簡略化のため「Knowledge bases for Amazon Bedrock」で作成したナレッジベースを使うことにします。

ナレッジベースの準備

ナレッジベースのデータソースとして、AWS公式ドキュメントの「Amazon S3ユーザーガイド」と「Amazon EFSユーザーガイド」を用います。

https://docs.aws.amazon.com/ja_jp/AmazonS3/latest/userguide/s3-userguide.pdf
https://docs.aws.amazon.com/ja_jp/efs/latest/ug/efs-ug.pdf

ダウンロードしたPDFファイルをS3バケットに配置して、ナレッジベースのデータソースに指定します。

※ 2024年4月のアップデートにより、ナレッジベースに複数のデータソースを指定できるようになっています。具体的な操作手順は下記ブログ記事を参照してください。

Step 1：ユーザー入力テキストを基に検索クエリーを生成する

それでは、RAGの処理を行うプログラムコードを作成していきます。

まず、ユーザー入力テキストから検索クエリーを生成するところまでです。

import boto3
import json

bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")
agent = boto3.client("bedrock-agent-runtime", region_name="us-east-1")

input_text = "S3とEFSの暗号化方式を比較してください。"


# Step 1: ユーザー入力テキストを基に検索クエリーを生成する

request_body = json.dumps({
    "message": input_text,
    "search_queries_only": True,
})

response = bedrock.invoke_model(
    body=request_body,
    contentType="application/json",
    accept="application/json",
    modelId="cohere.command-r-v1:0",
)

response_body = json.loads(response.get("body").read())
search_queries = response_body.get("search_queries")

リクエストボディのパラメーターsearch_queries_onlyにTrueをセットしてモデル (Command R) を呼び出すことにより、通常の「入力に対する応答生成」の動作でなく、「検索クエリー」を生成する動作が行われます。

レスポンスボディは以下のようになります。

{
    "chat_history": [],
    "finish_reason": "COMPLETE",
    "generation_id": "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX",
    "is_search_required": True,
    "response_id": "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX",
    "search_queries": [
        {
            "generation_id": "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX",
            "text": "S3 encryption"
        },
        {
            "generation_id": "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX",
            "text": "EFS encryption"
        }
    ],
    "text": ""
}

search_queriesにクエリー文字列が格納されています。今回のユーザー入力テキスト「S3とEFSの暗号化方式を比較してください」からは、「S3 encryption」「EFS encryption」という2つの検索クエリーが生成されました。

Step 2： Retrieverに対してクエリーを実行して検索結果を得る (Retrieve)

次に、Step 1で得られた「クエリー文字列」を使って、Retriever (今回はKnowledge bases for Amazon Bedrockのナレッジベース) に対して検索を行います。

# Step 2: Retrieverに対してクエリーを実行して検索結果を得る

documents = []

for query in search_queries:
    query_text = query.get("text")
    response = agent.retrieve(
        knowledgeBaseId="XXXXXXXXXX",
        retrievalQuery={"text": query_text},
        retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": 5}},
    )
    retrieval_results = response.get("retrievalResults")

    for result in retrieval_results:
        result_text = result.get("content").get("text")
        documents.append({"title": query_text, "snippet": result_text})

Knowledge bases for Amazon Bedrockでは、RAGを実行するretrieve_and_generateAPIの他に、検索 (Retrieve、つまりRAGの「R」) のみを実行するretrieveAPIが用意されているので、それを使います。

得られたクエリー文字列 (今回は2つ) のそれぞれを使って、ナレッジベースに対して検索を実行します。その際、検索結果を5つ取得するようにパラメーターで指定します。 ("numberOfResults": 5)

documentsという空の配列変数を用意しておき、検索結果を格納していきます。
検索結果1件につき、以下の情報を格納します。

title：検索で使用したクエリ文字列をセット
snippet：検索結果の内容をセット

Step 2を実行すると、documentsの中身は以下のようになることが期待されます。

[
    {
        "title": "S3 encryption",
        "snippet": "＜「S3 encryption」の検索結果・その1＞",
    },
    {
        "title": "S3 encryption",
        "snippet": "＜「S3 encryption」の検索結果・その2＞",
    },
      ：
    (省略)
      ：
    {
        "title": "EFS encryption",
        "snippet": "＜「EFS encryption」の検索結果・その1＞",
    },
    {
        "title": "EFS encryption",
        "snippet": "＜「EFS encryption」の検索結果・その2＞",
    },
      ：
    (省略)
      ：
]

Step 3：ユーザー入力テキストと検索結果を基に応答テキストを生成する (Generate)

最後に、Step 2の検索結果と、最初に与えられたユーザー入力テキストを使って、モデルに応答テキストを生成してもらいましょう。

# Step 3: ユーザー入力テキストと検索結果を基に応答テキストを生成する

request_body = json.dumps({
    "message": input_text,
    "documents": documents,
})

response = bedrock.invoke_model(
    body=request_body,
    contentType="application/json",
    accept="application/json",
    modelId="cohere.command-r-v1:0",
)

response_body = json.loads(response.get("body").read())
output_text = response_body.get("text")

print(output_text)

リクエストボディは以下のようになっています。

{
    "message": "S3とEFSの暗号化方式を比較してください。",
    "documents": [
        {
            "title": "S3 encryption",
            "snippet": "＜「S3 encryption」の検索結果・その1＞",
        },
          ：
        (省略)
          ：
        {
            "title": "EFS encryption",
            "snippet": "＜「EFS encryption」の検索結果・その1＞",
        },
          ：
        (省略)
          ：
    ]
}

このように、messageと併せてdocumentsを指定してモデル (Command R) を呼び出すと、モデルはdocumentsで与えられた情報に基づいてmessageの指示に回答しようとします。

レスポンスボディは以下のようになります。

{
    "chat_history": [
        {
            "message": "S3とEFSの暗号化方式を比較してください。",
            "role": "USER"
        },
        {
            "message": "Amazon S3では、オブジェクトはAmazon S3マネージドキー(SSE-S3)を使用してサーバー側の暗号化により自動的に暗号化されます。(以下略)",
            "role": "CHATBOT"
        }
    ],
    "citations": [
        {"document_ids": ["doc_0", "doc_1", "doc_3"], "start": 0, "end": 9, "text": "Amazon S3"},
        {"document_ids": ["doc_3"], "start": 19, "end": 35, "text": "Amazon S3マネージドキー"},
        {"document_ids": ["doc_1", "doc_3"], "start": 35, "end": 43, "text": "(SSE-S3)"},
        {"document_ids": ["doc_1", "doc_3"], "start": 48, "end": 57, "text": "サーバー側の暗号化"},
        {"document_ids": ["doc_3"], "start": 60, "end": 67, "text": "自動的に暗号化"},
        {"document_ids": ["doc_3"], "start": 72, "end": 104, "text": "オブジェクトのメタデータでは なく、オブジェクトデータのみが暗号化"},
        {"document_ids": ["doc_1", "doc_3"], "start": 172, "end": 207, "text": "、AWS KMSキーによる二層式サーバー側の暗号化（DSSE-KMS）"},
        {"document_ids": ["doc_1"], "start": 207, "end": 240, "text": "、または顧客提供のキーによるサーバー側の暗号化（SSE-C）です。"},
        {"document_ids": ["doc_0"], "start": 260, "end": 308, "text": "CopyObjectRequestにSSE-S3を指定することでサーバー側で暗号化することができ"},
        {"document_ids": ["doc_3"], "start": 321, "end": 365, "text": "ディスクに保存される前に暗号化され、オブジェクトをダウンロードするときに復号化されます。"},
        {"document_ids": ["doc_5", "doc_8", "doc_9"], "start": 366, "end": 376, "text": "Amazon EFS"},
        {"document_ids": ["doc_5", "doc_7"], "start": 379, "end": 399, "text": "転送中のデータの暗号化を有効にできます。"},
        {"document_ids": ["doc_7"], "start": 419, "end": 497, "text": "Amazon EFSマウントヘルパーを使用してファイルシステムをマウントするときにTransport Layer Security（TLS）を有効にします。"},
        {"document_ids": ["doc_5", "doc_8", "doc_9"], "start": 500, "end": 517, "text": "保管中のデータについても暗号化でき"},
        {"document_ids": ["doc_5"], "start": 518, "end": 595, "text": "AWS Management Console、AWS CLI、またはAmazon EFS APIまたはAWS SDKsのいずれかを使用してプログラムで作成"},
        {"document_ids": ["doc_5"], "start": 612, "end": 618, "text": " 暗号化が有効"},
        {"document_ids": ["doc_5"], "start": 644, "end": 664, "text": "暗号化の設定を変更することはできません。"}
    ],
    "documents": [
        {"id": "doc_0", "title": "S3 encryption", "snippet": "(省略)"},
        {"id": "doc_1", "title": "S3 encryption", "snippet": "(省略)"},
        {"id": "doc_3", "title": "S3 encryption", "snippet": "(省略)"},
        {"id": "doc_5", "title": "EFS encryption", "snippet": "(省略)"},
        {"id": "doc_7", "title": "EFS encryption", "snippet": "(省略)"},
        {"id": "doc_8", "title": "EFS encryption", "snippet": "(省略)"},
        {"id": "doc_9", "title": "EFS encryption", "snippet": "(省略)"},
    ],
    "finish_reason": "COMPLETE",
    "generation_id": "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX",
    "response_id": "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX",
    "text": "Amazon S3では、オブジェクトはAmazon S3マネージドキー(SSE-S3)を使用してサーバー側の暗号化により自動的に暗号化されます。(以下略)"
}

要素を上から見ていきましょう。

chat_historyには、ユーザーとLLMの会話履歴が格納されています。この時点では、ユーザーとLLMは1往復のやり取りしか行っていないため、chat_historyにも1往復分の会話履歴のみが格納されています。 (chat_historyの使い方については、別の機会に試してみたいと思います)

citationsには、応答テキストに含まれる引用についての情報が格納されています。 doc_0、doc_1、・・・は、この後に出てくるdocumentsつまりソースドキュメントの番号に対応しています。その他、引用の開始箇所・終了箇所などの情報が続きます。

documentsには、応答テキストの生成に用いられたソースドキュメントが格納されています。こうやって見ると、Step 2で取得した検索結果は「2つのクエリー × 5つの検索結果」で「計10個」あったはずですが、実際に参照されているのは7個のみであったようです。

そして、最後のtextには、応答テキストが格納されます。

こうして、3つのステップを経て得られた最終的な応答テキストは、以下のようになりました。

Amazon S3では、オブジェクトはAmazon S3マネージドキー(SSE-S3)を使用してサーバー側の暗号化により自動的に暗号化されます。オブジェクトのメタデータではなく、オブジェクトデータのみが暗号化AWS Key Management Service（AWS KMS）キーによるサーバー側の暗号化（SSE-KMS）、AWS KMSキーによる二層式サーバー側の暗号化（DSSE-KMS）、または顧客提供のキーによるサーバー側の暗号化（SSE-C）です。また、オブジェクトをアップロードする際、CopyObjectRequestにSSE-S3を指定することでサーバー側で暗号化することができ、この場合、オブジェクトはディスクに保存される前に暗号化され、オブジェクトをダウンロードするときに復号化されます。 Amazon EFSでは、転送中のデータの暗号化を有効にできます。転送中のデータの暗号化を有効にするには、Amazon EFSマウントヘルパーを使用してファイルシステムをマウントするときにTransport Layer Security（TLS）を有効にします。また、保管中のデータについても暗号化でき、AWS Management Console、AWS CLI、またはAmazon EFS APIまたはAWS SDKsのいずれかを使用してプログラムで作成されたファイルシステムは、作成時に暗号化が有効になります。しかし、一度作成されたファイルシステムの暗号化の設定を変更することはできません。

そのため、Amazon S3の暗号化方式はより柔軟で、オブジェクトデータとメタデータの両方を暗号化できる一方、Amazon EFSはファイルシステムの作成時にのみ設定でき、転送中のデータの暗号化については追加のソフトウェアのインストールが必要になります。