Creating text data for RAG from Wikipedia dump data

$ ls -l
-rw-rw-r--  1 kenji kenji 4189767963  7月 29 22:03 jawiki-20240720-pages-articles.xml.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

$ sudo apt install python3-pip
$ pip3 install wikiextractor

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

# Extract pages from jawiki dump data for astronomy-related categories (astrophysics and astronomy) and convert to text.
#
# This program runs with the following pre-prepared.
# jawiki dump data from AA under the directory given by "WIKI_PATH" into the DB directory
# The pages belonging to the category to be extracted are stored in files named "wiki_00-wiki_99" in the DB directory # under the directory given by "WIKI_PATH".
# The pages belonging to the category to be extracted are extracted using PetScan as "pageid" column in "INTEREST_PAGES".

import os
import glob
import csv

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

# Create a list of files under the target directory
# Sort the file list so that the page numbers are in ascending order.
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) + "/*"))

# Create a list of pages to be extracted
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:
        # Use only articles where namespace is null.　Otherwise, skip.
        # Other articles are subcategories, etc., 
        # and "pageid" is not guaranteed to be in ascending order under those subcategories.
        if line[3] != "":
            continue
        pageid_list.append(int(line[2]))
        count += 1

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

================================================================================
抽出する記事数:13842
================================================================================

import re

# Prepare a regular expression to extract id/url/title/body from the following structure.
"""
<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=".+"')

# Read the block given by the file path and
# Return a list of articles (ARTICLE).
# Remove line breaks.
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)

# Extract one article (ARTICLE) given in "doc" and return them.
# extract id/url/title/text (body) and return them.
# skip the body text because the article title is at the beginning ([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>')
    # There are cases where titles do not overlap.
    # In that case, do not skip the first title.
    if text[len(title):(len(title)+len(title))] == title:
        text = text[len(title):]
    return doc_id, url, text

import time

# Output articles whose "docid" in the Wikipedia dump data matches "pageid" belonging to the category prepared in advance.
# Use print() for output, not write(). In case of text output, print() is easier to handle.

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:
            # Always keep docid <= pageid.
            # In this case, pageid was not found.
            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
================================================================================