OpenSearchの最近のアップデートで、RAGに使えそうな機能が続々と追加されています。
- Hybrid search（2.11～）
- Reranking（2.12～）
- Text chunking（2.13～）
全部使ったら、最強のナレッジベースができるかもと思いまして、試してみました！
「ロングコンテキスト対応 自動チャンク分割 ハイブリッド&リランク検索」ナレッジベースができました
Amazon OpenSearch Serviceとは
Amazon OpenSearch ServiceはオープンソースのOpenSearchをAWSがマネージドなサービスとして提供しているものです。
Amazon OpenSearch ServiceとAmazon OpenSearch Serverlessがありますが、今回扱うのは OpenSearch Service の方です。
今回構築するもの
データ登録と検索のときの処理をパイプラインとして登録します。
-
データ登録時の動作
-
検索時の動作
環境構築
OpenSearch Serviceの構築
AWSのAmazon OpenSearch Serviceを構築します。
画面キャプチャは省きますが、大まかな設定項目はこちらです。
- インスタンスは
m6g.large.search - スタンバイやマルチAZは無効にして1インスタンス構成
- ネットワークは
パブリックアクセス -
きめ細かなアクセスコントロールを有効化 - アクセスポリシーは
きめ細かなアクセスコントロールのみを使用
Python環境の準備
OpenSearchのPythonクライアントであるopensearch-pyを使用します。
pip install boto3 opensearch-py
Pythonのコードで、OpenSearchクライアントを生成します。
きめ細かなアクセスコントロールで指定したIAM権限でのアクセスとなるため、認証情報をセットする必要があります。
from urllib.parse import urlparse
from boto3 import Session
from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth
domain_endpoint = (
"https://*****.ap-northeast-1.es.amazonaws.com"
)
region = "ap-northeast-1"
url = urlparse(domain_endpoint)
service = "es"
credentials = Session().get_credentials()
auth = AWSV4SignerAuth(credentials, region, service)
aos_client = OpenSearch(
hosts=[{"host": url.netloc, "port": url.port or 443}],
http_auth=auth,
use_ssl=True,
verify_certs=True,
connection_class=RequestsHttpConnection,
timeout=120,
)
info = aos_client.info()
print(f"{info['version']['distribution']}: {info['version']['number']}")
OpenSearch Serviceの接続に成功すると、以下の内容が出力されます。
opensearch: 2.13.0
MLモデル
OpenSearchでMLモデルを使う方法には大きく2種類あります。
- Local model : OpenSearch内部でMLモデルを動かす。OpenSearchが事前にトレーニグしたモデルやカスタムモデルが利用可能
- Externally hosted model : OpenSearch外部のMLモデルを呼び出す。SageMaker、Bedrock、OpenAI、Cohereが利用可能
今回は以下の2つのMLモデルを使用します。
| MLモデル | モデル実行形態 | 用途 |
|---|---|---|
| Titan Text Embeddings V2 | Externally hosted model（Bedrock） | セマンティック検索 |
| ms-marco-MiniLM-L-6-v2 | Local model | リランク |
Titan Text Embeddings V2の設定
BedrockのTitan Text Embeddings V2を使用する設定を行います。
-
Connectorを作成
OpenSearchと外部との接続用コネクターを作成します。
OpenSearchからBedrockにアクセスための権限はIAMロールで制御します。Bedrockへアクセス可能なIAMロールを作成します。
次にconnectorを作成します。
opensearch_iam_role = "arn:aws:iam::*****:role/*****" bedrock_titan_model_id = "amazon.titan-embed-text-v2:0" bedrock_titan_region = "us-east-1" body = { "name": "Amazon Bedrock Connector: embedding", "description": "The connector to the Bedrock Titan embedding model", "version": 1, "protocol": "aws_sigv4", "parameters": {"region": bedrock_titan_region, "service_name": "bedrock"}, "credential": {"roleArn": opensearch_iam_role}, "actions": [ { "action_type": "predict", "method": "POST", "url": f"https://bedrock-runtime.{bedrock_titan_region}.amazonaws.com/model/{bedrock_titan_model_id}/invoke", "headers": { "content-type": "application/json", "x-amz-content-sha256": "required", }, "request_body": '{ "inputText": "${parameters.inputText}" }', "pre_process_function": '\n StringBuilder builder = new StringBuilder();\n builder.append("\\"");\n String first = params.text_docs[0];\n builder.append(first);\n builder.append("\\"");\n def parameters = "{" +"\\"inputText\\":" + builder + "}";\n return "{" +"\\"parameters\\":" + parameters + "}";', "post_process_function": '\n def name = "sentence_embedding";\n def dataType = "FLOAT32";\n if (params.embedding == null || params.embedding.length == 0) {\n return params.message;\n }\n def shape = [params.embedding.length];\n def json = "{" +\n "\\"name\\":\\"" + name + "\\"," +\n "\\"data_type\\":\\"" + dataType + "\\"," +\n "\\"shape\\":" + shape + "," +\n "\\"data\\":" + params.embedding +\n "}";\n return json;\n ', } ], } response = aos_client.transport.perform_request( method="POST", url="/_plugins/_ml/connectors/_create", body=body ) titan_connector_id = response["connector_id"] -
MLモデルを登録
作成したコネクターを指定しモデルを登録します。
model_name = "titan-embedded-model" body = { "name": model_name, "function_name": "remote", "connector_id": titan_connector_id, } response = aos_client.transport.perform_request( method="POST", url="/_plugins/_ml/models/_register", body=body ) titan_model_id = response["model_id"]モデルIDはあとで利用します。
-
MLモデルをデプロイ
登録したモデルをデプロイします。
response = aos_client.transport.perform_request( method="POST", url=f"/_plugins/_ml/models/{titan_model_id}/_deploy" )
これでTitan Text Embeddings V2が利用できるようになりました。
ms-marco-MiniLM-L-6-v2の設定
Local modelでMLモデルを動作させる方法は、モデルの登録とデプロイの2ステップですが、前準備として設定変更が必要です。
-
1ノードで動作するように設定変更
OpenSearchのデフォルト設定では、MLモデルは専用ノードで動作させる必要があります。今回は1ノード構成のため、設定を変更します。
aos_client.cluster.put_settings( body={"transient": {"plugins.ml_commons.only_run_on_ml_node": False}} ) -
MLモデルを登録
MLモデルを登録します。
登録（_register）処理はタスクを作成してすぐに完了しますが、次のステップに進むにはタスクの完了を待つ必要があります。
body = { "name": "huggingface/cross-encoders/ms-marco-MiniLM-L-6-v2", "version": "1.0.2", "model_format": "TORCH_SCRIPT", } response = aos_client.transport.perform_request( method="POST", url="/_plugins/_ml/models/_register", body=body ) task_id = response["task_id"] response = aos_client.transport.perform_request( method="GET", url=f"/_plugins/_ml/tasks/{task_id}" ) while True: if response["state"] == "COMPLETED": break import time time.sleep(5) response = aos_client.transport.perform_request( method="GET", url=f"/_plugins/_ml/tasks/{task_id}" ) rerank_model_id = response["model_id"] -
MLモデルをデプロイ
タスクが完了したらデプロイします。
response = aos_client.transport.perform_request( method="POST", url=f"/_plugins/_ml/models/{rerank_model_id}/_deploy" )
これでms-marco-MiniLM-L-6-v2の設定は完了です。
Ingest pipeline
インデックスに登録する際の処理をパイプラインで定義したものがIngest pipelineです。
https://opensearch.org/docs/latest/ingest-pipelines/
今回は「チャンク分割」と「ベクトル化」を行います。
チャンク分割（text_chunking）
チャンク分割はtext_chunkingというプロセッサーを使って行います。
アルゴリズムが「fixed_token_length」（トークン数での分割）と「delimiter」（区切り文字での分割）の2種類から選択が可能です。
fixed_token_lengthを使用する際の指定方法はこのようになります。
{
"text_chunking": {
"algorithm": {
"fixed_token_length": {
"token_limit": 500,
"overlap_rate": 0.2,
"tokenizer": "standard",
}
},
"field_map": {"body": "body_chunk"},
}
}
field_mapでどのプロパティをチャンク分割するか（body）、分割したチャンクをどのプロパティに登録するか（body_chunk）を指定します。
ベクトル化（text_embedding）
ベクトル化はtext_embeddingというプロセッサーを使います。
ベクトル化にはBedrockの（Titan Text Embeddings V2）を使用しますのでモデルIDを指定します。
field_mapでベクトル化対象のプロパティ（body_chunk）とベクトルデータを格納するプロパティ（body_chunk_embedding）を指定します。
{
"text_embedding": {
"model_id": titan_model_id,
"field_map": {"body_chunk": "body_chunk_embedding"},
}
}
Ingest pipelineを作成
チャンク分割とベクトル化を行うIngest pipelineを作成します。
aos_client.ingest.put_pipeline(
id="text-chunking-embedding-ingest-pipeline",
body={
"description": "A text chunking and embedding ingest pipeline",
"processors": [
{
"text_chunking": {
"algorithm": {
"fixed_token_length": {
"token_limit": 500,
"overlap_rate": 0.2,
"tokenizer": "standard",
}
},
"field_map": {"body": "body_chunk"},
}
},
{
"text_embedding": {
"model_id": titan_model_id,
"field_map": {"body_chunk": "body_chunk_embedding"},
}
},
],
},
)
Search pipeline
Ingest pipelineはデータ登録時のパイプラインでしたが、Search pipelineは検索時に使用するパイプラインです。
Search pipelineではハイブリッド検索の「正規化」と検索後の「リランク」、リランク後に上位以外の「切り捨て」を行います。
Search pipelineで使用するプロセッサーには以下の3種類があります。
- Search request processors : 検索リクエストに対する処理を行う
- Search response processors : 検索結果に対する処理を行う
- Search phase results processors : 検索処理の途中で処理を行う
正規化（normalization-processor）
正規化は「Search phase results processors」のnormalization-processorプロセッサーで行います。
設定項目の詳細は把握できていませんが、ドキュメントに従い、設定します。
{
"normalization-processor": {
"normalization": {"technique": "min_max"},
"combination": {
"technique": "arithmetic_mean",
"parameters": {"weights": [0.3, 0.7]},
},
}
}
リランク（rerank）
リランクは「response_processors」のrerankで行います。
使用するMLモデルやリランク時に参照するプロパティ（body_cunk）を指定します。
{
"rerank": {
"ml_opensearch": {"model_id": rerank_model_id},
"context": {"document_fields": ["body_chunk"]},
}
}
切り捨て（truncate_hits）
リランクでは並び替えを行うだけなので、並び替えた後に上位だけを残します。切り捨てはtruncate_hitsを使用します。
response_processorsのrerankのあとに指定します。
{
"truncate_hits": {
"target_size": 3,
}
}
Search pipelineを作成
それではSearch pipelineを作成しましょう。
aos_client.search_pipeline.put(
id="hybrid-rerank-search-pipeline",
body={
"description": "Post processor for hybrid rerank search",
"phase_results_processors": [
{
"normalization-processor": {
"normalization": {"technique": "min_max"},
"combination": {
"technique": "arithmetic_mean",
"parameters": {"weights": [0.3, 0.7]},
},
}
}
],
"response_processors": [
{
"rerank": {
"ml_opensearch": {"model_id": rerank_model_id},
"context": {"document_fields": ["body_chunk"]},
}
},
{
"truncate_hits": {
"target_size": 3,
}
},
],
},
)
インデックスの作成
データを保持するインデックスを作成します。
settings項目
settings項目では以下のような設定を行います。
- データ登録時のデフォルトパイプラインを指定
- Indexに日本語用のアナライザーを設定
- ベクトル検索を行うため、k-NN pluginを有効化
"settings": {
"default_pipeline": "text-chunking-embedding-ingest-pipeline",
"index": {
"analysis": {
"analyzer": {
"custom_kuromoji_analyzer": {
"tokenizer": "kuromoji_tokenizer",
"filter": ["kuromoji_baseform", "ja_stop"],
"char_filter": ["icu_normalizer"],
}
}
},
"knn": True,
},
}
mappings項目
mappings項目では、プロパティの型を定義します。
| プロパティ名 | 型 | その他 |
|---|---|---|
| body | keyword | オリジナルの文字列を格納 |
| body_chunk | text | チャンク分割した文字列を格納。日本語アナライザーを指定 |
| body_chunk_embedding | knn_vector(1024次元) | ベクトル化した値を格納 |
keywordとtextはどちらも文字列を扱う型です。textはアナライザーによる分析が行われますが、keywordは分析を行わないので、完全一致検索のみに使用できます。
検索に使用するのはチャンク分割した文字列のみのため、このような設定としました。
インデックスを作成
インデックスを作成します。
aos_client.indices.create(
index=index_name,
body={
"settings": {
"default_pipeline": "text-chunking-embedding-ingest-pipeline",
"index": {
"analysis": {
"analyzer": {
"custom_kuromoji_analyzer": {
"tokenizer": "kuromoji_tokenizer",
"filter": ["kuromoji_baseform", "ja_stop"],
"char_filter": ["icu_normalizer"],
}
}
},
"knn": True,
},
},
"mappings": {
"properties": {
"body": {"type": "keyword"},
"body_chunk": {
"type": "text",
"analyzer": "custom_kuromoji_analyzer",
},
"body_chunk_embedding": {
"type": "nested",
"properties": {"knn": {"type": "knn_vector", "dimension": 1024}},
},
}
},
},
)
データ登録
検索するデータを登録します。
AWSのWhat's NewのRSSから最新ニュースを20件取ってきます。
import feedparser
import requests
from bs4 import BeautifulSoup
feed_url = "https://aws.amazon.com/jp/about-aws/whats-new/recent/feed/"
feeds = feedparser.parse(feed_url)
entries = feeds["entries"]
targets = []
for entry in entries[:20]: # 20件のみを対象とする
id = entry["id"]
title = entry["title"]
published = entry["published"]
link = entry["link"]
response = requests.get(link)
body = BeautifulSoup(response.text, features="html.parser").main.text
targets.append(
{"id": id, "title": title, "published": published, "link": link, "body": body}
)
target[0]の内容はこのようなものです。
{
"id": "e832c248432005f74b2af6e2fcc44ef29211ad5b",
"title": "Amazon Bedrock のナレッジベースでガードレールの設定が可能に",
"published": "Tue, 28 May 2024 09:06:46 +0000",
"link": "https://aws.amazon.com/jp/about-aws/whats-new/2024/05/knowledge-bases-amazon-bedrock-configure-guardrails/",
"body": "\n\n\n\n\n Amazon Bedrock のナレッジベースでガードレールの設定が可能に \n\n\n\n 投稿日: May 17, 2024 \n\n\nAmazon Bedrock のナレッジベース (KB) は、検索拡張生成 (RAG) のために基盤モデル (FM) を社内データソースに安全に接続することで、より関連性の高い正確な応答を提供します。Amazon Bedrockのガードレールがナレッジベースと統合されたことをお知らせいたします。ガードレールを使用すると、RAG アプリケーションの要件に合わせてカスタマイズされた保護手段と責任ある AI ポリシーを組み込むことができ、エンドユーザーエクスペリエンスの向上につながります。\n\n\n\n\nガードレールは、生成 AI アプリケーションとの望ましくない応答やインタラクションからユーザーを保護するための包括的なポリシーセットを提供します。まず、拒否するトピックのセットをカスタマイズすることで、アプリケーションのコンテキスト内で回避できます。次に、ヘイト、侮辱、性的、暴力、不正行為、扇動攻撃など、あらかじめ用意されている有害なカテゴリ別にコンテンツをフィルタリングできます。第三に、アプリケーション内でブロックする攻撃的で不適切な言葉を定義できます。最後に、機密情報 (個人を特定できる情報など) を含むユーザー入力をフィルタリングしたり、ユースケースに基づいてモデル回答内の機密情報を編集したりできます。ガードレールは、モデルに送信される入力だけでなく、基盤モデルによって生成されたコンテンツにも適用できます。\n\n\n\n\nKnowledge Bases 内のこの機能が、アジアパシフィック (シンガポール)、アジアパシフィック (シドニー)、アジアパシフィック (東京)、欧州 (フランクフルト)、米国東部 (バージニア北部)、米国西部 (オレゴン) の各リージョンで利用できるようになりました。詳細については、Knowledge Bases for Amazon Bedrock のドキュメントをご覧ください。開始するには、Amazon Bedrock コンソールにアクセスするか、RetrieveAndGenerate API を利用してください。\n\n\n\n\n »\n\n\n\n"
}
このデータをインデックスに登録します。
for target in targets:
aos_client.index(index=index_name, body=target)
データ検索
登録したデータを検索します。
登録内容の確認
想定通り登録されているか、1件取得してみます。
response = aos_client.search(
index=index_name,
body={
"query": {"match_all": {}},
},
)
print(response["hits"]["hits"][0]["_source"])
{
"body_chunk": [
"\n\n\n\n\n Amazon Bedrock のナレッジベースで推論パラメータの設定が可能に \n\n\n\n 投稿日: May 17, 2024 \n\n\nAmazon Bedrock のナレッジベース (KB) では、推論パラメーターを設定して、基礎モデル (FM) によって生成される応答のパーソナライズをより細かく制御できるようになりました。 \n\n\n\n\n今回のリリースでは、オプションで推論パラメーターを設定して、基礎モデルによって生成される応答のランダム性や長さなどのパラメーターを定義できます。温度や top-p などのいくつかの設定を調整することで、生成されるテキストのランダム性や多様性を制御できます。温度設定により、モデルが通常とは異なる単語や予期しない単語を選択する可能性が高まったり低くなったりします。温度の値が小さいほど、予想されるより一般的な単語が選択されます。top-p 設定では、モデルが考慮する単語オプションの数を制限します。この数を減らすと、考慮する単語の選択肢の数が少なくなり、出力がより一般的なものになります。\n\n\n\n\nランダム性と多様性に加えて、MaxTokens と stopsequence を通じて、基礎モデルの出力の長さを制限できます。MaxTokens 設定を使用して、生成された応答で返されるトークンの最小数または最大数を指定できます。最後に、stopsequence 設定を使用すると、モデルがそれ以上のトークンの生成を停止するように制御する役割を果たす文字列を設定できます。\n\n\n\n\nKnowledge Bases 内の推論パラメータ機能が、アジアパシフィック (シンガポール)、アジアパシフィック (シドニー)、アジアパシフィック (東京)、欧州 (フランクフルト)、米国東部 (バージニア北部)、米国西部 (オレゴン) の各リージョンで利用できるようになりました。詳細については、Knowledge Bases for Amazon Bedrock のドキュメントをご覧ください。開始",
"を停止するように制御する役割を果たす文字列を設定できます。\n\n\n\n\nKnowledge Bases 内の推論パラメータ機能が、アジアパシフィック (シンガポール)、アジアパシフィック (シドニー)、アジアパシフィック (東京)、欧州 (フランクフルト)、米国東部 (バージニア北部)、米国西部 (オレゴン) の各リージョンで利用できるようになりました。詳細については、Knowledge Bases for Amazon Bedrock のドキュメントをご覧ください。開始するには、Amazon Bedrock コンソールにアクセスするか、RetrieveAndGenerate API を利用してください。\n\n\n\n\n »\n\n\n\n"
],
"body_chunk_embedding": [
{
"knn": [
-0.094595894,
0.039916188,
0.03349132,
0.0185911,
-0.03198763,
...
-0.0051771216,
0.04016191,
-0.011557008
]
}
],
"link": "https://aws.amazon.com/jp/about-aws/whats-new/2024/05/knowledge-bases-amazon-bedrock-configure-parameters/",
"id": "93cca1181913559eee2df0b0525b24652140222e",
"published": "Tue, 28 May 2024 09:08:21 +0000",
"title": "Amazon Bedrock のナレッジベースで推論パラメータの設定が可能に",
"body": "\n\n\n\n\n Amazon Bedrock のナレッジベースで推論パラメータの設定が可能に \n\n\n\n 投稿日: May 17, 2024 \n\n\nAmazon Bedrock のナレッジベース (KB) では、推論パラメーターを設定して、基礎モデル (FM) によって生成される応答のパーソナライズをより細かく制御できるようになりました。 \n\n\n\n\n今回のリリースでは、オプションで推論パラメーターを設定して、基礎モデルによって生成される応答のランダム性や長さなどのパラメーターを定義できます。温度や top-p などのいくつかの設定を調整することで、生成されるテキストのランダム性や多様性を制御できます。温度設定により、モデルが通常とは異なる単語や予期しない単語を選択する可能性が高まったり低くなったりします。温度の値が小さいほど、予想されるより一般的な単語が選択されます。top-p 設定では、モデルが考慮する単語オプションの数を制限します。この数を減らすと、考慮する単語の選択肢の数が少なくなり、出力がより一般的なものになります。\n\n\n\n\nランダム性と多様性に加えて、MaxTokens と stopsequence を通じて、基礎モデルの出力の長さを制限できます。MaxTokens 設定を使用して、生成された応答で返されるトークンの最小数または最大数を指定できます。最後に、stopsequence 設定を使用すると、モデルがそれ以上のトークンの生成を停止するように制御する役割を果たす文字列を設定できます。\n\n\n\n\nKnowledge Bases 内の推論パラメータ機能が、アジアパシフィック (シンガポール)、アジアパシフィック (シドニー)、アジアパシフィック (東京)、欧州 (フランクフルト)、米国東部 (バージニア北部)、米国西部 (オレゴン) の各リージョンで利用できるようになりました。詳細については、Knowledge Bases for Amazon Bedrock のドキュメントをご覧ください。開始するには、Amazon Bedrock コンソールにアクセスするか、RetrieveAndGenerate API を利用してください。\n\n\n\n\n »\n\n\n\n"
}
以下のフォーマットで期待通り登録されています。
| プロパティ名 | 内容 |
|---|---|
| body_chunk | チャンク分割した文字列 |
| body_chunk_embedding |
body_chunkをベクトル化した値 |
| link | RSSのlink |
| id | RSSのid |
| published | RSSのpublished |
| title | RSSのtitle |
| body | ニュースの本文 |
キーワード検索
単純なキーワード検索をしてみましょう
response = aos_client.search(
index=index_name,
body={
"_source": {"exclude": ["body_chunk", "body_chunk_embedding"]},
"query": {"match": {"body_chunk": "ﾅﾚｯｼﾞﾍﾞｰｽ"}},
},
)
print(json.dumps(response["hits"]["hits"][0]["_source"], indent=2, ensure_ascii=False))
{
"link": "https://aws.amazon.com/jp/about-aws/whats-new/2024/05/knowledge-bases-amazon-bedrock-configure-guardrails/",
"id": "e832c248432005f74b2af6e2fcc44ef29211ad5b",
"published": "Tue, 28 May 2024 09:06:46 +0000",
"title": "Amazon Bedrock のナレッジベースでガードレールの設定が可能に",
"body": "\n\n\n\n\n Amazon Bedrock のナレッジベースでガードレールの設定が可能に \n\n\n\n 投稿日: May 17, 2024 \n\n\nAmazon Bedrock のナレッジベース (KB) は、....."
}
日本語用のアナライザーを設定しているので、半角カタカナと全角カタカナを同一の文字として検索ができます。
※ベクトルデータやチャンク分割した文字列は検索結果に不要なのでexcludeを指定します
セマンティック検索
次にセマンティック検索を行います。
チャンク分割後の文字列をベクトル化しているため、「配列の配列」としてベクトルデータが登録されています。
そのため、nestedというタイプの検索キーワードを使用します。
response = aos_client.search(
index=index_name,
body={
"_source": {"exclude": ["body_chunk", "body_chunk_embedding"]},
"query": {
"nested": {
"score_mode": "max",
"path": "body_chunk_embedding",
"query": {
"neural": {
"body_chunk_embedding.knn": {
"query_text": "Guardrails",
"model_id": titan_model_id,
}
}
},
}
},
},
)
print(json.dumps(response["hits"]["hits"][0]["_source"], indent=2, ensure_ascii=False))
{
"link": "https://aws.amazon.com/jp/about-aws/whats-new/2024/05/knowledge-bases-amazon-bedrock-configure-guardrails/",
"id": "e832c248432005f74b2af6e2fcc44ef29211ad5b",
"published": "Tue, 28 May 2024 09:06:46 +0000",
"title": "Amazon Bedrock のナレッジベースでガードレールの設定が可能に",
"body": "\n\n\n\n\n Amazon Bedrock のナレッジベースでガードレールの設定が可能に \n\n\n\n 投稿日: May 17, 2024 \n\n\nAmazon "
}
セマンティック検索なので、「Guardrails」というキーワードは存在しませんが、「ガードレール」が意味的に近いため検索できました。
ハイブリッドリランク検索
ハイブリッドとリランクを組み合わせた検索を行います。
これまでの検索と違い、Search pipelineを使用します。
だんだんJSONが複雑になってきましたが、paramsでSearch pipelineを指定しています。
また、リランク処理を行う際に「どの文字列との関連で並び替えを行うか」の対象文字列をextの部分で指定します。
query = "Guardrail ﾅﾚｯｼﾞﾍﾞｰｽ"
aos_client.search(
index=index_name,
body={
"_source": {"exclude": ["body_chunk", "body_chunk_embedding"]},
"query": {
"hybrid": {
"queries": [
{"match": {"body_chunk": {"query": query}}},
{
"nested": {
"score_mode": "max",
"path": "body_chunk_embedding",
"query": {
"neural": {
"body_chunk_embedding.knn": {
"query_text": query,
"model_id": titan_model_id,
}
}
},
}
},
],
}
},
"ext": {"rerank": {"query_context": {"query_text": query}}},
},
params={"search_pipeline": "hybrid-rerank-search-pipeline"},
)
print(json.dumps(response["hits"]["hits"][0]["_source"], indent=2, ensure_ascii=False))
{
"link": "https://aws.amazon.com/jp/about-aws/whats-new/2024/05/knowledge-bases-amazon-bedrock-configure-guardrails/",
"id": "e832c248432005f74b2af6e2fcc44ef29211ad5b",
"published": "Tue, 28 May 2024 09:06:46 +0000",
"title": "Amazon Bedrock のナレッジベースでガードレールの設定が可能に",
"body": "\n\n\n\n\n Amazon Bedrock のナレッジベースでガードレールの設定が可能に....."
}
RAG化
うまくできたように思うので、最後にBedrockのCohere Command Rと組み合わせてRAGにしてみましょう。
import boto3
import json
client = boto3.client("bedrock-runtime")
model_id = "cohere.command-r-v1:0"
Command Rを使って検索クエリーを生成する関数
def generate_search_queries(question: str):
body = json.dumps(
{
"message": question,
"search_queries_only": True,
}
)
response = client.invoke_model(modelId=model_id, body=body)
response_body = json.loads(response.get("body").read())
return response_body["search_queries"][0]["text"]
ハイブリッドリランク検索を行う関数
def hybrid_rerank_search(keyword_query: str, semanthic_query: str):
response = aos_client.search(
index=index_name,
body={
"_source": {"exclude": ["body_chunk", "body_chunk_embedding"]},
"query": {
"hybrid": {
"queries": [
{"match": {"body_chunk": {"query": keyword_query}}},
{
"nested": {
"score_mode": "max",
"path": "body_chunk_embedding",
"query": {
"neural": {
"body_chunk_embedding.knn": {
"query_text": semanthic_query,
"model_id": titan_model_id,
}
}
},
}
},
],
}
},
"ext": {"rerank": {"query_context": {"query_text": semanthic_query}}},
},
params={"search_pipeline": "hybrid-rerank-search-pipeline"},
)
return list(map(lambda x: x["_source"], response["hits"]["hits"]))
Command Rを使って回答を生成する関数
def generate_response(question: str, documents: dict):
body = json.dumps({"message": question, "documents": documents})
response = client.invoke_model(modelId=model_id, body=body)
response_body = json.loads(response.get("body").read())
return response_body
実行します！
question = "BedrockのKnowledgebaseってどんなサービス？"
search_query = generate_search_queries(question=question)
search_result = hybrid_rerank_search(
keyword_query=search_query,
semanthic_query=question,
)
response = generate_response(
question=question,
documents=search_result,
)
print(response["text"])
Amazon Bedrock のナレッジベース (KB) は、基盤モデル (FM) を社内データソースに安全に接続し、より関連性が高く、正確な応答を提供するサービスです。
最近のアップデートでは、ガードレールの設定や 推論パラメータの設定が可能となり、生成AIアプリケーションの保護や応答のパーソナライズに役立ちます。ガードレールは、AIアプリケーションとのインタラクションからユーザーを保護するため、望ましくない応答や不適切な言葉を避けるのに役立ちます。一方、推論パラメータは、基礎モデルによって生成される応答のランダム性や長さを制御でき、より細かく応答をカスタマイズできます。
それっぽい！！！
まとめ
色々組み合わせてナレッジベースができました。
副次的な効果ですが、自動チャンク分割のお陰で、「元の文字列」と「チャンク分割後の文字列」を両方保持したナレッジベースとなっています。そのため、分割したチャンク（ 子チャンク ）で検索した結果として、元の文字列（ 親チャンク ）の取得が可能です。
Claude 3やCommand Rは長いコンテキストを入力にできるので、親チャンクを知識としてプロンプトで渡すことで、精度の高い回答生成に期待できます 。