MENU
AWS

PineconeのReranking APIをさっそく使ってみた

内山浩佑

しつこく続いた咳がようやく落ち着いてきた内山です。

先日、Pinecone Inference APIのRerankingがpublic previewで公開されたようなので、さっそく試してみました。

https://docs.pinecone.io/release-notes/2024#features

Contents
  1. Pineconeとは
  2. Rerankingとは
  3. 検証用AWS環境
  4. AWS Marketplace経由でPineconeと連携
  5. KnowledgeBase for Amazon Bedrockの設定
  6. 実装
  7. 結果
  8. おわりに

Pineconeとは

https://www.pinecone.io/

Pineconeは、高性能なベクトル検索アプリケーションを簡単に構築できるフルマネージド型のクラウドネイティブなベクトルデータベースサービスです。主な特徴は以下のとおりです。

  • 数十億規模のベクトルデータでも高速検索が可能
  • データの追加・更新に合わせてインデックスがリアルタイムに更新される
  • ベクトル検索とメタデータフィルタの組み合わせで適切な結果を素早く取得可能
  • 運用の手間いらずで、導入・利用・拡張が容易

今回はベクトル検索とRerankingのAPIを利用することになります。

Rerankingとは

Rerankingとは、検索結果や文書の順位を再評価し、より適切な順番に並び替える手法です。主に以下の目的で活用されます。

  • 検索精度の向上
  • ユーザーの意図に適した結果の提供
  • 関連性の高い情報を上位に表示

実際のベンチマークでも、Rerankingすることによって精度が向上することが確認されているようです。(参考: https://www.llamaindex.ai/blog/boosting-rag-picking-the-best-embedding-reranker-models-42d079022e83)

最近では、主にRAG（Retrieval-Augmented Generation）の精度向上に活用されている技術です。

今回の検証でも、RAGを意識したものとなっています。

検証用AWS環境

今回の検証環境は、AWS上で構築しました。主に以下のサービスを使用しています。

  • Knowledge bases for Amazon Bedrock: Pineconeのベクトル検索
  • AWS Marketplace: Pineconeのサブスクリプション

AWS Marketplace経由でPineconeと連携

rerankエンドポイントは、Standardプラン以上に入っている必要があります。今回は AWS Marketplace経由でPineconeを使うことにしました。

AWSでは、Marketplaceを通じてPinecone serverlessをサブスクライブすることができます。このようにすることで、AWSアカウントと連携してPineconeを利用でき、利用料金もAWSの請求と一元化されるようになります。

AWS Marketplaceで、Pineconeをサブスクライブすると、以下のような画面になります。

Pay as you goという表記があるとおり、Pineconeの料金は使った分だけになります。

この画面からPineconeのサイトに行くと、アカウント作成画面が表示されるので、アカウントがない場合は作成しておきます。アカウントを作成すると、Pineconeの設定画面に入ることができます。

後ほどAPIキーを使用するので、以下のページに移動して、控えておきます。

KnowledgeBase for Amazon Bedrockの設定

KnowledgeBase for Amazon Bedrockは、Pineconeのベクトル検索を実施するために使用します。S3に情報のソースとなるデータを保存すると、KnowledgeBaseがそれを読み込んでベクトル化し、Pineconeのインデックスに保存するところも実施してくれます。ユーザーが入力した質問を元にPineconeのクエリを生成してくれます。

今回は、ソースとしてBedrockのユーザーガイドを利用しています。

https://docs.aws.amazon.com/pdfs/bedrock/latest/userguide/bedrock-ug.pdf

実装

検証用の実装では、RAGを行うプログラムにしました。Rerankingの様子は、LangSmithを使用して確認しました。

https://www.langchain.com/langsmith

RerankingのAPIの利用方法については、以下のページで確認できます。

https://docs.pinecone.io/guides/inference/rerank#2-rerank-documents

今回は、KnowledgeBaseと連携する処理を実装するために、LangChainを使用します。LangChainには、 0.1系と0.2系がありますが、今回は0.2系を使用しています。

LangChainのインストールは、以下のとおりです。

pip install langchain langchain_core langchain_aws

プロンプトテンプレートの設定は以下の通りです。

prompt = PromptTemplate.from_template("""
あなたは質問応答タスクのためのアシスタントです。
以下の提供された文脈を使用して質問に答えてください。
答えがわからない場合は、わからないと言ってください。

Question: {question} 

Context: {context} 

Answer:
""")

{question}には、ユーザの質問が埋め込まれます。{context}には、Rerankingされた結果が埋め込まれることになります。

最終的な回答を生成するモデルにはClaude 3 Sonnetを使用しています。

llm = ChatBedrock(
    model_id="anthropic.claude-3-sonnet-20240229-v1:0",
    model_kwargs=dict(temperature=0),
    region_name="us-east-1",
)

AmazonKnowledgeBasesRetrieverを使用して、KnowledgeBaseから結果を取得する実装は以下のとおりです。上位10件を取得するようにしています。

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id=os.environ["KNOWLEDGE_BASE_ID"],
    region_name="us-east-1",
    retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 10}},
)

