LlamaIndex の Reranker を試す

npaka

2024年4月3日 22:24

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

・LlamaIndex 0.10.14

1. Reranker

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

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

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

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

2. ドキュメントの準備

今回は、マンガペディアの「ぼっち・ざ・ろっく！」のドキュメントを用意しました。

・bocchi.txt

3. Rerankerの使用

Reranker付きの質問応答の実行手順は、次のとおりです。
クエリエンジンのポストプロセッサにRerankerを追加しています。

(1) パッケージのインストール。
Rerankerのパッケージ「llama-index-postprocessor-flag-embedding-reranker」「FlagEmbedding」もインストールします。

# パッケージのインストール
!pip install llama-index==0.10.14
!pip install llama-index-postprocessor-flag-embedding-reranker
!pip install FlagEmbedding

(2) 環境変数の準備。
左端の鍵アイコンで「OPENAI_API_KEY」を設定してからセルを実行してください。

import os
from google.colab import userdata

# 環境変数の準備 (左端の鍵アイコンでOPENAI_API_KEYを設定)
os.environ["OPENAI_API_KEY"] = userdata.get("OPENAI_API_KEY")

(3) ログレベルの設定。

import logging
import sys

# ログレベルの設定
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)

(4) Rerankerの準備。
今回は、多言語のRerankerモデル「BAAI/bge-reranker-v2-m3」を使います。top_n=5で関連性の高い5件に絞ります。

from llama_index.postprocessor.flag_embedding_reranker import FlagEmbeddingReranker

# Rerankerの準備
rerank = FlagEmbeddingReranker(
    model="BAAI/bge-reranker-v2-m3", 
    use_fp16=True, 
    top_n=5
)

(5) Colabにdataフォルダを作成してドキュメントを配置。
左端のフォルダアイコンでファイル一覧を表示し、右クリック「新しいフォルダ」でdataフォルダを作成し、ドキュメントをドラッグ&ドロップします。

(6) ドキュメントの読み込み。

from llama_index.core import SimpleDirectoryReader

# ドキュメントの読み込み
documents = SimpleDirectoryReader("data").load_data()

(7) インデックスの作成。

from llama_index.core import VectorStoreIndex

# インデックスの作成
index = VectorStoreIndex.from_documents(documents)

(8) クエリエンジンの作成。
ベクトル検索で10件を指定し、ポストプロセッサーにRerankerを追加しています。Rerankerでは10件を5件に絞っています。

# クエリエンジンの準備
query_engine = index.as_query_engine(
    similarity_top_k=10, 
    node_postprocessors=[rerank]
)

(9) 質問応答。

# 質問応答
response = query_engine.query("ぼっちちゃんの得意な楽器は？")
print(response)

後藤ひとりの得意な楽器はギターです。

(10) クエリに渡されたチャンクの確認。
Rerankの設定通り、関連性の高い5件に絞られてることがわかります。

# 取得したチャンクの確認
for node in response.source_nodes:
     print(node.get_text())
     print("----")

毎日練習したお陰でギターの腕前はプロ級になったが、...
----
ギターを担当している。髪をロングに伸ばし、明るい雰囲気を漂わせている。....
----
そのため過去に別のバンドに所属していたが、...
----
ひねくれ者なヨヨコは、...
----
ひとりがギターを弾いている姿を見て、...
----

この記事が気に入ったらサポートをしてみませんか？