RAG用途に使える、Wikipedia 日本語の embeddings とベクトル検索用の faiss index を作った
この記事は、情報検索・検索技術 Advent Calendar 2023の12月4日の記事である。
昨今のLLMの台頭により、外部情報を In-Context Learning として利用しLLMの生成結果の性能を高めることが可能な RAG(Retrieval Augmented Generation) の重要性の高まりを感じる。ただ、RAG を使ったシステムを構築してみようにも、データが少ないと面白みが少なかったりする。その為、Wikipedia 日本語の約550万文から簡単に検索可能でRAGの入力データとして使えるような embeddings と、素早い速度でベクトル検索できるような faiss 用の index を作成した。
例えば、Wikipedia から該当の文を検索する用途はこのように使える。
from datasets.download import DownloadManager
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import faiss
# wikipedia 日本語データセットのロード
wikija_dataset = load_dataset(
path="singletongue/wikipedia-utils",
name="passages-c400-jawiki-20230403",
split="train",
)
# faiss index のダウンロード
dm = DownloadManager()
index_local_path = dm.download(
f"https://huggingface.co/datasets/hotchpotch/wikipedia-passages-jawiki-embeddings/resolve/main/faiss_indexes/passages-c400-jawiki-20230403/multilingual-e5-small-passage/index_IVF2048_PQ96.faiss"
)
# faiss index のロード
faiss_index = faiss.read_index(index_local_path)
# embeddings へ変換するモデルのロード
model = SentenceTransformer("intfloat/multilingual-e5-small")
model.max_seq_length = 512
# embeddings へ変換
def to_emb(model, text, prefix="query: "):
return model.encode([prefix + text], normalize_embeddings=True)
emb = to_emb(model, "楽曲『約束はいらない』でデビューした、声優は誰?")
# faiss で検索して、関連 Top-5 を取り出す
TOP_K = 5
scores, indexes = faiss_index.search(emb, TOP_K)
for idx, (id, score) in enumerate(zip(indexes[0], scores[0])):
data = wikija_dataset[int(id)]
print((score, data["title"], data["text"][:100]))結果は以下。「誰？」という質問に対して、Top-3 に適切な人名が入っている結果となった。
(0.21018645, '約束はいらない', '「約束はいらない」(やくそくはいらない)は、坂本真綾のデビューシングル。')
(0.24241784, '約束はいらない', '坂本真綾の歌手デビュー作品。当時坂本はまだ無名の声優であったが、同曲がテーマソングとなったアニメ『天空のエスカフローネ』とともに知名度を上げることとなる。後に「指輪」が同アニメの劇場版映画の主題歌とな')
(0.2441598, '坂本真綾', '本格的な歌手活動は、1996年にテレビアニメ『天空のエスカフローネ』へ出演したことがきっかけで始めており、同作のオープニングテーマソングである「約束はいらない」(岩里祐穂作詞、菅野よう子作曲)をシング')
(0.26043588, 'Suara', '大阪府立豊中高等学校を経て大阪外国語大学でインドネシア語を専攻中にバンド・ユニットを組んで音楽活動を始めた。普段はお笑い番組が大好きなこともあってよく喋るほうだが、東京の仕事で標準語の喋りをする時は、')
(0.2690969, '早見沙織', '声優デビュー時より、数多くの主題歌やキャラクターソングを担当し、バラードからポップス、ヒットソングのカバー曲や英語の楽曲など、様々な曲を歌いこなす。2009年には吉田仁美とのユニット「blue dro')
このように HuggingFace dataset から Wikipedia 全文と検索用のインデクスのデータをロードできるため、難しいことをやらずとも簡単に利用が可能だ。検索に使っている faiss の index は省メモリで速く、550万件程度のベクトルなら大した計算機リソースがなくとも、CPUで高速に検索ができる。
また、HuggingFace Spaces に、RAGができるようなデモを作ったので Web-UI からRAG検索〜QAの応対を試すことができる。
RAG の In-Context として wikipeida 文を使う
続いて、先ほど質問「楽曲『約束はいらない』でデビューした、声優は誰？」の回答データを RAG の In-Context Learning として利用してみよう。まず ChatGPT 3.5 に対して質問してみると、以下の結果となった。
楽曲『約束はいらない』でデビューした、声優は誰？
「約束はいらない」は、声優の悠木碧（ゆうき あおい、英: Aoi Yuki）によって歌われた楽曲です。悠木碧は、日本の声優として知られており、多くのアニメ作品やゲームで声優として活躍しています。この曲は彼女のデビューシングルの一つとして知られています。
ChatGPT 3.5では間違った答えを出力してしまった。続いて先ほどの検索結果を Wikipedia の文を注入して生成、いわゆるIn-Context Learning として RAG の検索結果を使って実行してみる。
楽曲『約束はいらない』でデビューした、声優は誰？以下のデータにヒントが含まれる場合は参考にして答えよ。また以下のデータ自体が間違っている場合もあるため、その場合はデータを利用せずに答えよ。
(0.21018645, '約束はいらない', '「約束はいらない」(やくそくはいらない)は、坂本真綾のデビューシングル。') (0.24241784, '約束はいらない', '坂本真綾の歌手デビュー作品。当時坂本はまだ無名の声優であったが、同曲がテーマソングとなったアニメ『天空のエスカフローネ』とともに知名度を上げることとなる。後に「指輪」が同アニメの劇場版映画の主題歌とな') (0.2441598, '坂本真綾', '本格的な歌手活動は、1996年にテレビアニメ『天空のエスカフローネ』へ出演したことがきっかけで始めており、同作のオープニングテーマソングである「約束はいらない」(岩里祐穂作詞、菅野よう子作曲)をシング') (0.26043588, 'Suara', '大阪府立豊中高等学校を経て大阪外国語大学でインドネシア語を専攻中にバンド・ユニットを組んで音楽活動を始めた。普段はお笑い番組が大好きなこともあってよく喋るほうだが、東京の仕事で標準語の喋りをする時は、') (0.2690969, '早見沙織', '声優デビュー時より、数多くの主題歌やキャラクターソングを担当し、バラードからポップス、ヒットソングのカバー曲や英語の楽曲など、様々な曲を歌いこなす。2009年には吉田仁美とのユニット「blue dro')
提供されたデータによれば、楽曲「約束はいらない」は坂本真綾のデビューシングルであり、坂本真綾が歌手デビューした作品であることが確認できます。したがって、この楽曲でデビューした声優は坂本真綾（さかもと まあや）です。
きちんと回答することができた。試した内容はこちら。なおGPT4では、RAG で追加の情報を入れずとも最初から正解を回答できた。
wikipedia の文、passage (chunk) データをどう作るか
Wikipedia のデータを取得するには、全文のダンプデータ, CirrusSearchのダンプデータ, API を叩く方法等々がある。各々長所や短所があるのだが、綺麗なテキストを取得するには、APIを叩く方法が一番良いと個人的に思っている。ただ、APIを記事分叩くのが大変なので、最初はお手軽に使えるCirrusSearchダンプデータを chunk に分割して試していたりしていた。
そんな中、Masatoshi Suzukiさん作のsingletongue/wikipedia-utilsのデータに出会う。中身を見てみると記事タイトル、セクションタイトルも含む、綺麗なテキストデータとなっている。また passage も 200, 300, 400 chars と使いやすい文章長のデータとして公開されているので、こちらを元に embeddings を作成した。
文から embeddings への変換に何のモデルを使うか
文からの embeddings 変換モデルは利用用途によって性能(≒学習のさせ方)の方向性に種類があって、RAG のシステムにおいて重視されるのが質問に対してマッチ度が高くなるように学習された、 Retrieve (Reranking) タスクが強いモデルで、次点で似ている文章を探す性能が高い類似文章検索タスクが強いモデルであろう。
それらの性能が高そうなモデルのうち、無料で日本語モデルとして利用できるデータから、以下のモデルを embeddings に変換している。なお、OpenAI の embeddings API である ada-v2 は入力長tokensは長いというメリットがあるが、出力次元の大きさ(1536次元)や、変換したデータのライセンスの不明瞭さ、また何よりAPI利用料がかかるため除外した。
- multilingal-e5 シリーズ
- embeddings に変換する文章の先頭に、Retrieve用なら "passage: " を、それ以外なら "query: " をつけることを意識して使う必要があるが、かなり性能が高く感じる。
- pkshatech/GLuCoSE-base-ja
- cl-nagoya/sup-simcse-ja-base
これらを先ほどの passage-400 をこれらのモデルで embeddings に変換した。e5 対象文の先頭に "passage: " をつけるか "query: " をつけるかで、生成されるベクトルが別物になるので、各々の prefix をつけたデータで生成した embeddings を公開している。
embeddings に加え、検索用にそれらの faiss の index も公開している。index 作成のパラメータはこちらでの評価を参考にIVFのnlist=2048,PQの量子化オプションはベクトルの次元サイズの1/4(384なら96)とした。これらの各々のモデルで wikipedia 文の embeddings 検索がどのようになるのか、興味があればぜひ試してみてほしい。
変換に使った実装は、hotchpotch/wikipedia-passages-jawiki-embeddings-utilsで公開している。datasets_to_embs.pyではデータセットからembeddingsへの変換を、embs_to_faiss.pyではembeddingsからfaiss indexの実装である。
おわりに
本記事では、汎用的に試しやすいようなRAGのためのWikipedia日本語のembeddingsと検索インデクスの作成&紹介を行った。LLM が台頭してきた今年、検索のユースケースも大きく変わっていくであろうことから、検索技術の再発見・再定義が始まり出した年だと思っている。RAG(今回触れたのは、RAGの中でもベクトル検索という技術の1要素にすぎない)も含め、様々な検索ユースケースへの広がりを感じており、来年以降の検索技術の発展・活用も楽しみである。