Rerankingするための関数を定義します。Rerankingされた結果から上位3件を取得するようにしています。

def rerank_docs(docs):
    ret = pc.inference.rerank(
        model="bge-reranker-v2-m3",
        query=query,
        documents=[{"id": str(idx), "text": doc.page_content} for idx, doc in enumerate(docs)],
        top_n=3,
        return_documents=True,
    )
    return [docs[int(doc["document"]["id"])] for doc in ret.data]

Rerankingされた結果から最終的に空行区切りでプロンプトに埋め込むように整形する関数を定義します。

qa_chain = (
        {
            "context": retriever | rerank_docs | format_docs,
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
)

このchainに対して、クエリを投げて、結果を表示するようにします。

query = "Bedrockで使えるモデルは？"
print(qa_chain.invoke(query))

実行してみると、以下のような内容が出力されました。

提供された文脈から、Amazon Bedrockで使えるモデルは以下のようです。

- Cohere Command R
- Cohere Command R+
- Mistral AI Mistral Large 2 (24.07)
- Meta Llama 3.1 Instruct
- AI21 Labs Jurassic-2
- Amazon Titan Image Generator
- Amazon Titan Text 
- Amazon Titan Text Embeddings
- Anthropic Claude

また、カスタムモデルをインポートしてBedrockで使用したり、他のアカウントとカスタムモデルを共有することもできるようです。文脈からは具体的な使用可能モデルの一覧は確認できませんが、上記のモデルが利用可能であることがわかります。

結果

LangSmithを見に行くと、以下のように各chain毎に入力と出力が確認できるようになっています。

今回は、Pineconeのベクトル検索の結果が、どのようにRerankingされたかを確認したいので、rerank_docsの入出力を確認します。

以下のように、Reranking前の結果の中には、サンプルスクリプトの情報や目次の部分が上位に含まれてしまっていました。

);      const response = await client.send(command);    const models = response.modelSummaries;      console.log("Listing the available Bedrock foundation models:");      for (let model of models) {      console.log("=".repeat(42));      console.log(` Model: ${model.modelId}`);      console.log("-".repeat(42));      console.log(` Name: ${model.modelName}`);      console.log(` Provider: ${model.providerName}`);      console.log(` Model ARN: ${model.modelArn}`);      console.log(` Input modalities: ${model.inputModalities}`);      console.log(` Output modalities: ${model.outputModalities}`);      console.log(` Supported customizations: ${model.customizationsSupported}`);      console.log(` Supported inference types: ${model.inferenceTypesSupported}`);    Amazon Bedrock 1188    https://github.com/awsdocs/aws-doc-sdk-examples/tree/main/javascriptv3/example_code/bedrock#code-examples      Amazon Bedrock User Guide       console.log(` Lifecycle status: ${model.modelLifecycle.status}`);      console.log("=".repeat(42) + "\n");    }      const active = models.filter(      (m) => m.modelLifecycle.status === "ACTIVE",    ).length;    const legacy = models.filter(      (m) => m.modelLifecycle.status === "LEGACY",    ).length;      console.log(      `There are ${active} active
• Invoke Anthropic Claude on Amazon Bedrock using Bedrock's Converse API with a  response stream   • Invoke Anthropic Claude on Amazon Bedrock using the Invoke Model API   • Invoke Anthropic Claude models on Amazon Bedrock using the Invoke Model API with a  response stream   • A tool use demo illustrating how to connect AI models on Amazon Bedrock with a custom  tool or API   • Cohere Command for Amazon Bedrock Runtime using AWS SDKs   • Invoke Cohere Command on Amazon Bedrock using Bedrock's Converse API   • Invoke Cohere Command on Amazon Bedrock using Bedrock's Converse API with a  response stream   • Invoke Cohere Command R and R+ on Amazon Bedrock using the Invoke Model API   • Invoke Cohere Command on Amazon Bedrock using the Invoke Model API   • Invoke Cohere Command R and R+ on Amazon Bedrock using the Invoke Model API with a  response stream   • Invoke Cohere Command on Amazon Bedrock using the Invoke Model API with a response  stream   • A tool use demo illustrating how to connect AI models on Amazon Bedrock with a custom  tool or API   • Meta Llama for Amazon Bedrock Runtime using AWS SDKs   • Invoke Meta Llama on Amazon Bedrock using Bedrock's Converse API   • Invoke Meta

Reranking後は、余計な情報は省かれて、より関連のありそうなものだけが絞り込まれていました。

April 30, 2024   New models You can now use Cohere  Command R and Cohere  Command R+ models with  Amazon Bedrock.   April 29, 2024   1579    https://docs.aws.amazon.com/bedrock/latest/userguide/titan-models.html  https://docs.aws.amazon.com/bedrock/latest/userguide/br-studio.html  https://docs.aws.amazon.com/bedrock/latest/userguide/agents-alias-manage.html  https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-regions.html  https://docs.aws.amazon.com/general/latest/gr/bedrock.html  https://docs.aws.amazon.com/general/latest/gr/bedrock.html  https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-setup.html  https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html  https://docs.aws.amazon.com/bedrock/latest/userguide/pt-supported.html  https://docs.aws.amazon.com/bedrock/latest/userguide/pt-supported.html  https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html      Amazon Bedrock User Guide   New feature You can now import a custom  model into Amazon Bedrock.
August 1, 2024   New feature You can now share custom  models with other accounts in  Amazon Bedrock.   August 1, 2024   New managed policy Amazon Bedrock has added AmazonBedrockStudi  oPermissionsBounda  ry  to limit permissions of  the provisioned IAM principal  that the policy is attached to.   July 31, 2024   New model You can now use Mistral AI  Mistral Large 2 (24.07) model  with Amazon Bedrock.   July 24, 2024   1574    https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html  https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html  https://docs.aws.amazon.com/bedrock/latest/userguide/copy-model.html  https://docs.aws.amazon.com/bedrock/latest/userguide/share-model.html  https://docs.aws.amazon.com/bedrock/latest/userguide/security-iam-awsmanpol.html  https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html      Amazon Bedrock User Guide   New model You can now use Meta Llama  3.1 Instruct models with  Amazon Bedrock.   July 23, 2024   New feature You can now use Prompt  management and Prompt  flows with Amazon Bedrock  Studio.
Amazon Bedrock using Bedrock's Converse API   • Invoke AI21 Labs Jurassic-2 models on Amazon Bedrock using the Invoke Model API   • Amazon Titan Image Generator for Amazon Bedrock Runtime using AWS SDKs   • Invoke Amazon Titan Image on Amazon Bedrock to generate an image   • Amazon Titan Text for Amazon Bedrock Runtime using AWS SDKs   • Invoke Amazon Titan Text on Amazon Bedrock using Bedrock's Converse API   • Invoke Amazon Titan Text on Amazon Bedrock using Bedrock's Converse API with a  response stream   • Invoke Amazon Titan Text models on Amazon Bedrock using the Invoke Model API   • Invoke Amazon Titan Text models on Amazon Bedrock using the Invoke Model API with a  response stream   • Amazon Titan Text Embeddings for Amazon Bedrock Runtime using AWS SDKs   • Invoke Amazon Titan Text Embeddings on Amazon Bedrock   • Anthropic Claude for Amazon Bedrock Runtime using AWS SDKs 1181        Amazon Bedrock User Guide   • Invoke Anthropic Claude on Amazon Bedrock using Bedrock's Converse API   • Invoke Anthropic Claude on Amazon Bedrock using Bedrock's Converse API with a  response stream   • Invoke Anthropic Claude on Amazon Bedrock using the Invoke Model API   • Invoke Anthropic Claude models on Amazon Bedrock using the Invoke Model API

関連のある情報のみがLLMに渡されることによって、回答精度が高まりそうですね。

おわりに

今回の検証では、PineconeのReranking APIが有効に機能し、より関連性の高い検索結果を得ることができました。今後もこの技術を活用して、さらに精度の高いRAGを実現し、ユーザー体験の向上を目指していきたいと考えています。

AUTHOR
uchiko
uchiko
記事URLをコピーしました