LangChain で Cohere Reranker を試す
「LangChain」で「Cohere Reranker」を試したので、まとめました。
1. Cohere Reranker
RAGにおける「Reranker」は、取得したチャンクの中から、質問に対して最も関連性の高い情報を持つチャンクを選択する役割を担っています。
RAGの処理の流れ、次のとおりです。
「Reranker」を用いることで、検索された多数のドキュメントの中から、質問に対して最も有用な情報を持つドキュメントを効果的に選択することができ、RAGの生成する回答の品質向上に寄与します。
今回は、「LangChain」に統合されている「Cohere Reranker」を使います。
2. ドキュメントの準備
今回は、マンガペディアの「ぼっち・ざ・ろっく!」のドキュメントを用意しました。
・bocchi.txt
3. Rerankerの使用
Reranker付きの質問応答の実行手順は、次のとおりです。
ContextualCompressionRetrieverにRerankerを追加しています。
(1) パッケージのインストール。
「Cohere Reranker」が含まれているパッケージ「langchain_cohere」もインストールします。
# パッケージのインストール
!pip install langchain langchainhub
!pip install langchain_openai langchain_cohere
!pip install unstructured faiss-gpu
(2) 環境変数の準備。
左端の鍵アイコンで「OPENAI_API_KEY」と「COHERE_API_KEY」を設定してからセルを実行してください。
import os
from google.colab import userdata
# 環境変数の準備 (左端の鍵アイコンでOPENAI_API_KEYとCOHERE_API_KEYを設定)
os.environ["OPENAI_API_KEY"] = userdata.get("OPENAI_API_KEY")
os.environ["COHERE_API_KEY"] = userdata.get("COHERE_API_KEY")
(3) Colabにdataフォルダを作成してドキュメントを配置。
左端のフォルダアイコンでファイル一覧を表示し、右クリック「新しいフォルダ」でdataフォルダを作成し、ドキュメントをドラッグ&ドロップします。
(4) ドキュメントとベクトルストアとRetrieverの準備。
Retrieverではベクトルストアからベクトル検索で関連するチャンクを10個取得するよう設定しています。
from langchain.document_loaders import DirectoryLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
# ドキュメントの準備
loader = DirectoryLoader("./data/")
docs = loader.load()
documents = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200
).split_documents(docs)
# ベクトルストアの準備
vector = FAISS.from_documents(documents, OpenAIEmbeddings())
# Retrieverの準備
retriever = vector.as_retriever(
search_kwargs={"k": 10} # ベクトル検索で10個取得
)
(5) ContextualCompressionRetrieverの準備。
「ContextualCompressionRetriever」はコンテキストを圧縮するRetrieverです。base_compresserにCohereRerankを設定して、Rerankで3個取得するよう設定しています。
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_cohere import CohereRerank
# ContextualCompressionRetrieverの準備
compression_retriever = ContextualCompressionRetriever(
base_compressor=CohereRerank(top_n=3), # Rerankで3個取得
base_retriever=retriever
)
(6) クエリエンジンの作成。
ベクトル検索で10件を指定し、ポストプロセッサーにRerankerを追加しています。Rerankerでは10件を5件に絞っています。
from langchain.chat_models import ChatOpenAI
from langchain.schema.runnable import RunnablePassthrough
from langchain import hub
# チャットモデルの準備
chat_model = ChatOpenAI(
model_name="gpt-3.5-turbo",
temperature=0
)
# プロンプトテンプレートの準備
rag_prompt = hub.pull("rlm/rag-prompt")
# チェーンの準備
rag_chain = {
"context": compression_retriever,
"question": RunnablePassthrough()
} | rag_prompt | chat_model
(7) 質問応答。
# 質問応答
output = rag_chain.invoke("ぼっちちゃんの得意な楽器は?")
print(output.content)
4. LangSmithでの動作確認
LangSmithでRerankが動作していることを確認します。
(1) LangSmithの有効化。
左端の鍵アイコンで「LANGCHAIN_API_KEY」(LangSmithのAPIキー)を設定してからセルを実行してください。
import os
from uuid import uuid4
# 環境変数の準備
unique_id = uuid4().hex[0:8]
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = f"Tracing Walkthrough - {unique_id}"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGCHAIN_API_KEY"] = userdata.get("LANGCHAIN_API_KEY")
(2) LangSmithで内部処理を確認。
1つ目のRetrieverでは、ベクトル検索で10個のチャンクを取得しています。ドキュメント名を表示されています。待ち時間は0.13sです。
2つ目のRetrieverでは、Rerankerで3つのチャンクを取得しています。スコアも表示されています。待ち時間は0.51sです。