russell@unturf.com
f90b7c69f0
Drops in for any dated snapshot in dumps.wikimedia.org/archive — pages-articles.xml.bz2 and pages-meta-current.xml.bz2 (cur snapshot, single-revision-per-page) plus pages-meta-history.xml.bz2 (multi_revision=True). abstract.xml feed yields pre-distilled summaries at ~1/100th the chunk volume. Streams .bz2/.gz directly via iterparse with bounded memory; same shard / resume contract as the SQL source. Falls back to title-prefix namespace filtering when older export schemas omit the per-page <ns> element (e.g. enwiki 20101011). Makefile defaults target enwiki 20101011 (6.2 GB). Override with WP_XML_YEAR / WP_XML_MONTH / WP_XML_DATE / WP_XML_LANG to fetch any other archived snapshot.
320 lines
12 KiB
Python
320 lines
12 KiB
Python
"""MediaWiki XML 'pages-articles' / 'pages-meta-current' / 'abstract' source.
|
|
|
|
Phase IV (2006+) Wikipedia dumps switched from MySQL extended INSERT
|
|
syntax to a custom XML schema. These sources stream that XML and yield
|
|
Documents the same way the WikipediaSqlDump source yields rows from
|
|
the 2003 cur table — same URI scheme, same edge extraction, same
|
|
shard / resume contract — so downstream ingest, sharding, and resume
|
|
logic Just Works.
|
|
|
|
Two source classes:
|
|
|
|
* `WikipediaXmlDump` — pages-articles.xml.bz2 (cur snapshot, main NS) or
|
|
pages-meta-current.xml.bz2 (cur snapshot, all NS) or
|
|
pages-meta-history.xml.bz2 (full history; pass multi_revision=True).
|
|
One Document per page (or per revision in history mode).
|
|
|
|
* `WikipediaAbstractDump` — abstract.xml (uncompressed). Each `<doc>`
|
|
has <title>, <url>, <abstract>: short pre-distilled summaries of every
|
|
article. Useful as cheap retrieval targets for FTS5 queries.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import bz2
|
|
import gzip
|
|
import re
|
|
from pathlib import Path
|
|
from typing import IO, Iterator
|
|
from xml.etree.ElementTree import Element, iterparse
|
|
|
|
from aborist.document import Document, Edge
|
|
from aborist.source import Source
|
|
from aborist.sources.wikipedia import _extract_wikilinks
|
|
|
|
|
|
# MediaWiki export schema declares an xmlns; ElementTree returns tag names
|
|
# as `{http://www.mediawiki.org/xml/export-0.10/}page`. Strip that prefix
|
|
# off so the rest of the parser doesn't have to care which schema version
|
|
# any given dump uses.
|
|
_NS_RE = re.compile(r"^\{[^}]+\}")
|
|
|
|
# Body that begins with `#REDIRECT [[Target]]` is a redirect — skip.
|
|
# Real dumps also use the explicit <redirect title="..."/> element, but
|
|
# older or imported pages sometimes lack that and only carry the body
|
|
# directive; we belt-and-suspenders both.
|
|
_REDIRECT_RE = re.compile(r"\s*#REDIRECT\b", re.IGNORECASE)
|
|
|
|
NAMESPACE_MAIN = 0
|
|
|
|
|
|
# Title-prefix-encoded namespaces. Only used as a fallback filter when the
|
|
# export schema doesn't declare <ns> per page. Excludes Wiktionary-style
|
|
# language prefixes (those don't appear in enwiki main namespace anyway).
|
|
_NON_MAIN_NS_PREFIXES = frozenset({
|
|
"Talk", "User", "User talk", "Wikipedia", "Wikipedia talk",
|
|
"File", "File talk", "MediaWiki", "MediaWiki talk",
|
|
"Template", "Template talk", "Help", "Help talk",
|
|
"Category", "Category talk", "Portal", "Portal talk",
|
|
"Book", "Book talk", "Draft", "Draft talk",
|
|
"Module", "Module talk", "Image", "Image talk",
|
|
})
|
|
|
|
|
|
def _strip_ns(tag: str) -> str:
|
|
return _NS_RE.sub("", tag)
|
|
|
|
|
|
def _open_streaming(path: Path) -> IO[bytes]:
|
|
"""Open a file as a binary byte stream, transparently decompressing."""
|
|
p = str(path)
|
|
if p.endswith(".bz2"):
|
|
return bz2.open(p, "rb")
|
|
if p.endswith(".gz"):
|
|
return gzip.open(p, "rb")
|
|
return open(p, "rb")
|
|
|
|
|
|
class WikipediaXmlDump(Source):
|
|
"""Iterates a MediaWiki XML dump.
|
|
|
|
Phase IV format: `<mediawiki>` root containing one `<siteinfo>` plus
|
|
many `<page>` children. Each `<page>` has `<title>`, `<ns>`, `<id>`,
|
|
optional `<redirect>`, and one or more `<revision>` children.
|
|
|
|
For pages-articles or pages-meta-current dumps, each page has exactly
|
|
one `<revision>` (the current one). For pages-meta-history dumps each
|
|
page has many. By default we emit the LAST revision per page (matches
|
|
cur semantics); pass `multi_revision=True` to emit every revision
|
|
(matches old semantics, with `supersedes` chaining at the URI level
|
|
via aborist's existing prior-document detection).
|
|
"""
|
|
|
|
source_type = "wikipedia_xml"
|
|
|
|
def __init__(
|
|
self,
|
|
path: str | Path,
|
|
*,
|
|
namespace: int = NAMESPACE_MAIN,
|
|
base_uri: str = "https://en.wikipedia.org/wiki/",
|
|
shard: tuple[int, int] | None = None,
|
|
start_id: int = 0,
|
|
multi_revision: bool = False,
|
|
):
|
|
self.path = Path(path)
|
|
self.namespace = namespace
|
|
self.base_uri = base_uri
|
|
if shard is not None:
|
|
rank, total = shard
|
|
if not (0 <= rank < total) or total < 1:
|
|
raise ValueError(
|
|
f"invalid shard {shard}: need 0 <= rank < total, total >= 1"
|
|
)
|
|
self.shard = shard
|
|
# Resume support — same contract as WikipediaSqlDump.
|
|
self.start_id = start_id
|
|
self.last_id: int = start_id
|
|
self.multi_revision = multi_revision
|
|
|
|
def iter_documents(self) -> Iterator[Document]:
|
|
rank, total = (0, 1) if self.shard is None else self.shard
|
|
idx = 0
|
|
with _open_streaming(self.path) as f:
|
|
context = iter(iterparse(f, events=("start", "end")))
|
|
try:
|
|
_, root = next(context)
|
|
except StopIteration:
|
|
return
|
|
for event, elem in context:
|
|
if event != "end":
|
|
continue
|
|
tag = _strip_ns(elem.tag)
|
|
if tag != "page":
|
|
# Free siteinfo/etc immediately — no useful data downstream.
|
|
if tag in ("siteinfo", "namespaces", "namespace"):
|
|
elem.clear()
|
|
continue
|
|
# Stride filter applied at the page index. Every shard sees
|
|
# the full XML stream but only emits its share — matches the
|
|
# SQL source's behavior so ingest counters line up.
|
|
if idx % total != rank:
|
|
idx += 1
|
|
elem.clear()
|
|
root.clear()
|
|
continue
|
|
idx += 1
|
|
yield from self._page_to_docs(elem)
|
|
elem.clear()
|
|
root.clear()
|
|
|
|
def _page_to_docs(self, page: Element) -> Iterator[Document]:
|
|
title: str | None = None
|
|
ns_text: str | None = None
|
|
page_id: int = 0
|
|
is_redirect = False
|
|
revisions: list[Element] = []
|
|
for child in page:
|
|
ctag = _strip_ns(child.tag)
|
|
if ctag == "title":
|
|
title = (child.text or "").strip() or None
|
|
elif ctag == "ns":
|
|
ns_text = (child.text or "").strip()
|
|
elif ctag == "id":
|
|
try:
|
|
page_id = int((child.text or "0").strip())
|
|
except ValueError:
|
|
pass
|
|
elif ctag == "redirect":
|
|
# Self-closing element with title attr. Body may also start
|
|
# with #REDIRECT — both flag the same condition.
|
|
is_redirect = True
|
|
elif ctag == "revision":
|
|
revisions.append(child)
|
|
|
|
try:
|
|
ns_int = int(ns_text) if ns_text is not None else None
|
|
except ValueError:
|
|
ns_int = None
|
|
# Older export schemas (pre-0.4) had no <ns> child element under
|
|
# <page>; namespace was encoded into the title prefix only. The
|
|
# `pages-articles.xml.bz2` dump is documented to contain just the
|
|
# main namespace plus redirects, so trust the file when <ns> is
|
|
# absent. For files that include sister namespaces (e.g.
|
|
# `pages-meta-current.xml.bz2`), filter on the title prefix as
|
|
# a safety belt — any title starting `Talk:` / `User:` /
|
|
# `Wikipedia:` / `File:` / `Template:` / `Help:` / `Category:` /
|
|
# `Portal:` / `Book:` / `Draft:` / `MediaWiki:` is non-main.
|
|
if ns_int is None:
|
|
if title and ":" in title:
|
|
head = title.split(":", 1)[0]
|
|
if head in _NON_MAIN_NS_PREFIXES:
|
|
return
|
|
elif ns_int != self.namespace:
|
|
return
|
|
if is_redirect:
|
|
return
|
|
if not title or not revisions:
|
|
return
|
|
|
|
# Resume fast-forward.
|
|
if self.start_id and page_id <= self.start_id:
|
|
return
|
|
if page_id > self.last_id:
|
|
self.last_id = page_id
|
|
|
|
revs = revisions if self.multi_revision else revisions[-1:]
|
|
uri = self.base_uri + title.replace(" ", "_")
|
|
for rev in revs:
|
|
text = ""
|
|
timestamp = ""
|
|
rev_id = 0
|
|
for child in rev:
|
|
ctag = _strip_ns(child.tag)
|
|
if ctag == "text":
|
|
text = child.text or ""
|
|
elif ctag == "timestamp":
|
|
timestamp = (child.text or "").strip()
|
|
elif ctag == "id":
|
|
try:
|
|
rev_id = int((child.text or "0").strip())
|
|
except ValueError:
|
|
pass
|
|
if not text:
|
|
continue
|
|
if _REDIRECT_RE.match(text):
|
|
continue
|
|
edges = _extract_wikilinks(text, self.base_uri)
|
|
extra: dict = {"page_id": str(page_id)}
|
|
if timestamp:
|
|
extra["timestamp"] = timestamp
|
|
if rev_id:
|
|
extra["rev_id"] = str(rev_id)
|
|
yield Document(
|
|
uri=uri,
|
|
content=text,
|
|
source_type=self.source_type,
|
|
title=title,
|
|
edges=edges,
|
|
extra=extra,
|
|
)
|
|
|
|
|
|
class WikipediaAbstractDump(Source):
|
|
"""Iterates a MediaWiki `abstract.xml` dump.
|
|
|
|
Format:
|
|
<feed>
|
|
<doc>
|
|
<title>Wikipedia: Anarchism</title>
|
|
<url>http://en.wikipedia.org/wiki/Anarchism</url>
|
|
<abstract>Anarchism is a political philosophy...</abstract>
|
|
<links>...</links>
|
|
</doc>
|
|
...
|
|
</feed>
|
|
|
|
Each <doc> already contains a one-paragraph summary — pre-distilled
|
|
by the dump generator. Ingesting these alongside (or instead of) full
|
|
pages-articles gives FTS5 a high-signal corpus at ~1/100th the bytes.
|
|
"""
|
|
|
|
source_type = "wikipedia_abstract"
|
|
|
|
def __init__(self, path: str | Path):
|
|
self.path = Path(path)
|
|
# Resume contract: abstract docs have no stable numeric id, so
|
|
# start_id / last_id are no-ops here. Kept for interface symmetry.
|
|
self.start_id = 0
|
|
self.last_id = 0
|
|
|
|
def iter_documents(self) -> Iterator[Document]:
|
|
with _open_streaming(self.path) as f:
|
|
context = iter(iterparse(f, events=("start", "end")))
|
|
try:
|
|
_, root = next(context)
|
|
except StopIteration:
|
|
return
|
|
for event, elem in context:
|
|
if event != "end":
|
|
continue
|
|
tag = _strip_ns(elem.tag)
|
|
if tag != "doc":
|
|
continue
|
|
title = ""
|
|
url = ""
|
|
abstract = ""
|
|
links: list[str] = []
|
|
for child in elem:
|
|
ctag = _strip_ns(child.tag)
|
|
if ctag == "title":
|
|
title = (child.text or "").strip()
|
|
elif ctag == "url":
|
|
url = (child.text or "").strip()
|
|
elif ctag == "abstract":
|
|
abstract = (child.text or "").strip()
|
|
elif ctag == "links":
|
|
for sublink in child:
|
|
for sub in sublink:
|
|
if _strip_ns(sub.tag) == "link":
|
|
href = (sub.text or "").strip()
|
|
if href:
|
|
links.append(href)
|
|
elem.clear()
|
|
root.clear()
|
|
if not url or not abstract:
|
|
continue
|
|
# Abstract feed prepends "Wikipedia: " to titles. Strip it.
|
|
clean_title = title
|
|
if clean_title.startswith("Wikipedia: "):
|
|
clean_title = clean_title[len("Wikipedia: "):]
|
|
edges = [
|
|
Edge(edge_type="abstract_link", dst_uri=u) for u in links
|
|
]
|
|
yield Document(
|
|
uri=url,
|
|
content=abstract,
|
|
source_type=self.source_type,
|
|
title=clean_title or None,
|
|
edges=edges,
|
|
)
|