source/src/strata/data/wikipedia.py

"""Wikipedia shard streaming utilities for STRATA data pipelines."""

from __future__ import annotations

import glob
import hashlib
import random
from pathlib import Path
from typing import Iterator


def document_split(text: str, *, holdout_mod: int) -> str:
    """Deterministically route a document to ``"train"`` or ``"heldout"``.

    Uses a stable content hash (not Python's process-randomised ``hash``) so a
    document lands in the same bucket across runs and machines. ``holdout_mod``
    is the reciprocal held-out fraction (e.g. 20 -> ~5% held out). Documents with
    ``sha1(text) % holdout_mod == 0`` are held out; the rest are training. This
    guarantees the two splits are disjoint by construction.
    """

    digest = hashlib.sha1(text.encode("utf-8")).digest()
    bucket = int.from_bytes(digest[:8], "big") % holdout_mod
    return "heldout" if bucket == 0 else "train"


def iter_parquet_texts(files: list[str], *, min_chars: int) -> Iterator[str]:
    """Yield text documents from Parquet shards with a minimum character length."""

    import pyarrow.parquet as pq

    for path in files:
        parquet = pq.ParquetFile(path)
        for batch in parquet.iter_batches(batch_size=512, columns=["text"]):
            for text in batch.column("text").to_pylist():
                if text and len(text) >= min_chars:
                    yield text


def discover_wikipedia_shards(
    data_root: Path,
    languages: list[str],
    *,
    snapshot: str,
    seed: int,
) -> dict[str, list[str]]:
    """Return shuffled raw Wikipedia shard paths for each requested language."""

    rng = random.Random(seed)
    shards: dict[str, list[str]] = {}
    for language in languages:
        pattern = str(
            data_root / "raw" / "wikipedia" / f"{snapshot}.{language}" / "*.parquet"
        )
        files = sorted(glob.glob(pattern))
        if not files:
            raise FileNotFoundError(
                f"no Wikipedia shards for language {language!r} at {pattern}"
            )
        rng.shuffle(files)
        shards[language] = files
    return shards


def iter_wikipedia_documents(
    data_root: Path,
    languages: list[str],
    *,
    min_chars: int,
    seed: int,
    snapshot: str,
    holdout_mod: int | None = None,
    split: str = "train",
) -> Iterator[tuple[str, str]]:
    """Round-robin ``(language, text)`` stream over raw Wikipedia shards.

    When ``holdout_mod`` is set, only documents whose deterministic
    :func:`document_split` bucket equals ``split`` (``"train"`` or ``"heldout"``)
    are yielded, giving a reproducible train/held-out partition that is disjoint
    by construction. With ``holdout_mod=None`` every document is yielded
    (unchanged legacy behaviour).
    """

    if split not in ("train", "heldout"):
        raise ValueError(f"split must be 'train' or 'heldout', got {split!r}")
    if holdout_mod is not None and holdout_mod < 2:
        raise ValueError(f"holdout_mod must be >= 2, got {holdout_mod}")

    shards = discover_wikipedia_shards(
        data_root,
        languages,
        snapshot=snapshot,
        seed=seed,
    )
    generators = {
        language: iter_parquet_texts(files, min_chars=min_chars)
        for language, files in shards.items()
    }
    active = list(languages)
    while active:
        for language in list(active):
            try:
                text = next(generators[language])
            except StopIteration:
                active.remove(language)
                continue
            if holdout_mod is not None and document_split(text, holdout_mod=holdout_mod) != split:
                continue
            yield language, text