ocr-foundry 0.2.0__py3-none-any.whl

khocr_gen/corpus.py

@@ -0,0 +1,621 @@
+"""Corpus loading, filtering, and counting.
+
+Handles UTF-8 plain-text corpora and JSONL corpus manifests: streams lines,
+applies length filters after normalization, and supports parallel count-only
+scans.
+
+JSONL manifests (one JSON object per line) are auto-detected by file extension
+(``.jsonl``).  Each object must have a ``"text"`` key.  Optional keys:
+``"source_id"``, ``"source_group"`` (alias ``"doc_id"``), ``"language"``.
+"""
+
+from __future__ import annotations
+
+import json as _json
+import logging
+import multiprocessing as mp
+import os
+from collections import Counter
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from .errors import InputValidationError
+from .normalizer import NormalizerConfig, normalize
+from .text_support import line_separator_error
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+_log = logging.getLogger("khocr_gen.corpus")
+
+_AUTO_PARALLEL_MIN_BYTES = 16 * 1024 * 1024
+
+_SKIP = 0
+_TOO_SHORT = 1
+_TOO_LONG = 2
+_PASS = 3
+
+
+# ---------------------------------------------------------------------------
+# CorpusItem: str subclass carrying optional provenance metadata
+# ---------------------------------------------------------------------------
+
+
+class CorpusItem(str):
+    """A text string with optional provenance metadata from a JSONL manifest.
+
+    Behaves exactly like ``str`` in every context that expects a string value.
+    The extra attributes survive pickling (needed for multiprocessing workers)
+    because ``str`` subclasses keep a ``__dict__``.
+    """
+    source_id: str | None
+    source_group: str | None
+    language: str | None
+    script: str | None
+    domain: str | None
+    source_uri: str | None
+    license: str | None
+    retrieved_at: str | None
+    source_sha256: str | None
+    normalization_policy: str | None
+    consent_or_restrictions: str | None
+    raw_text: str | None
+
+    def __new__(
+        cls,
+        text: str,
+        source_id: str | None = None,
+        source_group: str | None = None,
+        language: str | None = None,
+        script: str | None = None,
+        domain: str | None = None,
+        source_uri: str | None = None,
+        license: str | None = None,
+        retrieved_at: str | None = None,
+        source_sha256: str | None = None,
+        normalization_policy: str | None = None,
+        consent_or_restrictions: str | None = None,
+        raw_text: str | None = None,
+        **kwargs: Any,
+    ) -> CorpusItem:
+        obj = super().__new__(cls, text)
+        obj.source_id = source_id
+        obj.source_group = source_group
+        obj.language = language
+        obj.script = script
+        obj.domain = domain
+        obj.source_uri = source_uri
+        obj.license = license
+        obj.retrieved_at = retrieved_at
+        obj.source_sha256 = source_sha256
+        obj.normalization_policy = normalization_policy
+        obj.consent_or_restrictions = consent_or_restrictions
+        obj.raw_text = raw_text
+        return obj
+
+    def __repr__(self) -> str:
+        parts = [repr(str(self))]
+        for field_name in (
+            "source_id",
+            "source_group",
+            "language",
+            "script",
+            "domain",
+            "source_uri",
+            "license",
+            "retrieved_at",
+            "source_sha256",
+            "normalization_policy",
+            "consent_or_restrictions",
+            "raw_text",
+        ):
+            val = getattr(self, field_name, None)
+            if val is not None:
+                parts.append(f"{field_name}={val!r}")
+        return f"CorpusItem({', '.join(parts)})"
+
+def _is_jsonl(path: str | Path) -> bool:
+    """True when *path* has a ``.jsonl`` extension (case-insensitive)."""
+    return str(path).lower().endswith(".jsonl")
+
+
+def _parse_jsonl_record(
+    raw_line: str,
+    line_number: int,
+    *,
+    strict: bool = True,
+) -> dict[str, Any] | None:
+    """Parse one JSONL line into a dict with at least a ``"text"`` key.
+
+    Returns *None* for blank/empty lines.  Raises on malformed JSON or
+    missing ``"text"`` when *strict* is True; logs a warning and returns
+    *None* when *strict* is False.
+    """
+    stripped = raw_line.strip()
+    if not stripped:
+        return None
+    try:
+        obj = _json.loads(stripped)
+    except _json.JSONDecodeError as exc:
+        msg = f"JSONL line {line_number}: invalid JSON: {exc}"
+        if strict:
+            raise InputValidationError(msg, code="jsonl_parse_error") from exc
+        _log.warning(msg)
+        return None
+    if not isinstance(obj, dict):
+        msg = f"JSONL line {line_number}: expected JSON object, got {type(obj).__name__}"
+        if strict:
+            raise InputValidationError(msg, code="jsonl_parse_error")
+        _log.warning(msg)
+        return None
+    if "text" not in obj:
+        msg = f"JSONL line {line_number}: missing required 'text' key"
+        if strict:
+            raise InputValidationError(msg, code="jsonl_parse_error")
+        _log.warning(msg)
+        return None
+    if not isinstance(obj["text"], str):
+        msg = f"JSONL line {line_number}: 'text' must be a string"
+        if strict:
+            raise InputValidationError(msg, code="jsonl_parse_error")
+        _log.warning(msg)
+        return None
+    # Accept doc_id as alias for source_group
+    if "doc_id" in obj and "source_group" not in obj:
+        obj["source_group"] = obj["doc_id"]
+    return obj
+
+
+def _classify_jsonl_text(
+    text: str,
+    min_length: int,
+    max_length: int,
+    normalizer: NormalizerConfig | None = None,
+) -> tuple[int, str | None]:
+    """Classify a ``"text"`` value from a JSONL record through the same
+    normalization and length-filter pipeline as plain-text lines."""
+    return _classify_raw_line(text, min_length, max_length, normalizer)
+
+
+def _empty_stats() -> dict[str, int]:
+    return {"total": 0, "passing": 0, "too_short": 0, "too_long": 0}
+
+
+def _merge_stats(base: dict[str, int], update: dict[str, int]) -> dict[str, int]:
+    for key, value in update.items():
+        base[key] += value
+    return base
+
+
+def _update_stats_for_status(stats: dict[str, int], status: int) -> None:
+    if status == _SKIP:
+        return
+    stats["total"] += 1
+    if status == _TOO_SHORT:
+        stats["too_short"] += 1
+    elif status == _TOO_LONG:
+        stats["too_long"] += 1
+    else:
+        stats["passing"] += 1
+
+
+def _classify_raw_line(
+    raw_line: str,
+    min_length: int,
+    max_length: int,
+    normalizer: NormalizerConfig | None = None,
+) -> tuple[int, str | None]:
+    # Zl/Zp survive strip() when embedded. Never accept them as one corpus line.
+    if line_separator_error(raw_line):
+        return _SKIP, None
+
+    raw = raw_line.strip()
+    if not raw:
+        return _SKIP, None
+
+    raw_len = len(raw)
+    if raw_len < min_length:
+        return _TOO_SHORT, None
+    if raw_len > max_length:
+        return _TOO_LONG, None
+
+    # Normalize (Khmer canonical reorder) before length check
+    line = normalize(raw, normalizer)
+    line_len = len(line)
+    if line_len < min_length:
+        return _TOO_SHORT, None
+    if line_len > max_length:
+        return _TOO_LONG, None
+
+    return _PASS, line
+
+
+def _count_lines(
+    lines, min_length: int, max_length: int, normalizer: NormalizerConfig | None = None
+) -> dict[str, int]:
+    stats = _empty_stats()
+    for raw_line in lines:
+        status, _ = _classify_raw_line(raw_line, min_length, max_length, normalizer)
+        _update_stats_for_status(stats, status)
+    return stats
+
+
+def _count_file_chunk(
+    path_str: str,
+    start: int,
+    end: int,
+    min_length: int,
+    max_length: int,
+    normalizer: NormalizerConfig | None = None,
+) -> dict[str, int]:
+    _UTF8_BOM = b"\xef\xbb\xbf"
+    with open(path_str, "rb") as fh:
+        if start != 0:
+            fh.seek(start - 1)
+            if fh.read(1) not in {b"\n", b"\r"}:
+                fh.readline()
+        else:
+            fh.seek(start)
+            maybe_bom = fh.read(len(_UTF8_BOM))
+            if not maybe_bom.startswith(_UTF8_BOM):
+                fh.seek(start)
+
+        stats = _empty_stats()
+        while True:
+            line_start = fh.tell()
+            if line_start >= end:
+                break
+            raw_line = fh.readline()
+            if not raw_line:
+                break
+            status, _ = _classify_raw_line(
+                raw_line.decode("utf-8", errors="replace"),
+                min_length=min_length,
+                max_length=max_length,
+                normalizer=normalizer,
+            )
+            _update_stats_for_status(stats, status)
+    return stats
+
+
+def _iter_chunk_ranges(path: Path, workers: int):
+    file_size = path.stat().st_size
+    chunk_size = max(1, file_size // workers)
+    start = 0
+    for worker_index in range(workers):
+        end = file_size if worker_index == workers - 1 else start + chunk_size
+        yield start, end
+        start = end
+
+
+def _resolve_worker_count(path: Path, workers: int) -> int:
+    requested = max(0, int(workers))
+    cpu_count = os.cpu_count() or 1
+    if requested == 1:
+        return 1
+    if requested > 1:
+        return min(requested, cpu_count)
+    if cpu_count <= 1:
+        return 1
+    try:
+        if path.stat().st_size < _AUTO_PARALLEL_MIN_BYTES:
+            return 1
+    except OSError:
+        return 1
+    return cpu_count
+
+
+def _count_corpus_serial(
+    path: Path,
+    min_length: int,
+    max_length: int,
+    normalizer: NormalizerConfig | None = None,
+) -> dict[str, int]:
+    stats = _empty_stats()
+    with path.open("r", encoding="utf-8", errors="replace") as fh:
+        batch_stats = _count_lines(fh, min_length, max_length, normalizer)
+        _merge_stats(stats, batch_stats)
+    return stats
+
+
+def load_corpus(
+    path: str | Path,
+    min_length: int = 1,
+    max_length: int = 260,
+    max_lines: int = 0,
+    normalizer: NormalizerConfig | None = None,
+) -> Generator[str | CorpusItem, None, None]:
+    """Stream lines from a corpus file, applying filters.
+
+    Auto-detects JSONL format by file extension (``.jsonl``).  For JSONL
+    files, each yielded value is a :class:`CorpusItem` carrying provenance
+    metadata.  For plain-text files, plain ``str`` values are yielded
+    (backward compatible).
+
+    Args:
+        path: Path to corpus text file (UTF-8, one string per line) or a
+            ``.jsonl`` manifest.
+        min_length: Skip lines shorter than this many characters (after normalization).
+        max_length: Skip lines longer than this many characters (after normalization).
+        max_lines: Stop after yielding this many lines (0 = unlimited).
+        normalizer: Normalization config (uses defaults if None).
+
+    Yields:
+        Filtered, normalized text strings (``str`` for .txt, ``CorpusItem``
+        for .jsonl).
+    """
+    path = Path(path)
+    if not path.exists():
+        raise InputValidationError(f"Corpus file not found: {path}", code="corpus_not_found")
+
+    if _is_jsonl(path):
+        yield from _load_corpus_jsonl(path, min_length, max_length, max_lines, normalizer)
+        return
+
+    yielded = 0
+    with path.open("r", encoding="utf-8", errors="replace") as fh:
+        for raw_line in fh:
+            status, line = _classify_raw_line(raw_line, min_length, max_length, normalizer)
+            if status != _PASS or line is None:
+                continue
+
+            yield line
+            yielded += 1
+            if max_lines > 0 and yielded >= max_lines:
+                break
+
+
+def _load_corpus_jsonl(
+    path: Path,
+    min_length: int,
+    max_length: int,
+    max_lines: int,
+    normalizer: NormalizerConfig | None,
+) -> Generator[CorpusItem, None, None]:
+    """Load a JSONL corpus manifest, yielding :class:`CorpusItem` values."""
+    yielded = 0
+    with path.open("r", encoding="utf-8-sig", errors="replace") as fh:
+        for line_number, raw_line in enumerate(fh, start=1):
+            record = _parse_jsonl_record(raw_line, line_number, strict=True)
+            if record is None:
+                continue
+            text = record["text"]
+            status, normalized = _classify_jsonl_text(
+                text, min_length, max_length, normalizer
+            )
+            if status != _PASS or normalized is None:
+                continue
+            item = CorpusItem(
+                normalized,
+                source_id=record.get("source_id"),
+                source_group=record.get("source_group"),
+                language=record.get("language"),
+                script=record.get("script"),
+                domain=record.get("domain"),
+                source_uri=record.get("source_uri"),
+                license=record.get("license"),
+                retrieved_at=record.get("retrieved_at"),
+                source_sha256=record.get("source_sha256"),
+                normalization_policy=record.get("normalization_policy"),
+                consent_or_restrictions=record.get("consent_or_restrictions"),
+                raw_text=record.get("raw_text", text),
+            )
+            yield item
+            yielded += 1
+            if max_lines > 0 and yielded >= max_lines:
+                break
+
+
+def load_corpus_with_stats(
+    path: str | Path,
+    min_length: int = 1,
+    max_length: int = 260,
+    max_lines: int = 0,
+    normalizer: NormalizerConfig | None = None,
+    *,
+    strict_jsonl: bool = False,
+) -> tuple[list[str | CorpusItem], dict[str, int]]:
+    """Load a corpus and return ``(items, stats)`` in one pass.
+
+    Shares normalization, length filtering and line-separator rejection with
+    :func:`load_corpus`, but also reports why records were dropped.  This is
+    used by layout generation, which needs explicit rejection counts.
+
+    Stats keys:
+        ``total``: non-blank records considered (after ``max_lines`` cap)
+        ``passing``: items returned
+        ``too_short`` / ``too_long``: rejected by length
+        ``malformed``: JSONL line that was not a valid object with a string ``text``
+        ``separator``: record contained U+2028/U+2029
+        ``blank``: blank line
+
+    ``strict_jsonl=False`` (the default) skips malformed JSONL records instead
+    of raising, so a single bad line does not abort the whole run.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise InputValidationError(f"Corpus file not found: {path}", code="corpus_not_found")
+
+    stats = {
+        "total": 0,
+        "passing": 0,
+        "too_short": 0,
+        "too_long": 0,
+        "empty": 0,
+        "has_line_separator": 0,
+        "malformed": 0,
+        "separator": 0,
+        "blank": 0,
+    }
+    items: list[str | CorpusItem] = []
+    is_jsonl = _is_jsonl(path)
+    encoding = "utf-8-sig" if is_jsonl else "utf-8"
+
+    with path.open("r", encoding=encoding, errors="replace") as fh:
+        for line_number, raw_line in enumerate(fh, start=1):
+            if not raw_line.strip():
+                stats["blank"] += 1
+                continue
+
+            if is_jsonl:
+                record = _parse_jsonl_record(raw_line, line_number, strict=strict_jsonl)
+                if record is None:
+                    stats["malformed"] += 1
+                    continue
+                text = record["text"]
+                metadata = {
+                    "source_id": record.get("source_id"),
+                    "source_group": record.get("source_group"),
+                    "language": record.get("language"),
+                    "script": record.get("script"),
+                    "domain": record.get("domain"),
+                    "source_uri": record.get("source_uri"),
+                    "license": record.get("license"),
+                    "retrieved_at": record.get("retrieved_at"),
+                    "source_sha256": record.get("source_sha256"),
+                    "normalization_policy": record.get("normalization_policy"),
+                    "consent_or_restrictions": record.get("consent_or_restrictions"),
+                    "raw_text": record.get("raw_text", text),
+                }
+            else:
+                text = raw_line
+                metadata = {}
+            stats["total"] += 1
+            if line_separator_error(text):
+                stats["separator"] += 1
+                stats["has_line_separator"] += 1
+                continue
+            # _classify_raw_line strips, length-filters and normalizes in one go.
+            status, normalized = _classify_raw_line(text, min_length, max_length, normalizer)
+            if status == _SKIP:
+                # Separators were already rejected above, so this is a
+                # whitespace-only text value (e.g. JSONL {"text": "   "}).
+                stats["blank"] += 1
+                stats["empty"] += 1
+                continue
+            if status == _TOO_SHORT:
+                stats["too_short"] += 1
+                continue
+            if status == _TOO_LONG:
+                stats["too_long"] += 1
+                continue
+            if normalized is None:
+                continue
+
+            if is_jsonl:
+                item: str | CorpusItem = CorpusItem(normalized, **metadata)
+            else:
+                item = normalized
+            items.append(item)
+            stats["passing"] += 1
+            if max_lines > 0 and stats["passing"] >= max_lines:
+                break
+
+    return items, stats
+
+
+def count_corpus(
+    path: str | Path,
+    min_length: int = 1,
+    max_length: int = 260,
+    workers: int = 0,
+    normalizer: NormalizerConfig | None = None,
+) -> dict[str, int]:
+    """Scan the full corpus and return filter statistics.
+
+    Auto-detects JSONL format by file extension.  JSONL counting always
+    runs serially (records are not byte-splittable).
+
+    Returns a dict with keys:
+        total, passing, too_short, too_long
+    """
+    path = Path(path)
+    if not path.exists():
+        raise InputValidationError(f"Corpus file not found: {path}", code="corpus_not_found")
+
+    if _is_jsonl(path):
+        return _count_corpus_jsonl(path, min_length, max_length, normalizer)
+
+    resolved_workers = _resolve_worker_count(path, workers)
+    if resolved_workers <= 1:
+        return _count_corpus_serial(path, min_length, max_length, normalizer)
+
+    stats = _empty_stats()
+    try:
+        with mp.Pool(processes=resolved_workers) as pool:
+            ranges = list(_iter_chunk_ranges(path, resolved_workers))
+            for batch_stats in pool.starmap(
+                _count_file_chunk,
+                [(str(path), s, e, min_length, max_length, normalizer) for s, e in ranges],
+                chunksize=1,
+            ):
+                _merge_stats(stats, batch_stats)
+    except (OSError, PermissionError):
+        return _count_corpus_serial(path, min_length, max_length, normalizer)
+
+    return stats
+
+
+def _count_corpus_jsonl(
+    path: Path,
+    min_length: int,
+    max_length: int,
+    normalizer: NormalizerConfig | None,
+) -> dict[str, int]:
+    """Count passing/short/long records in a JSONL corpus manifest."""
+    stats = _empty_stats()
+    with path.open("r", encoding="utf-8-sig", errors="replace") as fh:
+        for line_number, raw_line in enumerate(fh, start=1):
+            record = _parse_jsonl_record(raw_line, line_number, strict=True)
+            if record is None:
+                continue
+            text = record["text"]
+            status, _ = _classify_jsonl_text(text, min_length, max_length, normalizer)
+            _update_stats_for_status(stats, status)
+    return stats
+
+
+def char_frequencies(
+    path: str | Path,
+    min_length: int = 1,
+    max_length: int = 260,
+    max_lines: int = 0,
+    normalizer: NormalizerConfig | None = None,
+) -> Counter[str]:
+    """Count character occurrences across the filtered, normalized corpus.
+
+    Used to identify rare characters for oversampling: lines that contain
+    infrequently-seen characters can be rendered more often so the model
+    sees enough examples of them.
+    """
+    freq: Counter[str] = Counter()
+    for line in load_corpus(
+        path,
+        min_length=min_length,
+        max_length=max_length,
+        max_lines=max_lines,
+        normalizer=normalizer,
+    ):
+        freq.update(line)
+    return freq
+
+
+def rare_chars_from_frequencies(freq: Counter[str], percentile: float) -> set[str]:
+    """Return the least-frequent *percentile* percent of distinct characters in *freq*.
+
+    Args:
+        freq: Character occurrence counts, e.g. from :func:`char_frequencies`.
+        percentile: Value in (0, 100]. 5.0 selects the least-frequent 5% of
+            distinct characters seen in the corpus (by character type, not
+            by occurrence count).
+
+    Returns:
+        The set of characters considered "rare". Empty if *freq* is empty
+        or *percentile* is 0.
+    """
+    if not freq or percentile <= 0:
+        return set()
+    percentile = min(100.0, percentile)
+    ordered = sorted(freq.items(), key=lambda kv: kv[1])
+    cutoff = max(1, round(len(ordered) * percentile / 100.0))
+    return {ch for ch, _ in ordered[:cutoff]}