情報検索拡張型生成(Retrieval Augmented Generation、RAG)は、自身のドキュメントから情報を検索するための手法です。これは2段階のプロセスで構成されています:

1. ナレッジベースの作成: ドキュメントを小さな部分ドキュメントに分割し、各部分の埋め込みを計算します。これらの埋め込みと元のチャンクは通常ベクトルストアに保存されます。

2. 検索と生成: 新しいユーザークエリが入力されると、クエリの埋め込みが計算されます。保存されたチャンクの埋め込みとクエリの埋め込みを比較して、最も関連性の高いチャンクを検索します。検索されたチャンクはオリジナルのクエリに追加され、言語モデル(LLM)に入力されて最終的な応答が生成されます。

このセットアップは「検索拡張型生成」パイプラインと呼ばれ、LLMの生成がナレッジベースから検索された関連情報で拡張されています。

LangchainやLlamaIndexなどのフレームワークがRAGパイプラインの実装を簡単にしてくれますが、コアコンポーネントはPython、埋め込みモデル、LLMを使って構築することができます。このアプローチにより、基本概念をより深く理解でき、パイプラインのカスタマイズにも柔軟性が得られます。

検索拡張型生成(RAG)システムを構築する際、重要なステップの1つはインプットドキュメントを小さな、より管理しやすい部分に分割することです。この例では、ウィキペディア記事をパラグラフに分割しています。

このアプローチの理由は、パラグラフがしばしば論理的な情報単位を表しており、効果的に検索・利用できるためです。ドキュメントをこのように分割することで、最終的な応答に含める最も関連性の高いセクションを特定しやすくなります。

このステップのコードは以下の通りです:

# ドキュメントをパラグラフに分割
chunks = [paragraph.strip() for paragraph in text.split('\n') if paragraph.strip()]

ここでは、入力テキストを改行文字で分割してパラグラフを抽出し、前後の空白を取り除いています。

得られたchunksリストには個々のパラグラフが含まれており、これらを埋め込んでベクトルストアに保存できます。これにより、ユーザークエリに基づいて最も関連性の高いパラグラフを効率的に検索し、言語モデルによる最終的な応答に含めることができます。

ドキュメントチャンクとユーザークエリを埋め込むために、Sentence Transformersライブラリを使用します。このライブラリには高品質な埋め込みを生成できる事前学習モデルが用意されています。

最初に埋め込みモデルをロードします:

from sentence_transformers import SentenceTransformer
embedding_model = SentenceTransformer('all-mpnet-base-v2')

次に、ドキュメントチャンクの埋め込みを計算します:

chunk_embeddings = embedding_model.encode(chunks, normalize_embeddings=True)

ここでchunksはドキュメントからの文章チャンクのリストです。normalize_embeddings=Trueオプションは、後の類似度計算に重要な埋め込みの正規化を行います。

ユーザークエリの埋め込みを取得するには、以下のように実行します:

query_embedding = embedding_model.encode([user_query], normalize_embeddings=True)[0]

これで、ドキュメントチャンクとユーザークエリの埋め込みが得られ、検索ステップで使用できます。

ユーザークエリに基づいて最も関連性の高いチャンクを検索するには、クエリの埋め込みとドキュメントチャンクの埋め込みの間のコサイン類似度を計算する必要があります。手順は以下の通りです:

1. 同じ埋め込みモデルを使ってユーザークエリの埋め込みを計算する。
2. クエリの埋め込みとドキュメントチャンクの埋め込みの間の内積を計算する。これがコサイン類似度のスコアになる。
3. ドキュメントチャンクをコサイン類似度スコアに基づいてソートし、上位K個のチャンクを選択する。
4. 選択されたチャンクのテキストコンテンツを取得し、言語モデルの文脈として使用する。

主なステップは以下の通りです:

1. クエリの埋め込みを計算する
2. コサイン類似度スコアを計算する
3. 上位K個のチャンクを選択する
4. チャンクのテキストコンテンツを取得する

この単純なアプローチにより、ユーザークエリを言語モデルに入力する前に、ドキュメントコレクションから最も関連性の高い情報を効率的に検索できます。

検索されたチャンクとユーザークエリを使って応答を生成するには、まずこれらを含むプロンプトを作成します。その後、OpenAI GPT-4モデルを使ってこのプロンプトに基づいて応答を生成します。

コードは以下の通りです:

# ユーザークエリを取得
user_query = "What is the capital of France?"

# ユーザークエリに基づいて最も関連性の高いチャンクを検索
top_chunk_ids = [6, 8, 5]
top_chunks = [chunks[i] for i in top_chunk_ids]

# プロンプトを作成
prompt = "Use the following context to answer the question at the end. If you don't know the answer, say that you don't know and do not try to make up an answer.\n\nContext:\n"
for chunk in top_chunks:
    prompt += chunk + "\n\n"
prompt += f"\nQuestion: {user_query}"

# OpenAI GPT-4を使って応答を生成
response = openai_client.chat.create(
    model="gpt-4",
    messages=[
        {"role": "user", "content": prompt}
    ],
    max_tokens=1024,
    n=1,
    stop=None,
    temperature=0.7,
).choices[0].message.content

print(response)

このコードでは、まずユーザークエリに基づいて最も関連性の高い3つのチャンクを検索しています。次に、これらのチャンクとユーザークエリを含むプロンプトを作成し、OpenAI GPT-4モデルを使って応答を生成しています。

生成された応答は、検索されたチャンクの情報を活用して、ユーザークエリに対する簡潔で関連性の高い答えになります。

このチュートリアルでは、わずか10行のPythonコードでドキュメント検索システムを備えたチャットシステムを構築する方法を学びました。情報検索拡張型生成(RAG)パイプラインの主要コンポーネントについて、ナレッジベースの作成、ドキュメントの分割、埋め込み、検索、そして大規模言語モデルを使った生成について説明しました。

ここでのポイントは以下の通りです:

1. Langchainやllamaindexのような複雑なフレームワークを使わずに、PythonといくつかのライブラリだけでRAGパイプラインを実装できる。
2. ドキュメントの構造(パラグラフなど)に基づいて分割するのは、ほとんどのユースケースで簡単で効果的な戦略である。
3. ドキュメントチャンクとユーザークエリの埋め込みを計算し、類似度スコアを求めることで、ユーザークエリを拡張するための最も関連性の高い情報を検索できる。
4. 検索されたチャンクをユーザークエリに統合し、大規模言語モデルに入力することで、関連性の高い情報的な応答を生成できる。

この例は基礎を提供するものですが、より堅牢で高度なRAGシステムを構築する機会は多数あります。Langchainやllamaindexのようなフレームワークは、様々なベクトルストアや言語モデルと統合する際に役立ちます。ただし、Pythonの実装から始めることで、RAGパイプラインの中核概念とコンポーネントをより深く理解できます。

より高度なRAG手法に興味がある場合は、「RAG Beyond Basics」コースをチェックすることをおすすめします。そこではより複雑で本番環境に適したRAGシステムの構築について詳しく説明しています。