Knowledge base for Amazon BedrockがGAしたので、KendraではなくKnowledge base for Amazon Bedrockを使用したRAGをLangChainを使って実装してみます。
インデックスしてしまえばLangChainからはKendraとほぼ同じ使用感で使用でき、チャンクサイズの自由度が高いのと、Agents for Amazon Bedrockから「も」呼べるのが良いところかなと思います。あと料金ですかね
必要ライブラリのアップデート
boto3とLangChainを最新にします。
LangChainを動作確認したバージョンは0.0.345です。
pip install -U boto3
pip install -U langchain==0.0.345
Knowledge base の作成
マネコンのAmazon BedrockのOrchestrationのKnowledge baseのCreate knowledge baseから作成します。
名前を変えるかどうか程度でわりとデフォルトでいけるのですが、Set up data sourceでドキュメントが格納されたS3を指定します。それとAddional settings - optionalから、ベクトルDBのチャンクサイズとチャンク間でオーバーラップさせるパーセンテージを指定できます。チャンクサイズがどうあるべきかはユースケースとストアするドキュメント次第かなと思いますが、Claudeの入力可能トークンの大きさを頼ってとりあえず大きめにします（しました）。
おおまかに1000トークン弱でPDF1ページぐらいかな？という感じなので、この設定であれば3-4ページ分の大きさでチャンクを作ってくれるはずです。後はデフォルトで作成します（OpenSearchにベクトルDBが作成されます。不要になればマネコンのOpenSearchから削除してください）。
作成が終わったらSyncを実行する事でS3がクロールされてベクトル化されます。
私はテスト用に以下のPDFをS3に格納、Syncしました。
作成したKnowledge baseのKnowledge base overviewの右上から、Knowledge base IDをコピーしておきます。
RAGアプリの構築
まず最小のRAGを作ってみます。
knowledge_base_idには各自のKnowledge base IDを設定します。
実際には上述のチャンクサイズやベクトルDBからの取得件数はチューニング必要かなと思います。今の設定ではベクトルの類似度が高いチャンク（PDFの3-4ページごと）を10件、つまり、PDFから類似度が高いページを30-40ページ程度引っ張ってきて、参考ドキュメントとしてプロンプトに渡している動きになります。金額が嵩みそうなのでLLMはClaude Instantを使用しています。
from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.retrievers.bedrock import AmazonKnowledgeBasesRetriever
# クエリ内容を設定
user_input="bedrockで使用可能なモデルIDは"
#各自のKnowledge base ID
knowledge_base_id = "XXXXXXXXXX"
# ベクトルDBからの取得件数
retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}}
# LLMのモデルID
model_id="anthropic.claude-instant-v1" #Claude Instant
#model_id="anthropic.claude-v2:1" #Claude2.1
# 最大出力トークン数
model_kwargs={"max_tokens_to_sample": 1000} #Claude系の場合
# Retriever(KnowledgeBase)の定義
retriever = AmazonKnowledgeBasesRetriever(knowledge_base_id=knowledge_base_id,retrieval_config=retrieval_config)
# LLMの定義
llm = Bedrock(model_id=model_id,model_kwargs=model_kwargs)
# Chainの定義
qa = RetrievalQA.from_chain_type(retriever=retriever,llm=llm)
# chainの実行
result = qa.run(user_input)
print(result)
Amazon Bedrock で現在使用可能な主なモデルとそのモデルIDは以下の通りです。
- Titan Text G1-Express 8K (アマゾン): titan-text-express-v1:0:8キロ
- Titan Embeddings G1-Text (アマゾン): titan-embed-text-v1:2:8キロ
- Anthropic Claude V2 18K: anthropic.claude-v2:0:18k
- Anthropic Claude V2 100K: anthropic.claude-v2:0:100k
- Anthropic Claude Instant V1 100K: anthropic.claude-instant-v1:2:100メートル
- Stable Diffusion XL 1.0: stable-diffusion-xl-v1:0
- Meta Llama 2 Chat 13B: meta.llama2-13b-chat-v1:0:4キロバイト
これらのモデルIDは、プロビジョニングされたスループットを作成する際に使用します。
また、ベースモデルやカスタムモデルを指定する際の参考情報となります。
新しいモデルがサポートされる場合は、Amazon Bedrock ドキュメントやコンソールからその情報を
確認ください。
元のドキュメントの日本語が怪しめなのでRAG結果の日本語も怪しめですが、きちんと元PDFから取得されています。
プロンプトを指定する場合は以下のように書きます。プロンプトを指定した方が生成内容が安定するような気がします。
from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.retrievers.bedrock import AmazonKnowledgeBasesRetriever
from langchain.prompts import PromptTemplate
# クエリ内容を設定
user_input="bedrockで使用可能なモデルIDは"
#各自のKnowledge base ID
knowledge_base_id = "XXXXXXXXXX"
# ベクトルDBからの取得件数
retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}}
# LLMのモデルID
model_id="anthropic.claude-instant-v1" #Claude Instant
#model_id="anthropic.claude-v2:1" #Claude2.1
# 最大出力トークン数
model_kwargs={"max_tokens_to_sample": 1000} #Claude系の場合
# Retriever(KnowledgeBase)の定義
retriever = AmazonKnowledgeBasesRetriever(knowledge_base_id=knowledge_base_id,retrieval_config=retrieval_config)
# LLMの定義
llm = Bedrock(model_id=model_id,model_kwargs=model_kwargs)
# promptの定義
prompt_template = """
<documents>タグには参考文書が書かれています。
<documents>{context}</documents>
\n\nHuman: 上記参考文書を元に、<question>に対して1000文字以内で説明してください。言語の指定が無い場合は日本語で答えてください。但し固有名詞と思われる単語は英語のままとしてください。
もし<question>の内容が参考文書に無かった場合は「文書にありません」と答えてください。
<question>{question}</question>
\n\nAssistant:"""
PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
# Chainの定義
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(retriever=retriever,llm=llm,chain_type_kwargs=chain_type_kwargs)
# chainの実行
result = qa.run(user_input)
print(result)
Streamlitでガワを被せるとこんな感じです。
from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.retrievers.bedrock import AmazonKnowledgeBasesRetriever
from langchain.prompts import PromptTemplate
#各自のKnowledge base ID
knowledge_base_id = "XXXXXXXXXX"
# ベクトルDBからの取得件数
retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}}
# LLMのモデルID
model_id="anthropic.claude-instant-v1" #Claude Instant
#model_id="anthropic.claude-v2:1" #Claude2.1
# 最大出力トークン数
model_kwargs={"max_tokens_to_sample": 1000} #Claude系の場合
# Retriever(KnowledgeBase)の定義
retriever = AmazonKnowledgeBasesRetriever(knowledge_base_id=knowledge_base_id,retrieval_config=retrieval_config)
# LLMの定義
llm = Bedrock(model_id=model_id,model_kwargs=model_kwargs)
# promptの定義
prompt_template = """
<documents>タグには参考文書が書かれています。
<documents>{context}</documents>
\n\nHuman: 上記参考文書を元に、<question>に対して1000文字以内で説明してください。言語の指定が無い場合は日本語で答えてください。但し固有名詞と思われる単語は英語のままとしてください。
もし<question>の内容が参考文書に無かった場合は「文書にありません」と答えてください。
<question>{question}</question>
\n\nAssistant:"""
PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
# Chainの定義
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(retriever=retriever,llm=llm,chain_type_kwargs=chain_type_kwargs)
# Streamlitでガワを被せる
import streamlit as st
st.title("LangChainからのKnowledge base for Amazon Bedrock呼出し")
input_text = st.text_input("このテキストをBedrockに送信します")
send_button = st.button("送信")
if send_button:
# chainの実行
result = qa.run(input_text)
# 結果をStreamlitに出力
st.write(result)