src/korbyte/corpus.py

"""Streaming, deterministic, privacy-conscious corpus preparation."""

from __future__ import annotations

import hashlib
import json
import re
import sys
from collections.abc import Iterable, Iterator
from dataclasses import replace
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from datasets import load_dataset

from .config import DEFAULT_SEED, DEFAULT_SOURCES, Paths, SourceSpec

EMAIL_RE = re.compile(r"(?<![\w.+-])[\w.+-]+@[\w.-]+\.[A-Za-z]{2,}(?![\w.-])")
URL_RE = re.compile(r"(?i)\b(?:https?://|www\.)\S+")
LONG_NUMBER_RE = re.compile(r"(?<!\d)(?:\d[\s-]?){8,}\d(?!\d)")
CONTROL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
SPACE_RE = re.compile(r"[ \t\f\v]+")
HANGUL_RE = re.compile(r"[\u1100-\u11ff\u3130-\u318f\uac00-\ud7af]")
LATIN_RE = re.compile(r"[A-Za-z]")
LETTER_RE = re.compile(r"[^\W\d_]", re.UNICODE)


def sanitize_text(text: str) -> str:
    """Remove control bytes and redact obvious contact/identifier patterns."""

    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = CONTROL_RE.sub("", text)
    text = EMAIL_RE.sub("<EMAIL>", text)
    text = URL_RE.sub("<URL>", text)
    return LONG_NUMBER_RE.sub("<LONG_NUMBER>", text)


def script_ratio(text: str, target_script: str) -> float:
    """Return the target-script share among Unicode letters."""

    letters = len(LETTER_RE.findall(text))
    if not letters:
        return 0.0
    if target_script == "hangul":
        return len(HANGUL_RE.findall(text)) / letters
    if target_script == "latin":
        return len(LATIN_RE.findall(text)) / letters
    return 1.0


def iter_chunks(text: str, *, max_chars: int = 8_192) -> Iterator[str]:
    """Yield bounded lines, preferring paragraph and whitespace boundaries."""

    for paragraph in text.split("\n"):
        paragraph = SPACE_RE.sub(" ", paragraph).strip()
        if len(paragraph) < 24:
            continue
        start = 0
        while start < len(paragraph):
            end = min(start + max_chars, len(paragraph))
            if end < len(paragraph):
                boundary = paragraph.rfind(" ", start + max_chars // 2, end)
                if boundary > start:
                    end = boundary
            chunk = paragraph[start:end].strip()
            if len(chunk) >= 24:
                yield chunk
            start = max(end, start + 1)


def _source_iterator(spec: SourceSpec, *, seed: int) -> Iterable[dict[str, Any]]:
    dataset = load_dataset(
        spec.dataset_id,
        spec.config_name,
        split=spec.split,
        revision=spec.revision,
        streaming=True,
    )
    if spec.shuffle_buffer:
        dataset = dataset.shuffle(seed=seed, buffer_size=spec.shuffle_buffer)
    return dataset


def prepare_corpus(
    root: Path,
    *,
    scale: float = 1.0,
    seed: int = DEFAULT_SEED,
    sources: tuple[SourceSpec, ...] = DEFAULT_SOURCES,
) -> dict[str, Any]:
    """Stream public corpora into a filtered local training file and manifest."""

    if not 0 < scale <= 1:
        raise ValueError("scale must be in the interval (0, 1]")
    paths = Paths(root)
    paths.corpus.parent.mkdir(parents=True, exist_ok=True)
    paths.corpus_manifest.parent.mkdir(parents=True, exist_ok=True)

    seen: set[bytes] = set()
    source_results: list[dict[str, Any]] = []
    corpus_digest = hashlib.sha256()
    total_characters = 0
    total_lines = 0

    with paths.corpus.open("w", encoding="utf-8", newline="\n") as output:
        for source_index, original_spec in enumerate(sources):
            budget = max(100_000, round(original_spec.character_budget * scale))
            spec = replace(original_spec, character_budget=budget)
            accepted_characters = 0
            accepted_lines = 0
            inspected_documents = 0

            for row in _source_iterator(spec, seed=seed + source_index):
                inspected_documents += 1
                raw_text = row.get(spec.text_field)
                if not isinstance(raw_text, str):
                    continue
                for chunk in iter_chunks(sanitize_text(raw_text)):
                    if script_ratio(chunk, spec.target_script) < spec.min_target_script_ratio:
                        continue
                    digest = hashlib.blake2b(chunk.encode("utf-8"), digest_size=16).digest()
                    if digest in seen:
                        continue
                    seen.add(digest)
                    encoded = f"{chunk}\n".encode()
                    output.write(encoded.decode())
                    corpus_digest.update(encoded)
                    accepted_characters += len(chunk)
                    accepted_lines += 1
                    total_characters += len(chunk)
                    total_lines += 1
                    if accepted_characters >= budget:
                        break
                if accepted_characters >= budget:
                    break

            source_results.append(
                {
                    **spec.as_dict(),
                    "accepted_characters": accepted_characters,
                    "accepted_lines": accepted_lines,
                    "inspected_documents": inspected_documents,
                }
            )
            print(
                f"prepared {spec.key}: {accepted_characters:,} chars, {accepted_lines:,} lines",
                file=sys.stderr,
            )

    manifest = {
        "schema_version": 1,
        "created_at": datetime.now(UTC).isoformat(),
        "seed": seed,
        "scale": scale,
        "total_characters": total_characters,
        "total_lines": total_lines,
        "sha256": corpus_digest.hexdigest(),
        "sources": source_results,
        "redactions": ["email", "url", "long-number"],
        "redistributed_training_text": False,
    }
    paths.corpus_manifest.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return manifest