ksbyte/data_prep.py

"""Offline data preparation: raw corpus -> normalized -> packed byte shards.

Run ONCE per (dataset, normalization-policy). Idempotent: if the shards and
meta.json already exist and `cfg.rebuild_data` is False, it is a no-op.

Output artifacts in `cfg.data_dir`:
  train.bin / val.bin / test.bin   uint16 token streams (0..255 bytes, 256=BOS,
                                   257=EOS); each document is  BOS + utf8 + EOS.
  meta.json                        vocab/specials, per-split token counts,
                                   bytes_per_word (for word-ppl), a frozen copy
                                   of the normalization policy, and provenance.

Source selection:
  * cfg.local_text_file set  -> read that UTF-8 file (newline-delimited docs).
  * else                     -> stream cfg.hf_dataset via `datasets`.

Splitting is deterministic and content-hash based (like ks_diacritizer), so the
same row always lands in the same split regardless of order or machine.
"""
from __future__ import annotations

import hashlib
import json
import os
from typing import Dict, Iterator, List, Optional, Tuple

import numpy as np

from .config import BOS_ID, EOS_ID, VOCAB_SIZE, ByteLMConfig
from .logging_utils import get_logger
from .metrics import ks_ratio as compute_ks_ratio
from .normalize import LMNormalizer, NormConfig, run_guards

SPLITS = ("train", "val", "test")


# --------------------------- source iteration ------------------------------- #
def _split_long_text(text: str, max_chars: int) -> Iterator[str]:
    """Yield document-sized pieces from HF rows that may be whole-corpus blobs.

    The published corpus currently arrives as a few very large strings rather
    than millions of rows. We split first on Kashmiri/Urdu sentence punctuation
    and then, for pathological long sentences, on word boundaries. This mirrors
    the local newline-doc path while keeping every emitted document under the
    normal `max_chars` guard.
    """
    text = text.strip()
    if not text:
        return
    if len(text) <= max_chars:
        yield text
        return

    buf: List[str] = []
    cur = ""
    boundaries = {"۔", "؟", "!", "."}
    for ch in text:
        cur += ch
        if ch in boundaries:
            sent = cur.strip()
            cur = ""
            if not sent:
                continue
            if sum(len(x) + 1 for x in buf) + len(sent) > max_chars and buf:
                yield " ".join(buf).strip()
                buf = []
            if len(sent) <= max_chars:
                buf.append(sent)
            else:
                words = sent.split()
                chunk = []
                n = 0
                for w in words:
                    if chunk and n + 1 + len(w) > max_chars:
                        yield " ".join(chunk)
                        chunk, n = [], 0
                    chunk.append(w)
                    n += len(w) + 1
                if chunk:
                    yield " ".join(chunk)
    tail = cur.strip()
    if tail:
        if sum(len(x) + 1 for x in buf) + len(tail) > max_chars and buf:
            yield " ".join(buf).strip()
            buf = []
        buf.append(tail)
    if buf:
        yield " ".join(buf).strip()


def _pick_text_column(columns: List[str], sample_row: dict) -> str:
    """Heuristic: the column whose sample value is the longest string."""
    best, best_len = None, -1
    for c in columns:
        v = sample_row.get(c)
        if isinstance(v, str) and len(v) > best_len:
            best, best_len = c, len(v)
    if best is None:
        raise ValueError(f"no string column found among {columns}")
    return best


