WikipediaダンプデータからRAG向けテキストデータ作成

 August 13, 2024に投稿  |  Kenji Arai  • 翻訳：en

モチベーション

LangChainを使ってRAGを試しているのだが、確認用に使用するデータに何を使おうかと考えていたところ、wikipediaのダンプデータを使うことにした。全体ではボリュームを大きいので、自分の興味のある天文関係のカテゴリーのデータを使うことにした。

ここでは、wikipediaダンプデータから特定のカテゴリーのデータのみを取り出す一連の手順をまとめた。

情報源

  1. Index of /jawiki/ 日本語のwikipedia dumpのトップページ。今回自分は、「20240720」ディレクトリ配下のデータを使用した。
  2. Wikipediaの特定カテゴリの記事のみを取得する 自分が実施しようと考えていたことが掲載されているページ。大いに参考にさせてもらった、感謝。
  3. Wikipediaの特定カテゴリ以下の記事だけを取得する（サブカテゴリの取得） 最初自分もMySQL（MariaDB）にwikipediaデータ、カテゴリ、ページ情報を格納して、検索することを試そうとした。mariadbのコンテナをdocker-compose.ymlで起動し、jawikiという名前のデータベースを作成し、category, categorylinks, pageテーブルまでは作成したが、検索用SQL作成するあたりで挫けた。
  4. attardi/wikiextractor wikiダンプデータを整形するツールを紹介しているページ。
  5. Python3.7のWikiExtractor3.0.4で起こるImportError wikiダンプデータ整形ツールでエラーが発生した場合の対処が紹介されている。
  6. Wikipedia:PetScan wikipediaの記事カテゴリー（およびサブカテゴリ）を検索し、条件に合った記事を情報（タイトルやページID）を得るツールのページ。

手順

jawikiダンプデータのダウンロード

情報源 1.で示した、~/jawiki/20240720/配下の「jawiki-20240720-pages-articles.xml.bz2」を作業用のディレクトリにダウンロードする。

$ ls -l
-rw-rw-r--  1 kenji kenji 4189767963  7月 29 22:03 jawiki-20240720-pages-articles.xml.bz2

ダンプデータを整形

情報源4.から得た「WikiExtractor.py」を使って、ダウンロードしたbz2ファイルを整形する。

$ python3 WikiExtractor.py -b 500K -o jawiki jawiki-20240720-pages-articles.xml.bz2

次のようなエラーが発生した。

Traceback (most recent call last):
  File "/ext/nfs/workspace/wikipedia/WikiExtractor.py", line 66, in <module>
    from .extract import Extractor, ignoreTag, define_template, acceptedNamespaces
ImportError: attempted relative import with no known parent package

ここで、ネットで検索したところ情報源5.がヒットしたので、次のように必要なパッケージをインストールした。

$ sudo apt install python3-pip
$ pip3 install wikiextractor

改めて、次のコマンドにより、jawikiディレクトリに「AA」から「DB」ディレクトリが作成され、その配下にwiki_00からwiki_99まで100個のテキストファイルが格納されている。「DB」配下のみwiki_28まで。

$ python3 -m wikiextractor.WikiExtractor -b 500K -o jawiki jawiki-20240720-pages-articles.xml.bz2

自分の環境では、約28分掛かった。

興味あるカテゴリーのページIDを得る 〜 PetScan

情報源6.を開いて、左上の「PetScanを開く」で開かれるツール。今回は「カテゴリ」と「出力」タブのみを使用する。

考え方

カテゴリーとしては、「天体物理学」と「天文学」を考えていたのだが、カテゴリ深度（サブカテゴリの深さ）をどの程度にするかは、得られるタイトルを見て決めた。天体物理学はカテゴリ深度：4、天文学はカテゴリ深度：2とした。

検索条件

上記の考え方をもとに、次のように設定した。

「カテゴリ」タグでは、言語：ja、カテゴリ：天体物理学|4 天文学|2、組み合わせ：ユニオン（和集合）

複数カテゴリを選ぶ時は、カテゴリ深度で指定はせず、カテゴリで、各カテゴリ毎に指定する。 「出力」タブで、出力形式：CSVを選ぶ。

「カテゴリ」タブの左下の「実行」をクリックすると、「ダウンロード.csv」がダウンロードされる。

自分は、これをastronomy.csvにrenameした。

テキスト化プログラム

ここまでで、抽出すべきページIDはastronomy.csvの「pageid」列に格納されており、テキストデータは、「jawiki」ディレクトリ配下に格納されている。

以下に紹介するプログラムは、以下の2つの仮定で作られている。

  • 「jakiki」配下の「AA」から「DB」に格納されているwiki_??ファイル中のdocidは、昇順となっている。
  • 「astronomy.csv」の「pageid」列も昇順となっている。
