Amazon BedrockでRAGに特化したモデル Command R/R+を試す

Tsuyoshi Yamasaki

2024年6月27日 04:02

はじめに

2024年4月､Amazon BedrockでCohere Command R/R+が利用可能になりました｡

AWSのリリース文によるとこのように書かれています｡

ハルシネーションを軽減するための引用を用いた高度な検索拡張生成 (RAG)、複雑なビジネスタスクを自動化するための複数ステップのツールの使用など、長いコンテキストのタスク向けに最適化されています。

https://aws.amazon.com/jp/about-aws/whats-new/2024/04/cohere-command-r-r-plus-amazon-bedrock/

Command R/R+はRAGに強く､ハルシネーションを軽減する仕組みがあるようです｡RAGの他にToolという機能があるようですが､ここではRAGを試してみます｡
具体的な実装例は､末尾にあるQiitaのリンク先を参照してください｡

Command R/R+のRAG

Cohereの公式ドキュメントによると､Command R/R+のRAG処理は3つのステップがあると書かれています｡

Three steps of RAG
The RAG workflow generally consists of 3 steps:
- Generating search queries for finding relevant documents.
- Fetching relevant documents from an external data source using the generated search queries.
- Generating a response with inline citations using the fetched documents.

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

つまり､このようなステップを踏みます｡

関連するドキュメントを見つけるための検索クエリを生成
生成された検索クエリを使用して、外部データソースから関連ドキュメントを取得
取得したドキュメントを使用してインライン引用を含む応答を生成

3つのステップのうち1番目と3番目にCommand R/R+のAPIを利用します｡2番目は､今回はKnowledge Bases for Amazon Bedrockから取得してみます｡

Step 1: 関連するドキュメントを見つけるための検索クエリを生成

以下は､Amazon BedrockのユーザーガイドにあるCohere Command R and Command R+ modelsのリクエストパラメータを抜粋したものです｡最初のステップでは､search_queries_onlyパラメータに "true"を設定します｡これにより､レスポンスボディのsearch_queriesに検索クエリが格納されます｡

bedrock = boto3.client(service_name='bedrock-runtime', region_name="us-east-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["search_queries"]

検索クエリとは具体的にどのようなものでしょうか｡たとえば､｢Command R/R+とはどのようなモデルですか?｣という質問をモデルに渡すと､｢Command R/R+ AI language model 特徴機能｣のようにいくつかの単語をスペースで区切った内容が検索クエリとして返ってきます｡Googleなどの検索エンジンでなにかを検索するときに入力する文字列と似ています｡

他のモデルで同様の検索クエリを生成しようとすると､プロンプトを工夫したり､レスポンスを加工して検索クエリとして整えるなど手間がかかるかもしれません｡Command R/R+は検索クエリを生成するための単純なオプションが用意されている点で､他のモデルにない特徴を備えています｡

Step 2: 生成された検索クエリを使用して、外部データソースから関連ドキュメントを取得

Step1で生成した検索クエリを使用して外部データソース (ここでは､Knowledge Bases for Amazon Bedrock)から関連ドキュメントを取得します｡

for query in search_queries:
    query_text = query["text"]
    response = agent.retrieve(
        knowledgeBaseId=knowledge_base_id,
        retrievalQuery={
            "text": query_text
        },
        retrievalConfiguration={
            "vectorSearchConfiguration": {
                "numberOfResults": 5
            }
        },
    )

    retrieval_results = response['retrievalResults']

    for result in retrieval_results:
        documents.append(
            {
                "title": query_text,
                "snippet": result['content']['text'],
            }
        )

for文にあるsearch_queriesは､Step1で生成した検索クエリが格納されています｡たとえば､Step1で､｢Command R/R+ AI language model 特徴機能｣｢Command R/R+ モデル自然言語処理 NLP｣という2種類の検索クエリが返ってきた場合､Knowledge Baseに対して検索を2回実行します｡retrievalConfigurationで"numberOfResults": 5としているので､2 * 5 = 10件の検索結果を得ます｡そして､documentsに検索結果を格納します｡

Amazon Bedrockのドキュメントによると､documentsは"A list of texts that the model can cite to generate a more accurate reply. (より正確な返信を生成するためにモデルが引用できるテキストのリスト)"と書かれています｡モデルはここに格納されているテキストを引用して最終的な回答を生成します｡良以下はdocumentsのフォーマット例です｡厳密に項目が定められているわけではないため､モデルの回答生成を手助けする情報としてURLやAuthorなど様々なメタデータを追加することができます｡

"documents": [
    {
        "title": "Tall penguins", 
        "snippet": "Emperor penguins are the tallest."
    },
    {
        "title": "Penguin habitats",
        "snippet": "Emperor penguins only live in Antarctica."
    }
]

Knowledge Bases for Amazon Bedrockはメタデータフィルタリングをサポートしており､ソースデータ毎にメタデータを設定することでretrieveの際にそれらを取得することができます｡

Step3: 取得したドキュメントを使用してインライン引用を含む応答を生成

Step2で作成したdocumentsに加え､ユーザーが入力した質問(input_text)とシステムプロンプトでリクエストボディを作成してモデルを呼び出します｡Step1と異なり､search_queries_onlyはfalseを設定します｡デフォルトがOffなので､search_queries_onlyは省略可能です｡

system_instruction = """
あなたは親切で知識豊富なチャットアシスタントです。
あなたにはユーザーが知りたい情報に関連する複数のドキュメントの抜粋が提供されます｡
これらの情報をもとに、ユーザーの質問に対する回答を提供してください。
"""

request_body = json.dumps({
    "preamble": system_instruction,
    "message": input_text,
    "documents": documents,
    "search_queries_only": false,
})


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)

モデルはdocumentsの内容に基づいて､質問に対する回答と引用情報を出力します｡引用情報は以下のようなレスポンスで示されます｡

{
  "text": "The tallest penguins, Emperor penguins, live in Antarctica.",  
  "citations": [  
    {"start": 22, "end": 38, "text": "Emperor penguins", "document_ids": ["doc_0"]},  
    {"start": 48, "end": 59, "text": "Antarctica.", "document_ids": ["doc_1"]}  
  ],  
  "documents": [  
    {"id": "doc_0", "title": "Tall penguins", "snippet": "Emperor penguins are the tallest.", "url": ""},  
    {"id": "doc_1", "title": "Penguin habitats", "snippet": "Emperor penguins only live in Antarctica.", "url": ""}  
  ] 
}

実装次第で以下のような整形が可能です｡

The tallest penguins, Emperor penguins[0], live in Antarctica[1].
[0] Tall penguins , http://…
[1] Penguin habitats , http://…

実装例

上記の内容をもとに､コマンドラインベースの回答生成やSlackチャットボットを実装しました｡Command R/R+を利用する際の参考になれば幸いです｡

参考情報

今回の内容は以下のページを参考にしました｡その他､実装の際に参考となったサイトについては､Qiitaの記事にリンクを掲載しています｡

まとめ

Command R/R+を使ってRAGを試してみました。search_queries_onlyオプションの設定のみで容易に検索クエリが生成可能なため､Advanced RAGにも応用できそうです｡検索クエリー生成はCommand R/R+を使い､応答テキスト生成はモデルの性能や得意分野を考慮して他のモデルを使うなどの組み合わせも考えられます｡その場合､たとえばbotocore clientのretrieve_and_generateではレスポンスで引用情報が返されるためStep3の処理に代用できそうです｡

この記事が気に入ったらサポートをしてみませんか？