LangChain で Cohere Reranker を試す

「LangChain」で「Cohere Reranker」を試したので、まとめました。

1. Cohere Reranker

RAGにおける「Reranker」は、取得したチャンクの中から、質問に対して最も関連性の高い情報を持つチャンクを選択する役割を担っています。

RAGの処理の流れ、次のとおりです。

(1) 質問を入力
(2) 質問に関連するチャンクを取得
(3) Rerankerは各チャンクと質問の関連性をスコア付けし、最も関連性の高いチャンクを選択
(4) 選択されたチャンクで、質問に対する回答を生成

「Reranker」を用いることで、検索された多数のドキュメントの中から、質問に対して最も有用な情報を持つドキュメントを効果的に選択することができ、RAGの生成する回答の品質向上に寄与します。

今回は、「LangChain」に統合されている「Cohere Reranker」を使います。

Cohere Reranker | 🦜️🔗 LangChain

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です。

OpenAI GPT-4V／ChatGPT／GPTs 人工知能プログラミング実践入門 | 株式会社ボーンデジタル