File size: 3,669 Bytes
7c5e40e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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
|