111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
import bz2
|
|
import xml.etree.ElementTree as ET
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Generator, Dict, Any
|
|
|
|
def extract_text_from_wikipedia_dump(bz2_path: str, output_dir: str, limit: int = None) -> int:
|
|
"""Extract text from Wikipedia XML dump"""
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
count = 0
|
|
with bz2.BZ2File(bz2_path, 'r') as f:
|
|
context = ET.iterparse(f, events=('start', 'end'))
|
|
context = iter(context)
|
|
|
|
event, root = next(context)
|
|
|
|
for event, elem in context:
|
|
if event == 'end' and elem.tag.endswith('page'):
|
|
doc = parse_page(elem)
|
|
if doc and doc.get('text'):
|
|
filename = f"{output_dir}/wiki_{doc['id']}.json"
|
|
with open(filename, 'w', encoding='utf-8') as out:
|
|
json.dump(doc, out, ensure_ascii=False)
|
|
|
|
count += 1
|
|
if count % 1000 == 0:
|
|
print(f"Processed {count} documents")
|
|
|
|
if limit and count >= limit:
|
|
break
|
|
|
|
root.clear()
|
|
|
|
return count
|
|
|
|
def parse_page(page_elem) -> Dict[str, Any]:
|
|
"""Parse a single Wikipedia page"""
|
|
ns = {'wiki': 'http://www.mediawiki.org/xml/export-0.10/'}
|
|
|
|
title_elem = page_elem.find('wiki:title', ns) or page_elem.find('.//title')
|
|
page_id_elem = page_elem.find('wiki:id', ns) or page_elem.find('.//id')
|
|
revision_elem = page_elem.find('.//revision')
|
|
|
|
if revision_elem is None:
|
|
return None
|
|
|
|
text_elem = revision_elem.find('wiki:text', ns) or revision_elem.find('.//text')
|
|
|
|
if title_elem is None or text_elem is None:
|
|
return None
|
|
|
|
title = title_elem.text or ""
|
|
page_id = page_id_elem.text if page_id_elem is not None else ""
|
|
text = text_elem.text or ""
|
|
|
|
# Skip redirects and special pages
|
|
if text.startswith('#REDIRECT') or title.startswith('Wikipedia:'):
|
|
return None
|
|
|
|
# Basic text cleaning
|
|
text = clean_wikipedia_markup(text)
|
|
|
|
if len(text.split()) < 50:
|
|
return None
|
|
|
|
return {
|
|
"id": f"wiki_{page_id}",
|
|
"title": title,
|
|
"text": text,
|
|
"source": "wikipedia",
|
|
"lang": "ru" # Change for English
|
|
}
|
|
|
|
def clean_wikipedia_markup(text: str) -> str:
|
|
"""Remove Wikipedia markup"""
|
|
import re
|
|
|
|
# Remove wiki links
|
|
text = re.sub(r'\[\[([^\|]*\|)?([^\]]*)\]\]', r'\2', text)
|
|
# Remove external links
|
|
text = re.sub(r'\[http[^\s]+ ([^\]]*)\]', r'\1', text)
|
|
# Remove templates
|
|
text = re.sub(r'\{\{[^}]*\}\}', '', text)
|
|
# Remove categories
|
|
text = re.sub(r'\[\[Category:[^\]]*\]\]', '', text)
|
|
# Remove refs
|
|
text = re.sub(r'<ref[^>]*>[^<]*</ref>', '', text)
|
|
# Remove HTML tags
|
|
text = re.sub(r'<[^>]*>', '', text)
|
|
# Clean whitespace
|
|
text = re.sub(r'\n\n+', '\n', text)
|
|
text = re.sub(r' +', ' ', text)
|
|
|
|
return text.strip()
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) < 3:
|
|
print("Usage: python parse_wikipedia.py <input_bz2> <output_dir> [limit]")
|
|
sys.exit(1)
|
|
|
|
input_file = sys.argv[1]
|
|
output_dir = sys.argv[2]
|
|
limit = int(sys.argv[3]) if len(sys.argv) > 3 else None
|
|
|
|
count = extract_text_from_wikipedia_dump(input_file, output_dir, limit)
|
|
print(f"Successfully processed {count} documents")
|