RankGPTを日本語検索データセットで試してみる

まとめ

リランキングの手法の一つして、GPTを用いたリランキング手法を日本語の検索のデータセットで試してみました。text-embedding-3-large単体での検索結果をGPT4でリランキングしてみましたが、あまり検索精度は改善されませんでした。

元ネタ：Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents

この論文で、LLMによる新たなリランキング手法が、GPT-4を用いた場合に、従来手法を上回る性能を発揮したと報告されています。そのリランキング手法は、permutation generationと命され、パッセージの配列、及び、その配列をクエリとの関連性に基づき並び替える指示をLLMにプロンプトで与える手法です。

permutation generationのプロンプト（参考論文より）

検証内容

OpenAIの埋め込みモデル3種類（3-large, 3-small, ada-002）について、文章検索に埋め込みを利用し検索精度を調べました。また、埋め込みモデル単体の場合の検索精度と、GPTによるリランキングを行った後の検索精度を比較し、この手法の効果を検証しました。
評価データには、尼崎市のFAQデータセットから作成した検索データセットを用いています。計748件のクエリに対し、クエリと関連したFAQのアンサー（複数個あり）を検索データセットの正解データとしています。評価指標には上位5件のヒット率とMRRを用いました。

検証結果

使用したモデル（[]内の数字は埋め込みの次元）・リランキングの条件ごとに、上位5件のヒット率とMRRで検索精度を評価しました。リランキングを行う際は、事前に埋め込みモデルを検索に利用して上位10件のパッセージを取得し、その10件に対してリランキングを行っています。評価指標の計算には、Pythonのランキング評価用パッケージのranxを利用しています。

モデル名及びリランキングの有無	hit_rate@5	mrr@5
text-embedding-3-large [3072]	0.773	0.620
text-embedding-3-small [1536]	0.663	0.500
text-embedding-ada-002 [1536]	0.691	0.561
text-embedding-3-large [3072] + GPT3.5 ReRanking	0.774	0.619
text-embedding-3-small [1536] + GPT3.5 ReRanking	0.672	0.531

まず、上3つを見てみると3-largeモデルが頭一つ抜けた成績を修めています。予想通りでした。また、ada-002の方が3-smallより良い成績を修めていました。後者の結果は、ベンチマークの結果からすると予想外でした。OpenAIのブログポストによると、3-small[1536]の方がada-002よりMTEBで僅かに良い性能を示したとされており、また、JapaneseEmbeddingEvalでも、MIRACLの日本語データセットで3-small（おそらく[1536]）の方がada-002より高いスコアを出しています。言語やデータセットや評価指標の特徴に起因して、優れた結果を出すモデルが3-smallとada-002で逆転するのかなと推察します。

次に、下二つと上二つのデータ、つまり、GPT3.5によるリランキング有り無しの違いを見てみます。3-largeではリランキングの効果はほぼ見られませんでした。3-smallに対しては、リランキングの効果がわずかにみられ、少し検索精度が向上していました。初回検索の精度が高いと、リランキングの効果はあまり見られないのかもしれません。

LLMの性能を上げれば、初回検索の精度が高くてもリランキングの効果が確認できると予想し、GPT3.5よりも強力なGPT4で、3-largeモデルによる検索結果をリランキングし、評価指標の値を次の表にまとめました。APIコストを抑えるため、評価セットのデータを100件に絞って評価しています。結果を見ると、上位5件のヒット率に変化はほぼありませんでしたが、MRRはヒット率に比べて大きく上昇しました。ヒット率は、上位検索結果中の正解データの有無できまり、MRRは検索結果中の正解データの順位を加味して決まるので、平均的にみて、上位5件トータルの質は変わらなかったものの、その中の並び方は少し改善したのだと考えられます。とはいえ、しょっぱいですね。。。

モデル名及びリランキングの有無	hit_rate@5 (only 100 data)	mrr@5 (only 100 data)
text-embedding-3-large [3072]	0.69	0.56
text-embedding-3-large [3072] + GPT4 ReRanking	0.7	0.60

試せていませんが、他リランキング手法についても、text-embedding-3-large単体での検索結果をどの程度改善するのか調べたいですね。

実装方法・コード

以下、使用したコードの一部をまとめます。類似したコードは省いています。

インポート

import pandas as pd
import os
import re
from pathlib import Path
from langchain_openai import OpenAIEmbeddings
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_core.documents.base import Document
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv
load_dotenv()

データセットのロード

testset_path = "./data/testset.txt"
testset_df = pd.read_csv(testset_path, sep='\t', header=None).fillna("None")
testset_df.columns = ['Query', 'A_ID', 'B_ID', 'C_ID']
testset_df.head()

answers_path = './data/answers_in_Amagasaki.txt'
answers_df = pd.read_csv(answers_path, sep='\t', header=None, names=['ID', 'Answer'])
answers = answers_df["Answer"].values.tolist()
answers_df.head()

埋め込みモデルを搭載したVectorStoreの作成

3-small, ada-002モデルも同様に作成しました。

store = LocalFileStore("./cache/")
docs = [Document(page_content=answer, metadata={"id": str(i)}) for i, answer in enumerate(answers)]
embeddings_large = OpenAIEmbeddings(model="text-embedding-3-large")
cached_embedder_large = CacheBackedEmbeddings.from_bytes_store(
    embeddings_large, store, namespace=embeddings_large.model
)
db_large_path = "./chroma_db_large"
if not Path(db_large_path).exists():
    db_large = Chroma.from_documents(docs, cached_embedder_large, persist_directory=db_large_path)