ヘッダ
# jawikiのダンプデータから、天文関連のカテゴリー（天体物理学、天文学）のページを抽出し、テキスト化する。
#
# 次のことが事前に準備された状態で、このプログラムは動作する。
# jawikiのタンプデータは、"WIKI_PATH"で与えられたディレクトリ配下のAAからDBディレクトリ内に内に
# wiki_00-wiki_99という名前のファイルに分割して格納されている。
# 抽出すべきカテゴリーに属するページは、PetScanを使って、"INTEREST_PAGES"の"pageid"列として抽出されている。
ファイルリスト、ページIDリストを作成
import os
import glob
import csv

WIKI_PATH = "jawiki"
INTEREST_PAGES = "astronomy.csv"
EXTRACT = "textdb"

# 対象ディレクトリ配下のファイルリストを作成
# ページ番号が昇順になるように、ファイルリストはソートしておく。
file_list = []
dir_list = sorted(os.listdir(WIKI_PATH))
for d in dir_list:
    file_list += sorted(glob.glob(os.path.join(WIKI_PATH, d) + "/*"))

# 抽出すべきページのリストを作成
pageid_list = []
with open(INTEREST_PAGES, encoding='utf-8') as f:
    csvreader = csv.reader(f)
    header = next(csvreader)  # skip header
    count = 0
    for line in csvreader:
        # namespaceがNullの記事のみを使う。　それ以外はスキップする。
        # それ以外はサブカテゴリー等であり、その配下については、"pageid"が昇順が保証されない。
        if line[3] != "":
            continue
        pageid_list.append(int(line[2]))
        count += 1

print("="*80)
print("抽出する記事数:{}".format(count))
print("="*80)

実行結果は次のとおり。

================================================================================
抽出する記事数:13842
================================================================================
記事を抽出する関数
import re

# 次のような構造からid/url/title/本文を抽出するための、正規表現を準備。
"""
<doc id="5" url="https://ja.wikipedia.org/wiki?curid=5" title="アンパサンド">
アンパサンド

アンパサンド（&amp;, ）は、並立助詞「…と…」を意味する記号である。・・・・

</doc>
"""

doc_re = re.compile(r'<doc id=.+?</doc>')
head_re = re.compile(r'<doc id=.+?">')
id_re = re.compile(r'id="\d+"')
url_re = re.compile(r'url=".+?"')
title_re = re.compile(r'title=".+"')

# ファイルパスで与えられたブロックを読み込み
# 記事（article）のリストを返す。
# 改行は取り除く。
def get_article_list(file):
    f = open(file, 'r', encoding='utf-8')
    block_data = f.read()
    block_data = block_data.replace('\n', '')
    return doc_re.findall(block_data)

# "doc"で与えられた一つの記事（article)から
# id/url/title/text(本文）を抽出し、それらを返す。
# 本文は、記事タイトルが冒頭にあるので、それを飛ばす（[len(title):])。
def get_article(doc):
    head = head_re.search(doc)
    id_tag = id_re.search(head.group())
    doc_id = re.search(r'\d+', id_tag.group())
    doc_id = int(doc_id.group())
    url_tag = url_re.search(head.group())
    url = url_tag.group()[len("url="):].replace('"', '')
    title_tag = title_re.search(head.group())
    doc_title = re.search(r'".+"',title_tag.group())
    title = doc_title.group().replace('"', '')
#    text = doc.replace(head.group(), '').rstrip('</doc>')[len(title):]
    text = doc.replace(head.group(), '').rstrip('</doc>')
    # タイトルが重複しないケースがある。
    # その場合、1つ目のタイトルをスキップすることはしない。
    if text[len(title):(len(title)+len(title))] == title:
        text = text[len(title):]
    return doc_id, url, text
メインループ
import time

# Wikipediaダンプデータ中の"docid"が事前に準備したカテゴリに属する"pageid"と一致した記事を出力する。
# 出力は、write()ではなく、print()を使う。テキスト出力の場合、print()の方が扱いやすい。

start_time = time.time()

ff = open(EXTRACT, 'w', encoding='utf-8')

no = 0
pageid = pageid_list[no]

id_list = []
no_page_list = []

for file in file_list:
    for page in get_article_list(file):
        docid, url, text = get_article(page)
        id_list.append(docid)
        if docid < pageid:
            continue
        if docid == pageid:
            print(text, file=ff)  # write text into file
            no += 1
            pageid = pageid_list[no]
        if docid > pageid:
            # 常に、docid <= pageid が保たれていること。
            # この場合は、pageidが見つからなかった。
            no_page_list.append(pageid)
            no +=1
            pageid = pageid_list[no]

ff.close()

end_time = time.time()
processing_time = end_time - start_time

print("processing_time(sec): ", processing_time)
print("="*80)
print("全記事数:{}\t抽出した記事数:{}\t未検出記事数:{}".format(len(id_list), no, len(no_page_list)))
print("="*80)

実行結果は次のとおり。

processing_time(sec):  70.98843002319336
================================================================================
全記事数:2304095	抽出した記事数:13837	未検出記事数:1
================================================================================

まとめ

今回作成した「textdb」という名前のテキストファイルを使って、今後RAGのベクトルDBに格納していく。