151 lines
6.2 KiB
Python
151 lines
6.2 KiB
Python
#!/usr/bin/env python3
|
||
"""Extract clean Russian sentences from an official Wikisource XML dump."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import bz2
|
||
import heapq
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import sys
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from xml.etree import ElementTree
|
||
|
||
|
||
SENTENCE_BOUNDARY = re.compile(r"(?<=[.!?…])\s+")
|
||
WHITESPACE = re.compile(r"\s+")
|
||
RUSSIAN_LETTER = re.compile(r"[А-Яа-яЁё]")
|
||
ANY_LETTER = re.compile(r"[^\W\d_]", re.UNICODE)
|
||
BAD_MARKUP = ("{{", "}}", "[[", "]]", "http://", "https://", "<math", "</", "{|", "|}")
|
||
|
||
|
||
def local_name(tag: str) -> str:
|
||
return tag.rsplit("}", 1)[-1]
|
||
|
||
|
||
def descendant_text(element, name: str) -> str:
|
||
for child in element.iter():
|
||
if local_name(child.tag) == name:
|
||
return child.text or ""
|
||
return ""
|
||
|
||
|
||
def clean_page_text(source: str) -> str:
|
||
import mwparserfromhell
|
||
|
||
source = re.sub(r"<ref\b[^>]*>.*?</ref>|<ref\b[^>]*/>", " ", source, flags=re.IGNORECASE | re.DOTALL)
|
||
source = re.sub(r"\{\|.*?\|\}", " ", source, flags=re.DOTALL)
|
||
plain = mwparserfromhell.parse(source).strip_code(normalize=True, collapse=True)
|
||
return WHITESPACE.sub(" ", plain).strip()
|
||
|
||
|
||
def acceptable(sentence: str) -> bool:
|
||
if not 40 <= len(sentence) <= 240 or any(marker in sentence for marker in BAD_MARKUP):
|
||
return False
|
||
if sentence[-1:] not in ".!?…":
|
||
return False
|
||
words = sentence.split()
|
||
if not 6 <= len(words) <= 45:
|
||
return False
|
||
letters = ANY_LETTER.findall(sentence)
|
||
if not letters:
|
||
return False
|
||
russian = RUSSIAN_LETTER.findall(sentence)
|
||
return len(russian) / len(letters) >= 0.75
|
||
|
||
|
||
def extract(args: argparse.Namespace) -> None:
|
||
temporary = args.output.with_suffix(args.output.suffix + ".part")
|
||
if args.output.exists() or args.metadata.exists() or temporary.exists():
|
||
raise FileExistsError("Refusing to overwrite output, metadata, or partial output")
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
seen: set[bytes] = set()
|
||
sample_heap: list[tuple[int, bytes, str]] = []
|
||
pages = 0
|
||
accepted = 0
|
||
candidates = 0
|
||
with bz2.open(args.dump, "rb") as source, temporary.open("w", encoding="utf-8", newline="\n") as output:
|
||
for _, element in ElementTree.iterparse(source, events=("end",)):
|
||
if local_name(element.tag) != "page":
|
||
continue
|
||
pages += 1
|
||
namespace = descendant_text(element, "ns")
|
||
title = descendant_text(element, "title")
|
||
wiki_text = descendant_text(element, "text")
|
||
if namespace == "0" and wiki_text and not title.startswith(("Категория:", "Шаблон:", "Справка:")):
|
||
for sentence in SENTENCE_BOUNDARY.split(clean_page_text(wiki_text)):
|
||
# Leading dashes carry dialogue/prosody information and must survive extraction.
|
||
sentence = sentence.strip(" \t\r\n")
|
||
if not acceptable(sentence):
|
||
continue
|
||
digest = hashlib.blake2b(sentence.encode("utf-8"), digest_size=16).digest()
|
||
candidates += 1
|
||
if args.sampling == "first":
|
||
if digest in seen:
|
||
continue
|
||
seen.add(digest)
|
||
output.write(sentence + "\n")
|
||
accepted += 1
|
||
if accepted % 10_000 == 0:
|
||
print(f"sentences={accepted} pages={pages}", flush=True)
|
||
if accepted >= args.max_sentences:
|
||
break
|
||
else:
|
||
score = int.from_bytes(digest, "big")
|
||
if digest in seen:
|
||
continue
|
||
if len(sample_heap) < args.max_sentences:
|
||
heapq.heappush(sample_heap, (-score, digest, sentence)); seen.add(digest)
|
||
elif score < -sample_heap[0][0]:
|
||
_, removed_digest, _ = heapq.heapreplace(sample_heap, (-score, digest, sentence))
|
||
seen.remove(removed_digest); seen.add(digest)
|
||
element.clear()
|
||
if args.sampling == "first" and accepted >= args.max_sentences:
|
||
break
|
||
if pages % 10_000 == 0:
|
||
print(f"pages={pages} candidates={candidates} retained={len(sample_heap)}", flush=True)
|
||
if args.sampling == "hash":
|
||
selected = sorted((-negative_score, sentence) for negative_score, _, sentence in sample_heap)
|
||
for _, sentence in selected:
|
||
output.write(sentence + "\n")
|
||
accepted = len(selected)
|
||
temporary.replace(args.output)
|
||
metadata = {
|
||
"source_url": args.source_url,
|
||
"dump_file": args.dump.name,
|
||
"dump_bytes": args.dump.stat().st_size,
|
||
"pages_scanned": pages,
|
||
"candidate_sentences": candidates,
|
||
"sentences": accepted,
|
||
"output_bytes": args.output.stat().st_size,
|
||
"created_utc": datetime.now(timezone.utc).isoformat(),
|
||
"sampling": args.sampling,
|
||
"filters": {"min_chars": 40, "max_chars": 240, "min_words": 6, "max_words": 45, "min_russian_letter_ratio": 0.75},
|
||
}
|
||
args.metadata.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||
print(json.dumps(metadata, ensure_ascii=False, indent=2))
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--dump", type=Path, required=True)
|
||
parser.add_argument("--output", type=Path, required=True)
|
||
parser.add_argument("--metadata", type=Path, required=True)
|
||
parser.add_argument("--max-sentences", type=int, default=200_000)
|
||
parser.add_argument(
|
||
"--sampling", choices=("first", "hash"), default="hash",
|
||
help="hash scans the complete dump and keeps a deterministic min-hash sample",
|
||
)
|
||
parser.add_argument(
|
||
"--source-url",
|
||
default="https://dumps.wikimedia.org/ruwikisource/latest/ruwikisource-latest-pages-articles-multistream.xml.bz2",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(extract(parse_args()) or 0)
|