def _iter_raw_docs(cfg: ByteLMConfig, logger) -> Iterator[Tuple[str, Optional[float]]]:
    """Yield (raw_text, ks_ratio_or_None) documents from the configured source."""
    if cfg.local_text_file:
        logger.info(f"reading local text file: {cfg.local_text_file}")
        if not os.path.exists(cfg.local_text_file):
            raise FileNotFoundError(cfg.local_text_file)
        with open(cfg.local_text_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    yield line, None
        return

    try:
        from datasets import load_dataset
    except ImportError as e:
        raise ImportError(
            "`datasets` is required for HF loading. `pip install datasets`, "
            "or set cfg.local_text_file to a local .txt corpus."
        ) from e

    logger.info(f"loading HF dataset: {cfg.hf_dataset} (rev={cfg.hf_revision})")
    ds = load_dataset(cfg.hf_dataset, revision=cfg.hf_revision)
    split_name = "train" if "train" in ds else list(ds.keys())[0]
    ds = ds[split_name]
    text_col = cfg.text_col
    if text_col == "auto":
        text_col = _pick_text_column(ds.column_names, ds[0])
        logger.info(f"auto-detected text column: {text_col!r}")
    has_ksr = "ks_ratio" in ds.column_names
    for row in ds:
        text = row.get(text_col)
        if isinstance(text, str):
            # Some HF exports pack many newline-delimited documents into a few
            # giant rows/files. Treat those lines like the local .txt path so
            # max_chars filtering does not drop the whole corpus as one row.
            for line in text.splitlines():
                for doc in _split_long_text(line, cfg.max_chars):
                    yield doc, (float(row["ks_ratio"]) if has_ksr else None)


# ------------------------------ encoding ------------------------------------ #
def _encode_doc(text: str) -> np.ndarray:
    """utf-8 bytes wrapped with BOS/EOS, as uint16."""
    raw = text.encode("utf-8")
    arr = np.empty(len(raw) + 2, dtype=np.uint16)
    arr[0] = BOS_ID
    arr[1:-1] = np.frombuffer(raw, dtype=np.uint8)
    arr[-1] = EOS_ID
    return arr


def _split_of(text: str, cfg: ByteLMConfig) -> str:
    """Deterministic content-hash split assignment."""
    h = hashlib.sha1(f"{cfg.split_seed}:{text}".encode("utf-8")).digest()
    frac = int.from_bytes(h[:8], "big") / float(1 << 64)
    if frac < cfg.val_frac:
        return "val"
    if frac < cfg.val_frac + cfg.test_frac:
        return "test"
    return "train"


# ------------------------------ entrypoint ---------------------------------- #
def _meta_path(cfg: ByteLMConfig) -> str:
    return os.path.join(cfg.data_dir, "meta.json")


def shards_exist(cfg: ByteLMConfig) -> bool:
    if not os.path.exists(_meta_path(cfg)):
        return False
    return all(os.path.exists(os.path.join(cfg.data_dir, f"{s}.bin")) for s in SPLITS)


def load_meta(cfg: ByteLMConfig) -> dict:
    with open(_meta_path(cfg), "r", encoding="utf-8") as f:
        return json.load(f)


def prepare_data(cfg: ByteLMConfig, logger=None) -> dict:
    """Build (or reuse) the packed byte shards. Returns the meta dict."""
    logger = logger or get_logger("ksbyte.data_prep", cfg.run_dir)
    cfg.validate()

    if shards_exist(cfg) and not cfg.rebuild_data:
        meta = load_meta(cfg)
        logger.info(f"data shards already present in {cfg.data_dir} "
                    f"(train={meta['counts']['train']:,} tokens) — skipping prep")
        return meta

    # Guard the vendored normalizer BEFORE touching any text.
    guard_report = run_guards(strict=True)
    logger.info(f"normalizer guard OK ({guard_report['protected_count']} protected letters; "
                f"tatweel->{guard_report['tatweel_maps_to']}, zwnj(raw)->{guard_report['zwnj_maps_to']})")

    normalizer = LMNormalizer(NormConfig(
        zwnj_policy=cfg.zwnj_policy,
        digit_policy=cfg.digit_policy,
        remove_diacritics=cfg.remove_diacritics,
    ))

    os.makedirs(cfg.data_dir, exist_ok=True)
    buffers: Dict[str, List[np.ndarray]] = {s: [] for s in SPLITS}
    seen_hashes: set = set()
    stats = {"raw": 0, "kept": 0, "dropped_short": 0, "dropped_ksr": 0,
             "dropped_dup": 0, "content_bytes": 0, "words": 0}

    for raw_text, ksr in _iter_raw_docs(cfg, logger):
        stats["raw"] += 1
        text = normalizer(raw_text)
        if not (cfg.min_chars <= len(text) <= cfg.max_chars):
            stats["dropped_short"] += 1
            continue
        ratio = ksr if ksr is not None else compute_ks_ratio(text)
        if not cfg.keep_mixed_script and ratio < cfg.min_ks_ratio:
            stats["dropped_ksr"] += 1
            continue
        if cfg.dedup:
            hh = hashlib.sha1(text.encode("utf-8")).digest()
            if hh in seen_hashes:
                stats["dropped_dup"] += 1
                continue
            seen_hashes.add(hh)

        split = _split_of(text, cfg)
        buffers[split].append(_encode_doc(text))
        stats["kept"] += 1
        stats["content_bytes"] += len(text.encode("utf-8"))
        stats["words"] += len(text.split())
        if stats["raw"] % 50_000 == 0:
            logger.info(f"  processed {stats['raw']:,} raw docs (kept {stats['kept']:,})")

    counts = {}
    for split in SPLITS:
        arrs = buffers[split]
        stream = (np.concatenate(arrs) if arrs else np.empty(0, dtype=np.uint16))
        path = os.path.join(cfg.data_dir, f"{split}.bin")
        stream.tofile(path)
        counts[split] = int(stream.size)
        logger.info(f"  wrote {split}.bin: {stream.size:,} tokens -> {path}")

    if counts["train"] == 0:
        raise RuntimeError("train split is empty after preparation — check filters/source")

    bytes_per_word = (stats["content_bytes"] / stats["words"]) if stats["words"] else float("nan")
    meta = {
        "vocab_size": VOCAB_SIZE,
        "bos_id": BOS_ID,
        "eos_id": EOS_ID,
        "counts": counts,
        "bytes_per_word": bytes_per_word,
        "stats": stats,
        "normalization": {"zwnj_policy": cfg.zwnj_policy,
                          "digit_policy": cfg.digit_policy,
                          "remove_diacritics": cfg.remove_diacritics},
        "source": (cfg.local_text_file or cfg.hf_dataset),
        "dtype": "uint16",
    }
    with open(_meta_path(cfg), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)
    logger.info(f"data prep done: kept {stats['kept']:,}/{stats['raw']:,} docs, "
                f"bytes/word={bytes_per_word:.2f}, meta -> {_meta_path(cfg)}")
    return meta


if __name__ == "__main__":
    import argparse

    p = argparse.ArgumentParser(description="Prepare ks_byte_lm data shards")
    p.add_argument("--local_text_file", default=None)
    p.add_argument("--data_dir", default="data")
    p.add_argument("--rebuild", action="store_true")
    args = p.parse_args()
    cfg = ByteLMConfig().merge({
        "local_text_file": args.local_text_file,
        "data_dir": args.data_dir,
        "rebuild_data": args.rebuild,
    })
    prepare_data(cfg)