else:
    db_large = Chroma(persist_directory=db_large_path, embedding_function=cached_embedder_large)

ranx用の正解データ・検索結果データの収集

データセットを整形し正解データを用意するための関数

クエリとのFAQのアンサー関連度を3段階に分けていますが、ヒット率とMRRの計算には、各指標の定義から反映されないはずです。

def convert_df_to_dict(df, level_to_int):
    qrels_dict = {}
    # DataFrameの各行に対して処理
    for index, row in df.iterrows():
        sub_dict = {}
        # 各関連性レベル（A、B、C）について処理
        for level in level_to_int.keys():
            id_column = f'{level}_ID'
            if row[id_column] is not None and row[id_column] != 'None':
                # スペース区切りのIDを分割し、各IDに対してサブ辞書に追加
                for doc_id in row[id_column].split(' '):
                    sub_dict[f'd_{doc_id}'] = level_to_int[level]
        # クエリをキーとしてメイン辞書に追加
        qrels_dict[f'q_{index+1}'] = sub_dict

    return qrels_dict

# 関連性レベルを数値化するための辞書
level_to_int = {'A': 3, 'B': 2, 'C': 1}

# DataFrameを指定されたフォーマットの辞書に変換
qrels_dict = convert_df_to_dict(testset_df, level_to_int)

検索結果データの取得及び整形のための関数

def run_test(base_retriever):
    run_dict = {}
    for i, query in enumerate(testset_df["Query"].values.tolist(), start=1):
        docs = base_retriever.get_relevant_documents(query)
        run_dict[f'q_{i}'] = {}
        for j, doc in enumerate(docs):
          doc_id = doc.metadata["id"]
          score = len(docs) - j
          run_dict[f'q_{i}']= run_dict[f'q_{i}'] | {f"d_{doc_id}": score}
    return run_dict

検索結果の取得

埋め込みモデルを用いたベクトル検索とその評価

base_retriever_large = db_large.as_retriever(search_kwargs={"k": 10})
run_dict_large = run_test(base_retriever_large)

qrels = Qrels(qrels_dict)
run_large = Run(run_dict_large)
evaluate(qrels, run_large, ["hit_rate@5", "mrr@5", "ndcg@5"])

ベクトル検索の結果のリランキングとその評価

items_large = run_dict_to_item(run_dict_large)
rerank_items_large = rerank(items_large)
run_dict_large_rerank35 = item_to_run_dict(rerank_items_large)

qrels = Qrels(qrels_dict)
run_large = Run(run_dict_large_rerank35)
evaluate(qrels, run_large, ["hit_rate@5", "mrr@5", "ndcg@5"])

リランキングは、元ネタの参考論文著者たちが提供しているモジュール RankGPTを使い、実施しています。

from rank_gpt import permutation_pipeline
def rerank(items, rank_start=0, rank_end=10, model_name="gpt-3.5-turbo"):
    new_items = []
    for item in items:
        new_item = permutation_pipeline(item, rank_start=rank_start, rank_end=rank_end, model_name=model_name, api_key=os.environ["OPENAI_API_KEY"])
        new_items.append(new_item)
    return new_items

permutation_pipeline内部で、create_permutation_instructionという関数が呼ばれているのですが、日本語対応していないので、日本語用に少し書き換えて使っています。リランキング対象のパッセージについて、1パッセージ当たりの読み込む長さを制限する箇所があり、そこを日本語用にアレンジしました。

def create_permutation_instruction(item=None, rank_start=0, rank_end=100, model_name='gpt-3.5-turbo'):
    query = item['query']
    num = len(item['hits'][rank_start: rank_end])

    # max_length = 300 # For Engish
    max_length = 1000  # For Japanese

    messages = get_prefix_prompt(query, num)
    rank = 0
    for hit in item['hits'][rank_start: rank_end]:
        rank += 1
        content = hit['content']
        content = content.replace('Title: Content: ', '')
        content = content.strip()
        # For Japanese should cut by character: content = content[:int(max_length)]
        # content = ' '.join(content.split()[:int(max_length)])
        content = content[:max_length]
        messages.append({'role': 'user', 'content': f"[{rank}] {content}"})
        messages.append({'role': 'assistant', 'content': f'Received passage [{rank}].'})
    messages.append({'role': 'user', 'content': get_post_prompt(query, num)})

    return messages

ranx用のインプット（dict）とrank_gptのインプット（item）を変換するコードは以下です。

def run_dict_to_item(run_dict):
    queries = testset_df["Query"].values.tolist()
    items = []
    for q_id_str, related_docs in run_dict.items():
        query_id = int(re.findall(r'\d+', q_id_str)[0])
        query = {"query_id": query_id, "query": queries[query_id - 1]}
        hited_related_docs = []
        for d_id_str, _ in related_docs.items():
            doc_id = int(re.findall(r'\d+', d_id_str)[0])
            content = answers_df[answers_df["ID"].astype(int) == doc_id]["Answer"].values[0]
            hited_related_docs.append({"id": doc_id, "content": content})
        hits = {"hits": hited_related_docs}
        item = query | hits
        items.append(item)
    return items
    
def item_to_run_dict(items):
    run_dict = {}
    for item in items:
        query_id = item["query_id"]
        run_dict[f'q_{query_id}'] = {}
        for j, doc in enumerate(item["hits"]):
            doc_id = doc["id"]
            score = len(item["hits"]) - j
            run_dict[f'q_{query_id}']= run_dict[f'q_{query_id}'] | {f"d_{doc_id}": score}
    return run_dict

参考

検証方法の参考にさせていただきました。
・Ahogrammerさんブログより 文書検索におけるリランキングの効果を検証する